diff --git a/drivers/pg/driver.go b/drivers/pg/driver.go index 5b8adb1b..db757f57 100644 --- a/drivers/pg/driver.go +++ b/drivers/pg/driver.go @@ -3,6 +3,7 @@ package pg import ( "context" "fmt" + "strings" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" @@ -186,3 +187,127 @@ func (s *Driver) OptimizeStorage(ctx context.Context) error { return optimizeStorage(ctx, conn) } + +// resolveKindIDs maps kinds to their integer IDs, refreshing the schema cache once on a miss. It returns the resolved +// IDs alongside any kinds that remain undefined after the refresh, so callers can decide whether an unresolved kind is +// a tolerable no-op (include predicates) or must fail closed (exclude predicates). +func (s *Driver) resolveKindIDs(ctx context.Context, kinds graph.Kinds) ([]int16, graph.Kinds, error) { + if len(kinds) == 0 { + return nil, nil, nil + } + + s.lock.RLock() + if kindIDs, missingKinds := s.mapKinds(kinds); len(missingKinds) == 0 { + s.lock.RUnlock() + return kindIDs, nil, nil + } + s.lock.RUnlock() + + s.lock.Lock() + defer s.lock.Unlock() + + if err := s.Fetch(ctx); err != nil { + return nil, nil, err + } + + kindIDs, missingKinds := s.mapKinds(kinds) + return kindIDs, missingKinds, nil +} + +// DeleteNodesByKinds performs a server-side, set-based delete of nodes using the kind_ids GIN index instead of +// streaming node IDs through the application. A node is deleted when its kind_ids overlap includeAny (or, when +// includeAny is empty, for every node) and do not overlap excludeAny. Deleting nodes fires the delete_node_edges +// trigger, cascading the attached edge deletes. +// +// includeAny is mapped to kind IDs tolerantly: include kinds that are not defined in the database map to no IDs and +// therefore match no nodes, so a request that targets only undefined kinds is a safe no-op rather than an accidental +// full delete. excludeAny is mapped fail-closed: if any exclude kind is undefined the delete is refused, because +// silently dropping an exclusion would widen the delete and could remove protected nodes (e.g. an unresolved +// MigrationData would turn a guarded wipe into an unguarded delete from node). +func (s *Driver) DeleteNodesByKinds(ctx context.Context, includeAny graph.Kinds, excludeAny graph.Kinds) error { + includeIDs, _, err := s.resolveKindIDs(ctx, includeAny) + if err != nil { + return err + } + + excludeIDs, excludeMissing, err := s.resolveKindIDs(ctx, excludeAny) + if err != nil { + return err + } + if len(excludeMissing) > 0 { + return fmt.Errorf("cannot exclude undefined kinds from node delete: %v", excludeMissing) + } + + statement, arguments := buildNodeDeleteStatement(len(includeAny) > 0, includeIDs, excludeIDs) + + conn, err := s.pool.Acquire(ctx) + if err != nil { + return fmt.Errorf("acquire connection for node delete: %w", err) + } + defer conn.Release() + + if _, err := conn.Exec(ctx, statement, arguments...); err != nil { + return fmt.Errorf("%s: %w", statement, err) + } + + return nil +} + +// buildNodeDeleteStatement renders the node delete statement and its positional arguments for the given resolved kind +// IDs. The include predicate is emitted whenever an include filter was requested (includeRequested), even if includeIDs +// is empty, so that targeting only undefined kinds matches no nodes. The exclude predicate is emitted only when +// excludeIDs is non-empty, so an unresolved exclusion can never widen the delete into an unguarded wipe. +func buildNodeDeleteStatement(includeRequested bool, includeIDs []int16, excludeIDs []int16) (string, []any) { + var ( + predicates []string + arguments []any + ) + + if includeRequested { + arguments = append(arguments, includeIDs) + predicates = append(predicates, fmt.Sprintf("kind_ids operator (pg_catalog.&&) $%d::int2[]", len(arguments))) + } + + if len(excludeIDs) > 0 { + arguments = append(arguments, excludeIDs) + predicates = append(predicates, fmt.Sprintf("not (kind_ids operator (pg_catalog.&&) $%d::int2[])", len(arguments))) + } + + statement := "delete from node" + if len(predicates) > 0 { + statement += " where " + strings.Join(predicates, " and ") + } + + return statement, arguments +} + +// DeleteRelationshipsByKinds performs a server-side, set-based delete of relationships whose kind_id matches any of +// the given kinds, using the edge_kind_id_id_start_id_end_id_index covering index instead of streaming relationship +// IDs through the application. +// +// kinds are mapped to kind IDs tolerantly: kinds that are not defined in the database map to no IDs. An empty kinds +// argument, or one that maps entirely to undefined kinds, deletes nothing rather than every relationship. +func (s *Driver) DeleteRelationshipsByKinds(ctx context.Context, kinds graph.Kinds) error { + if len(kinds) == 0 { + return nil + } + + kindIDs, _, err := s.resolveKindIDs(ctx, kinds) + if err != nil { + return err + } + + const statement = "delete from edge where kind_id = any($1::int2[])" + + conn, err := s.pool.Acquire(ctx) + if err != nil { + return fmt.Errorf("acquire connection for relationship delete: %w", err) + } + defer conn.Release() + + if _, err := conn.Exec(ctx, statement, kindIDs); err != nil { + return fmt.Errorf("%s: %w", statement, err) + } + + return nil +} diff --git a/drivers/pg/driver_test.go b/drivers/pg/driver_test.go new file mode 100644 index 00000000..65c30285 --- /dev/null +++ b/drivers/pg/driver_test.go @@ -0,0 +1,100 @@ +package pg + +import ( + "context" + "testing" + + "github.com/specterops/dawgs/graph" + "github.com/stretchr/testify/require" +) + +// TestBuildNodeDeleteStatement covers the statement/argument construction for DeleteNodesByKinds, including the guard +// that prevents an unresolved exclusion from widening the delete into an unguarded wipe. +func TestBuildNodeDeleteStatement(t *testing.T) { + var ( + includeIDs = []int16{1, 2} + excludeIDs = []int16{9} + ) + + t.Run("no filters deletes all nodes", func(t *testing.T) { + statement, arguments := buildNodeDeleteStatement(false, nil, nil) + require.Equal(t, "delete from node", statement) + require.Empty(t, arguments) + }) + + t.Run("include only", func(t *testing.T) { + statement, arguments := buildNodeDeleteStatement(true, includeIDs, nil) + require.Equal(t, "delete from node where kind_ids operator (pg_catalog.&&) $1::int2[]", statement) + require.Equal(t, []any{includeIDs}, arguments) + }) + + t.Run("exclude only", func(t *testing.T) { + statement, arguments := buildNodeDeleteStatement(false, nil, excludeIDs) + require.Equal(t, "delete from node where not (kind_ids operator (pg_catalog.&&) $1::int2[])", statement) + require.Equal(t, []any{excludeIDs}, arguments) + }) + + t.Run("include and exclude are positionally numbered", func(t *testing.T) { + statement, arguments := buildNodeDeleteStatement(true, includeIDs, excludeIDs) + require.Equal(t, "delete from node where kind_ids operator (pg_catalog.&&) $1::int2[] and not (kind_ids operator (pg_catalog.&&) $2::int2[])", statement) + require.Equal(t, []any{includeIDs, excludeIDs}, arguments) + }) + + t.Run("empty excludeIDs cannot widen the delete", func(t *testing.T) { + // A requested-but-unresolved exclusion must never emit a not(... && '{}') clause that matches every row. + statement, arguments := buildNodeDeleteStatement(false, nil, []int16{}) + require.Equal(t, "delete from node", statement) + require.Empty(t, arguments) + + statement, arguments = buildNodeDeleteStatement(true, includeIDs, []int16{}) + require.Equal(t, "delete from node where kind_ids operator (pg_catalog.&&) $1::int2[]", statement) + require.Equal(t, []any{includeIDs}, arguments) + }) + + t.Run("include requested with empty IDs is a tolerant no-op predicate", func(t *testing.T) { + statement, arguments := buildNodeDeleteStatement(true, []int16{}, nil) + require.Equal(t, "delete from node where kind_ids operator (pg_catalog.&&) $1::int2[]", statement) + require.Equal(t, []any{[]int16{}}, arguments) + }) +} + +// TestResolveKindIDsDefinedFastPath exercises the cache-hit path of resolveKindIDs, which resolves defined kinds +// without touching the database. The cache-miss/refresh and fail-closed exclude paths require a live pool and are +// covered by the integration suite. +func TestResolveKindIDsDefinedFastPath(t *testing.T) { + ctx := context.Background() + + driver := &Driver{SchemaManager: NewSchemaManager(nil, 0)} + + var ( + userKind = graph.StringKind("User") + groupKind = graph.StringKind("Group") + ) + driver.kindsByID[userKind] = 1 + driver.kindsByID[groupKind] = 2 + + t.Run("defined kinds resolve with no missing", func(t *testing.T) { + ids, missing, err := driver.resolveKindIDs(ctx, graph.Kinds{userKind, groupKind}) + require.NoError(t, err) + require.Empty(t, missing) + require.ElementsMatch(t, []int16{1, 2}, ids) + }) + + t.Run("empty kinds short-circuit", func(t *testing.T) { + ids, missing, err := driver.resolveKindIDs(ctx, nil) + require.NoError(t, err) + require.Nil(t, ids) + require.Nil(t, missing) + }) +} + +// TestDeleteRelationshipsByKindsEmptyIsNoop verifies that an empty kinds request returns before acquiring a +// connection, so it is a safe no-op rather than deleting every relationship. +func TestDeleteRelationshipsByKindsEmptyIsNoop(t *testing.T) { + ctx := context.Background() + + driver := &Driver{SchemaManager: NewSchemaManager(nil, 0)} + + require.NoError(t, driver.DeleteRelationshipsByKinds(ctx, nil)) + require.NoError(t, driver.DeleteRelationshipsByKinds(ctx, graph.Kinds{})) +} diff --git a/integration/pgsql_delete_by_kind_test.go b/integration/pgsql_delete_by_kind_test.go new file mode 100644 index 00000000..917db51b --- /dev/null +++ b/integration/pgsql_delete_by_kind_test.go @@ -0,0 +1,166 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +//go:build manual_integration + +package integration + +import ( + "context" + "os" + "testing" + + "github.com/specterops/dawgs/drivers/pg" + "github.com/specterops/dawgs/graph" +) + +// nodesByKindDeleter mirrors the capability the BloodHound delete path detects on the PostgreSQL driver. +type nodesByKindDeleter interface { + DeleteNodesByKinds(ctx context.Context, includeAny graph.Kinds, excludeAny graph.Kinds) error +} + +// TestPostgreSQLDeleteNodesByKinds verifies the server-side, set-based node delete: includeAny restricts the delete to +// nodes carrying one of the listed kinds, excludeAny protects nodes carrying one of the listed kinds, undefined include +// kinds are a safe no-op while undefined exclude kinds fail closed, and deleting nodes cascades incident edges. +func TestPostgreSQLDeleteNodesByKinds(t *testing.T) { + connStr := os.Getenv("CONNECTION_STRING") + if connStr == "" { + t.Skip("CONNECTION_STRING env var is not set") + } + + driver, err := DriverFromConnectionString(connStr) + if err != nil { + t.Fatalf("failed to detect driver: %v", err) + } + if driver != pg.DriverName { + t.Skip("CONNECTION_STRING is not a PostgreSQL connection string") + } + + var ( + kindA = graph.StringKind("DeleteByKindA") + kindB = graph.StringKind("DeleteByKindB") + edgeKind = graph.StringKind("DeleteByKindEdge") + missing = graph.StringKind("DeleteByKindMissing") + db, ctx = SetupDBWithKinds(t, CleanupGraph, graph.Kinds{kindA, kindB}, graph.Kinds{edgeKind}) + ) + + deleter, hasCapability := graph.AsDriver[nodesByKindDeleter](db) + if !hasCapability { + t.Fatal("PostgreSQL driver does not implement DeleteNodesByKinds") + } + + // fixture creates two kindA nodes, two kindB nodes, and an A->B edge that must cascade when its start is deleted. + createFixture := func() { + if err := db.WriteTransaction(ctx, func(tx graph.Transaction) error { + a0, err := tx.CreateNode(graph.NewProperties(), kindA) + if err != nil { + return err + } + if _, err := tx.CreateNode(graph.NewProperties(), kindA); err != nil { + return err + } + + b0, err := tx.CreateNode(graph.NewProperties(), kindB) + if err != nil { + return err + } + if _, err := tx.CreateNode(graph.NewProperties(), kindB); err != nil { + return err + } + + _, err = tx.CreateRelationshipByIDs(a0.ID, b0.ID, edgeKind, graph.NewProperties()) + return err + }); err != nil { + t.Fatalf("failed to create delete-by-kind fixture: %v", err) + } + } + + t.Run("includeAny deletes matching kinds and cascades edges", func(t *testing.T) { + createFixture() + + if err := deleter.DeleteNodesByKinds(ctx, graph.Kinds{kindA}, nil); err != nil { + t.Fatalf("DeleteNodesByKinds(include kindA) failed: %v", err) + } + + if count := countByCypher(t, ctx, db, "MATCH (n:DeleteByKindA) RETURN count(n)"); count != 0 { + t.Fatalf("kindA node count: got %d, want 0", count) + } + if count := countByCypher(t, ctx, db, "MATCH (n:DeleteByKindB) RETURN count(n)"); count != 2 { + t.Fatalf("kindB node count: got %d, want 2", count) + } + if count := countByCypher(t, ctx, db, "MATCH ()-[r:DeleteByKindEdge]->() RETURN count(r)"); count != 0 { + t.Fatalf("edge count after cascade: got %d, want 0", count) + } + + cleanupAll(t, ctx, deleter) + }) + + t.Run("excludeAny protects matching kinds", func(t *testing.T) { + createFixture() + + // Delete every node except those carrying kindB. + if err := deleter.DeleteNodesByKinds(ctx, nil, graph.Kinds{kindB}); err != nil { + t.Fatalf("DeleteNodesByKinds(exclude kindB) failed: %v", err) + } + + if count := countByCypher(t, ctx, db, "MATCH (n:DeleteByKindA) RETURN count(n)"); count != 0 { + t.Fatalf("kindA node count: got %d, want 0", count) + } + if count := countByCypher(t, ctx, db, "MATCH (n:DeleteByKindB) RETURN count(n)"); count != 2 { + t.Fatalf("kindB node count: got %d, want 2", count) + } + + cleanupAll(t, ctx, deleter) + }) + + t.Run("undefined include kinds are a safe no-op", func(t *testing.T) { + createFixture() + + if err := deleter.DeleteNodesByKinds(ctx, graph.Kinds{missing}, nil); err != nil { + t.Fatalf("DeleteNodesByKinds(include missing) failed: %v", err) + } + + if count := countByCypher(t, ctx, db, "MATCH (n) RETURN count(n)"); count != 4 { + t.Fatalf("node count after no-op delete: got %d, want 4", count) + } + + cleanupAll(t, ctx, deleter) + }) + + t.Run("undefined exclude kinds fail closed and delete nothing", func(t *testing.T) { + createFixture() + + // An unresolved exclusion would otherwise collapse to an unguarded delete; the driver must refuse instead. + if err := deleter.DeleteNodesByKinds(ctx, nil, graph.Kinds{missing}); err == nil { + t.Fatal("DeleteNodesByKinds(exclude missing) succeeded, want error") + } + + if count := countByCypher(t, ctx, db, "MATCH (n) RETURN count(n)"); count != 4 { + t.Fatalf("node count after failed delete: got %d, want 4", count) + } + + cleanupAll(t, ctx, deleter) + }) +} + +// cleanupAll removes every node between subtests so each starts from an empty graph. +func cleanupAll(t *testing.T, ctx context.Context, deleter nodesByKindDeleter) { + t.Helper() + + if err := deleter.DeleteNodesByKinds(ctx, nil, nil); err != nil { + t.Fatalf("failed to clean up nodes between subtests: %v", err) + } +} diff --git a/integration/pgsql_delete_relationships_by_kind_test.go b/integration/pgsql_delete_relationships_by_kind_test.go new file mode 100644 index 00000000..716502de --- /dev/null +++ b/integration/pgsql_delete_relationships_by_kind_test.go @@ -0,0 +1,147 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +//go:build manual_integration + +package integration + +import ( + "context" + "os" + "testing" + + "github.com/specterops/dawgs/drivers/pg" + "github.com/specterops/dawgs/graph" +) + +// relationshipsByKindDeleter mirrors the capability the BloodHound delete path detects on the PostgreSQL driver. +type relationshipsByKindDeleter interface { + DeleteRelationshipsByKinds(ctx context.Context, kinds graph.Kinds) error +} + +// TestPostgreSQLDeleteRelationshipsByKinds verifies the server-side, set-based relationship delete: the listed kinds +// restrict the delete to relationships carrying one of those kinds, nodes are left intact, multiple kinds are unioned, +// undefined kinds are a safe no-op, and an empty request deletes nothing. +func TestPostgreSQLDeleteRelationshipsByKinds(t *testing.T) { + connStr := os.Getenv("CONNECTION_STRING") + if connStr == "" { + t.Skip("CONNECTION_STRING env var is not set") + } + + driver, err := DriverFromConnectionString(connStr) + if err != nil { + t.Fatalf("failed to detect driver: %v", err) + } + if driver != pg.DriverName { + t.Skip("CONNECTION_STRING is not a PostgreSQL connection string") + } + + var ( + nodeKind = graph.StringKind("RelByKindNode") + edgeKindA = graph.StringKind("RelByKindEdgeA") + edgeKindB = graph.StringKind("RelByKindEdgeB") + missing = graph.StringKind("RelByKindMissing") + db, ctx = SetupDBWithKinds(t, CleanupGraph, graph.Kinds{nodeKind}, graph.Kinds{edgeKindA, edgeKindB}) + ) + + deleter, hasCapability := graph.AsDriver[relationshipsByKindDeleter](db) + if !hasCapability { + t.Fatal("PostgreSQL driver does not implement DeleteRelationshipsByKinds") + } + + // fixture creates three nodes joined by two edgeKindA edges and one edgeKindB edge. Nodes must survive every delete. + createFixture := func() { + if err := db.WriteTransaction(ctx, func(tx graph.Transaction) error { + n0, err := tx.CreateNode(graph.NewProperties(), nodeKind) + if err != nil { + return err + } + n1, err := tx.CreateNode(graph.NewProperties(), nodeKind) + if err != nil { + return err + } + n2, err := tx.CreateNode(graph.NewProperties(), nodeKind) + if err != nil { + return err + } + + if _, err := tx.CreateRelationshipByIDs(n0.ID, n1.ID, edgeKindA, graph.NewProperties()); err != nil { + return err + } + if _, err := tx.CreateRelationshipByIDs(n1.ID, n2.ID, edgeKindA, graph.NewProperties()); err != nil { + return err + } + _, err = tx.CreateRelationshipByIDs(n0.ID, n2.ID, edgeKindB, graph.NewProperties()) + return err + }); err != nil { + t.Fatalf("failed to create delete-relationships-by-kind fixture: %v", err) + } + } + + t.Run("deletes relationships of the given kind and leaves nodes", func(t *testing.T) { + createFixture() + + if err := deleter.DeleteRelationshipsByKinds(ctx, graph.Kinds{edgeKindA}); err != nil { + t.Fatalf("DeleteRelationshipsByKinds(edgeKindA) failed: %v", err) + } + + if count := countByCypher(t, ctx, db, "MATCH ()-[r:RelByKindEdgeA]->() RETURN count(r)"); count != 0 { + t.Fatalf("edgeKindA count: got %d, want 0", count) + } + if count := countByCypher(t, ctx, db, "MATCH ()-[r:RelByKindEdgeB]->() RETURN count(r)"); count != 1 { + t.Fatalf("edgeKindB count: got %d, want 1", count) + } + if count := countByCypher(t, ctx, db, "MATCH (n:RelByKindNode) RETURN count(n)"); count != 3 { + t.Fatalf("node count: got %d, want 3", count) + } + + ClearGraph(t, db, ctx) + }) + + t.Run("multiple kinds delete every matching relationship", func(t *testing.T) { + createFixture() + + if err := deleter.DeleteRelationshipsByKinds(ctx, graph.Kinds{edgeKindA, edgeKindB}); err != nil { + t.Fatalf("DeleteRelationshipsByKinds(edgeKindA, edgeKindB) failed: %v", err) + } + + if count := countByCypher(t, ctx, db, "MATCH ()-[r]->() RETURN count(r)"); count != 0 { + t.Fatalf("edge count: got %d, want 0", count) + } + if count := countByCypher(t, ctx, db, "MATCH (n:RelByKindNode) RETURN count(n)"); count != 3 { + t.Fatalf("node count: got %d, want 3", count) + } + + ClearGraph(t, db, ctx) + }) + + t.Run("undefined and empty kinds are a safe no-op", func(t *testing.T) { + createFixture() + + if err := deleter.DeleteRelationshipsByKinds(ctx, graph.Kinds{missing}); err != nil { + t.Fatalf("DeleteRelationshipsByKinds(missing) failed: %v", err) + } + if err := deleter.DeleteRelationshipsByKinds(ctx, nil); err != nil { + t.Fatalf("DeleteRelationshipsByKinds(nil) failed: %v", err) + } + + if count := countByCypher(t, ctx, db, "MATCH ()-[r]->() RETURN count(r)"); count != 3 { + t.Fatalf("edge count after no-op delete: got %d, want 3", count) + } + + ClearGraph(t, db, ctx) + }) +}