Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 21 additions & 11 deletions cmd/api/src/database/migration/extensions/ad_graph_schema.sql
Original file line number Diff line number Diff line change
Expand Up @@ -15,25 +15,35 @@
-- SPDX-License-Identifier: Apache-2.0
-- Code generated by Cuelang code gen. DO NOT EDIT!
-- Cuelang source: github.com/specterops/bloodhound/-/tree/main/packages/cue/schemas/
CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS void AS $$
CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS SMALLINT AS $$
DECLARE
kind_id SMALLINT;
BEGIN
IF NOT EXISTS (SELECT id FROM kind WHERE kind.name = node_kind_name) THEN
INSERT INTO kind (name) VALUES (node_kind_name);
INSERT INTO kind (name) VALUES (node_kind_name) RETURNING id INTO kind_id;
ELSE
SELECT id FROM kind WHERE kind.name = node_kind_name INTO kind_id;
END IF;

RETURN kind_id;
END $$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION genscript_upsert_source_kind(kind_name TEXT) RETURNS void AS $$
CREATE OR REPLACE FUNCTION genscript_upsert_source_kind(kind_name TEXT) RETURNS SMALLINT AS $$
DECLARE
retrieved_kind_id SMALLINT;
source_kind_id SMALLINT;
BEGIN
SELECT k.id INTO retrieved_kind_id FROM kind k WHERE k.name = kind_name;
IF retrieved_kind_id IS NULL THEN
INSERT INTO kind (name) VALUES (kind_name)
RETURNING id INTO retrieved_kind_id;
SELECT genscript_upsert_kind(kind_name) INTO retrieved_kind_id;
END IF;
IF NOT EXISTS (SELECT sk.id FROM source_kinds sk WHERE sk.kind_id = retrieved_kind_id) THEN
INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id);
INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id) RETURNING id INTO source_kind_id;
ELSE
SELECT sk.id FROM source_kinds sk WHERE sk.kind_id = retrieved_kind_id INTO source_kind_id;
END IF;

RETURN source_kind_id;
END $$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION genscript_upsert_schema_node_kind(v_extension_id INT, v_kind_name VARCHAR(256), v_display_name TEXT, v_description TEXT, v_is_display_kind BOOLEAN, v_icon TEXT, v_icon_color TEXT) RETURNS void AS $$
Expand All @@ -42,7 +52,7 @@ DECLARE
BEGIN
SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_kind_name;
IF retrieved_kind_id IS NULL THEN
RAISE EXCEPTION 'couldn''t find matching kind_id';
SELECT genscript_upsert_kind(v_kind_name) INTO retrieved_kind_id;
END IF;

IF NOT EXISTS (SELECT id FROM schema_node_kinds nk WHERE nk.kind_id = retrieved_kind_id) THEN
Expand All @@ -59,7 +69,7 @@ DECLARE
BEGIN
SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_kind_name;
IF retrieved_kind_id IS NULL THEN
RAISE EXCEPTION 'couldn''t find matching kind_id';
SELECT genscript_upsert_kind(v_kind_name) INTO retrieved_kind_id;
END IF;

IF NOT EXISTS (SELECT id FROM schema_relationship_kinds ek WHERE ek.kind_id = retrieved_kind_id) THEN
Expand All @@ -78,12 +88,12 @@ DECLARE
BEGIN
SELECT id INTO retrieved_environment_kind_id FROM kind WHERE name = v_environment_kind_name;
IF retrieved_environment_kind_id IS NULL THEN
RAISE EXCEPTION 'couldn''t find matching kind_id';
SELECT genscript_upsert_kind(v_environment_kind_name) INTO retrieved_environment_kind_id;
END IF;

SELECT sk.id INTO retrieved_source_kind_id FROM source_kinds sk JOIN kind k ON sk.kind_id = k.id WHERE k.name = v_source_kind_name;
IF retrieved_source_kind_id IS NULL THEN
RAISE EXCEPTION 'couldn''t find matching kind_id';
SELECT genscript_upsert_source_kind(v_source_kind_name) INTO retrieved_source_kind_id;
END IF;

IF NOT EXISTS (SELECT id FROM schema_environments se WHERE se.schema_extension_id = v_extension_id) THEN
Expand All @@ -102,7 +112,7 @@ DECLARE
BEGIN
SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_principal_kind_name;
IF retrieved_kind_id IS NULL THEN
RAISE EXCEPTION 'couldn''t find matching kind_id';
SELECT genscript_upsert_kind(v_principal_kind_name) INTO retrieved_kind_id;
END IF;

IF NOT EXISTS (SELECT 1 FROM schema_environments_principal_kinds pk WHERE pk.principal_kind = retrieved_kind_id) THEN
Expand Down
32 changes: 21 additions & 11 deletions cmd/api/src/database/migration/extensions/az_graph_schema.sql
Original file line number Diff line number Diff line change
Expand Up @@ -15,25 +15,35 @@
-- SPDX-License-Identifier: Apache-2.0
-- Code generated by Cuelang code gen. DO NOT EDIT!
-- Cuelang source: github.com/specterops/bloodhound/-/tree/main/packages/cue/schemas/
CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS void AS $$
CREATE OR REPLACE FUNCTION genscript_upsert_kind(node_kind_name TEXT) RETURNS SMALLINT AS $$
DECLARE
kind_id SMALLINT;
BEGIN
IF NOT EXISTS (SELECT id FROM kind WHERE kind.name = node_kind_name) THEN
INSERT INTO kind (name) VALUES (node_kind_name);
INSERT INTO kind (name) VALUES (node_kind_name) RETURNING id INTO kind_id;
ELSE
SELECT id FROM kind WHERE kind.name = node_kind_name INTO kind_id;
END IF;

RETURN kind_id;
END $$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION genscript_upsert_source_kind(kind_name TEXT) RETURNS void AS $$
CREATE OR REPLACE FUNCTION genscript_upsert_source_kind(kind_name TEXT) RETURNS SMALLINT AS $$
DECLARE
retrieved_kind_id SMALLINT;
source_kind_id SMALLINT;
BEGIN
SELECT k.id INTO retrieved_kind_id FROM kind k WHERE k.name = kind_name;
IF retrieved_kind_id IS NULL THEN
INSERT INTO kind (name) VALUES (kind_name)
RETURNING id INTO retrieved_kind_id;
SELECT genscript_upsert_kind(kind_name) INTO retrieved_kind_id;
END IF;
IF NOT EXISTS (SELECT sk.id FROM source_kinds sk WHERE sk.kind_id = retrieved_kind_id) THEN
INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id);
INSERT INTO source_kinds (kind_id) VALUES (retrieved_kind_id) RETURNING id INTO source_kind_id;
ELSE
SELECT sk.id FROM source_kinds sk WHERE sk.kind_id = retrieved_kind_id INTO source_kind_id;
END IF;

RETURN source_kind_id;
END $$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION genscript_upsert_schema_node_kind(v_extension_id INT, v_kind_name VARCHAR(256), v_display_name TEXT, v_description TEXT, v_is_display_kind BOOLEAN, v_icon TEXT, v_icon_color TEXT) RETURNS void AS $$
Expand All @@ -42,7 +52,7 @@ DECLARE
BEGIN
SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_kind_name;
IF retrieved_kind_id IS NULL THEN
RAISE EXCEPTION 'couldn''t find matching kind_id';
SELECT genscript_upsert_kind(v_kind_name) INTO retrieved_kind_id;
END IF;

IF NOT EXISTS (SELECT id FROM schema_node_kinds nk WHERE nk.kind_id = retrieved_kind_id) THEN
Expand All @@ -59,7 +69,7 @@ DECLARE
BEGIN
SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_kind_name;
IF retrieved_kind_id IS NULL THEN
RAISE EXCEPTION 'couldn''t find matching kind_id';
SELECT genscript_upsert_kind(v_kind_name) INTO retrieved_kind_id;
END IF;

IF NOT EXISTS (SELECT id FROM schema_relationship_kinds ek WHERE ek.kind_id = retrieved_kind_id) THEN
Expand All @@ -78,12 +88,12 @@ DECLARE
BEGIN
SELECT id INTO retrieved_environment_kind_id FROM kind WHERE name = v_environment_kind_name;
IF retrieved_environment_kind_id IS NULL THEN
RAISE EXCEPTION 'couldn''t find matching kind_id';
SELECT genscript_upsert_kind(v_environment_kind_name) INTO retrieved_environment_kind_id;
END IF;

