From f732fabcef20f12bd67f369a185354d3a8ae83a7 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Fri, 23 Jan 2026 07:43:31 -0800 Subject: [PATCH 01/22] add kinds method --- cmd/api/src/database/kind.go | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/cmd/api/src/database/kind.go b/cmd/api/src/database/kind.go index 0487f28375d..17aaa2249c0 100644 --- a/cmd/api/src/database/kind.go +++ b/cmd/api/src/database/kind.go @@ -23,6 +23,7 @@ import ( type Kind interface { GetKindByName(ctx context.Context, name string) (model.Kind, error) + GetKindById(ctx context.Context, id int32) (model.Kind, error) } func (s *BloodhoundDB) GetKindByName(ctx context.Context, name string) (model.Kind, error) { @@ -45,3 +46,24 @@ func (s *BloodhoundDB) GetKindByName(ctx context.Context, name string) (model.Ki return kind, nil } + +func (s *BloodhoundDB) GetKindById(ctx context.Context, id int32) (model.Kind, error) { + const query = ` + SELECT id, name + FROM kind + WHERE id = $1; + ` + + var kind model.Kind + result := s.db.WithContext(ctx).Raw(query, id).Scan(&kind) + + if result.Error != nil { + return model.Kind{}, result.Error + } + + if result.RowsAffected == 0 || kind.ID == 0 { + return model.Kind{}, ErrNotFound + } + + return kind, nil +} From b3b3d115385f679f68c2c64bc25cae09c5c07052 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Fri, 23 Jan 2026 08:57:50 -0800 Subject: [PATCH 02/22] new db method --- cmd/api/src/database/graphschema.go | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/cmd/api/src/database/graphschema.go b/cmd/api/src/database/graphschema.go index 57a6fd7357c..cc32381add5 100644 --- a/cmd/api/src/database/graphschema.go +++ b/cmd/api/src/database/graphschema.go @@ -61,6 +61,7 @@ type OpenGraphSchema interface { CreateSchemaRelationshipFinding(ctx context.Context, extensionId int32, relationshipKindId int32, environmentId int32, name string, displayName string) (model.SchemaRelationshipFinding, error) GetSchemaRelationshipFindingById(ctx context.Context, findingId int32) (model.SchemaRelationshipFinding, error) GetSchemaRelationshipFindingByName(ctx context.Context, name string) (model.SchemaRelationshipFinding, error) + GetSchemaRelationshipFindingsByEnvironmentId(ctx context.Context, environmentId int32) ([]model.SchemaRelationshipFinding, error) DeleteSchemaRelationshipFinding(ctx context.Context, findingId int32) error CreateRemediation(ctx context.Context, findingId int32, shortDescription string, longDescription string, shortRemediation string, longRemediation string) (model.Remediation, error) @@ -720,6 +721,22 @@ func (s *BloodhoundDB) GetSchemaRelationshipFindingByName(ctx context.Context, n return finding, nil } +// GetSchemaRelationshipFindingsByEnvironmentId - retrieves all schema relationship findings for a given environment. +func (s *BloodhoundDB) GetSchemaRelationshipFindingsByEnvironmentId(ctx context.Context, environmentId int32) ([]model.SchemaRelationshipFinding, error) { + var findings []model.SchemaRelationshipFinding + var finding model.SchemaRelationshipFinding + + if result := s.db.WithContext(ctx).Raw(fmt.Sprintf(` + SELECT id, schema_extension_id, relationship_kind_id, environment_id, name, display_name, created_at + FROM %s WHERE environment_id = ?`, + finding.TableName()), + environmentId).Scan(&findings); result.Error != nil { + return nil, CheckError(result) + } + + return findings, nil +} + // DeleteSchemaRelationshipFinding - deletes a schema relationship finding by id. func (s *BloodhoundDB) DeleteSchemaRelationshipFinding(ctx context.Context, findingId int32) error { var finding model.SchemaRelationshipFinding From cd9ea4f6268d9366cf0717e6402dcf0dd600780c Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Fri, 23 Jan 2026 10:56:22 -0800 Subject: [PATCH 03/22] working on og-tool --- cmd/api/src/database/sourcekinds.go | 34 +++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/cmd/api/src/database/sourcekinds.go b/cmd/api/src/database/sourcekinds.go index 6aa76b8a24a..29307b450f9 100644 --- a/cmd/api/src/database/sourcekinds.go +++ b/cmd/api/src/database/sourcekinds.go @@ -29,6 +29,7 @@ type SourceKindsData interface { DeactivateSourceKindsByName(ctx context.Context, kinds graph.Kinds) error RegisterSourceKind(ctx context.Context) func(sourceKind graph.Kind) error GetSourceKindByName(ctx context.Context, name string) (SourceKind, error) + GetSourceKindById(ctx context.Context, id int32) (SourceKind, error) } // RegisterSourceKind returns a function that inserts a source kind by name, @@ -127,6 +128,39 @@ func (s *BloodhoundDB) GetSourceKindByName(ctx context.Context, name string) (So return kind, nil } +func (s *BloodhoundDB) GetSourceKindById(ctx context.Context, id int32) (SourceKind, error) { + const query = ` + SELECT id, name, active + FROM source_kinds + WHERE id = $1 AND active = true; + ` + + type rawSourceKind struct { + ID int + Name string + Active bool + } + + var raw rawSourceKind + result := s.db.WithContext(ctx).Raw(query, id).Scan(&raw) + + if result.Error != nil { + return SourceKind{}, result.Error + } + + if result.RowsAffected == 0 || raw.ID == 0 { + return SourceKind{}, ErrNotFound + } + + kind := SourceKind{ + ID: raw.ID, + Name: graph.StringKind(raw.Name), + Active: raw.Active, + } + + return kind, nil +} + func (s *BloodhoundDB) DeactivateSourceKindsByName(ctx context.Context, kinds graph.Kinds) error { if len(kinds) == 0 { return nil From 61c3098a8481ff9ab47471df35ae2e70f941123f Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Tue, 27 Jan 2026 10:05:29 -0800 Subject: [PATCH 04/22] mocks --- cmd/api/src/database/mocks/db.go | 45 ++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/cmd/api/src/database/mocks/db.go b/cmd/api/src/database/mocks/db.go index 27d97625997..4f0c0f1e6aa 100644 --- a/cmd/api/src/database/mocks/db.go +++ b/cmd/api/src/database/mocks/db.go @@ -1992,6 +1992,21 @@ func (mr *MockDatabaseMockRecorder) GetInstallation(ctx any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInstallation", reflect.TypeOf((*MockDatabase)(nil).GetInstallation), ctx) } +// GetKindById mocks base method. +func (m *MockDatabase) GetKindById(ctx context.Context, id int32) (model.Kind, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetKindById", ctx, id) + ret0, _ := ret[0].(model.Kind) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetKindById indicates an expected call of GetKindById. +func (mr *MockDatabaseMockRecorder) GetKindById(ctx, id any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetKindById", reflect.TypeOf((*MockDatabase)(nil).GetKindById), ctx, id) +} + // GetKindByName mocks base method. func (m *MockDatabase) GetKindByName(ctx context.Context, name string) (model.Kind, error) { m.ctrl.T.Helper() @@ -2277,6 +2292,21 @@ func (mr *MockDatabaseMockRecorder) GetSchemaRelationshipFindingByName(ctx, name return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSchemaRelationshipFindingByName", reflect.TypeOf((*MockDatabase)(nil).GetSchemaRelationshipFindingByName), ctx, name) } +// GetSchemaRelationshipFindingsByEnvironmentId mocks base method. +func (m *MockDatabase) GetSchemaRelationshipFindingsByEnvironmentId(ctx context.Context, environmentId int32) ([]model.SchemaRelationshipFinding, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetSchemaRelationshipFindingsByEnvironmentId", ctx, environmentId) + ret0, _ := ret[0].([]model.SchemaRelationshipFinding) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetSchemaRelationshipFindingsByEnvironmentId indicates an expected call of GetSchemaRelationshipFindingsByEnvironmentId. +func (mr *MockDatabaseMockRecorder) GetSchemaRelationshipFindingsByEnvironmentId(ctx, environmentId any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSchemaRelationshipFindingsByEnvironmentId", reflect.TypeOf((*MockDatabase)(nil).GetSchemaRelationshipFindingsByEnvironmentId), ctx, environmentId) +} + // GetScopeForSavedQuery mocks base method. func (m *MockDatabase) GetScopeForSavedQuery(ctx context.Context, queryID int64, userID uuid.UUID) (database.SavedQueryScopeMap, error) { m.ctrl.T.Helper() @@ -2363,6 +2393,21 @@ func (mr *MockDatabaseMockRecorder) GetSharedSavedQueries(ctx, userID any) *gomo return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSharedSavedQueries", reflect.TypeOf((*MockDatabase)(nil).GetSharedSavedQueries), ctx, userID) } +// GetSourceKindById mocks base method. +func (m *MockDatabase) GetSourceKindById(ctx context.Context, id int32) (database.SourceKind, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetSourceKindById", ctx, id) + ret0, _ := ret[0].(database.SourceKind) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetSourceKindById indicates an expected call of GetSourceKindById. +func (mr *MockDatabaseMockRecorder) GetSourceKindById(ctx, id any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSourceKindById", reflect.TypeOf((*MockDatabase)(nil).GetSourceKindById), ctx, id) +} + // GetSourceKindByName mocks base method. func (m *MockDatabase) GetSourceKindByName(ctx context.Context, name string) (database.SourceKind, error) { m.ctrl.T.Helper() From 58fcf096922499b5bc45ab4431ebe7c43c5d0e12 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Thu, 29 Jan 2026 09:48:18 -0800 Subject: [PATCH 05/22] generate metatrees double loop extensions -> environments --- cmd/api/src/database/graphschema.go | 45 +++++++++++++++++++++++++++++ cmd/api/src/database/mocks/db.go | 15 ++++++++++ 2 files changed, 60 insertions(+) diff --git a/cmd/api/src/database/graphschema.go b/cmd/api/src/database/graphschema.go index 5abc3be3b7d..a355b03f4bc 100644 --- a/cmd/api/src/database/graphschema.go +++ b/cmd/api/src/database/graphschema.go @@ -56,6 +56,7 @@ type OpenGraphSchema interface { GetEnvironmentByKinds(ctx context.Context, environmentKindId, sourceKindId int32) (model.SchemaEnvironment, error) GetEnvironmentById(ctx context.Context, environmentId int32) (model.SchemaEnvironment, error) GetEnvironments(ctx context.Context) ([]model.SchemaEnvironment, error) + GetEnvironmentsFiltered(ctx context.Context, filters model.Filters) ([]model.SchemaEnvironment, error) DeleteEnvironment(ctx context.Context, environmentId int32) error CreateSchemaRelationshipFinding(ctx context.Context, extensionId int32, relationshipKindId int32, environmentId int32, name string, displayName string) (model.SchemaRelationshipFinding, error) @@ -624,6 +625,50 @@ func (s *BloodhoundDB) GetEnvironments(ctx context.Context) ([]model.SchemaEnvir return result, nil } +// GetEnvironmentsFiltered - retrieves schema environments filtered by the given criteria. +// Common use case: filter by schema_extension_id to get all environments for a specific extension. +// Example: filters := model.Filters{"schema_extension_id": []model.Filter{{Operator: model.Equals, Value: "1"}}} +func (s *BloodhoundDB) GetEnvironmentsFiltered(ctx context.Context, filters model.Filters) ([]model.SchemaEnvironment, error) { + var result []model.SchemaEnvironment + + sqlFilter, err := buildSQLFilter(filters) + if err != nil { + return nil, err + } + + whereClause := "" + if sqlFilter.sqlString != "" { + whereClause = fmt.Sprintf("WHERE %s", sqlFilter.sqlString) + } + + query := fmt.Sprintf(` + SELECT + se.id, + se.schema_extension_id, + ext.display_name as schema_extension_display_name, + se.environment_kind_id, + k.name as environment_kind_name, + se.source_kind_id, + se.created_at, + se.updated_at, + se.deleted_at + FROM schema_environments se + INNER JOIN kind k ON se.environment_kind_id = k.id + INNER JOIN schema_extensions ext ON se.schema_extension_id = ext.id + %s + ORDER BY se.id`, whereClause) + + if err := CheckError(s.db.WithContext(ctx).Raw(query, sqlFilter.params...).Scan(&result)); err != nil { + return nil, err + } + + if result == nil { + result = []model.SchemaEnvironment{} + } + + return result, nil +} + // GetEnvironmentByKinds - retrieves an environment by its environment kind and source kind. func (s *BloodhoundDB) GetEnvironmentByKinds(ctx context.Context, environmentKindId, sourceKindId int32) (model.SchemaEnvironment, error) { var env model.SchemaEnvironment diff --git a/cmd/api/src/database/mocks/db.go b/cmd/api/src/database/mocks/db.go index af1c6492840..aa38811d881 100644 --- a/cmd/api/src/database/mocks/db.go +++ b/cmd/api/src/database/mocks/db.go @@ -1762,6 +1762,21 @@ func (mr *MockDatabaseMockRecorder) GetEnvironments(ctx any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnvironments", reflect.TypeOf((*MockDatabase)(nil).GetEnvironments), ctx) } +// GetEnvironmentsFiltered mocks base method. +func (m *MockDatabase) GetEnvironmentsFiltered(ctx context.Context, filters model.Filters) ([]model.SchemaEnvironment, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetEnvironmentsFiltered", ctx, filters) + ret0, _ := ret[0].([]model.SchemaEnvironment) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetEnvironmentsFiltered indicates an expected call of GetEnvironmentsFiltered. +func (mr *MockDatabaseMockRecorder) GetEnvironmentsFiltered(ctx, filters any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnvironmentsFiltered", reflect.TypeOf((*MockDatabase)(nil).GetEnvironmentsFiltered), ctx, filters) +} + // GetFlag mocks base method. func (m *MockDatabase) GetFlag(ctx context.Context, id int32) (appcfg.FeatureFlag, error) { m.ctrl.T.Helper() From 1e14852c9d103134c18a7e1a7ae83ce49cc7d30b Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Mon, 2 Feb 2026 11:46:41 -0800 Subject: [PATCH 06/22] uppercase environmetn_id on nodes --- cmd/api/src/services/graphify/convertors.go | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/cmd/api/src/services/graphify/convertors.go b/cmd/api/src/services/graphify/convertors.go index 1acab6af23a..b6c5e7830c3 100644 --- a/cmd/api/src/services/graphify/convertors.go +++ b/cmd/api/src/services/graphify/convertors.go @@ -52,6 +52,13 @@ func ConvertGenericNode(entity ein.GenericNode, converted *ConvertedData) error } } + // Uppercase environment_id if present + if rawEnvID, ok := node.PropertyMap["environment_id"]; ok { + if envID, ok := rawEnvID.(string); ok { + node.PropertyMap["environment_id"] = strings.ToUpper(envID) + } + } + // the first element in node.Labels determines which icon the UI renders for the node. // it is critical to specify this information because a node can have up to 3 kinds. if len(node.Labels) > 0 { From eddcd85bf0e0c7daff9399f7d43caf1cc463035f Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Mon, 2 Feb 2026 16:11:41 -0800 Subject: [PATCH 07/22] just generate --- cmd/ui/src/ducks/global/types.ts | 2 +- cmd/ui/src/ducks/graph/graphutils.ts | 2 +- .../bh-shared-ui/src/components/AppIcon/Icons/index.ts | 2 +- .../bh-shared-ui/src/components/DropdownSelector/index.ts | 2 +- .../bh-shared-ui/src/components/ExploreTable/index.ts | 2 +- packages/javascript/bh-shared-ui/src/components/index.ts | 8 ++++---- packages/javascript/bh-shared-ui/src/hooks/index.ts | 2 +- .../src/providers/NotificationProvider/index.ts | 2 +- .../src/views/Explore/ExploreSearch/SavedQueries/index.ts | 2 +- .../javascript/bh-shared-ui/src/views/Explore/index.ts | 2 +- .../bh-shared-ui/src/views/PrivilegeZones/index.tsx | 4 ++-- 11 files changed, 15 insertions(+), 15 deletions(-) diff --git a/cmd/ui/src/ducks/global/types.ts b/cmd/ui/src/ducks/global/types.ts index c6a25d5170b..aed7ea493a1 100644 --- a/cmd/ui/src/ducks/global/types.ts +++ b/cmd/ui/src/ducks/global/types.ts @@ -37,9 +37,9 @@ export { GLOBAL_CLOSE_SNACKBAR, GLOBAL_FETCH_ASSET_GROUPS, GLOBAL_REMOVE_SNACKBAR, + GLOBAL_SET_ASSET_GROUPS, GLOBAL_SET_ASSET_GROUP_EDIT, GLOBAL_SET_ASSET_GROUP_INDEX, - GLOBAL_SET_ASSET_GROUPS, GLOBAL_SET_AUTO_RUN_QUERIES, GLOBAL_SET_DARK_MODE, GLOBAL_SET_DOMAIN, diff --git a/cmd/ui/src/ducks/graph/graphutils.ts b/cmd/ui/src/ducks/graph/graphutils.ts index b444568363f..a408d92aeb3 100644 --- a/cmd/ui/src/ducks/graph/graphutils.ts +++ b/cmd/ui/src/ducks/graph/graphutils.ts @@ -246,6 +246,7 @@ const setFontIcons = (data: Items): void => { }; export { + ICONS, applyRelWidths, findRootId, findRootRelsIds, @@ -255,7 +256,6 @@ export { getLinksIndex, getNodesIndex, handleLabels, - ICONS, setFontIcons, withLinkImact, }; diff --git a/packages/javascript/bh-shared-ui/src/components/AppIcon/Icons/index.ts b/packages/javascript/bh-shared-ui/src/components/AppIcon/Icons/index.ts index 066fe8a8571..69c1814d96d 100644 --- a/packages/javascript/bh-shared-ui/src/components/AppIcon/Icons/index.ts +++ b/packages/javascript/bh-shared-ui/src/components/AppIcon/Icons/index.ts @@ -15,11 +15,11 @@ // SPDX-License-Identifier: Apache-2.0 export * from './AttackPaths'; -export * from './BarChart'; export * from './BHCELogo'; export * from './BHCELogoFull'; export * from './BHELogo'; export * from './BHELogoFull'; +export * from './BarChart'; export * from './CalendarDay'; export * from './CaretDown'; export * from './CaretUp'; diff --git a/packages/javascript/bh-shared-ui/src/components/DropdownSelector/index.ts b/packages/javascript/bh-shared-ui/src/components/DropdownSelector/index.ts index 4088f5582d1..75127953753 100644 --- a/packages/javascript/bh-shared-ui/src/components/DropdownSelector/index.ts +++ b/packages/javascript/bh-shared-ui/src/components/DropdownSelector/index.ts @@ -14,7 +14,7 @@ // // SPDX-License-Identifier: Apache-2.0 -export * from './constants'; export { default as DropdownSelector } from './DropdownSelector'; export { default as DropdownTrigger } from './DropdownTrigger'; +export * from './constants'; export * from './types'; diff --git a/packages/javascript/bh-shared-ui/src/components/ExploreTable/index.ts b/packages/javascript/bh-shared-ui/src/components/ExploreTable/index.ts index 47437f67541..d28626e4956 100644 --- a/packages/javascript/bh-shared-ui/src/components/ExploreTable/index.ts +++ b/packages/javascript/bh-shared-ui/src/components/ExploreTable/index.ts @@ -14,6 +14,6 @@ // // SPDX-License-Identifier: Apache-2.0 -export * from './explore-table-utils'; export { default as ExploreTable } from './ExploreTable'; export type { ManageColumnsComboBoxOption } from './ManageColumnsComboBox/ManageColumnsComboBox'; +export * from './explore-table-utils'; diff --git a/packages/javascript/bh-shared-ui/src/components/index.ts b/packages/javascript/bh-shared-ui/src/components/index.ts index edd8688e636..7e641bdfbf1 100644 --- a/packages/javascript/bh-shared-ui/src/components/index.ts +++ b/packages/javascript/bh-shared-ui/src/components/index.ts @@ -118,6 +118,10 @@ export * from './PasswordResetForm'; export { default as PasswordResetForm } from './PasswordResetForm'; export * from './PrebuiltSearchList'; export { default as PrebuiltSearchList } from './PrebuiltSearchList'; +export * from './SSOProviderInfoPanel'; +export { default as SSOProviderInfoPanel } from './SSOProviderInfoPanel'; +export * from './SSOProviderTable'; +export { default as SSOProviderTable } from './SSOProviderTable'; export * from './SearchCurrentNodes'; export { default as SearchCurrentNodes } from './SearchCurrentNodes'; export * from './SearchInput'; @@ -128,10 +132,6 @@ export * from './SetupKeyDialog'; export { default as SetupKeyDialog } from './SetupKeyDialog'; export * from './SimpleEnvironmentSelector'; export * from './SourceKindsCheckboxes'; -export * from './SSOProviderInfoPanel'; -export { default as SSOProviderInfoPanel } from './SSOProviderInfoPanel'; -export * from './SSOProviderTable'; -export { default as SSOProviderTable } from './SSOProviderTable'; export * from './StatusIndicator'; export * from './TextWithFallback'; export { default as TextWithFallback } from './TextWithFallback'; diff --git a/packages/javascript/bh-shared-ui/src/hooks/index.ts b/packages/javascript/bh-shared-ui/src/hooks/index.ts index cb2c584c55f..7a285b0df1f 100644 --- a/packages/javascript/bh-shared-ui/src/hooks/index.ts +++ b/packages/javascript/bh-shared-ui/src/hooks/index.ts @@ -46,10 +46,10 @@ export * from './useMatchingPaths'; export * from './useMountEffect'; export * from './useObjectState'; export { default as useOnClickOutside } from './useOnClickOutside'; +export * from './usePZParams'; export * from './usePermissions'; export * from './usePrebuiltQueries'; export * from './usePreviousValue'; -export * from './usePZParams'; export { default as useRoleBasedFiltering } from './useRoleBasedFiltering'; export * from './useSavedQueries'; export * from './useSearch'; diff --git a/packages/javascript/bh-shared-ui/src/providers/NotificationProvider/index.ts b/packages/javascript/bh-shared-ui/src/providers/NotificationProvider/index.ts index 2485013da7b..06309632267 100644 --- a/packages/javascript/bh-shared-ui/src/providers/NotificationProvider/index.ts +++ b/packages/javascript/bh-shared-ui/src/providers/NotificationProvider/index.ts @@ -14,6 +14,6 @@ // // SPDX-License-Identifier: Apache-2.0 +export { default as NotificationsProvider } from './NotificationsProvider'; export * from './hooks'; export * from './model'; -export { default as NotificationsProvider } from './NotificationsProvider'; diff --git a/packages/javascript/bh-shared-ui/src/views/Explore/ExploreSearch/SavedQueries/index.ts b/packages/javascript/bh-shared-ui/src/views/Explore/ExploreSearch/SavedQueries/index.ts index c89053c320c..6f294f90f24 100644 --- a/packages/javascript/bh-shared-ui/src/views/Explore/ExploreSearch/SavedQueries/index.ts +++ b/packages/javascript/bh-shared-ui/src/views/Explore/ExploreSearch/SavedQueries/index.ts @@ -18,8 +18,8 @@ export { default as ConfirmUpdateQueryDialog } from './ConfirmUpdateQueryDialog' export { default as CypherSearchMessage } from './CypherSearchMessage'; export { default as ImportQueryDialog } from './ImportQueryDialog'; export { default as QuerySearchFilter } from './QuerySearchFilter'; -export { default as SavedQueryPermissions } from './SavedQueryPermissions'; export { default as SaveQueryActionMenu } from './SaveQueryActionMenu'; export { default as SaveQueryDialog } from './SaveQueryDialog'; +export { default as SavedQueryPermissions } from './SavedQueryPermissions'; export { default as TagToZoneLabel } from './TagToZoneLabel'; export { default as TagToZoneLabelDialog } from './TagToZoneLabelDialog'; diff --git a/packages/javascript/bh-shared-ui/src/views/Explore/index.ts b/packages/javascript/bh-shared-ui/src/views/Explore/index.ts index f9ccbf319b0..73f1c169105 100644 --- a/packages/javascript/bh-shared-ui/src/views/Explore/index.ts +++ b/packages/javascript/bh-shared-ui/src/views/Explore/index.ts @@ -20,6 +20,6 @@ export { default as ContextMenuPrivilegeZonesEnabled } from './ContextMenu/Conte export { default as CopyMenuItem } from './ContextMenu/CopyMenuItem'; export * from './EdgeInfo'; export * from './ExploreSearch'; -export * from './fragments'; export * from './InfoStyles'; +export * from './fragments'; export * from './providers'; diff --git a/packages/javascript/bh-shared-ui/src/views/PrivilegeZones/index.tsx b/packages/javascript/bh-shared-ui/src/views/PrivilegeZones/index.tsx index 1f1e01ed088..93d58c3d881 100644 --- a/packages/javascript/bh-shared-ui/src/views/PrivilegeZones/index.tsx +++ b/packages/javascript/bh-shared-ui/src/views/PrivilegeZones/index.tsx @@ -18,9 +18,9 @@ import PrivilegeZones from './PrivilegeZones'; export { EntityRulesInformation, SelectedDetails } from './Details'; export * from './Filters'; -export { CreateRuleButtonLink, EditRuleButtonLink, EditTagButtonLink, PageDescription, ZonesLink } from './fragments'; export * from './PrivilegeZonesContext'; -export * from './utils'; export * from './ZoneIcon'; +export { CreateRuleButtonLink, EditRuleButtonLink, EditTagButtonLink, PageDescription, ZonesLink } from './fragments'; +export * from './utils'; export default PrivilegeZones; From 309111e61e9d8197aaa8dacd908f4772eef01ac1 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Tue, 3 Feb 2026 05:45:21 -0800 Subject: [PATCH 08/22] Tenant is incorrect --- cmd/api/src/database/migration/extensions/az_graph_schema.sql | 3 +-- packages/go/schemagen/generator/sql.go | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/cmd/api/src/database/migration/extensions/az_graph_schema.sql b/cmd/api/src/database/migration/extensions/az_graph_schema.sql index 7625a81fa4e..b9dd7d39761 100644 --- a/cmd/api/src/database/migration/extensions/az_graph_schema.sql +++ b/cmd/api/src/database/migration/extensions/az_graph_schema.sql @@ -262,8 +262,7 @@ BEGIN PERFORM genscript_upsert_schema_relationship_kind(extension_id, 'AZRoleApprover', '', true); PERFORM genscript_upsert_source_kind('AZBase'); - PERFORM genscript_upsert_kind('Tenant'); - SELECT genscript_upsert_schema_environments(extension_id, 'Tenant', 'AZBase') INTO environment_id; + SELECT genscript_upsert_schema_environments(extension_id, 'AZTenant', 'AZBase') INTO environment_id; PERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'AZUser'); PERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'AZVM'); PERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'AZServicePrincipal'); diff --git a/packages/go/schemagen/generator/sql.go b/packages/go/schemagen/generator/sql.go index 29e7e770553..bf8255b9059 100644 --- a/packages/go/schemagen/generator/sql.go +++ b/packages/go/schemagen/generator/sql.go @@ -365,8 +365,7 @@ func GenerateADSpecifics(sb io.StringWriter) { func GenerateAZSpecifics(sb io.StringWriter) { sb.WriteString("\tPERFORM genscript_upsert_source_kind('AZBase');\n") - sb.WriteString("\tPERFORM genscript_upsert_kind('Tenant');\n") - sb.WriteString("\tSELECT genscript_upsert_schema_environments(extension_id, 'Tenant', 'AZBase') INTO environment_id;\n") + sb.WriteString("\tSELECT genscript_upsert_schema_environments(extension_id, 'AZTenant', 'AZBase') INTO environment_id;\n") sb.WriteString("\tPERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'AZUser');\n") sb.WriteString("\tPERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'AZVM');\n") sb.WriteString("\tPERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'AZServicePrincipal');\n") From 912af8be5f4029fc003c1cf7eec0739c2399f6ff Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Tue, 3 Feb 2026 06:16:20 -0800 Subject: [PATCH 09/22] just prepare for codereview --- cmd/ui/src/ducks/global/types.ts | 2 +- cmd/ui/src/ducks/graph/graphutils.ts | 2 +- .../bh-shared-ui/src/components/AppIcon/Icons/index.ts | 2 +- .../bh-shared-ui/src/components/DropdownSelector/index.ts | 2 +- .../bh-shared-ui/src/components/ExploreTable/index.ts | 2 +- packages/javascript/bh-shared-ui/src/components/index.ts | 8 ++++---- packages/javascript/bh-shared-ui/src/hooks/index.ts | 2 +- .../src/providers/NotificationProvider/index.ts | 2 +- .../src/views/Explore/ExploreSearch/SavedQueries/index.ts | 2 +- .../javascript/bh-shared-ui/src/views/Explore/index.ts | 2 +- .../bh-shared-ui/src/views/PrivilegeZones/index.tsx | 4 ++-- 11 files changed, 15 insertions(+), 15 deletions(-) diff --git a/cmd/ui/src/ducks/global/types.ts b/cmd/ui/src/ducks/global/types.ts index aed7ea493a1..c6a25d5170b 100644 --- a/cmd/ui/src/ducks/global/types.ts +++ b/cmd/ui/src/ducks/global/types.ts @@ -37,9 +37,9 @@ export { GLOBAL_CLOSE_SNACKBAR, GLOBAL_FETCH_ASSET_GROUPS, GLOBAL_REMOVE_SNACKBAR, - GLOBAL_SET_ASSET_GROUPS, GLOBAL_SET_ASSET_GROUP_EDIT, GLOBAL_SET_ASSET_GROUP_INDEX, + GLOBAL_SET_ASSET_GROUPS, GLOBAL_SET_AUTO_RUN_QUERIES, GLOBAL_SET_DARK_MODE, GLOBAL_SET_DOMAIN, diff --git a/cmd/ui/src/ducks/graph/graphutils.ts b/cmd/ui/src/ducks/graph/graphutils.ts index a408d92aeb3..b444568363f 100644 --- a/cmd/ui/src/ducks/graph/graphutils.ts +++ b/cmd/ui/src/ducks/graph/graphutils.ts @@ -246,7 +246,6 @@ const setFontIcons = (data: Items): void => { }; export { - ICONS, applyRelWidths, findRootId, findRootRelsIds, @@ -256,6 +255,7 @@ export { getLinksIndex, getNodesIndex, handleLabels, + ICONS, setFontIcons, withLinkImact, }; diff --git a/packages/javascript/bh-shared-ui/src/components/AppIcon/Icons/index.ts b/packages/javascript/bh-shared-ui/src/components/AppIcon/Icons/index.ts index 69c1814d96d..066fe8a8571 100644 --- a/packages/javascript/bh-shared-ui/src/components/AppIcon/Icons/index.ts +++ b/packages/javascript/bh-shared-ui/src/components/AppIcon/Icons/index.ts @@ -15,11 +15,11 @@ // SPDX-License-Identifier: Apache-2.0 export * from './AttackPaths'; +export * from './BarChart'; export * from './BHCELogo'; export * from './BHCELogoFull'; export * from './BHELogo'; export * from './BHELogoFull'; -export * from './BarChart'; export * from './CalendarDay'; export * from './CaretDown'; export * from './CaretUp'; diff --git a/packages/javascript/bh-shared-ui/src/components/DropdownSelector/index.ts b/packages/javascript/bh-shared-ui/src/components/DropdownSelector/index.ts index 75127953753..4088f5582d1 100644 --- a/packages/javascript/bh-shared-ui/src/components/DropdownSelector/index.ts +++ b/packages/javascript/bh-shared-ui/src/components/DropdownSelector/index.ts @@ -14,7 +14,7 @@ // // SPDX-License-Identifier: Apache-2.0 +export * from './constants'; export { default as DropdownSelector } from './DropdownSelector'; export { default as DropdownTrigger } from './DropdownTrigger'; -export * from './constants'; export * from './types'; diff --git a/packages/javascript/bh-shared-ui/src/components/ExploreTable/index.ts b/packages/javascript/bh-shared-ui/src/components/ExploreTable/index.ts index d28626e4956..47437f67541 100644 --- a/packages/javascript/bh-shared-ui/src/components/ExploreTable/index.ts +++ b/packages/javascript/bh-shared-ui/src/components/ExploreTable/index.ts @@ -14,6 +14,6 @@ // // SPDX-License-Identifier: Apache-2.0 +export * from './explore-table-utils'; export { default as ExploreTable } from './ExploreTable'; export type { ManageColumnsComboBoxOption } from './ManageColumnsComboBox/ManageColumnsComboBox'; -export * from './explore-table-utils'; diff --git a/packages/javascript/bh-shared-ui/src/components/index.ts b/packages/javascript/bh-shared-ui/src/components/index.ts index 7e641bdfbf1..edd8688e636 100644 --- a/packages/javascript/bh-shared-ui/src/components/index.ts +++ b/packages/javascript/bh-shared-ui/src/components/index.ts @@ -118,10 +118,6 @@ export * from './PasswordResetForm'; export { default as PasswordResetForm } from './PasswordResetForm'; export * from './PrebuiltSearchList'; export { default as PrebuiltSearchList } from './PrebuiltSearchList'; -export * from './SSOProviderInfoPanel'; -export { default as SSOProviderInfoPanel } from './SSOProviderInfoPanel'; -export * from './SSOProviderTable'; -export { default as SSOProviderTable } from './SSOProviderTable'; export * from './SearchCurrentNodes'; export { default as SearchCurrentNodes } from './SearchCurrentNodes'; export * from './SearchInput'; @@ -132,6 +128,10 @@ export * from './SetupKeyDialog'; export { default as SetupKeyDialog } from './SetupKeyDialog'; export * from './SimpleEnvironmentSelector'; export * from './SourceKindsCheckboxes'; +export * from './SSOProviderInfoPanel'; +export { default as SSOProviderInfoPanel } from './SSOProviderInfoPanel'; +export * from './SSOProviderTable'; +export { default as SSOProviderTable } from './SSOProviderTable'; export * from './StatusIndicator'; export * from './TextWithFallback'; export { default as TextWithFallback } from './TextWithFallback'; diff --git a/packages/javascript/bh-shared-ui/src/hooks/index.ts b/packages/javascript/bh-shared-ui/src/hooks/index.ts index 7a285b0df1f..cb2c584c55f 100644 --- a/packages/javascript/bh-shared-ui/src/hooks/index.ts +++ b/packages/javascript/bh-shared-ui/src/hooks/index.ts @@ -46,10 +46,10 @@ export * from './useMatchingPaths'; export * from './useMountEffect'; export * from './useObjectState'; export { default as useOnClickOutside } from './useOnClickOutside'; -export * from './usePZParams'; export * from './usePermissions'; export * from './usePrebuiltQueries'; export * from './usePreviousValue'; +export * from './usePZParams'; export { default as useRoleBasedFiltering } from './useRoleBasedFiltering'; export * from './useSavedQueries'; export * from './useSearch'; diff --git a/packages/javascript/bh-shared-ui/src/providers/NotificationProvider/index.ts b/packages/javascript/bh-shared-ui/src/providers/NotificationProvider/index.ts index 06309632267..2485013da7b 100644 --- a/packages/javascript/bh-shared-ui/src/providers/NotificationProvider/index.ts +++ b/packages/javascript/bh-shared-ui/src/providers/NotificationProvider/index.ts @@ -14,6 +14,6 @@ // // SPDX-License-Identifier: Apache-2.0 -export { default as NotificationsProvider } from './NotificationsProvider'; export * from './hooks'; export * from './model'; +export { default as NotificationsProvider } from './NotificationsProvider'; diff --git a/packages/javascript/bh-shared-ui/src/views/Explore/ExploreSearch/SavedQueries/index.ts b/packages/javascript/bh-shared-ui/src/views/Explore/ExploreSearch/SavedQueries/index.ts index 6f294f90f24..c89053c320c 100644 --- a/packages/javascript/bh-shared-ui/src/views/Explore/ExploreSearch/SavedQueries/index.ts +++ b/packages/javascript/bh-shared-ui/src/views/Explore/ExploreSearch/SavedQueries/index.ts @@ -18,8 +18,8 @@ export { default as ConfirmUpdateQueryDialog } from './ConfirmUpdateQueryDialog' export { default as CypherSearchMessage } from './CypherSearchMessage'; export { default as ImportQueryDialog } from './ImportQueryDialog'; export { default as QuerySearchFilter } from './QuerySearchFilter'; +export { default as SavedQueryPermissions } from './SavedQueryPermissions'; export { default as SaveQueryActionMenu } from './SaveQueryActionMenu'; export { default as SaveQueryDialog } from './SaveQueryDialog'; -export { default as SavedQueryPermissions } from './SavedQueryPermissions'; export { default as TagToZoneLabel } from './TagToZoneLabel'; export { default as TagToZoneLabelDialog } from './TagToZoneLabelDialog'; diff --git a/packages/javascript/bh-shared-ui/src/views/Explore/index.ts b/packages/javascript/bh-shared-ui/src/views/Explore/index.ts index 73f1c169105..f9ccbf319b0 100644 --- a/packages/javascript/bh-shared-ui/src/views/Explore/index.ts +++ b/packages/javascript/bh-shared-ui/src/views/Explore/index.ts @@ -20,6 +20,6 @@ export { default as ContextMenuPrivilegeZonesEnabled } from './ContextMenu/Conte export { default as CopyMenuItem } from './ContextMenu/CopyMenuItem'; export * from './EdgeInfo'; export * from './ExploreSearch'; -export * from './InfoStyles'; export * from './fragments'; +export * from './InfoStyles'; export * from './providers'; diff --git a/packages/javascript/bh-shared-ui/src/views/PrivilegeZones/index.tsx b/packages/javascript/bh-shared-ui/src/views/PrivilegeZones/index.tsx index 93d58c3d881..1f1e01ed088 100644 --- a/packages/javascript/bh-shared-ui/src/views/PrivilegeZones/index.tsx +++ b/packages/javascript/bh-shared-ui/src/views/PrivilegeZones/index.tsx @@ -18,9 +18,9 @@ import PrivilegeZones from './PrivilegeZones'; export { EntityRulesInformation, SelectedDetails } from './Details'; export * from './Filters'; -export * from './PrivilegeZonesContext'; -export * from './ZoneIcon'; export { CreateRuleButtonLink, EditRuleButtonLink, EditTagButtonLink, PageDescription, ZonesLink } from './fragments'; +export * from './PrivilegeZonesContext'; export * from './utils'; +export * from './ZoneIcon'; export default PrivilegeZones; From 2f0f74ffb8c97701d02c995af169e02cbb3a748e Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Tue, 3 Feb 2026 07:06:28 -0800 Subject: [PATCH 10/22] consolidate getenv methods --- cmd/api/src/database/graphschema.go | 75 ++++++++++------------------- 1 file changed, 26 insertions(+), 49 deletions(-) diff --git a/cmd/api/src/database/graphschema.go b/cmd/api/src/database/graphschema.go index ada6dc8b5b1..cd6125cc174 100644 --- a/cmd/api/src/database/graphschema.go +++ b/cmd/api/src/database/graphschema.go @@ -595,40 +595,10 @@ func (s *BloodhoundDB) CreateEnvironment(ctx context.Context, extensionId int32, return schemaEnvironment, nil } -// GetEnvironments - retrieves list of schema environments. -func (s *BloodhoundDB) GetEnvironments(ctx context.Context) ([]model.SchemaEnvironment, error) { - var result []model.SchemaEnvironment - - query := ` - SELECT - se.id, - se.schema_extension_id, - ext.display_name as schema_extension_display_name, - se.environment_kind_id, - k.name as environment_kind_name, - se.source_kind_id, - se.created_at, - se.updated_at, - se.deleted_at - FROM schema_environments se - INNER JOIN kind k ON se.environment_kind_id = k.id - INNER JOIN schema_extensions ext ON se.schema_extension_id = ext.id - ORDER BY se.id` - - if err := CheckError(s.db.WithContext(ctx).Raw(query).Scan(&result)); err != nil { - return nil, err - } - - if result == nil { - result = []model.SchemaEnvironment{} - } - - return result, nil -} - // GetEnvironmentsFiltered - retrieves schema environments filtered by the given criteria. +// This is the core implementation that all other GetEnvironment* methods delegate to. // Common use case: filter by schema_extension_id to get all environments for a specific extension. -// Example: filters := model.Filters{"schema_extension_id": []model.Filter{{Operator: model.Equals, Value: "1"}}} +// Example: filters := model.Filters{"se.schema_extension_id": []model.Filter{{Operator: model.Equals, Value: "1"}}} func (s *BloodhoundDB) GetEnvironmentsFiltered(ctx context.Context, filters model.Filters) ([]model.SchemaEnvironment, error) { var result []model.SchemaEnvironment @@ -670,37 +640,44 @@ func (s *BloodhoundDB) GetEnvironmentsFiltered(ctx context.Context, filters mode return result, nil } +// GetEnvironments - retrieves all schema environments. +func (s *BloodhoundDB) GetEnvironments(ctx context.Context) ([]model.SchemaEnvironment, error) { + return s.GetEnvironmentsFiltered(ctx, model.Filters{}) +} + // GetEnvironmentByKinds - retrieves an environment by its environment kind and source kind. func (s *BloodhoundDB) GetEnvironmentByKinds(ctx context.Context, environmentKindId, sourceKindId int32) (model.SchemaEnvironment, error) { - var env model.SchemaEnvironment + filters := model.Filters{ + "se.environment_kind_id": []model.Filter{{Operator: model.Equals, Value: fmt.Sprintf("%d", environmentKindId)}}, + "se.source_kind_id": []model.Filter{{Operator: model.Equals, Value: fmt.Sprintf("%d", sourceKindId)}}, + } - if result := s.db.WithContext(ctx).Raw( - "SELECT * FROM schema_environments WHERE environment_kind_id = ? AND source_kind_id = ? AND deleted_at IS NULL", - environmentKindId, sourceKindId, - ).Scan(&env); result.Error != nil { - return model.SchemaEnvironment{}, CheckError(result) - } else if result.RowsAffected == 0 { + envs, err := s.GetEnvironmentsFiltered(ctx, filters) + if err != nil { + return model.SchemaEnvironment{}, err + } + if len(envs) == 0 { return model.SchemaEnvironment{}, ErrNotFound } - return env, nil + return envs[0], nil } // GetEnvironmentById - retrieves a schema environment by id. func (s *BloodhoundDB) GetEnvironmentById(ctx context.Context, environmentId int32) (model.SchemaEnvironment, error) { - var schemaEnvironment model.SchemaEnvironment + filters := model.Filters{ + "se.id": []model.Filter{{Operator: model.Equals, Value: fmt.Sprintf("%d", environmentId)}}, + } - if result := s.db.WithContext(ctx).Raw(fmt.Sprintf(` - SELECT id, schema_extension_id, environment_kind_id, source_kind_id, created_at, updated_at, deleted_at - FROM %s WHERE id = ?`, - schemaEnvironment.TableName()), - environmentId).Scan(&schemaEnvironment); result.Error != nil { - return model.SchemaEnvironment{}, CheckError(result) - } else if result.RowsAffected == 0 { + envs, err := s.GetEnvironmentsFiltered(ctx, filters) + if err != nil { + return model.SchemaEnvironment{}, err + } + if len(envs) == 0 { return model.SchemaEnvironment{}, ErrNotFound } - return schemaEnvironment, nil + return envs[0], nil } // DeleteEnvironment - deletes a schema environment by id. From fcf0fd0d5ec2156fe26202cffe23d5895a3aec7c Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Tue, 3 Feb 2026 07:17:29 -0800 Subject: [PATCH 11/22] single helper --- cmd/api/src/database/graphschema.go | 88 ++++++++++++++++++++--------- 1 file changed, 60 insertions(+), 28 deletions(-) diff --git a/cmd/api/src/database/graphschema.go b/cmd/api/src/database/graphschema.go index cd6125cc174..0e4333363b5 100644 --- a/cmd/api/src/database/graphschema.go +++ b/cmd/api/src/database/graphschema.go @@ -711,54 +711,86 @@ func (s *BloodhoundDB) CreateSchemaRelationshipFinding(ctx context.Context, exte return finding, nil } +// getSchemaRelationshipFindingsFiltered - retrieves schema relationship findings filtered by the given criteria. +// This is the core implementation that all other GetSchemaRelationshipFinding* methods delegate to. +func (s *BloodhoundDB) getSchemaRelationshipFindingsFiltered(ctx context.Context, filters model.Filters) ([]model.SchemaRelationshipFinding, error) { + var result []model.SchemaRelationshipFinding + + sqlFilter, err := buildSQLFilter(filters) + if err != nil { + return nil, err + } + + whereClause := "" + if sqlFilter.sqlString != "" { + whereClause = fmt.Sprintf("WHERE %s", sqlFilter.sqlString) + } + + query := fmt.Sprintf(` + SELECT + srf.id, + srf.schema_extension_id, + srf.relationship_kind_id, + srf.environment_id, + srf.name, + srf.display_name, + srf.created_at + FROM schema_relationship_findings srf + %s + ORDER BY srf.id`, whereClause) + + if err := CheckError(s.db.WithContext(ctx).Raw(query, sqlFilter.params...).Scan(&result)); err != nil { + return nil, err + } + + if result == nil { + result = []model.SchemaRelationshipFinding{} + } + + return result, nil +} + // GetSchemaRelationshipFindingById - retrieves a schema relationship finding by id. func (s *BloodhoundDB) GetSchemaRelationshipFindingById(ctx context.Context, findingId int32) (model.SchemaRelationshipFinding, error) { - var finding model.SchemaRelationshipFinding + filters := model.Filters{ + "srf.id": []model.Filter{{Operator: model.Equals, Value: fmt.Sprintf("%d", findingId)}}, + } - if result := s.db.WithContext(ctx).Raw(fmt.Sprintf(` - SELECT id, schema_extension_id, relationship_kind_id, environment_id, name, display_name, created_at - FROM %s WHERE id = ?`, - finding.TableName()), - findingId).Scan(&finding); result.Error != nil { - return model.SchemaRelationshipFinding{}, CheckError(result) - } else if result.RowsAffected == 0 { + findings, err := s.getSchemaRelationshipFindingsFiltered(ctx, filters) + if err != nil { + return model.SchemaRelationshipFinding{}, err + } + if len(findings) == 0 { return model.SchemaRelationshipFinding{}, ErrNotFound } - return finding, nil + return findings[0], nil } // GetSchemaRelationshipFindingByName - retrieves a schema relationship finding by finding name. func (s *BloodhoundDB) GetSchemaRelationshipFindingByName(ctx context.Context, name string) (model.SchemaRelationshipFinding, error) { - var finding model.SchemaRelationshipFinding + filters := model.Filters{ + "srf.name": []model.Filter{{Operator: model.Equals, Value: name}}, + } - if result := s.db.WithContext(ctx).Raw(fmt.Sprintf(` - SELECT id, schema_extension_id, relationship_kind_id, environment_id, name, display_name, created_at - FROM %s WHERE name = ?`, - finding.TableName()), - name).Scan(&finding); result.Error != nil { - return model.SchemaRelationshipFinding{}, CheckError(result) - } else if result.RowsAffected == 0 { + findings, err := s.getSchemaRelationshipFindingsFiltered(ctx, filters) + if err != nil { + return model.SchemaRelationshipFinding{}, err + } + if len(findings) == 0 { return model.SchemaRelationshipFinding{}, ErrNotFound } - return finding, nil + return findings[0], nil } // GetSchemaRelationshipFindingsByEnvironmentId - retrieves all schema relationship findings for a given environment. func (s *BloodhoundDB) GetSchemaRelationshipFindingsByEnvironmentId(ctx context.Context, environmentId int32) ([]model.SchemaRelationshipFinding, error) { - var findings []model.SchemaRelationshipFinding - var finding model.SchemaRelationshipFinding - - if result := s.db.WithContext(ctx).Raw(fmt.Sprintf(` - SELECT id, schema_extension_id, relationship_kind_id, environment_id, name, display_name, created_at - FROM %s WHERE environment_id = ?`, - finding.TableName()), - environmentId).Scan(&findings); result.Error != nil { - return nil, CheckError(result) + filters := model.Filters{ + "srf.environment_id": []model.Filter{{Operator: model.Equals, Value: fmt.Sprintf("%d", environmentId)}}, } - return findings, nil + return s.getSchemaRelationshipFindingsFiltered(ctx, filters) } // DeleteSchemaRelationshipFinding - deletes a schema relationship finding by id. From d83a5f9aebc15da53fa2de5a224f78ec4c31c72f Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Fri, 6 Feb 2026 10:00:20 -0800 Subject: [PATCH 12/22] remove raise exceptions -> genscript upsert kind --- .../database/migration/extensions/ad_graph_schema.sql | 10 +++++----- .../database/migration/extensions/az_graph_schema.sql | 9 +++++---- packages/go/schemagen/generator/sql.go | 11 +++++------ 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/cmd/api/src/database/migration/extensions/ad_graph_schema.sql b/cmd/api/src/database/migration/extensions/ad_graph_schema.sql index c930a9f0aa6..e92abafb206 100644 --- a/cmd/api/src/database/migration/extensions/ad_graph_schema.sql +++ b/cmd/api/src/database/migration/extensions/ad_graph_schema.sql @@ -71,14 +71,16 @@ DECLARE BEGIN SELECT id INTO retreived_environment_kind_id FROM kind WHERE name = v_environment_kind_name; IF retreived_environment_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + PERFORM genscript_upsert_kind(v_environment_kind_name); + SELECT id INTO retreived_environment_kind_id FROM kind WHERE name = v_environment_kind_name; END IF; SELECT id INTO retreived_source_kind_id FROM source_kinds WHERE name = v_source_kind_name; IF retreived_source_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + PERFORM genscript_upsert_source_kind(v_source_kind_name); + SELECT id INTO retreived_source_kind_id FROM source_kinds WHERE name = v_source_kind_name; END IF; - + IF NOT EXISTS (SELECT id FROM schema_environments se WHERE se.schema_extension_id = v_extension_id) THEN INSERT INTO schema_environments (schema_extension_id, environment_kind_id, source_kind_id) VALUES (v_extension_id, retreived_environment_kind_id, retreived_source_kind_id) RETURNING id INTO schema_environment_id; ELSE @@ -327,8 +329,6 @@ BEGIN PERFORM genscript_upsert_schema_relationship_kind(extension_id, 'HasTrustKeys', '', true); PERFORM genscript_upsert_schema_relationship_kind(extension_id, 'ProtectAdminGroups', '', false); - PERFORM genscript_upsert_source_kind('Base'); - PERFORM genscript_upsert_kind('Domain'); SELECT genscript_upsert_schema_environments(extension_id, 'Domain', 'Base') INTO environment_id; PERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'User'); PERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'Computer'); diff --git a/cmd/api/src/database/migration/extensions/az_graph_schema.sql b/cmd/api/src/database/migration/extensions/az_graph_schema.sql index b9dd7d39761..2275e16c4f8 100644 --- a/cmd/api/src/database/migration/extensions/az_graph_schema.sql +++ b/cmd/api/src/database/migration/extensions/az_graph_schema.sql @@ -71,14 +71,16 @@ DECLARE BEGIN SELECT id INTO retreived_environment_kind_id FROM kind WHERE name = v_environment_kind_name; IF retreived_environment_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + PERFORM genscript_upsert_kind(v_environment_kind_name); + SELECT id INTO retreived_environment_kind_id FROM kind WHERE name = v_environment_kind_name; END IF; SELECT id INTO retreived_source_kind_id FROM source_kinds WHERE name = v_source_kind_name; IF retreived_source_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + PERFORM genscript_upsert_source_kind(v_source_kind_name); + SELECT id INTO retreived_source_kind_id FROM source_kinds WHERE name = v_source_kind_name; END IF; - + IF NOT EXISTS (SELECT id FROM schema_environments se WHERE se.schema_extension_id = v_extension_id) THEN INSERT INTO schema_environments (schema_extension_id, environment_kind_id, source_kind_id) VALUES (v_extension_id, retreived_environment_kind_id, retreived_source_kind_id) RETURNING id INTO schema_environment_id; ELSE @@ -261,7 +263,6 @@ BEGIN PERFORM genscript_upsert_schema_relationship_kind(extension_id, 'AZRoleEligible', '', true); PERFORM genscript_upsert_schema_relationship_kind(extension_id, 'AZRoleApprover', '', true); - PERFORM genscript_upsert_source_kind('AZBase'); SELECT genscript_upsert_schema_environments(extension_id, 'AZTenant', 'AZBase') INTO environment_id; PERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'AZUser'); PERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'AZVM'); diff --git a/packages/go/schemagen/generator/sql.go b/packages/go/schemagen/generator/sql.go index bf8255b9059..d5e5871b810 100644 --- a/packages/go/schemagen/generator/sql.go +++ b/packages/go/schemagen/generator/sql.go @@ -241,14 +241,16 @@ DECLARE BEGIN SELECT id INTO retreived_environment_kind_id FROM kind WHERE name = v_environment_kind_name; IF retreived_environment_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + PERFORM genscript_upsert_kind(v_environment_kind_name); + SELECT id INTO retreived_environment_kind_id FROM kind WHERE name = v_environment_kind_name; END IF; SELECT id INTO retreived_source_kind_id FROM source_kinds WHERE name = v_source_kind_name; IF retreived_source_kind_id IS NULL THEN - RAISE EXCEPTION 'couldn''t find matching kind_id'; + PERFORM genscript_upsert_source_kind(v_source_kind_name); + SELECT id INTO retreived_source_kind_id FROM source_kinds WHERE name = v_source_kind_name; END IF; - + IF NOT EXISTS (SELECT id FROM schema_environments se WHERE se.schema_extension_id = v_extension_id) THEN INSERT INTO schema_environments (schema_extension_id, environment_kind_id, source_kind_id) VALUES (v_extension_id, retreived_environment_kind_id, retreived_source_kind_id) RETURNING id INTO schema_environment_id; ELSE @@ -356,15 +358,12 @@ DROP FUNCTION IF EXISTS genscript_upsert_schema_environments_principal_kinds;`) } func GenerateADSpecifics(sb io.StringWriter) { - sb.WriteString("\tPERFORM genscript_upsert_source_kind('Base');\n") - sb.WriteString("\tPERFORM genscript_upsert_kind('Domain');\n") sb.WriteString("\tSELECT genscript_upsert_schema_environments(extension_id, 'Domain', 'Base') INTO environment_id;\n") sb.WriteString("\tPERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'User');\n") sb.WriteString("\tPERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'Computer');\n") } func GenerateAZSpecifics(sb io.StringWriter) { - sb.WriteString("\tPERFORM genscript_upsert_source_kind('AZBase');\n") sb.WriteString("\tSELECT genscript_upsert_schema_environments(extension_id, 'AZTenant', 'AZBase') INTO environment_id;\n") sb.WriteString("\tPERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'AZUser');\n") sb.WriteString("\tPERFORM genscript_upsert_schema_environments_principal_kinds(environment_id, 'AZVM');\n") From 974e4834f5439e36596d30aecdf3018838e7d094 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Fri, 6 Feb 2026 10:24:54 -0800 Subject: [PATCH 13/22] better type for sourcekind methods --- cmd/api/src/api/v2/database_wipe.go | 2 +- cmd/api/src/api/v2/kinds.go | 2 +- cmd/api/src/daemons/datapipe/pipeline.go | 2 +- cmd/api/src/database/sourcekinds.go | 16 ++++--- .../database/sourcekinds_integration_test.go | 44 +++++++++---------- 5 files changed, 35 insertions(+), 31 deletions(-) diff --git a/cmd/api/src/api/v2/database_wipe.go b/cmd/api/src/api/v2/database_wipe.go index 605f0a030f2..3347d3b3a4c 100644 --- a/cmd/api/src/api/v2/database_wipe.go +++ b/cmd/api/src/api/v2/database_wipe.go @@ -279,7 +279,7 @@ func (s Resources) BuildDeleteRequest(ctx context.Context, userID string, payloa found := false for _, sk := range sourceKinds { if sk.ID == id { - requestedKinds = append(requestedKinds, sk.Name) + requestedKinds = append(requestedKinds, sk.ToKind()) found = true break } diff --git a/cmd/api/src/api/v2/kinds.go b/cmd/api/src/api/v2/kinds.go index f60538268e7..7ed32945c3c 100644 --- a/cmd/api/src/api/v2/kinds.go +++ b/cmd/api/src/api/v2/kinds.go @@ -59,7 +59,7 @@ func (s Resources) ListSourceKinds(response http.ResponseWriter, request *http.R } else { // inject 0, Sourceless into the payload. We don't track this as an official kind // but it will facilitate delete requests for data that isn't associated with a kind. - kinds = append(kinds, database.SourceKind{ID: 0, Name: graph.StringKind("Sourceless")}) + kinds = append(kinds, database.SourceKind{ID: 0, Name: "Sourceless"}) api.WriteBasicResponse(request.Context(), ListSourceKindsResponse{Kinds: kinds}, http.StatusOK, response) } } diff --git a/cmd/api/src/daemons/datapipe/pipeline.go b/cmd/api/src/daemons/datapipe/pipeline.go index 835bedf4ccf..189d9fcdc93 100644 --- a/cmd/api/src/daemons/datapipe/pipeline.go +++ b/cmd/api/src/daemons/datapipe/pipeline.go @@ -127,7 +127,7 @@ func PurgeGraphData( func extractKindNames(sourceKinds []database.SourceKind) graph.Kinds { var kinds graph.Kinds for _, k := range sourceKinds { - kinds = append(kinds, k.Name) + kinds = append(kinds, k.ToKind()) } return kinds } diff --git a/cmd/api/src/database/sourcekinds.go b/cmd/api/src/database/sourcekinds.go index 29307b450f9..ca4e53e7d44 100644 --- a/cmd/api/src/database/sourcekinds.go +++ b/cmd/api/src/database/sourcekinds.go @@ -58,9 +58,13 @@ func (s *BloodhoundDB) RegisterSourceKind(ctx context.Context) func(sourceKind g } type SourceKind struct { - ID int `json:"id"` - Name graph.Kind `json:"name"` - Active bool `json:"active"` + ID int `json:"id"` + Name string `json:"name"` + Active bool `json:"active"` +} + +func (s *SourceKind) ToKind() graph.Kind { + return graph.StringKind(s.Name) } func (s *BloodhoundDB) GetSourceKinds(ctx context.Context) ([]SourceKind, error) { @@ -87,7 +91,7 @@ func (s *BloodhoundDB) GetSourceKinds(ctx context.Context) ([]SourceKind, error) for i, k := range kinds { out[i] = SourceKind{ ID: k.ID, - Name: graph.StringKind(k.Name), + Name: k.Name, Active: k.Active, } } @@ -121,7 +125,7 @@ func (s *BloodhoundDB) GetSourceKindByName(ctx context.Context, name string) (So kind := SourceKind{ ID: raw.ID, - Name: graph.StringKind(raw.Name), + Name: raw.Name, Active: raw.Active, } @@ -154,7 +158,7 @@ func (s *BloodhoundDB) GetSourceKindById(ctx context.Context, id int32) (SourceK kind := SourceKind{ ID: raw.ID, - Name: graph.StringKind(raw.Name), + Name: raw.Name, Active: raw.Active, } diff --git a/cmd/api/src/database/sourcekinds_integration_test.go b/cmd/api/src/database/sourcekinds_integration_test.go index 0d8931d504b..a3fe95c2bd4 100644 --- a/cmd/api/src/database/sourcekinds_integration_test.go +++ b/cmd/api/src/database/sourcekinds_integration_test.go @@ -56,12 +56,12 @@ func TestRegisterSourceKind(t *testing.T) { sourceKinds: []database.SourceKind{ { ID: 2, - Name: graph.StringKind("AZBase"), + Name: "AZBase", Active: true, }, { ID: 1, - Name: graph.StringKind("Base"), + Name: "Base", Active: true, }, }, @@ -80,17 +80,17 @@ func TestRegisterSourceKind(t *testing.T) { sourceKinds: []database.SourceKind{ { ID: 2, - Name: graph.StringKind("AZBase"), + Name: "AZBase", Active: true, }, { ID: 1, - Name: graph.StringKind("Base"), + Name: "Base", Active: true, }, { ID: 3, - Name: graph.StringKind("harnessEdge.Kind"), + Name: "harnessEdge.Kind", Active: true, }, }, @@ -118,17 +118,17 @@ func TestRegisterSourceKind(t *testing.T) { sourceKinds: []database.SourceKind{ { ID: 2, - Name: graph.StringKind("AZBase"), + Name: "AZBase", Active: true, }, { ID: 1, - Name: graph.StringKind("Base"), + Name: "Base", Active: true, }, { ID: 3, - Name: graph.StringKind("Kind"), + Name: "Kind", Active: true, }, }, @@ -178,12 +178,12 @@ func TestGetSourceKinds(t *testing.T) { sourceKinds: []database.SourceKind{ { ID: 2, - Name: graph.StringKind("AZBase"), + Name: "AZBase", Active: true, }, { ID: 1, - Name: graph.StringKind("Base"), + Name: "Base", Active: true, }, }, @@ -234,7 +234,7 @@ func TestGetSourceKindByName(t *testing.T) { // simply testing the default returned source_kinds sourceKind: database.SourceKind{ ID: 2, - Name: graph.StringKind("AZBase"), + Name: "AZBase", Active: true, }, }, @@ -283,12 +283,12 @@ func TestDeactivateSourceKindsByName(t *testing.T) { sourceKinds: []database.SourceKind{ { ID: 2, - Name: graph.StringKind("AZBase"), + Name: "AZBase", Active: true, }, { ID: 1, - Name: graph.StringKind("Base"), + Name: "Base", Active: true, }, }, @@ -307,12 +307,12 @@ func TestDeactivateSourceKindsByName(t *testing.T) { sourceKinds: []database.SourceKind{ { ID: 2, - Name: graph.StringKind("AZBase"), + Name: "AZBase", Active: true, }, { ID: 1, - Name: graph.StringKind("Base"), + Name: "Base", Active: true, }, }, @@ -331,12 +331,12 @@ func TestDeactivateSourceKindsByName(t *testing.T) { sourceKinds: []database.SourceKind{ { ID: 2, - Name: graph.StringKind("AZBase"), + Name: "AZBase", Active: true, }, { ID: 1, - Name: graph.StringKind("Base"), + Name: "Base", Active: true, }, }, @@ -365,17 +365,17 @@ func TestDeactivateSourceKindsByName(t *testing.T) { sourceKinds: []database.SourceKind{ { ID: 4, - Name: graph.StringKind("AnotherKind"), + Name: "AnotherKind", Active: true, }, { ID: 2, - Name: graph.StringKind("AZBase"), + Name: "AZBase", Active: true, }, { ID: 1, - Name: graph.StringKind("Base"), + Name: "Base", Active: true, }, }, @@ -404,12 +404,12 @@ func TestDeactivateSourceKindsByName(t *testing.T) { sourceKinds: []database.SourceKind{ { ID: 2, - Name: graph.StringKind("AZBase"), + Name: "AZBase", Active: true, }, { ID: 1, - Name: graph.StringKind("Base"), + Name: "Base", Active: true, }, }, From de91119c20b636b3f0ee6ca4e11fb80b2fa589a1 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Fri, 6 Feb 2026 10:45:24 -0800 Subject: [PATCH 14/22] drop compound constraint on schema_environments --- cmd/api/src/database/graphschema.go | 19 ------------ .../database/migration/migrations/v8.7.0.sql | 31 +++++++++++++++++++ cmd/api/src/database/mocks/db.go | 15 --------- .../src/database/upsert_schema_environment.go | 2 +- cmd/api/src/database/upsert_schema_finding.go | 4 +-- 5 files changed, 34 insertions(+), 37 deletions(-) create mode 100644 cmd/api/src/database/migration/migrations/v8.7.0.sql diff --git a/cmd/api/src/database/graphschema.go b/cmd/api/src/database/graphschema.go index 0e4333363b5..de3ab05fb00 100644 --- a/cmd/api/src/database/graphschema.go +++ b/cmd/api/src/database/graphschema.go @@ -53,7 +53,6 @@ type OpenGraphSchema interface { GetGraphSchemaRelationshipKindsWithSchemaName(ctx context.Context, filters model.Filters, sort model.Sort, skip, limit int) (model.GraphSchemaRelationshipKindsWithNamedSchema, int, error) CreateEnvironment(ctx context.Context, extensionId int32, environmentKindId int32, sourceKindId int32) (model.SchemaEnvironment, error) - GetEnvironmentByKinds(ctx context.Context, environmentKindId, sourceKindId int32) (model.SchemaEnvironment, error) GetEnvironmentById(ctx context.Context, environmentId int32) (model.SchemaEnvironment, error) GetEnvironments(ctx context.Context) ([]model.SchemaEnvironment, error) GetEnvironmentsFiltered(ctx context.Context, filters model.Filters) ([]model.SchemaEnvironment, error) @@ -645,24 +644,6 @@ func (s *BloodhoundDB) GetEnvironments(ctx context.Context) ([]model.SchemaEnvir return s.GetEnvironmentsFiltered(ctx, model.Filters{}) } -// GetEnvironmentByKinds - retrieves an environment by its environment kind and source kind. -func (s *BloodhoundDB) GetEnvironmentByKinds(ctx context.Context, environmentKindId, sourceKindId int32) (model.SchemaEnvironment, error) { - filters := model.Filters{ - "se.environment_kind_id": []model.Filter{{Operator: model.Equals, Value: fmt.Sprintf("%d", environmentKindId)}}, - "se.source_kind_id": []model.Filter{{Operator: model.Equals, Value: fmt.Sprintf("%d", sourceKindId)}}, - } - - envs, err := s.GetEnvironmentsFiltered(ctx, filters) - if err != nil { - return model.SchemaEnvironment{}, err - } - if len(envs) == 0 { - return model.SchemaEnvironment{}, ErrNotFound - } - - return envs[0], nil -} - // GetEnvironmentById - retrieves a schema environment by id. func (s *BloodhoundDB) GetEnvironmentById(ctx context.Context, environmentId int32) (model.SchemaEnvironment, error) { filters := model.Filters{ diff --git a/cmd/api/src/database/migration/migrations/v8.7.0.sql b/cmd/api/src/database/migration/migrations/v8.7.0.sql new file mode 100644 index 00000000000..e4120590d09 --- /dev/null +++ b/cmd/api/src/database/migration/migrations/v8.7.0.sql @@ -0,0 +1,31 @@ +-- 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 +-- Drop the compound unique constraint on schema_environments (environment_kind_id, source_kind_id) +-- and add a unique constraint on just environment_kind_id +ALTER TABLE IF EXISTS schema_environments + DROP CONSTRAINT IF EXISTS schema_environments_environment_kind_id_source_kind_id_key; + +DO $$ + BEGIN + IF NOT EXISTS ( + SELECT 1 + FROM pg_constraint + WHERE conname = 'schema_environments_environment_kind_id_key' + ) THEN + ALTER TABLE schema_environments + ADD CONSTRAINT schema_environments_environment_kind_id_key UNIQUE (environment_kind_id); + END IF; + END$$; diff --git a/cmd/api/src/database/mocks/db.go b/cmd/api/src/database/mocks/db.go index e0fb306a8f5..e3502607994 100644 --- a/cmd/api/src/database/mocks/db.go +++ b/cmd/api/src/database/mocks/db.go @@ -1717,21 +1717,6 @@ func (mr *MockDatabaseMockRecorder) GetEnvironmentById(ctx, environmentId any) * return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnvironmentById", reflect.TypeOf((*MockDatabase)(nil).GetEnvironmentById), ctx, environmentId) } -// GetEnvironmentByKinds mocks base method. -func (m *MockDatabase) GetEnvironmentByKinds(ctx context.Context, environmentKindId, sourceKindId int32) (model.SchemaEnvironment, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetEnvironmentByKinds", ctx, environmentKindId, sourceKindId) - ret0, _ := ret[0].(model.SchemaEnvironment) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// GetEnvironmentByKinds indicates an expected call of GetEnvironmentByKinds. -func (mr *MockDatabaseMockRecorder) GetEnvironmentByKinds(ctx, environmentKindId, sourceKindId any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnvironmentByKinds", reflect.TypeOf((*MockDatabase)(nil).GetEnvironmentByKinds), ctx, environmentKindId, sourceKindId) -} - // GetEnvironmentTargetedAccessControlForUser mocks base method. func (m *MockDatabase) GetEnvironmentTargetedAccessControlForUser(ctx context.Context, user model.User) ([]model.EnvironmentTargetedAccessControl, error) { m.ctrl.T.Helper() diff --git a/cmd/api/src/database/upsert_schema_environment.go b/cmd/api/src/database/upsert_schema_environment.go index f266623d9df..4fbc3ed8ca6 100644 --- a/cmd/api/src/database/upsert_schema_environment.go +++ b/cmd/api/src/database/upsert_schema_environment.go @@ -119,7 +119,7 @@ func (s *BloodhoundDB) validateAndTranslatePrincipalKinds(ctx context.Context, p // The unique constraint on (environment_kind_id, source_kind_id) of the Schema Environment table ensures no // duplicate pairs exist, enabling this upsert logic. func (s *BloodhoundDB) replaceSchemaEnvironment(ctx context.Context, graphSchema model.SchemaEnvironment) (int32, error) { - if existing, err := s.GetEnvironmentByKinds(ctx, graphSchema.EnvironmentKindId, graphSchema.SourceKindId); err != nil && !errors.Is(err, ErrNotFound) { + if existing, err := s.GetEnvironmentById(ctx, graphSchema.EnvironmentKindId); err != nil && !errors.Is(err, ErrNotFound) { return 0, fmt.Errorf("error retrieving schema environment: %w", err) } else if !errors.Is(err, ErrNotFound) { // Environment exists - delete it first diff --git a/cmd/api/src/database/upsert_schema_finding.go b/cmd/api/src/database/upsert_schema_finding.go index 9bca51a1271..e9940261034 100644 --- a/cmd/api/src/database/upsert_schema_finding.go +++ b/cmd/api/src/database/upsert_schema_finding.go @@ -36,14 +36,14 @@ func (s *BloodhoundDB) UpsertFinding(ctx context.Context, extensionId int32, sou return model.SchemaRelationshipFinding{}, err } - sourceKindId, err := s.validateAndTranslateSourceKind(ctx, sourceKindName) + _, err = s.validateAndTranslateSourceKind(ctx, sourceKindName) if err != nil { return model.SchemaRelationshipFinding{}, err } // The unique constraint on (environment_kind_id, source_kind_id) of the Schema Environment table ensures no // duplicate pairs exist, enabling this logic. - environment, err := s.GetEnvironmentByKinds(ctx, environmentKindId, sourceKindId) + environment, err := s.GetEnvironmentById(ctx, environmentKindId) if err != nil { return model.SchemaRelationshipFinding{}, err } From 60e539e0fb1fdedd3dc96513ade36c7dbe66a877 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Fri, 6 Feb 2026 11:04:29 -0800 Subject: [PATCH 15/22] fix logic for getenvbykindid --- cmd/api/src/database/graphschema.go | 8 ++++---- .../src/database/graphschema_integration_test.go | 16 ++++++++-------- cmd/api/src/database/mocks/db.go | 12 ++++++------ .../src/database/upsert_schema_environment.go | 8 ++++---- cmd/api/src/database/upsert_schema_finding.go | 6 +++--- 5 files changed, 25 insertions(+), 25 deletions(-) diff --git a/cmd/api/src/database/graphschema.go b/cmd/api/src/database/graphschema.go index de3ab05fb00..af312d543a2 100644 --- a/cmd/api/src/database/graphschema.go +++ b/cmd/api/src/database/graphschema.go @@ -53,7 +53,7 @@ type OpenGraphSchema interface { GetGraphSchemaRelationshipKindsWithSchemaName(ctx context.Context, filters model.Filters, sort model.Sort, skip, limit int) (model.GraphSchemaRelationshipKindsWithNamedSchema, int, error) CreateEnvironment(ctx context.Context, extensionId int32, environmentKindId int32, sourceKindId int32) (model.SchemaEnvironment, error) - GetEnvironmentById(ctx context.Context, environmentId int32) (model.SchemaEnvironment, error) + GetEnvironmentByEnvironmentKindId(ctx context.Context, environmentKindId int32) (model.SchemaEnvironment, error) GetEnvironments(ctx context.Context) ([]model.SchemaEnvironment, error) GetEnvironmentsFiltered(ctx context.Context, filters model.Filters) ([]model.SchemaEnvironment, error) DeleteEnvironment(ctx context.Context, environmentId int32) error @@ -644,10 +644,10 @@ func (s *BloodhoundDB) GetEnvironments(ctx context.Context) ([]model.SchemaEnvir return s.GetEnvironmentsFiltered(ctx, model.Filters{}) } -// GetEnvironmentById - retrieves a schema environment by id. -func (s *BloodhoundDB) GetEnvironmentById(ctx context.Context, environmentId int32) (model.SchemaEnvironment, error) { +// GetEnvironmentByEnvironmentKindId - retrieves a schema environment by environment_kind_id. +func (s *BloodhoundDB) GetEnvironmentByEnvironmentKindId(ctx context.Context, environmentKindId int32) (model.SchemaEnvironment, error) { filters := model.Filters{ - "se.id": []model.Filter{{Operator: model.Equals, Value: fmt.Sprintf("%d", environmentId)}}, + "se.environment_kind_id": []model.Filter{{Operator: model.Equals, Value: fmt.Sprintf("%d", environmentKindId)}}, } envs, err := s.GetEnvironmentsFiltered(ctx, filters) diff --git a/cmd/api/src/database/graphschema_integration_test.go b/cmd/api/src/database/graphschema_integration_test.go index 020f22147b6..0a433fcb2ed 100644 --- a/cmd/api/src/database/graphschema_integration_test.go +++ b/cmd/api/src/database/graphschema_integration_test.go @@ -1453,7 +1453,7 @@ func TestGetSchemaEnvironmentById(t *testing.T) { testSuite := testCase.setup() defer teardownIntegrationTestSuite(t, &testSuite) - got, err := testSuite.BHDatabase.GetEnvironmentById(testSuite.Context, testCase.args.environmentId) + got, err := testSuite.BHDatabase.GetEnvironmentByEnvironmentKindId(testSuite.Context, testCase.args.environmentId) if testCase.want.err != nil { assert.ErrorIs(t, err, testCase.want.err) } else { @@ -1532,7 +1532,7 @@ func TestDeleteSchemaEnvironment(t *testing.T) { assert.NoError(t, err) // Verify deletion by trying to get the environment - _, err = testSuite.BHDatabase.GetEnvironmentById(testSuite.Context, testCase.args.environmentId) + _, err = testSuite.BHDatabase.GetEnvironmentByEnvironmentKindId(testSuite.Context, testCase.args.environmentId) assert.ErrorIs(t, err, database.ErrNotFound) } }) @@ -1560,11 +1560,11 @@ func TestTransaction_SchemaEnvironment(t *testing.T) { require.NoError(t, err) // Verify both environments were created - env1, err := testSuite.BHDatabase.GetEnvironmentById(testSuite.Context, 1) + env1, err := testSuite.BHDatabase.GetEnvironmentByEnvironmentKindId(testSuite.Context, 1) require.NoError(t, err) assert.Equal(t, int32(1), env1.EnvironmentKindId) - env2, err := testSuite.BHDatabase.GetEnvironmentById(testSuite.Context, 2) + env2, err := testSuite.BHDatabase.GetEnvironmentByEnvironmentKindId(testSuite.Context, 2) require.NoError(t, err) assert.Equal(t, int32(2), env2.EnvironmentKindId) }) @@ -1589,7 +1589,7 @@ func TestTransaction_SchemaEnvironment(t *testing.T) { require.ErrorIs(t, err, expectedErr) // Verify the environment was NOT created (rolled back) - _, err = testSuite.BHDatabase.GetEnvironmentById(testSuite.Context, 1) + _, err = testSuite.BHDatabase.GetEnvironmentByEnvironmentKindId(testSuite.Context, 1) assert.ErrorIs(t, err, database.ErrNotFound) }) @@ -1614,7 +1614,7 @@ func TestTransaction_SchemaEnvironment(t *testing.T) { require.Error(t, err) // Verify the first environment was NOT created (rolled back due to second failure) - _, err = testSuite.BHDatabase.GetEnvironmentById(testSuite.Context, 1) + _, err = testSuite.BHDatabase.GetEnvironmentByEnvironmentKindId(testSuite.Context, 1) assert.ErrorIs(t, err, database.ErrNotFound) }) @@ -1637,7 +1637,7 @@ func TestTransaction_SchemaEnvironment(t *testing.T) { require.NoError(t, err) // Verify the environment does not exist - _, err = testSuite.BHDatabase.GetEnvironmentById(testSuite.Context, 1) + _, err = testSuite.BHDatabase.GetEnvironmentByEnvironmentKindId(testSuite.Context, 1) assert.ErrorIs(t, err, database.ErrNotFound) }) } @@ -2488,7 +2488,7 @@ func TestDeleteSchemaExtension_CascadeDeletesAllDependents(t *testing.T) { _, err = testSuite.BHDatabase.GetGraphSchemaRelationshipKindById(testSuite.Context, relationshipKind.ID) assert.ErrorIs(t, err, database.ErrNotFound) - _, err = testSuite.BHDatabase.GetEnvironmentById(testSuite.Context, environment.ID) + _, err = testSuite.BHDatabase.GetEnvironmentByEnvironmentKindId(testSuite.Context, environment.ID) assert.ErrorIs(t, err, database.ErrNotFound) _, err = testSuite.BHDatabase.GetSchemaRelationshipFindingById(testSuite.Context, relationshipFinding.ID) diff --git a/cmd/api/src/database/mocks/db.go b/cmd/api/src/database/mocks/db.go index e3502607994..7d8e3a472a1 100644 --- a/cmd/api/src/database/mocks/db.go +++ b/cmd/api/src/database/mocks/db.go @@ -1702,19 +1702,19 @@ func (mr *MockDatabaseMockRecorder) GetDatapipeStatus(ctx any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDatapipeStatus", reflect.TypeOf((*MockDatabase)(nil).GetDatapipeStatus), ctx) } -// GetEnvironmentById mocks base method. -func (m *MockDatabase) GetEnvironmentById(ctx context.Context, environmentId int32) (model.SchemaEnvironment, error) { +// GetEnvironmentByEnvironmentKindId mocks base method. +func (m *MockDatabase) GetEnvironmentByEnvironmentKindId(ctx context.Context, environmentKindId int32) (model.SchemaEnvironment, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetEnvironmentById", ctx, environmentId) + ret := m.ctrl.Call(m, "GetEnvironmentByEnvironmentKindId", ctx, environmentKindId) ret0, _ := ret[0].(model.SchemaEnvironment) ret1, _ := ret[1].(error) return ret0, ret1 } -// GetEnvironmentById indicates an expected call of GetEnvironmentById. -func (mr *MockDatabaseMockRecorder) GetEnvironmentById(ctx, environmentId any) *gomock.Call { +// GetEnvironmentByEnvironmentKindId indicates an expected call of GetEnvironmentByEnvironmentKindId. +func (mr *MockDatabaseMockRecorder) GetEnvironmentByEnvironmentKindId(ctx, environmentKindId any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnvironmentById", reflect.TypeOf((*MockDatabase)(nil).GetEnvironmentById), ctx, environmentId) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnvironmentByEnvironmentKindId", reflect.TypeOf((*MockDatabase)(nil).GetEnvironmentByEnvironmentKindId), ctx, environmentKindId) } // GetEnvironmentTargetedAccessControlForUser mocks base method. diff --git a/cmd/api/src/database/upsert_schema_environment.go b/cmd/api/src/database/upsert_schema_environment.go index 4fbc3ed8ca6..42d2293863e 100644 --- a/cmd/api/src/database/upsert_schema_environment.go +++ b/cmd/api/src/database/upsert_schema_environment.go @@ -115,11 +115,11 @@ func (s *BloodhoundDB) validateAndTranslatePrincipalKinds(ctx context.Context, p } // replaceSchemaEnvironment creates or updates a schema environment. -// If an environment with the given kinds exists, it deletes it first before creating the new one. -// The unique constraint on (environment_kind_id, source_kind_id) of the Schema Environment table ensures no -// duplicate pairs exist, enabling this upsert logic. +// If an environment with the given environment_kind_id exists, it deletes it first before creating the new one. +// The unique constraint on environment_kind_id of the Schema Environment table ensures no +// duplicates exist, enabling this upsert logic. func (s *BloodhoundDB) replaceSchemaEnvironment(ctx context.Context, graphSchema model.SchemaEnvironment) (int32, error) { - if existing, err := s.GetEnvironmentById(ctx, graphSchema.EnvironmentKindId); err != nil && !errors.Is(err, ErrNotFound) { + if existing, err := s.GetEnvironmentByEnvironmentKindId(ctx, graphSchema.EnvironmentKindId); err != nil && !errors.Is(err, ErrNotFound) { return 0, fmt.Errorf("error retrieving schema environment: %w", err) } else if !errors.Is(err, ErrNotFound) { // Environment exists - delete it first diff --git a/cmd/api/src/database/upsert_schema_finding.go b/cmd/api/src/database/upsert_schema_finding.go index e9940261034..2f15e6c93ad 100644 --- a/cmd/api/src/database/upsert_schema_finding.go +++ b/cmd/api/src/database/upsert_schema_finding.go @@ -41,9 +41,9 @@ func (s *BloodhoundDB) UpsertFinding(ctx context.Context, extensionId int32, sou return model.SchemaRelationshipFinding{}, err } - // The unique constraint on (environment_kind_id, source_kind_id) of the Schema Environment table ensures no - // duplicate pairs exist, enabling this logic. - environment, err := s.GetEnvironmentById(ctx, environmentKindId) + // The unique constraint on environment_kind_id of the Schema Environment table ensures no + // duplicates exist, enabling this logic. + environment, err := s.GetEnvironmentByEnvironmentKindId(ctx, environmentKindId) if err != nil { return model.SchemaRelationshipFinding{}, err } From 717fe7d28cc52e4b7c2fb13a569a57cc93e15a49 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Mon, 9 Feb 2026 11:53:34 -0800 Subject: [PATCH 16/22] embed tierID --- packages/go/analysis/tiering/helpers.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/go/analysis/tiering/helpers.go b/packages/go/analysis/tiering/helpers.go index 556c4ea7b74..d1275929ba3 100644 --- a/packages/go/analysis/tiering/helpers.go +++ b/packages/go/analysis/tiering/helpers.go @@ -24,6 +24,7 @@ import ( ) type SearchTierNodesCtx struct { + TierID int // Tier ID for findings processing IsTierZero bool PrimaryTierKind graph.Kind SearchTierNodes graph.Criteria @@ -32,8 +33,9 @@ type SearchTierNodesCtx struct { SearchPrimaryTierNodesRel graph.Criteria } -func NewSearchTierNodesCtx(tieringEnabled bool, isTierZero bool, primaryKind graph.Kind, tierKinds ...graph.Kind) SearchTierNodesCtx { +func NewSearchTierNodesCtx(tieringEnabled bool, isTierZero bool, tierID int, primaryKind graph.Kind, tierKinds ...graph.Kind) SearchTierNodesCtx { return SearchTierNodesCtx{ + TierID: tierID, IsTierZero: isTierZero, PrimaryTierKind: primaryKind, SearchTierNodesRel: searchTierNodesRel(tieringEnabled, tierKinds...), From 4b91907822d9866b463512c46354a63f13389209 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Mon, 9 Feb 2026 13:21:51 -0800 Subject: [PATCH 17/22] add back in GetEnvironmentById --- cmd/api/src/database/graphschema.go | 18 ++++ cmd/api/src/database/mocks/db.go | 15 ++++ ...psert_schema_extension_integration_test.go | 87 ++++++------------- 3 files changed, 61 insertions(+), 59 deletions(-) diff --git a/cmd/api/src/database/graphschema.go b/cmd/api/src/database/graphschema.go index f097a3dab9f..c09eeab55be 100644 --- a/cmd/api/src/database/graphschema.go +++ b/cmd/api/src/database/graphschema.go @@ -54,6 +54,7 @@ type OpenGraphSchema interface { CreateEnvironment(ctx context.Context, extensionId int32, environmentKindId int32, sourceKindId int32) (model.SchemaEnvironment, error) GetEnvironmentByEnvironmentKindId(ctx context.Context, environmentKindId int32) (model.SchemaEnvironment, error) + GetEnvironmentById(ctx context.Context, environmentId int32) (model.SchemaEnvironment, error) GetEnvironments(ctx context.Context) ([]model.SchemaEnvironment, error) GetEnvironmentsFiltered(ctx context.Context, filters model.Filters) ([]model.SchemaEnvironment, error) DeleteEnvironment(ctx context.Context, environmentId int32) error @@ -673,6 +674,23 @@ func (s *BloodhoundDB) GetEnvironmentsByExtensionId(ctx context.Context, extensi } +// GetEnvironmentById - retrieves a schema environment by id. +func (s *BloodhoundDB) GetEnvironmentById(ctx context.Context, environmentId int32) (model.SchemaEnvironment, error) { + var schemaEnvironment model.SchemaEnvironment + + if result := s.db.WithContext(ctx).Raw(fmt.Sprintf(` + SELECT id, schema_extension_id, environment_kind_id, source_kind_id, created_at, updated_at, deleted_at + FROM %s WHERE id = ?`, + schemaEnvironment.TableName()), + environmentId).Scan(&schemaEnvironment); result.Error != nil { + return model.SchemaEnvironment{}, CheckError(result) + } else if result.RowsAffected == 0 { + return model.SchemaEnvironment{}, ErrNotFound + } + + return schemaEnvironment, nil +} + // GetEnvironmentByEnvironmentKindId - retrieves a schema environment by environment_kind_id. func (s *BloodhoundDB) GetEnvironmentByEnvironmentKindId(ctx context.Context, environmentKindId int32) (model.SchemaEnvironment, error) { filters := model.Filters{ diff --git a/cmd/api/src/database/mocks/db.go b/cmd/api/src/database/mocks/db.go index 7d8e3a472a1..cc0bef93df1 100644 --- a/cmd/api/src/database/mocks/db.go +++ b/cmd/api/src/database/mocks/db.go @@ -1717,6 +1717,21 @@ func (mr *MockDatabaseMockRecorder) GetEnvironmentByEnvironmentKindId(ctx, envir return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnvironmentByEnvironmentKindId", reflect.TypeOf((*MockDatabase)(nil).GetEnvironmentByEnvironmentKindId), ctx, environmentKindId) } +// GetEnvironmentById mocks base method. +func (m *MockDatabase) GetEnvironmentById(ctx context.Context, environmentId int32) (model.SchemaEnvironment, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetEnvironmentById", ctx, environmentId) + ret0, _ := ret[0].(model.SchemaEnvironment) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetEnvironmentById indicates an expected call of GetEnvironmentById. +func (mr *MockDatabaseMockRecorder) GetEnvironmentById(ctx, environmentId any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnvironmentById", reflect.TypeOf((*MockDatabase)(nil).GetEnvironmentById), ctx, environmentId) +} + // GetEnvironmentTargetedAccessControlForUser mocks base method. func (m *MockDatabase) GetEnvironmentTargetedAccessControlForUser(ctx context.Context, user model.User) ([]model.EnvironmentTargetedAccessControl, error) { m.ctrl.T.Helper() diff --git a/cmd/api/src/database/upsert_schema_extension_integration_test.go b/cmd/api/src/database/upsert_schema_extension_integration_test.go index d4b3df2366d..43fab7ff12b 100644 --- a/cmd/api/src/database/upsert_schema_extension_integration_test.go +++ b/cmd/api/src/database/upsert_schema_extension_integration_test.go @@ -1151,74 +1151,43 @@ func getAndCompareGraphExtension(t *testing.T, testContext context.Context, db * require.NoError(t, err) require.Containsf(t, want.EnvironmentsInput[idx].PrincipalKinds, dawgsPrincipalKind.Name, "PrincipalKind - Name mismatch") } - - // Test Findings - gotSchemaRelationshipFinding, err = db.GetSchemaRelationshipFindingsBySchemaExtensionId(testContext, gotGraphExtension.ID) - require.NoError(t, err) - - require.Equalf(t, len(want.RelationshipFindingsInput), len(gotSchemaRelationshipFinding), "mismatched number of findings") - for i, finding := range gotSchemaRelationshipFinding { - // Finding - require.Greater(t, finding.ID, int32(0)) - require.Equalf(t, gotGraphExtension.ID, finding.SchemaExtensionId, "RelationshipFindingInput - graph schema extension id should be greater than 0") - - dawgsFindingRelationshipKind, err = db.GetKindById(testContext, finding.RelationshipKindId) - require.NoError(t, err) - require.Equalf(t, want.RelationshipFindingsInput[i].RelationshipKindName, dawgsFindingRelationshipKind.Name, "RelationshipFindingInput - relationship kind name mismatch") - - // TODO: fix this - require.Equalf(t, want.RelationshipFindingsInput[i].EnvironmentKindName, gotEnvironment.EnvironmentKindName, "RelationshipFindingInput - environment kind name mismatch") - - require.Equalf(t, want.RelationshipFindingsInput[i].Name, finding.Name, "RelationshipFindingInput - name mismatch") - require.Equalf(t, want.RelationshipFindingsInput[i].DisplayName, finding.DisplayName, "RelationshipFindingInput - display name mismatch") - - // Remediation - gotRemediation, err = db.GetRemediationByFindingId(testContext, finding.ID) - require.NoError(t, err) - - require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.ShortRemediation, gotRemediation.ShortRemediation, "Remediation - short_remediation mismatch") - require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.LongRemediation, gotRemediation.LongRemediation, "Remediation - long_remediation mismatch") - require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.ShortDescription, gotRemediation.ShortDescription, "Remediation - short_description mismatch") - require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.LongDescription, gotRemediation.LongDescription, "Remediation - long_description mismatch") - - } } - // // Test Findings - // gotSchemaRelationshipFinding, err = db.GetSchemaRelationshipFindingsBySchemaExtensionId(testContext, gotGraphExtension.ID) - // require.NoError(t, err) + // Test Findings + gotSchemaRelationshipFinding, err = db.GetSchemaRelationshipFindingsBySchemaExtensionId(testContext, gotGraphExtension.ID) + require.NoError(t, err) - // require.Equalf(t, len(want.RelationshipFindingsInput), len(gotSchemaRelationshipFinding), "mismatched number of findings") - // for i, finding := range gotSchemaRelationshipFinding { - // // Finding - // require.Greater(t, finding.ID, int32(0)) - // require.Equalf(t, gotGraphExtension.ID, finding.SchemaExtensionId, "RelationshipFindingInput - graph schema extension id should be greater than 0") + require.Equalf(t, len(want.RelationshipFindingsInput), len(gotSchemaRelationshipFinding), "mismatched number of findings") + for i, finding := range gotSchemaRelationshipFinding { + // Finding + require.Greater(t, finding.ID, int32(0)) + require.Equalf(t, gotGraphExtension.ID, finding.SchemaExtensionId, "RelationshipFindingInput - graph schema extension id should be greater than 0") - // dawgsFindingRelationshipKind, err = db.GetKindById(testContext, finding.RelationshipKindId) - // require.NoError(t, err) - // require.Equalf(t, want.RelationshipFindingsInput[i].RelationshipKindName, dawgsFindingRelationshipKind.Name, "RelationshipFindingInput - relationship kind name mismatch") + dawgsFindingRelationshipKind, err = db.GetKindById(testContext, finding.RelationshipKindId) + require.NoError(t, err) + require.Equalf(t, want.RelationshipFindingsInput[i].RelationshipKindName, dawgsFindingRelationshipKind.Name, "RelationshipFindingInput - relationship kind name mismatch") - // // TODO: fix this - // // was: - // findingEnvironment, err = db.GetEnvironmentById(testContext, finding.EnvironmentId) - // require.NoError(t, err) - // dawgsFindingEnvironmentKind, err = db.GetKindById(testContext, findingEnvironment.EnvironmentKindId) - // require.NoError(t, err) - // require.Equalf(t, want.RelationshipFindingsInput[i].EnvironmentKindName, dawgsFindingEnvironmentKind.Name, "RelationshipFindingInput - environment kind name mismatch") + // TODO: fix this + // was: + findingEnvironment, err := db.GetEnvironmentById(testContext, finding.EnvironmentId) + require.NoError(t, err) + dawgsFindingEnvironmentKind, err := db.GetKindById(testContext, findingEnvironment.EnvironmentKindId) + require.NoError(t, err) + require.Equalf(t, want.RelationshipFindingsInput[i].EnvironmentKindName, dawgsFindingEnvironmentKind.Name, "RelationshipFindingInput - environment kind name mismatch") - // require.Equalf(t, want.RelationshipFindingsInput[i].Name, finding.Name, "RelationshipFindingInput - name mismatch") - // require.Equalf(t, want.RelationshipFindingsInput[i].DisplayName, finding.DisplayName, "RelationshipFindingInput - display name mismatch") + require.Equalf(t, want.RelationshipFindingsInput[i].Name, finding.Name, "RelationshipFindingInput - name mismatch") + require.Equalf(t, want.RelationshipFindingsInput[i].DisplayName, finding.DisplayName, "RelationshipFindingInput - display name mismatch") - // // Remediation - // gotRemediation, err = db.GetRemediationByFindingId(testContext, finding.ID) - // require.NoError(t, err) + // Remediation + gotRemediation, err = db.GetRemediationByFindingId(testContext, finding.ID) + require.NoError(t, err) - // require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.ShortRemediation, gotRemediation.ShortRemediation, "Remediation - short_remediation mismatch") - // require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.LongRemediation, gotRemediation.LongRemediation, "Remediation - long_remediation mismatch") - // require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.ShortDescription, gotRemediation.ShortDescription, "Remediation - short_description mismatch") - // require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.LongDescription, gotRemediation.LongDescription, "Remediation - long_description mismatch") + require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.ShortRemediation, gotRemediation.ShortRemediation, "Remediation - short_remediation mismatch") + require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.LongRemediation, gotRemediation.LongRemediation, "Remediation - long_remediation mismatch") + require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.ShortDescription, gotRemediation.ShortDescription, "Remediation - short_description mismatch") + require.Equalf(t, want.RelationshipFindingsInput[i].RemediationInput.LongDescription, gotRemediation.LongDescription, "Remediation - long_description mismatch") - // } + } return gotGraphExtension.ID } From fc2af668c2890b0332245dfa62003f2aa3ee4332 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Mon, 9 Feb 2026 13:38:52 -0800 Subject: [PATCH 18/22] sourcekind cleanup --- cmd/api/src/database/sourcekinds.go | 81 ++++++----------------------- 1 file changed, 15 insertions(+), 66 deletions(-) diff --git a/cmd/api/src/database/sourcekinds.go b/cmd/api/src/database/sourcekinds.go index b0d6662f572..c8d8ec13b02 100644 --- a/cmd/api/src/database/sourcekinds.go +++ b/cmd/api/src/database/sourcekinds.go @@ -75,28 +75,12 @@ func (s *BloodhoundDB) GetSourceKinds(ctx context.Context) ([]SourceKind, error) ORDER BY name ASC; ` - type rawSourceKind struct { - ID int - Name string - Active bool - } - - var kinds []rawSourceKind - result := s.db.WithContext(ctx).Raw(query).Scan(&kinds) - if err := result.Error; err != nil { + var kinds []SourceKind + if err := s.db.WithContext(ctx).Raw(query).Scan(&kinds).Error; err != nil { return nil, fmt.Errorf("failed to fetch source kinds: %w", err) } - out := make([]SourceKind, len(kinds)) - for i, k := range kinds { - out[i] = SourceKind{ - ID: k.ID, - Name: k.Name, - Active: k.Active, - } - } - - return out, nil + return kinds, nil } func (s *BloodhoundDB) GetSourceKindByName(ctx context.Context, name string) (SourceKind, error) { @@ -106,30 +90,18 @@ func (s *BloodhoundDB) GetSourceKindByName(ctx context.Context, name string) (So WHERE name = $1 AND active = true; ` - type rawSourceKind struct { - ID int - Name string - Active bool - } - - var raw rawSourceKind - result := s.db.WithContext(ctx).Raw(query, name).Scan(&raw) + var sourceKind SourceKind + result := s.db.WithContext(ctx).Raw(query, name).Scan(&sourceKind) if result.Error != nil { return SourceKind{}, result.Error } - if result.RowsAffected == 0 || raw.ID == 0 { + if result.RowsAffected == 0 || sourceKind.ID == 0 { return SourceKind{}, ErrNotFound } - kind := SourceKind{ - ID: raw.ID, - Name: raw.Name, - Active: raw.Active, - } - - return kind, nil + return sourceKind, nil } func (s *BloodhoundDB) GetSourceKindById(ctx context.Context, id int32) (SourceKind, error) { @@ -139,30 +111,18 @@ func (s *BloodhoundDB) GetSourceKindById(ctx context.Context, id int32) (SourceK WHERE id = $1 AND active = true; ` - type rawSourceKind struct { - ID int - Name string - Active bool - } - - var raw rawSourceKind - result := s.db.WithContext(ctx).Raw(query, id).Scan(&raw) + var sourceKind SourceKind + result := s.db.WithContext(ctx).Raw(query, id).Scan(&sourceKind) if result.Error != nil { return SourceKind{}, result.Error } - if result.RowsAffected == 0 || raw.ID == 0 { + if result.RowsAffected == 0 || sourceKind.ID == 0 { return SourceKind{}, ErrNotFound } - kind := SourceKind{ - ID: raw.ID, - Name: raw.Name, - Active: raw.Active, - } - - return kind, nil + return sourceKind, nil } func (s *BloodhoundDB) GetSourceKindByID(ctx context.Context, id int) (SourceKind, error) { @@ -171,30 +131,19 @@ func (s *BloodhoundDB) GetSourceKindByID(ctx context.Context, id int) (SourceKin FROM source_kinds WHERE id = $1 AND active = true; ` - type rawSourceKind struct { - ID int - Name string - Active bool - } - var raw rawSourceKind - result := s.db.WithContext(ctx).Raw(query, id).Scan(&raw) + var sourceKind SourceKind + result := s.db.WithContext(ctx).Raw(query, id).Scan(&sourceKind) if result.Error != nil { return SourceKind{}, result.Error } - if result.RowsAffected == 0 || raw.ID == 0 { + if result.RowsAffected == 0 || sourceKind.ID == 0 { return SourceKind{}, ErrNotFound } - kind := SourceKind{ - ID: raw.ID, - Name: raw.Name, - Active: raw.Active, - } - - return kind, nil + return sourceKind, nil } func (s *BloodhoundDB) DeactivateSourceKindsByName(ctx context.Context, kinds graph.Kinds) error { From d9b15a4f05788f114de72c030d2af7e3e6e303ad Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Tue, 10 Feb 2026 11:13:38 -0800 Subject: [PATCH 19/22] getsourcekindbyids variadic --- cmd/api/src/database/mocks/db.go | 35 ++++++++++++++++------------- cmd/api/src/database/sourcekinds.go | 24 +++++++++----------- 2 files changed, 31 insertions(+), 28 deletions(-) diff --git a/cmd/api/src/database/mocks/db.go b/cmd/api/src/database/mocks/db.go index cc0bef93df1..7589b172c1a 100644 --- a/cmd/api/src/database/mocks/db.go +++ b/cmd/api/src/database/mocks/db.go @@ -2423,21 +2423,6 @@ func (mr *MockDatabaseMockRecorder) GetSharedSavedQueries(ctx, userID any) *gomo return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSharedSavedQueries", reflect.TypeOf((*MockDatabase)(nil).GetSharedSavedQueries), ctx, userID) } -// GetSourceKindById mocks base method. -func (m *MockDatabase) GetSourceKindById(ctx context.Context, id int32) (database.SourceKind, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetSourceKindById", ctx, id) - ret0, _ := ret[0].(database.SourceKind) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// GetSourceKindById indicates an expected call of GetSourceKindById. -func (mr *MockDatabaseMockRecorder) GetSourceKindById(ctx, id any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSourceKindById", reflect.TypeOf((*MockDatabase)(nil).GetSourceKindById), ctx, id) -} - // GetSourceKindByName mocks base method. func (m *MockDatabase) GetSourceKindByName(ctx context.Context, name string) (database.SourceKind, error) { m.ctrl.T.Helper() @@ -2468,6 +2453,26 @@ func (mr *MockDatabaseMockRecorder) GetSourceKinds(ctx any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSourceKinds", reflect.TypeOf((*MockDatabase)(nil).GetSourceKinds), ctx) } +// GetSourceKindsByIds mocks base method. +func (m *MockDatabase) GetSourceKindsByIds(ctx context.Context, ids ...int32) ([]database.SourceKind, error) { + m.ctrl.T.Helper() + varargs := []any{ctx} + for _, a := range ids { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "GetSourceKindsByIds", varargs...) + ret0, _ := ret[0].([]database.SourceKind) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetSourceKindsByIds indicates an expected call of GetSourceKindsByIds. +func (mr *MockDatabaseMockRecorder) GetSourceKindsByIds(ctx any, ids ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx}, ids...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSourceKindsByIds", reflect.TypeOf((*MockDatabase)(nil).GetSourceKindsByIds), varargs...) +} + // GetTimeRangedAssetGroupCollections mocks base method. func (m *MockDatabase) GetTimeRangedAssetGroupCollections(ctx context.Context, assetGroupID int32, from, to int64, order string) (model.AssetGroupCollections, error) { m.ctrl.T.Helper() diff --git a/cmd/api/src/database/sourcekinds.go b/cmd/api/src/database/sourcekinds.go index c8d8ec13b02..0eadae13c47 100644 --- a/cmd/api/src/database/sourcekinds.go +++ b/cmd/api/src/database/sourcekinds.go @@ -29,7 +29,7 @@ type SourceKindsData interface { DeactivateSourceKindsByName(ctx context.Context, kinds graph.Kinds) error RegisterSourceKind(ctx context.Context) func(sourceKind graph.Kind) error GetSourceKindByName(ctx context.Context, name string) (SourceKind, error) - GetSourceKindById(ctx context.Context, id int32) (SourceKind, error) + GetSourceKindsByIds(ctx context.Context, ids ...int32) ([]SourceKind, error) } // RegisterSourceKind returns a function that inserts a source kind by name, @@ -104,25 +104,23 @@ func (s *BloodhoundDB) GetSourceKindByName(ctx context.Context, name string) (So return sourceKind, nil } -func (s *BloodhoundDB) GetSourceKindById(ctx context.Context, id int32) (SourceKind, error) { +func (s *BloodhoundDB) GetSourceKindsByIds(ctx context.Context, ids ...int32) ([]SourceKind, error) { + if len(ids) == 0 { + return []SourceKind{}, nil + } + const query = ` SELECT id, name, active FROM source_kinds - WHERE id = $1 AND active = true; + WHERE id = ANY(?) AND active = true; ` - var sourceKind SourceKind - result := s.db.WithContext(ctx).Raw(query, id).Scan(&sourceKind) - - if result.Error != nil { - return SourceKind{}, result.Error - } - - if result.RowsAffected == 0 || sourceKind.ID == 0 { - return SourceKind{}, ErrNotFound + var sourceKinds []SourceKind + if err := s.db.WithContext(ctx).Raw(query, pq.Array(ids)).Scan(&sourceKinds).Error; err != nil { + return nil, err } - return sourceKind, nil + return sourceKinds, nil } func (s *BloodhoundDB) GetSourceKindByID(ctx context.Context, id int) (SourceKind, error) { From 4785d4585a73973ac47ec00a2a25c5ce60c82d9a Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Tue, 10 Feb 2026 11:36:32 -0800 Subject: [PATCH 20/22] getkindsbyids --- cmd/api/src/database/kind.go | 25 +++++++------ cmd/api/src/database/kind_integration_test.go | 36 ++++++++++++++++--- cmd/api/src/database/mocks/db.go | 31 +++++++++------- ...psert_schema_extension_integration_test.go | 34 ++++++++++-------- 4 files changed, 83 insertions(+), 43 deletions(-) diff --git a/cmd/api/src/database/kind.go b/cmd/api/src/database/kind.go index 17aaa2249c0..fa253765b6d 100644 --- a/cmd/api/src/database/kind.go +++ b/cmd/api/src/database/kind.go @@ -18,12 +18,13 @@ package database import ( "context" + "github.com/lib/pq" "github.com/specterops/bloodhound/cmd/api/src/model" ) type Kind interface { GetKindByName(ctx context.Context, name string) (model.Kind, error) - GetKindById(ctx context.Context, id int32) (model.Kind, error) + GetKindsByIds(ctx context.Context, ids ...int32) ([]model.Kind, error) } func (s *BloodhoundDB) GetKindByName(ctx context.Context, name string) (model.Kind, error) { @@ -47,23 +48,25 @@ func (s *BloodhoundDB) GetKindByName(ctx context.Context, name string) (model.Ki return kind, nil } -func (s *BloodhoundDB) GetKindById(ctx context.Context, id int32) (model.Kind, error) { +func (s *BloodhoundDB) GetKindsByIds(ctx context.Context, ids ...int32) ([]model.Kind, error) { + if len(ids) == 0 { + return []model.Kind{}, nil + } + const query = ` SELECT id, name FROM kind - WHERE id = $1; + WHERE id = ANY(?); ` - var kind model.Kind - result := s.db.WithContext(ctx).Raw(query, id).Scan(&kind) - - if result.Error != nil { - return model.Kind{}, result.Error + var kinds []model.Kind + if err := s.db.WithContext(ctx).Raw(query, pq.Array(ids)).Scan(&kinds).Error; err != nil { + return nil, err } - if result.RowsAffected == 0 || kind.ID == 0 { - return model.Kind{}, ErrNotFound + if len(kinds) != len(ids) { + return nil, ErrNotFound } - return kind, nil + return kinds, nil } diff --git a/cmd/api/src/database/kind_integration_test.go b/cmd/api/src/database/kind_integration_test.go index dea63dada97..17b117181d7 100644 --- a/cmd/api/src/database/kind_integration_test.go +++ b/cmd/api/src/database/kind_integration_test.go @@ -129,16 +129,44 @@ func TestGetKindByID(t *testing.T) { var ( err error createdKind model.Kind - got model.Kind + got []model.Kind ) createdKind = tt.setup(t) - if got, err = testSuite.BHDatabase.GetKindById(testSuite.Context, createdKind.ID); tt.want.err != nil { + if got, err = testSuite.BHDatabase.GetKindsByIds(testSuite.Context, createdKind.ID); tt.want.err != nil { assert.EqualError(t, err, tt.want.err.Error()) } else { assert.NoError(t, err) - assert.Equal(t, tt.want.kind.Name, got.Name) - assert.Greater(t, got.ID, int32(0)) + assert.Equal(t, tt.want.kind.Name, got[0].Name) + assert.Greater(t, got[0].ID, int32(0)) } }) } } + +func TestGetKindsByIds_MultipleKinds(t *testing.T) { + testSuite := setupIntegrationTestSuite(t) + defer teardownIntegrationTestSuite(t, &testSuite) + + // Create two kinds + var kind1, kind2 model.Kind + result := testSuite.DB.WithContext(testSuite.Context).Raw(` + INSERT INTO kind (name) + VALUES ('Test_Kind_One') + RETURNING id, name;`).Scan(&kind1) + require.NoError(t, result.Error) + + result = testSuite.DB.WithContext(testSuite.Context).Raw(` + INSERT INTO kind (name) + VALUES ('Test_Kind_Two') + RETURNING id, name;`).Scan(&kind2) + require.NoError(t, result.Error) + + // Fetch both kinds by their IDs + kinds, err := testSuite.BHDatabase.GetKindsByIds(testSuite.Context, kind1.ID, kind2.ID) + require.NoError(t, err) + assert.Len(t, kinds, 2) + + // Verify both kinds are returned (order not guaranteed) + names := []string{kinds[0].Name, kinds[1].Name} + require.ElementsMatch(t, []string{"Test_Kind_One", "Test_Kind_Two"}, names) +} diff --git a/cmd/api/src/database/mocks/db.go b/cmd/api/src/database/mocks/db.go index 7589b172c1a..7d2aaf3cf16 100644 --- a/cmd/api/src/database/mocks/db.go +++ b/cmd/api/src/database/mocks/db.go @@ -2007,34 +2007,39 @@ func (mr *MockDatabaseMockRecorder) GetInstallation(ctx any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetInstallation", reflect.TypeOf((*MockDatabase)(nil).GetInstallation), ctx) } -// GetKindById mocks base method. -func (m *MockDatabase) GetKindById(ctx context.Context, id int32) (model.Kind, error) { +// GetKindByName mocks base method. +func (m *MockDatabase) GetKindByName(ctx context.Context, name string) (model.Kind, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetKindById", ctx, id) + ret := m.ctrl.Call(m, "GetKindByName", ctx, name) ret0, _ := ret[0].(model.Kind) ret1, _ := ret[1].(error) return ret0, ret1 } -// GetKindById indicates an expected call of GetKindById. -func (mr *MockDatabaseMockRecorder) GetKindById(ctx, id any) *gomock.Call { +// GetKindByName indicates an expected call of GetKindByName. +func (mr *MockDatabaseMockRecorder) GetKindByName(ctx, name any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetKindById", reflect.TypeOf((*MockDatabase)(nil).GetKindById), ctx, id) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetKindByName", reflect.TypeOf((*MockDatabase)(nil).GetKindByName), ctx, name) } -// GetKindByName mocks base method. -func (m *MockDatabase) GetKindByName(ctx context.Context, name string) (model.Kind, error) { +// GetKindsByIds mocks base method. +func (m *MockDatabase) GetKindsByIds(ctx context.Context, ids ...int32) ([]model.Kind, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetKindByName", ctx, name) - ret0, _ := ret[0].(model.Kind) + varargs := []any{ctx} + for _, a := range ids { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "GetKindsByIds", varargs...) + ret0, _ := ret[0].([]model.Kind) ret1, _ := ret[1].(error) return ret0, ret1 } -// GetKindByName indicates an expected call of GetKindByName. -func (mr *MockDatabaseMockRecorder) GetKindByName(ctx, name any) *gomock.Call { +// GetKindsByIds indicates an expected call of GetKindsByIds. +func (mr *MockDatabaseMockRecorder) GetKindsByIds(ctx any, ids ...any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetKindByName", reflect.TypeOf((*MockDatabase)(nil).GetKindByName), ctx, name) + varargs := append([]any{ctx}, ids...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetKindsByIds", reflect.TypeOf((*MockDatabase)(nil).GetKindsByIds), varargs...) } // GetLatestAssetGroupCollection mocks base method. diff --git a/cmd/api/src/database/upsert_schema_extension_integration_test.go b/cmd/api/src/database/upsert_schema_extension_integration_test.go index 43fab7ff12b..662aec9e1a4 100644 --- a/cmd/api/src/database/upsert_schema_extension_integration_test.go +++ b/cmd/api/src/database/upsert_schema_extension_integration_test.go @@ -1068,16 +1068,18 @@ func getAndCompareGraphExtension(t *testing.T, testContext context.Context, db * SetOperator: model.FilterAnd, } - gotNodeKinds model.GraphSchemaNodeKinds - gotRelationshipKinds model.GraphSchemaRelationshipKinds - gotProperties model.GraphSchemaProperties - gotSchemaEnvironments []model.SchemaEnvironment - gotPrincipalKinds model.SchemaEnvironmentPrincipalKinds - sourceKind database.SourceKind - dawgsPrincipalKind model.Kind - dawgsFindingRelationshipKind model.Kind - gotSchemaRelationshipFinding []model.SchemaRelationshipFinding - gotRemediation model.Remediation + gotNodeKinds model.GraphSchemaNodeKinds + gotRelationshipKinds model.GraphSchemaRelationshipKinds + gotProperties model.GraphSchemaProperties + gotSchemaEnvironments []model.SchemaEnvironment + gotPrincipalKinds model.SchemaEnvironmentPrincipalKinds + sourceKind database.SourceKind + dawgsPrincipalKinds []model.Kind + dawgsPrincipalKind model.Kind + dawgsFindingRelationshipKinds []model.Kind + dawgsFindingRelationshipKind model.Kind + gotSchemaRelationshipFinding []model.SchemaRelationshipFinding + gotRemediation model.Remediation ) // Test Node Kinds @@ -1147,7 +1149,8 @@ func getAndCompareGraphExtension(t *testing.T, testContext context.Context, db * require.Equalf(t, len(want.EnvironmentsInput[idx].PrincipalKinds), len(gotPrincipalKinds), "PrincipalKinds - count mismatch") for _, gotPrincipalKind := range gotPrincipalKinds { require.Equalf(t, gotEnvironment.ID, gotPrincipalKind.EnvironmentId, "PrincipalKind - EnvironmentId is invalid") - dawgsPrincipalKind, err = db.GetKindById(testContext, gotPrincipalKind.PrincipalKind) + dawgsPrincipalKinds, err = db.GetKindsByIds(testContext, gotPrincipalKind.PrincipalKind) + dawgsPrincipalKind = dawgsPrincipalKinds[0] require.NoError(t, err) require.Containsf(t, want.EnvironmentsInput[idx].PrincipalKinds, dawgsPrincipalKind.Name, "PrincipalKind - Name mismatch") } @@ -1163,16 +1166,17 @@ func getAndCompareGraphExtension(t *testing.T, testContext context.Context, db * require.Greater(t, finding.ID, int32(0)) require.Equalf(t, gotGraphExtension.ID, finding.SchemaExtensionId, "RelationshipFindingInput - graph schema extension id should be greater than 0") - dawgsFindingRelationshipKind, err = db.GetKindById(testContext, finding.RelationshipKindId) + dawgsFindingRelationshipKinds, err = db.GetKindsByIds(testContext, finding.RelationshipKindId) + dawgsFindingRelationshipKind = dawgsFindingRelationshipKinds[0] require.NoError(t, err) require.Equalf(t, want.RelationshipFindingsInput[i].RelationshipKindName, dawgsFindingRelationshipKind.Name, "RelationshipFindingInput - relationship kind name mismatch") - // TODO: fix this - // was: findingEnvironment, err := db.GetEnvironmentById(testContext, finding.EnvironmentId) require.NoError(t, err) - dawgsFindingEnvironmentKind, err := db.GetKindById(testContext, findingEnvironment.EnvironmentKindId) + + dawgsFindingEnvironmentKinds, err := db.GetKindsByIds(testContext, findingEnvironment.EnvironmentKindId) require.NoError(t, err) + dawgsFindingEnvironmentKind := dawgsFindingEnvironmentKinds[0] require.Equalf(t, want.RelationshipFindingsInput[i].EnvironmentKindName, dawgsFindingEnvironmentKind.Name, "RelationshipFindingInput - environment kind name mismatch") require.Equalf(t, want.RelationshipFindingsInput[i].Name, finding.Name, "RelationshipFindingInput - name mismatch") From 6b09cc03f9ed3cd24eb98f19b17d8c6c6660e364 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Tue, 10 Feb 2026 12:02:02 -0800 Subject: [PATCH 21/22] sourcekind integration test --- cmd/api/src/database/sourcekinds.go | 4 ++ .../database/sourcekinds_integration_test.go | 52 +++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/cmd/api/src/database/sourcekinds.go b/cmd/api/src/database/sourcekinds.go index 0eadae13c47..978ddb5799f 100644 --- a/cmd/api/src/database/sourcekinds.go +++ b/cmd/api/src/database/sourcekinds.go @@ -120,6 +120,10 @@ func (s *BloodhoundDB) GetSourceKindsByIds(ctx context.Context, ids ...int32) ([ return nil, err } + if len(sourceKinds) != len(ids) { + return nil, ErrNotFound + } + return sourceKinds, nil } diff --git a/cmd/api/src/database/sourcekinds_integration_test.go b/cmd/api/src/database/sourcekinds_integration_test.go index d7e3937d1fd..d328f298c05 100644 --- a/cmd/api/src/database/sourcekinds_integration_test.go +++ b/cmd/api/src/database/sourcekinds_integration_test.go @@ -495,3 +495,55 @@ func TestDeactivateSourceKindsByName(t *testing.T) { }) } } + +func TestGetSourceKindsByIds(t *testing.T) { + testSuite := setupIntegrationTestSuite(t) + defer teardownIntegrationTestSuite(t, &testSuite) + + t.Run("not found", func(t *testing.T) { + _, err := testSuite.BHDatabase.GetSourceKindsByIds(testSuite.Context, 9999) + require.ErrorIs(t, err, database.ErrNotFound) + }) + + t.Run("single kind", func(t *testing.T) { + // Create a source kind + var sourceKind database.SourceKind + result := testSuite.DB.WithContext(testSuite.Context).Raw(` + INSERT INTO source_kinds (name, active) + VALUES ('TestSourceKindOne', true) + RETURNING id, name, active;`).Scan(&sourceKind) + require.NoError(t, result.Error) + + // Fetch it by ID + kinds, err := testSuite.BHDatabase.GetSourceKindsByIds(testSuite.Context, int32(sourceKind.ID)) + require.NoError(t, err) + require.Len(t, kinds, 1) + assert.Equal(t, "TestSourceKindOne", kinds[0].Name) + assert.True(t, kinds[0].Active) + }) + + t.Run("multiple kinds", func(t *testing.T) { + // Create two source kinds + var kind1, kind2 database.SourceKind + result := testSuite.DB.WithContext(testSuite.Context).Raw(` + INSERT INTO source_kinds (name, active) + VALUES ('TestSourceKindTwo', true) + RETURNING id, name, active;`).Scan(&kind1) + require.NoError(t, result.Error) + + result = testSuite.DB.WithContext(testSuite.Context).Raw(` + INSERT INTO source_kinds (name, active) + VALUES ('TestSourceKindThree', true) + RETURNING id, name, active;`).Scan(&kind2) + require.NoError(t, result.Error) + + // Fetch both by their IDs + kinds, err := testSuite.BHDatabase.GetSourceKindsByIds(testSuite.Context, int32(kind1.ID), int32(kind2.ID)) + require.NoError(t, err) + require.Len(t, kinds, 2) + + // Verify both kinds are returned (order not guaranteed) + names := []string{kinds[0].Name, kinds[1].Name} + assert.ElementsMatch(t, []string{"TestSourceKindTwo", "TestSourceKindThree"}, names) + }) +} From d2cc8818ae941dc2f4959fae1333c70a96b8d9c1 Mon Sep 17 00:00:00 2001 From: Brandon Shearin Date: Tue, 10 Feb 2026 12:34:15 -0800 Subject: [PATCH 22/22] purpose-built func for getting rels by extID --- cmd/api/src/database/graphschema.go | 19 +++++++++++++++++++ cmd/api/src/database/mocks/db.go | 15 +++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/cmd/api/src/database/graphschema.go b/cmd/api/src/database/graphschema.go index c09eeab55be..fc787811d70 100644 --- a/cmd/api/src/database/graphschema.go +++ b/cmd/api/src/database/graphschema.go @@ -47,6 +47,7 @@ type OpenGraphSchema interface { CreateGraphSchemaRelationshipKind(ctx context.Context, name string, schemaExtensionId int32, description string, isTraversable bool) (model.GraphSchemaRelationshipKind, error) GetGraphSchemaRelationshipKinds(ctx context.Context, filters model.Filters, sort model.Sort, skip, limit int) (model.GraphSchemaRelationshipKinds, int, error) GetGraphSchemaRelationshipKindById(ctx context.Context, schemaRelationshipKindId int32) (model.GraphSchemaRelationshipKind, error) + GetTraversableRelationshipKindsByExtensionID(ctx context.Context, extensionID int32) (model.GraphSchemaRelationshipKinds, error) UpdateGraphSchemaRelationshipKind(ctx context.Context, schemaRelationshipKind model.GraphSchemaRelationshipKind) (model.GraphSchemaRelationshipKind, error) DeleteGraphSchemaRelationshipKind(ctx context.Context, schemaRelationshipKindId int32) error @@ -496,6 +497,24 @@ func (s *BloodhoundDB) GetGraphSchemaRelationshipKinds(ctx context.Context, rela } } +// GetTraversableRelationshipKindsByExtensionID returns all traversable relationship kinds for a given schema extension. +// This is a purpose-built query for the analysis pipeline that needs traversable edges for graph traversal. +func (s *BloodhoundDB) GetTraversableRelationshipKindsByExtensionID(ctx context.Context, extensionID int32) (model.GraphSchemaRelationshipKinds, error) { + const query = ` + SELECT rk.id, k.name, rk.schema_extension_id, rk.description, rk.is_traversable, + rk.created_at, rk.updated_at, rk.deleted_at + FROM schema_relationship_kinds rk + JOIN kind k ON rk.kind_id = k.id + WHERE rk.schema_extension_id = $1 AND rk.is_traversable = true + ` + + var kinds model.GraphSchemaRelationshipKinds + if result := s.db.WithContext(ctx).Raw(query, extensionID).Scan(&kinds); result.Error != nil { + return nil, CheckError(result) + } + return kinds, nil +} + func (s *BloodhoundDB) GetGraphSchemaRelationshipKindsWithSchemaName(ctx context.Context, relationshipKindFilters model.Filters, sort model.Sort, skip, limit int) (model.GraphSchemaRelationshipKindsWithNamedSchema, int, error) { var ( schemaRelationshipKinds = model.GraphSchemaRelationshipKindsWithNamedSchema{} diff --git a/cmd/api/src/database/mocks/db.go b/cmd/api/src/database/mocks/db.go index 7d2aaf3cf16..42d6a3a01e7 100644 --- a/cmd/api/src/database/mocks/db.go +++ b/cmd/api/src/database/mocks/db.go @@ -2493,6 +2493,21 @@ func (mr *MockDatabaseMockRecorder) GetTimeRangedAssetGroupCollections(ctx, asse return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTimeRangedAssetGroupCollections", reflect.TypeOf((*MockDatabase)(nil).GetTimeRangedAssetGroupCollections), ctx, assetGroupID, from, to, order) } +// GetTraversableRelationshipKindsByExtensionID mocks base method. +func (m *MockDatabase) GetTraversableRelationshipKindsByExtensionID(ctx context.Context, extensionID int32) (model.GraphSchemaRelationshipKinds, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetTraversableRelationshipKindsByExtensionID", ctx, extensionID) + ret0, _ := ret[0].(model.GraphSchemaRelationshipKinds) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetTraversableRelationshipKindsByExtensionID indicates an expected call of GetTraversableRelationshipKindsByExtensionID. +func (mr *MockDatabaseMockRecorder) GetTraversableRelationshipKindsByExtensionID(ctx, extensionID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTraversableRelationshipKindsByExtensionID", reflect.TypeOf((*MockDatabase)(nil).GetTraversableRelationshipKindsByExtensionID), ctx, extensionID) +} + // GetUser mocks base method. func (m *MockDatabase) GetUser(ctx context.Context, id uuid.UUID) (model.User, error) { m.ctrl.T.Helper()