diff --git a/libs/apps/generator/generator.go b/libs/apps/generator/generator.go index 3f99d80f900..3d38bed46b7 100644 --- a/libs/apps/generator/generator.go +++ b/libs/apps/generator/generator.go @@ -410,6 +410,12 @@ var appResourceSpecs = map[string]appResourceSpec{ staticFields: [][2]string{{"securable_type", "FUNCTION"}}, permission: "EXECUTE", }, + "model_service": { + yamlKey: "uc_securable", + varFields: [][2]string{{"name", "securable_full_name"}}, + staticFields: [][2]string{{"securable_type", "MODEL_SERVICE"}}, + permission: "EXECUTE", + }, "uc_connection": { yamlKey: "uc_securable", varFields: [][2]string{{"id", "securable_full_name"}}, diff --git a/libs/apps/generator/generator_test.go b/libs/apps/generator/generator_test.go index d7a1de85cc8..a9c63e3e273 100644 --- a/libs/apps/generator/generator_test.go +++ b/libs/apps/generator/generator_test.go @@ -246,6 +246,7 @@ func TestGenerateBundleResourcesDefaultPermissions(t *testing.T) { {"database", "CAN_CONNECT_AND_CREATE"}, {"volume", "READ_VOLUME"}, {"uc_function", "EXECUTE"}, + {"model_service", "EXECUTE"}, {"uc_connection", "USE_CONNECTION"}, {"genie_space", "CAN_VIEW"}, {"vector_search_index", "SELECT"}, @@ -466,6 +467,17 @@ func TestGenerateResourceYAMLAllTypes(t *testing.T) { "permission: EXECUTE", }, }, + { + name: "model_service maps to uc_securable MODEL_SERVICE", + resource: manifest.Resource{Type: "model_service", Alias: "Model Service", ResourceKey: "model", Permission: "EXECUTE"}, + expectContains: []string{ + "- name: model", + "uc_securable:", + "securable_full_name: ${var.model_name}", + "securable_type: MODEL_SERVICE", + "permission: EXECUTE", + }, + }, { name: "uc_connection maps to uc_securable CONNECTION", resource: manifest.Resource{Type: "uc_connection", Alias: "UC Connection", ResourceKey: "conn", Permission: "USE_CONNECTION"}, diff --git a/libs/apps/prompt/listers.go b/libs/apps/prompt/listers.go index 4bc2b1e301e..de8c3f1a404 100644 --- a/libs/apps/prompt/listers.go +++ b/libs/apps/prompt/listers.go @@ -155,6 +155,26 @@ func ListFunctionsInSchema(ctx context.Context, catalogName, schemaName string) return capResults(out), nil } +// ListModelServicesInSchema returns model services in catalog.schema, with IDs as catalog.schema.name. +func ListModelServicesInSchema(ctx context.Context, catalogName, schemaName string) ([]ListItem, error) { + w, err := workspaceClient(ctx) + if err != nil { + return nil, err + } + svcs, err := w.AiGateway.ListModelServicesAll(ctx, catalog.ListModelServicesRequest{ + Parent: "schemas/" + catalogName + "." + schemaName, + }) + if err != nil { + return nil, err + } + out := make([]ListItem, 0, min(len(svcs), maxListResults)) + for _, s := range svcs { + fullName := strings.TrimPrefix(s.Name, "model-services/") + out = append(out, ListItem{ID: fullName, Label: fullName}) + } + return capResults(out), nil +} + // listDatabasesResponse is the response from the /databases endpoint. type listDatabasesResponse struct { Databases []struct { diff --git a/libs/apps/prompt/listers_test.go b/libs/apps/prompt/listers_test.go new file mode 100644 index 00000000000..71049848dcc --- /dev/null +++ b/libs/apps/prompt/listers_test.go @@ -0,0 +1,32 @@ +package prompt_test + +import ( + "testing" + + "github.com/databricks/cli/libs/apps/prompt" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/catalog" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestListModelServicesInSchema(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + + m.GetMockAiGatewayAPI().EXPECT(). + ListModelServicesAll(ctx, catalog.ListModelServicesRequest{Parent: "schemas/main.default"}). + Return([]catalog.ModelService{ + {Name: "model-services/main.default.chat"}, + {Name: "model-services/main.default.embed"}, + }, nil) + + items, err := prompt.ListModelServicesInSchema(ctx, "main", "default") + require.NoError(t, err) + assert.Equal(t, []prompt.ListItem{ + {ID: "main.default.chat", Label: "main.default.chat"}, + {ID: "main.default.embed", Label: "main.default.embed"}, + }, items) +} diff --git a/libs/apps/prompt/prefetch.go b/libs/apps/prompt/prefetch.go index d3f38bade08..e2a2b017f19 100644 --- a/libs/apps/prompt/prefetch.go +++ b/libs/apps/prompt/prefetch.go @@ -26,6 +26,7 @@ const ( cacheKeyCatalogs = "_catalogs" cacheKeyDatabaseInstances = "_database_instances" cacheKeyPostgresProjects = "_postgres_projects" + cacheKeySystemAIModels = "_system_ai_model_services" ) // firstStepPrefetch maps multi-step resource types to the cache key and @@ -34,10 +35,11 @@ var firstStepPrefetch = map[string]struct { cacheKey string ctor pagedConstructor }{ - ResourceTypeVolume: {cacheKeyCatalogs, ListCatalogs}, - ResourceTypeUCFunction: {cacheKeyCatalogs, ListCatalogs}, - ResourceTypeDatabase: {cacheKeyDatabaseInstances, ListDatabaseInstances}, - ResourceTypePostgres: {cacheKeyPostgresProjects, ListPostgresProjects}, + ResourceTypeVolume: {cacheKeyCatalogs, ListCatalogs}, + ResourceTypeUCFunction: {cacheKeyCatalogs, ListCatalogs}, + ResourceTypeModelService: {cacheKeySystemAIModels, listSystemAIModelServices}, + ResourceTypeDatabase: {cacheKeyDatabaseInstances, ListDatabaseInstances}, + ResourceTypePostgres: {cacheKeyPostgresProjects, ListPostgresProjects}, } // PrefetchResources kicks off a background goroutine for every resource type diff --git a/libs/apps/prompt/prompt.go b/libs/apps/prompt/prompt.go index 6bbbaaa2d28..915bc6b4ec3 100644 --- a/libs/apps/prompt/prompt.go +++ b/libs/apps/prompt/prompt.go @@ -17,6 +17,7 @@ import ( "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/listing" "github.com/databricks/databricks-sdk-go/service/apps" "github.com/databricks/databricks-sdk-go/service/postgres" @@ -852,6 +853,39 @@ func PromptForUCFunction(ctx context.Context, r manifest.Resource, required bool return promptUCResource(ctx, r, required, "UC Function", "Fetching functions...", ListFunctionsInSchema) } +const browseCatalogsID = "__browse_catalogs__" + +// listSystemAIModelServices lists system.ai model services plus the browse entry; no access yields only the latter. +func listSystemAIModelServices(ctx context.Context) (*PagedFetcher, error) { + items, err := ListModelServicesInSchema(ctx, "system", "ai") + if errors.Is(err, apierr.ErrPermissionDenied) || errors.Is(err, apierr.ErrNotFound) { + items, err = nil, nil + } + if err != nil { + return nil, err + } + return &PagedFetcher{Items: append(items, ListItem{ID: browseCatalogsID, Label: "Browse other catalogs..."})}, nil +} + +// PromptForModelService opens on system.ai, since paging to the system catalog is slow in large workspaces. +func PromptForModelService(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { + f, err := getFetcherByKey(ctx, cacheKeySystemAIModels, "Fetching model services...", listSystemAIModelServices) + if err != nil { + return nil, err + } + value := browseCatalogsID + if len(f.Items) > 1 { + value, err = PromptFromList(ctx, resourceTitle("Select Model Service", r), "", f.Items, required) + if err != nil { + return nil, err + } + } + if value == browseCatalogsID { + return promptUCResource(ctx, r, required, "Model Service", "Fetching model services...", ListModelServicesInSchema) + } + return singleValueResult(r, value), nil +} + // PromptForUCConnection shows a picker for UC connections. func PromptForUCConnection(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { title := resourceTitle("Select UC Connection", r) diff --git a/libs/apps/prompt/prompt_test.go b/libs/apps/prompt/prompt_test.go index 239fb852cfa..9e77eb76aca 100644 --- a/libs/apps/prompt/prompt_test.go +++ b/libs/apps/prompt/prompt_test.go @@ -9,7 +9,11 @@ import ( "time" "github.com/databricks/cli/libs/apps/manifest" + "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/catalog" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -524,3 +528,30 @@ func TestPrintSuccessNotInPlace(t *testing.T) { assert.Contains(t, got, "cd my-app") assert.Contains(t, got, "npm run dev") } + +func TestListSystemAIModelServices(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + m.GetMockAiGatewayAPI().EXPECT(). + ListModelServicesAll(ctx, catalog.ListModelServicesRequest{Parent: "schemas/system.ai"}). + Return([]catalog.ModelService{{Name: "model-services/system.ai.chat"}}, nil) + + f, err := listSystemAIModelServices(ctx) + require.NoError(t, err) + assert.Equal(t, []ListItem{ + {ID: "system.ai.chat", Label: "system.ai.chat"}, + {ID: browseCatalogsID, Label: "Browse other catalogs..."}, + }, f.Items) +} + +func TestListSystemAIModelServicesWithoutAccess(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + m.GetMockAiGatewayAPI().EXPECT(). + ListModelServicesAll(ctx, catalog.ListModelServicesRequest{Parent: "schemas/system.ai"}). + Return(nil, apierr.ErrPermissionDenied) + + f, err := listSystemAIModelServices(ctx) + require.NoError(t, err) + assert.Equal(t, []ListItem{{ID: browseCatalogsID, Label: "Browse other catalogs..."}}, f.Items) +} diff --git a/libs/apps/prompt/resource_registry.go b/libs/apps/prompt/resource_registry.go index e6846ad2f2b..8744aa5594c 100644 --- a/libs/apps/prompt/resource_registry.go +++ b/libs/apps/prompt/resource_registry.go @@ -16,6 +16,7 @@ const ( ResourceTypeVectorSearchIndex = "vector_search_index" ResourceTypeUCFunction = "uc_function" ResourceTypeUCConnection = "uc_connection" + ResourceTypeModelService = "model_service" ResourceTypeDatabase = "database" // Lakebase Provisioned (V1) ResourceTypePostgres = "postgres" // Lakebase Autoscaling (V2) ResourceTypeGenieSpace = "genie_space" @@ -62,6 +63,8 @@ func GetPromptFunc(resourceType string) (PromptResourceFunc, bool) { return PromptForVectorSearchIndex, true case ResourceTypeUCFunction: return PromptForUCFunction, true + case ResourceTypeModelService: + return PromptForModelService, true case ResourceTypeUCConnection: return PromptForUCConnection, true case ResourceTypeDatabase: