feat(OpenGraph): Generate AD and AZ extension schemas and apply on startup - BED-6721

cmd/api/src/bootstrap/server.go

@@ -72,6 +72,14 @@ func MigrateDB(ctx context.Context, cfg config.Configuration, db database.Databa
 	return CreateDefaultAdmin(ctx, cfg, db, defaultAdminFunc)
 }
 
+func PopulateExtensionData(ctx context.Context, db database.Database) error {
+	if err := db.PopulateExtensionData(ctx); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func CreateDefaultAdmin(ctx context.Context, cfg config.Configuration, db database.Database, defaultAdminFunction func() (config.DefaultAdminConfiguration, error)) error {
 	var (
 		secretDigester = cfg.Crypto.Argon2.NewDigester()

cmd/api/src/daemons/changelog/ingestion_integration_test.go

@@ -116,6 +116,7 @@ func setupIntegrationTest(t *testing.T) IntegrationTestSuite {
 
 	db := database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver())
 	require.NoError(t, db.Migrate(ctx))
+	require.NoError(t, db.PopulateExtensionData(ctx))
 
 	return IntegrationTestSuite{
 		Context:      ctx,

cmd/api/src/daemons/datapipe/datapipe_integration_test.go

@@ -93,6 +93,9 @@ func setupIntegrationTestSuite(t *testing.T, fixturesPath string) IntegrationTes
 	err = graphDB.AssertSchema(ctx, graphschema.DefaultGraphSchema())
 	require.NoError(t, err)
 
+	err = db.PopulateExtensionData(ctx)
+	require.NoError(t, err)
+
 	ingestSchema, err := upload.LoadIngestSchema()
 	require.NoError(t, err)
 

cmd/api/src/database/database_integration_test.go

@@ -59,6 +59,9 @@ func setupIntegrationTestSuite(t *testing.T) IntegrationTestSuite {
 	err = db.Migrate(ctx)
 	require.NoError(t, err)
 
+	err = db.PopulateExtensionData(ctx)
+	require.NoError(t, err)
+
 	// #endregion
 
 	return IntegrationTestSuite{

cmd/api/src/database/db.go

@@ -34,6 +34,7 @@ import (
 	"github.com/specterops/bloodhound/cmd/api/src/services/agi"
 	"github.com/specterops/bloodhound/cmd/api/src/services/dataquality"
 	"github.com/specterops/bloodhound/cmd/api/src/services/upload"
+	"github.com/specterops/bloodhound/packages/go/bhlog/attr"
 	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
 )
@@ -99,6 +100,7 @@ type Database interface {
 
 	Wipe(ctx context.Context) error
 	Migrate(ctx context.Context) error
+	PopulateExtensionData(ctx context.Context) error
 	CreateInstallation(ctx context.Context) (model.Installation, error)
 	GetInstallation(ctx context.Context) (model.Installation, error)
 	HasInstallation(ctx context.Context) (bool, error)
@@ -280,7 +282,16 @@ func (s *BloodhoundDB) Wipe(ctx context.Context) error {
 func (s *BloodhoundDB) Migrate(ctx context.Context) error {
 	// Run the migrator
 	if err := migration.NewMigrator(s.db.WithContext(ctx)).ExecuteStepwiseMigrations(); err != nil {
-		slog.ErrorContext(ctx, fmt.Sprintf("Error during SQL database migration phase: %v", err))
+		slog.ErrorContext(ctx, "Error during SQL database migration phase", attr.Error(err))
+		return err
+	}
+
+	return nil
+}
+
+func (s *BloodhoundDB) PopulateExtensionData(ctx context.Context) error {
+	if err := migration.NewMigrator(s.db.WithContext(ctx)).ExecuteExtensionDataPopulation(); err != nil {
+		slog.ErrorContext(ctx, "Error during extensions data population phase", attr.Error(err))
 		return err
 	}
 

cmd/api/src/database/graphschema.go

@@ -223,7 +223,7 @@ func (s *BloodhoundDB) CreateGraphSchemaNodeKind(ctx context.Context, name strin
 		if strings.Contains(result.Error.Error(), DuplicateKeyValueErrorString) {
 			return model.GraphSchemaNodeKind{}, fmt.Errorf("%w: %v", ErrDuplicateSchemaNodeKindName, result.Error)
 		}
-		return model.GraphSchemaNodeKind{}, CheckError(result)
+		return model.GraphSchemaNodeKind{}, result.Error
 	}
 	return schemaNodeKind, nil
 }
@@ -232,11 +232,38 @@ func (s *BloodhoundDB) CreateGraphSchemaNodeKind(ctx context.Context, name strin
 // populated with data, as well as an integer indicating the total number of rows returned by the query (excluding any given pagination).
 func (s *BloodhoundDB) GetGraphSchemaNodeKinds(ctx context.Context, filters model.Filters, sort model.Sort, skip, limit int) (model.GraphSchemaNodeKinds, int, error) {
 	var (
-		schemaNodeKinds = model.GraphSchemaNodeKinds{}
-		totalRowCount   int
+		schemaNodeKinds        = model.GraphSchemaNodeKinds{}
+		totalRowCount          int
+		tableIdentifiedFilters = make(model.Filters, len(filters))
+		tableIdentifiedSort    = make(model.Sort, len(sort))
 	)
 
-	if filterAndPagination, err := parseFiltersAndPagination(filters, sort, skip, limit); err != nil {
+	// add table identifiers to filtering and sorting columns ensuring we don't return an ambiguous column error
+	for column, filter := range filters {
+		if column == "name" {
+			tableIdentifiedFilters[fmt.Sprintf("%s.%s", "k", column)] = filter
+			delete(filters, column)
+		} else {
+			tableIdentifiedFilters[fmt.Sprintf("%s.%s", "nk", column)] = filter
+			delete(filters, column)
+		}
+	}
+
+	for idx, sortItem := range sort {
+		if sort[idx].Column == "name" {
+			tableIdentifiedSort[idx] = model.SortItem{
+				Direction: sortItem.Direction,
+				Column:    fmt.Sprintf("%s.%s", "k", sort[idx].Column),
+			}
+		} else {
+			tableIdentifiedSort[idx] = model.SortItem{
+				Direction: sortItem.Direction,
+				Column:    fmt.Sprintf("%s.%s", "nk", sort[idx].Column),
+			}
+		}
+	}
+
+	if filterAndPagination, err := parseFiltersAndPagination(tableIdentifiedFilters, tableIdentifiedSort, skip, limit); err != nil {
 		return schemaNodeKinds, 0, err
 	} else {
 		sqlStr := fmt.Sprintf(`SELECT nk.id, k.name, nk.schema_extension_id, nk.display_name, nk.description, 
@@ -450,11 +477,38 @@ func (s *BloodhoundDB) CreateGraphSchemaEdgeKind(ctx context.Context, name strin
 // populated with data, as well as an integer indicating the total number of rows returned by the query (excluding any given pagination).
 func (s *BloodhoundDB) GetGraphSchemaEdgeKinds(ctx context.Context, edgeKindFilters model.Filters, sort model.Sort, skip, limit int) (model.GraphSchemaEdgeKinds, int, error) {
 	var (
-		schemaEdgeKinds = model.GraphSchemaEdgeKinds{}
-		totalRowCount   int
+		schemaEdgeKinds        = model.GraphSchemaEdgeKinds{}
+		totalRowCount          int
+		tableIdentifiedFilters = make(model.Filters, len(edgeKindFilters))
+		tableIdentifiedSort    = make(model.Sort, len(sort))
 	)
 
-	if filterAndPagination, err := parseFiltersAndPagination(edgeKindFilters, sort, skip, limit); err != nil {
+	// add table identifiers to filtering and sorting columns ensuring we don't return an ambiguous column error
+	for column, filters := range edgeKindFilters {
+		if column == "name" {
+			tableIdentifiedFilters[fmt.Sprintf("%s.%s", "k", column)] = filters
+			delete(edgeKindFilters, column)
+		} else {
+			tableIdentifiedFilters[fmt.Sprintf("%s.%s", "ek", column)] = filters
+			delete(edgeKindFilters, column)
+		}
+	}
+
+	for idx, sortItem := range sort {
+		if sort[idx].Column == "name" {
+			tableIdentifiedSort[idx] = model.SortItem{
+				Direction: sortItem.Direction,
+				Column:    fmt.Sprintf("%s.%s", "k", sort[idx].Column),
+			}
+		} else {
+			tableIdentifiedSort[idx] = model.SortItem{
+				Direction: sortItem.Direction,
+				Column:    fmt.Sprintf("%s.%s", "ek", sort[idx].Column),
+			}
+		}
+	}
+
+	if filterAndPagination, err := parseFiltersAndPagination(tableIdentifiedFilters, tableIdentifiedSort, skip, limit); err != nil {
 		return schemaEdgeKinds, 0, err
 	} else {
 		sqlStr := fmt.Sprintf(`SELECT ek.id, k.name, ek.schema_extension_id, ek.description, ek.is_traversable, 

cmd/api/src/database/graphschema_integration_test.go

@@ -458,7 +458,7 @@ func TestDatabase_GraphSchemaNodeKind_CRUD(t *testing.T) {
 	t.Run("fail - return error for filtering on non-existent column", func(t *testing.T) {
 		_, _, err = testSuite.BHDatabase.GetGraphSchemaNodeKinds(testSuite.Context,
 			model.Filters{"nonexistentcolumn": []model.Filter{{Operator: model.Equals, Value: "blah", SetOperator: model.FilterAnd}}}, model.Sort{}, 0, 0)
-		require.EqualError(t, err, "ERROR: column \"nonexistentcolumn\" does not exist (SQLSTATE 42703)")
+		require.EqualError(t, err, "ERROR: column nk.nonexistentcolumn does not exist (SQLSTATE 42703)")
 	})
 
 	// UPDATE
@@ -896,7 +896,7 @@ func TestDatabase_GraphSchemaEdgeKind_CRUD(t *testing.T) {
 	t.Run("fail - return error for filtering on non-existent column", func(t *testing.T) {
 		_, _, err = testSuite.BHDatabase.GetGraphSchemaEdgeKinds(testSuite.Context,
 			model.Filters{"nonexistentcolumn": []model.Filter{{Operator: model.Equals, Value: "blah", SetOperator: model.FilterAnd}}}, model.Sort{}, 0, 0)
-		require.EqualError(t, err, "ERROR: column \"nonexistentcolumn\" does not exist (SQLSTATE 42703)")
+		require.EqualError(t, err, "ERROR: column ek.nonexistentcolumn does not exist (SQLSTATE 42703)")
 	})
 
 	// UPDATE