feat: Add Extension to Saved Queries - BED-9321

cmd/api/src/api/v2/saved_queries.go

@@ -477,7 +477,7 @@ func (s Resources) CreateSavedQuery(response http.ResponseWriter, request *http.
 		api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, err.Error(), request), response)
 	} else if createRequest.Name == "" || createRequest.Query == "" {
 		api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, "the name and/or query field is empty", request), response)
-	} else if savedQuery, err := s.DB.CreateSavedQuery(request.Context(), user.ID, createRequest.Name, createRequest.Query, createRequest.Description); err != nil {
+	} else if savedQuery, err := s.DB.CreateSavedQuery(request.Context(), user.ID, createRequest.Name, createRequest.Query, createRequest.Description, nil, nil); err != nil {
 		if strings.Contains(err.Error(), "duplicate key value violates unique constraint") {
 			api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, "duplicate name for saved query: please choose a different name", request), response)
 		} else {

cmd/api/src/api/v2/saved_queries_test.go

@@ -189,7 +189,7 @@ func TestResources_CreateSavedQuery_DuplicateName(t *testing.T) {
 
 	req.Header.Set(headers.ContentType.String(), mediatypes.ApplicationJson.String())
 
-	mockDB.EXPECT().CreateSavedQuery(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(model.SavedQuery{}, fmt.Errorf("duplicate key value violates unique constraint \"idx_saved_queries_composite_index\""))
+	mockDB.EXPECT().CreateSavedQuery(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(model.SavedQuery{}, fmt.Errorf("duplicate key value violates unique constraint \"idx_saved_queries_composite_index\""))
 
 	router := mux.NewRouter()
 	router.HandleFunc(endpoint, resources.CreateSavedQuery).Methods("POST")
@@ -232,7 +232,7 @@ func TestResources_CreateSavedQuery_CreateFailure(t *testing.T) {
 
 	req.Header.Set(headers.ContentType.String(), mediatypes.ApplicationJson.String())
 
-	mockDB.EXPECT().CreateSavedQuery(gomock.Any(), userId, payload["name"], payload["query"], payload["description"]).Return(model.SavedQuery{}, fmt.Errorf("foo"))
+	mockDB.EXPECT().CreateSavedQuery(gomock.Any(), userId, payload["name"], payload["query"], payload["description"], gomock.Any(), gomock.Any()).Return(model.SavedQuery{}, fmt.Errorf("foo"))
 
 	router := mux.NewRouter()
 	router.HandleFunc(endpoint, resources.CreateSavedQuery).Methods("POST")
@@ -275,7 +275,7 @@ func TestResources_CreateSavedQuery(t *testing.T) {
 
 	req.Header.Set(headers.ContentType.String(), mediatypes.ApplicationJson.String())
 
-	mockDB.EXPECT().CreateSavedQuery(gomock.Any(), userId, payload["name"], payload["query"], payload["description"]).Return(model.SavedQuery{
+	mockDB.EXPECT().CreateSavedQuery(gomock.Any(), userId, payload["name"], payload["query"], payload["description"], gomock.Any(), gomock.Any()).Return(model.SavedQuery{
 		UserID:      userId.String(),
 		Name:        fmt.Sprintf("%v", payload["name"]),
 		Query:       fmt.Sprintf("%v", payload["query"]),

cmd/api/src/database/migration/migrations/20260901110807_v9_add_schema_extension_id_to_saved_queries.sql

@@ -0,0 +1,57 @@
+-- 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
+-- +goose Up
+ALTER TABLE saved_queries
+    ADD COLUMN IF NOT EXISTS schema_extension_id INTEGER REFERENCES schema_extensions (id) ON DELETE CASCADE,
+    ADD COLUMN IF NOT EXISTS query_key TEXT
+        CONSTRAINT chk_saved_queries_extension_shape
+        CHECK (
+            (schema_extension_id IS NULL AND query_key IS NULL AND user_id <> '00000000-0000-0000-0000-000000000000')
+            OR
+            (schema_extension_id IS NOT NULL AND query_key IS NOT NULL AND user_id = '00000000-0000-0000-0000-000000000000')
+        );
+
+CREATE INDEX IF NOT EXISTS idx_saved_queries_schema_extension_id
+    ON saved_queries (schema_extension_id) WHERE schema_extension_id IS NOT NULL;
+
+DROP INDEX IF EXISTS idx_saved_queries_composite_index;
+
+CREATE UNIQUE INDEX IF NOT EXISTS idx_saved_queries_user_id_name
+    ON saved_queries (user_id, name) WHERE schema_extension_id IS NULL;
+CREATE UNIQUE INDEX IF NOT EXISTS idx_saved_queries_schema_extension_id_name
+    ON saved_queries (schema_extension_id, name) WHERE schema_extension_id IS NOT NULL;
+
+CREATE UNIQUE INDEX IF NOT EXISTS idx_saved_queries_extension_query_key
+    ON saved_queries (schema_extension_id, query_key) WHERE schema_extension_id IS NOT NULL;
+
+
+-- +goose Down
+
+DELETE FROM saved_queries WHERE schema_extension_id IS NOT NULL;
+
+DROP INDEX IF EXISTS idx_saved_queries_user_id_name;
+DROP INDEX IF EXISTS idx_saved_queries_schema_extension_id_name;
+
+CREATE UNIQUE INDEX IF NOT EXISTS idx_saved_queries_composite_index
+    ON saved_queries USING btree (user_id, name);
+
+DROP INDEX IF EXISTS idx_saved_queries_schema_extension_id;
+
+DROP INDEX IF EXISTS idx_saved_queries_extension_query_key;
+
+ALTER TABLE saved_queries
+    DROP COLUMN IF EXISTS query_key,
+    DROP COLUMN IF EXISTS schema_extension_id;

cmd/api/src/database/saved_queries.go

@@ -30,7 +30,7 @@ import (
 type SavedQueriesData interface {
 	GetSavedQuery(ctx context.Context, savedQueryID int64) (model.SavedQuery, error)
 	ListSavedQueries(ctx context.Context, scope string, userID uuid.UUID, order string, filter model.SQLFilter, skip, limit int) ([]model.ScopedSavedQuery, int, error)
-	CreateSavedQuery(ctx context.Context, userID uuid.UUID, name string, query string, description string) (model.SavedQuery, error)
+	CreateSavedQuery(ctx context.Context, userID uuid.UUID, name string, query string, description string, schemaExtensionID *int32, queryKey *string) (model.SavedQuery, error)
 	UpdateSavedQuery(ctx context.Context, savedQuery model.SavedQuery) (model.SavedQuery, error)
 	DeleteSavedQuery(ctx context.Context, savedQueryID int64) error
 	SavedQueryBelongsToUser(ctx context.Context, userID uuid.UUID, savedQueryID int64) (bool, error)
@@ -97,19 +97,23 @@ func (s *BloodhoundDB) ListSavedQueries(ctx context.Context, scope string, userI
 	return queries, int(count), CheckError(result)
 }
 
-func (s *BloodhoundDB) CreateSavedQuery(ctx context.Context, userID uuid.UUID, name string, query string, description string) (model.SavedQuery, error) {
+func (s *BloodhoundDB) CreateSavedQuery(ctx context.Context, userID uuid.UUID, name string, query string, description string, schemaExtensionID *int32, queryKey *string) (model.SavedQuery, error) {
 	savedQuery := model.SavedQuery{
-		UserID:      userID.String(),
-		Name:        name,
-		Query:       query,
-		Description: description,
+		UserID:            userID.String(),
+		Name:              name,
+		Query:             query,
+		Description:       description,
+		SchemaExtensionID: schemaExtensionID,
+		QueryKey:          queryKey,
 	}
 
-	return savedQuery, CheckError(s.db.WithContext(ctx).Create(&savedQuery))
+	result := s.db.WithContext(ctx).Create(&savedQuery)
+	return savedQuery, CheckError(result)
 }
 
 func (s *BloodhoundDB) UpdateSavedQuery(ctx context.Context, savedQuery model.SavedQuery) (model.SavedQuery, error) {
-	return savedQuery, CheckError(s.db.WithContext(ctx).Save(&savedQuery))
+	result := s.db.WithContext(ctx).Save(&savedQuery)
+	return savedQuery, CheckError(result)
 }
 
 func (s *BloodhoundDB) DeleteSavedQuery(ctx context.Context, savedQueryID int64) error {
@@ -127,18 +131,14 @@ func (s *BloodhoundDB) SavedQueryBelongsToUser(ctx context.Context, userID uuid.
 // GetSharedSavedQueries returns all the saved queries that the given userID has access to, including global queries
 func (s *BloodhoundDB) GetSharedSavedQueries(ctx context.Context, userID uuid.UUID) (model.SavedQueries, error) {
 	savedQueries := model.SavedQueries{}
-
 	result := s.db.WithContext(ctx).Select("saved_queries.*").Joins("JOIN saved_queries_permissions sqp ON sqp.query_id = saved_queries.id").Where("sqp.shared_to_user_id = ? ", userID).Find(&savedQueries)
-
 	return savedQueries, CheckError(result)
 }
 
 // GetPublicSavedQueries returns all the queries that were shared publicly
 func (s *BloodhoundDB) GetPublicSavedQueries(ctx context.Context) (model.SavedQueries, error) {
 	savedQueries := model.SavedQueries{}
-
 	result := s.db.WithContext(ctx).Select("saved_queries.*").Joins("JOIN saved_queries_permissions sqp ON sqp.query_id = saved_queries.id").Where("sqp.public = true").Find(&savedQueries)
-
 	return savedQueries, CheckError(result)
 }