fix: BED-5080 - move improper int4 casts to int8 (#989)

SpecterOps · Dec 3, 2024 · 79cc344 · 79cc344
1 parent e8ea43f
commit 79cc344
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/cmd/api/src/api/tools/pg.go b/cmd/api/src/api/tools/pg.go
@@ -236,6 +236,8 @@ func (s *PGMigrator) SwitchPostgreSQL(response http.ResponseWriter, request *htt
 		api.WriteJSONResponse(request.Context(), map[string]any{
 			"error": fmt.Errorf("failed connecting to PostgreSQL: %w", err),
 		}, http.StatusInternalServerError, response)
+	} else if err := pgDB.AssertSchema(request.Context(), s.graphSchema); err != nil {
+		log.Errorf("Unable to assert graph schema in PostgreSQL: %v", err)
 	} else if err := SetGraphDriver(request.Context(), s.cfg, pg.DriverName); err != nil {
 		api.WriteJSONResponse(request.Context(), map[string]any{
 			"error": fmt.Errorf("failed updating graph database driver preferences: %w", err),

diff --git a/packages/go/dawgs/drivers/pg/query/format.go b/packages/go/dawgs/drivers/pg/query/format.go
@@ -122,7 +122,7 @@ func FormatNodeUpsert(graphTarget model.Graph, identityProperties []string) stri
 	return join(
 		"insert into ", graphTarget.Partitions.Node.Name, " as n ",
 		"(graph_id, kind_ids, properties) ",
-		"select $1::int4, unnest($2::text[])::int2[], unnest($3::jsonb[]) ",
+		"select $1, unnest($2::text[])::int2[], unnest($3::jsonb[]) ",
 		formatConflictMatcher(identityProperties, "id, graph_id"),
 		"do update set properties = n.properties || excluded.properties, kind_ids = uniq(sort(n.kind_ids || excluded.kind_ids)) ",
 		"returning id;",
@@ -132,7 +132,7 @@ func FormatNodeUpsert(graphTarget model.Graph, identityProperties []string) stri
 func FormatRelationshipPartitionUpsert(graphTarget model.Graph, identityProperties []string) string {
 	return join("insert into ", graphTarget.Partitions.Edge.Name, " as e ",
 		"(graph_id, start_id, end_id, kind_id, properties) ",
-		"select $1::int4, unnest($2::int4[]), unnest($3::int4[]), unnest($4::int2[]), unnest($5::jsonb[]) ",
+		"select $1, unnest($2::int8[]), unnest($3::int8[]), unnest($4::int2[]), unnest($5::jsonb[]) ",
 		formatConflictMatcher(identityProperties, "graph_id, start_id, end_id, kind_id"),
 		"do update set properties = e.properties || excluded.properties;",
 	)

diff --git a/packages/go/dawgs/drivers/pg/statements.go b/packages/go/dawgs/drivers/pg/statements.go
@@ -19,7 +19,7 @@ package pg
 const (
 	createNodeStatement               = `insert into node (graph_id, kind_ids, properties) values (@graph_id, @kind_ids, @properties) returning (id, kind_ids, properties)::nodeComposite;`
 	createNodeWithoutIDBatchStatement = `insert into node (graph_id, kind_ids, properties) select $1, unnest($2::text[])::int2[], unnest($3::jsonb[])`
-	createNodeWithIDBatchStatement    = `insert into node (graph_id, id, kind_ids, properties) select $1, unnest($2::int4[]), unnest($3::text[])::int2[], unnest($4::jsonb[])`
+	createNodeWithIDBatchStatement    = `insert into node (graph_id, id, kind_ids, properties) select $1, unnest($2::int8[]), unnest($3::text[])::int2[], unnest($4::jsonb[])`
 	deleteNodeWithIDStatement         = `delete from node where node.id = any($1)`
 
 	createEdgeStatement = `insert into edge (graph_id, start_id, end_id, kind_id, properties) values (@graph_id, @start_id, @end_id, @kind_id, @properties) returning (id, start_id, end_id, kind_id, properties)::edgeComposite;`
@@ -28,7 +28,7 @@ const (
 	//	     Azure post-processing. This was done because Azure post will submit the same creation request hundreds of
 	// 		 times for the same edge. In PostgreSQL this results in a constraint violation. For now this is best-effort
 	//		 until Azure post-processing can be refactored.
-	createEdgeBatchStatement  = `insert into edge as e (graph_id, start_id, end_id, kind_id, properties) select $1::int4, unnest($2::int4[]), unnest($3::int4[]), unnest($4::int2[]), unnest($5::jsonb[]) on conflict (graph_id, start_id, end_id, kind_id) do update set properties = e.properties || excluded.properties;`
+	createEdgeBatchStatement  = `insert into edge as e (graph_id, start_id, end_id, kind_id, properties) select $1, unnest($2::int8[]), unnest($3::int8[]), unnest($4::int2[]), unnest($5::jsonb[]) on conflict (graph_id, start_id, end_id, kind_id) do update set properties = e.properties || excluded.properties;`
 	deleteEdgeWithIDStatement = `delete from edge as e where e.id = any($1)`
 
 	edgePropertySetOnlyStatement      = `update edge set properties = properties || $1::jsonb where edge.id = $2`