SELECT sk.id INTO retrieved_source_kind_id FROM source_kinds sk JOIN kind k ON sk.kind_id = k.id WHERE k.name = v_source_kind_name;
IF retrieved_source_kind_id IS NULL THEN
RAISE EXCEPTION 'couldn''t find matching kind_id';
SELECT genscript_upsert_source_kind(v_source_kind_name) INTO retrieved_source_kind_id;
END IF;

IF NOT EXISTS (SELECT id FROM schema_environments se WHERE se.schema_extension_id = v_extension_id) THEN
Expand All @@ -102,7 +112,7 @@ DECLARE
BEGIN
SELECT id INTO retrieved_kind_id FROM kind WHERE name = v_principal_kind_name;
IF retrieved_kind_id IS NULL THEN
RAISE EXCEPTION 'couldn''t find matching kind_id';
SELECT genscript_upsert_kind(v_principal_kind_name) INTO retrieved_kind_id;
END IF;

IF NOT EXISTS (SELECT 1 FROM schema_environments_principal_kinds pk WHERE pk.principal_kind = retrieved_kind_id) THEN
Expand Down
190 changes: 190 additions & 0 deletions cmd/api/src/database/migration/extensions_integration_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,190 @@
// 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 integration

package migration_test

import (
"context"
"fmt"
"net/url"
"strings"
"testing"

"github.com/peterldowns/pgtestdb"
"gorm.io/gorm"

"github.com/specterops/bloodhound/cmd/api/src/auth"
"github.com/specterops/bloodhound/cmd/api/src/database"
"github.com/specterops/bloodhound/cmd/api/src/model"
"github.com/specterops/bloodhound/cmd/api/src/test/integration/utils"
"github.com/stretchr/testify/require"
)

type IntegrationTestSuite struct {
Context context.Context
BHDatabase *database.BloodhoundDB
DB *gorm.DB
}

func setupIntegrationTestSuite(t *testing.T) IntegrationTestSuite {
t.Helper()

var (
ctx = context.Background()
connConf = pgtestdb.Custom(t, getPostgresConfig(t), pgtestdb.NoopMigrator{})
gormDB *gorm.DB
db *database.BloodhoundDB
err error
)

// #region Setup for dbs

gormDB, err = database.OpenDatabase(connConf.URL())
require.NoError(t, err)

db = database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver())

err = db.Migrate(ctx)
require.NoError(t, err)

err = db.PopulateExtensionData(ctx)
require.NoError(t, err)

// #endregion

return IntegrationTestSuite{
Context: ctx,
BHDatabase: db,
DB: gormDB,
}
}

// getPostgresConfig reads key/value pairs from the default integration
// config file and creates a pgtestdb configuration object.
func getPostgresConfig(t *testing.T) pgtestdb.Config {
t.Helper()

config, err := utils.LoadIntegrationTestConfig()
require.NoError(t, err)

environmentMap := make(map[string]string)
for _, entry := range strings.Fields(config.Database.Connection) {
if parts := strings.SplitN(entry, "=", 2); len(parts) == 2 {
environmentMap[parts[0]] = parts[1]
}
}

if strings.HasPrefix(environmentMap["host"], "/") {
return pgtestdb.Config{
DriverName: "pgx",
User: environmentMap["user"],
Password: environmentMap["password"],
Database: environmentMap["dbname"],
Options: fmt.Sprintf("host=%s", url.PathEscape(environmentMap["host"])),
TestRole: &pgtestdb.Role{
Username: environmentMap["user"],
Password: environmentMap["password"],
Capabilities: "NOSUPERUSER NOCREATEROLE",
},
}
}

return pgtestdb.Config{
DriverName: "pgx",
Host: environmentMap["host"],
Port: environmentMap["port"],
User: environmentMap["user"],
Password: environmentMap["password"],
Database: environmentMap["dbname"],
Options: "sslmode=disable",
ForceTerminateConnections: true,
}
}

func (s *IntegrationTestSuite) teardownIntegrationTestSuite(t *testing.T) {
t.Helper()

if s.BHDatabase != nil {
s.BHDatabase.Close(s.Context)
}
}

func TestExtensions_GetOnStartExtensionData(t *testing.T) {
var (
testSuite = setupIntegrationTestSuite(t)
)
Comment on lines +127 to +130
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Assert the exact extension names (AD, AZ), not only record count.

The current checks verify count and per-row shape, but they don’t explicitly enforce the expected extension name set. Add a final name-set assertion to make the regression test stricter.

✅ Tighten expectations
 func TestExtensions_GetOnStartExtensionData(t *testing.T) {
 	var (
-		testSuite = setupIntegrationTestSuite(t)
+		testSuite      = setupIntegrationTestSuite(t)
+		extensionNames = make([]string, 0, 2)
 	)
@@
 	require.NoError(t, err)

 	require.Equal(t, 2, totalRecords)
+	require.Len(t, extensions, 2)

 	for _, extension := range extensions {
+		extensionNames = append(extensionNames, extension.Name)
 		require.True(t, extension.IsBuiltin, "All extensions should be marked as built-in")
@@
 		validateEnvironmentKind(t, extension.Name, environmentKind.Name)
 	}
+
+	require.ElementsMatch(t, []string{"AD", "AZ"}, extensionNames)

 }

Also applies to: 142-164

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@cmd/api/src/database/migration/extensions_integration_test.go` around lines
127 - 130, The test TestExtensions_GetOnStartExtensionData currently only checks
record count and per-row fields; update it (and the similar block around the
function starting at lines 142–164) to explicitly assert the exact set of
extension names {"AD","AZ"} is present. After extracting the rows (e.g., the
slice/variable built from the query in TestExtensions_GetOnStartExtensionData
and the other test), collect the name values and add an assertion that the names
match the expected set (use the test framework's set/element matching helper
such as require.ElementsMatch or assert.ElementsMatch) so the test fails if any
name is missing or unexpected. Ensure this assertion is placed after the per-row
shape checks and references the same row/name variables used in those tests.


defer testSuite.teardownIntegrationTestSuite(t)

err := testSuite.BHDatabase.PopulateExtensionData(testSuite.Context)
require.NoError(t, err)

// Validate Both Schema Extensions Exist
extensions, totalRecords, err := testSuite.BHDatabase.GetGraphSchemaExtensions(testSuite.Context, model.Filters{}, model.Sort{}, 0, 0)

require.NoError(t, err)

require.Equal(t, 2, totalRecords)

for _, extension := range extensions {
require.True(t, extension.IsBuiltin, "All extensions should be marked as built-in")
// Validate Schema Environments Exist
schemaEnvironments, err := testSuite.BHDatabase.GetEnvironmentsByExtensionId(testSuite.Context, extension.ID)
require.NoError(t, err)

// There should only be one schema environment per built-in extension
require.Len(t, schemaEnvironments, 1)
schemaEnvironment := schemaEnvironments[0]

// Validate Source Kinds Exist
sourceKind, err := testSuite.BHDatabase.GetSourceKindByID(testSuite.Context, int(schemaEnvironment.SourceKindId))
require.NoError(t, err)
require.NotNil(t, sourceKind)
validateSourceKind(t, extension.Name, sourceKind.Name.String())

// Validate Environment Kinds Exist
environmentKind, err := testSuite.BHDatabase.GetKindById(testSuite.Context, schemaEnvironment.EnvironmentKindId)
require.NoError(t, err)
validateEnvironmentKind(t, extension.Name, environmentKind.Name)
}

}

func validateSourceKind(t *testing.T, extensionName, sourceKindName string) {
t.Helper()
switch extensionName {
case "AD":
require.Equal(t, "Base", sourceKindName)
case "AZ":
require.Equal(t, "AZBase", sourceKindName)
default:
t.Fatalf("Invalid extension name %s", extensionName)
}
}

func validateEnvironmentKind(t *testing.T, extensionName, environmentKindName string) {
t.Helper()
switch extensionName {
case "AD":
require.Equal(t, "Domain", environmentKindName)
case "AZ":
require.Equal(t, "AZTenant", environmentKindName)
default:
t.Fatalf("Invalid extension name %s", extensionName)
}
}
Loading