diff --git a/baton_capabilities.json b/baton_capabilities.json index 91985cec..ed513611 100644 --- a/baton_capabilities.json +++ b/baton_capabilities.json @@ -12,6 +12,17 @@ ], "permissions": {} }, + { + "resourceType": { + "id": "catalog", + "displayName": "Catalog" + }, + "capabilities": [ + "CAPABILITY_SYNC", + "CAPABILITY_PROVISION" + ], + "permissions": {} + }, { "resourceType": { "id": "group", @@ -40,6 +51,17 @@ ], "permissions": {} }, + { + "resourceType": { + "id": "schema", + "displayName": "Schema" + }, + "capabilities": [ + "CAPABILITY_SYNC", + "CAPABILITY_PROVISION" + ], + "permissions": {} + }, { "resourceType": { "id": "service_principal", @@ -54,6 +76,17 @@ ], "permissions": {} }, + { + "resourceType": { + "id": "table", + "displayName": "Table" + }, + "capabilities": [ + "CAPABILITY_SYNC", + "CAPABILITY_PROVISION" + ], + "permissions": {} + }, { "resourceType": { "id": "user", diff --git a/config_schema.json b/config_schema.json index 2e162e4f..fcbfdaf5 100644 --- a/config_schema.json +++ b/config_schema.json @@ -163,6 +163,18 @@ "displayName": "Exclude Workspaces", "description": "Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces.", "stringSliceField": {} + }, + { + "name": "sync-unity-catalog", + "displayName": "Sync Unity Catalog", + "description": "Sync Unity Catalog catalogs and schemas (and their permission grants) for access reviews.", + "boolField": {} + }, + { + "name": "sync-unity-catalog-tables", + "displayName": "Sync Unity Catalog Tables", + "description": "Also sync Unity Catalog tables and their grants. Requires sync-unity-catalog. May be very large in big metastores.", + "boolField": {} } ], "constraints": [ @@ -181,6 +193,15 @@ "secondaryFieldNames": [ "workspaces" ] + }, + { + "kind": "CONSTRAINT_KIND_DEPENDENT_ON", + "fieldNames": [ + "sync-unity-catalog-tables" + ], + "secondaryFieldNames": [ + "sync-unity-catalog" + ] } ], "displayName": "Databricks", @@ -198,7 +219,9 @@ "hostname", "account-hostname", "workspaces", - "databricks-exclude-workspaces" + "databricks-exclude-workspaces", + "sync-unity-catalog", + "sync-unity-catalog-tables" ], "default": true } diff --git a/docs/connector.mdx b/docs/connector.mdx index ba8ddee1..7fe7d2da 100644 --- a/docs/connector.mdx +++ b/docs/connector.mdx @@ -11,9 +11,12 @@ sidebarTitle: "Databricks" | Resource | Sync | Provision | | :--- | :--- | :--- | | Accounts | | | +| Catalogs | | | | Groups | | | | Roles | | | +| Schemas | | | | Service principals | | | +| Tables | | | | Users | | | | Workspaces | | | @@ -219,6 +222,14 @@ stringData: # Mutually exclusive with BATON_WORKSPACES — set one or the other, never both. BATON_DATABRICKS_EXCLUDE_WORKSPACES: + # Optional: sync Unity Catalog catalogs and schemas (and their grants) for + # data access reviews. Off by default. + # BATON_SYNC_UNITY_CATALOG: true + + # Optional: also sync Unity Catalog tables and their grants. Requires + # BATON_SYNC_UNITY_CATALOG. Can be very large in big metastores. + # BATON_SYNC_UNITY_CATALOG_TABLES: true + # Optional: include if you want C1 to provision access using this connector BATON_PROVISIONING: true ``` diff --git a/pkg/config/conf.gen.go b/pkg/config/conf.gen.go index 6ca5facf..e8fdadea 100644 --- a/pkg/config/conf.gen.go +++ b/pkg/config/conf.gen.go @@ -13,6 +13,8 @@ type Databricks struct { WorkspaceTokens []string `mapstructure:"workspace-tokens"` BaseUrl string `mapstructure:"base-url"` DatabricksExcludeWorkspaces []string `mapstructure:"databricks-exclude-workspaces"` + SyncUnityCatalog bool `mapstructure:"sync-unity-catalog"` + SyncUnityCatalogTables bool `mapstructure:"sync-unity-catalog-tables"` } func (c *Databricks) findFieldByTag(tagValue string) (any, bool) { diff --git a/pkg/config/config.go b/pkg/config/config.go index d6a24c85..b661e9b1 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -70,6 +70,18 @@ var ( field.WithDescription("Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces."), field.WithDisplayName("Exclude Workspaces"), ) + SyncUnityCatalogField = field.BoolField( + "sync-unity-catalog", + field.WithDescription("Sync Unity Catalog catalogs and schemas (and their permission grants) for access reviews."), + field.WithDefaultValue(false), + field.WithDisplayName("Sync Unity Catalog"), + ) + SyncUnityCatalogTablesField = field.BoolField( + "sync-unity-catalog-tables", + field.WithDescription("Also sync Unity Catalog tables and their grants. Requires sync-unity-catalog. May be very large in big metastores."), + field.WithDefaultValue(false), + field.WithDisplayName("Sync Unity Catalog Tables"), + ) configFields = []field.SchemaField{ AccountHostnameField, AccountIdField, @@ -80,6 +92,8 @@ var ( WorkspaceTokensField, BaseURLField, ExcludeWorkspacesField, + SyncUnityCatalogField, + SyncUnityCatalogTablesField, } ) @@ -92,6 +106,7 @@ var Config = field.NewConfiguration( field.WithConstraints( field.FieldsMutuallyExclusive(WorkspacesField, ExcludeWorkspacesField), field.FieldsDependentOn([]field.SchemaField{WorkspaceTokensField}, []field.SchemaField{WorkspacesField}), + field.FieldsDependentOn([]field.SchemaField{SyncUnityCatalogTablesField}, []field.SchemaField{SyncUnityCatalogField}), ), field.WithFieldGroups([]field.SchemaFieldGroup{ { @@ -101,6 +116,7 @@ var Config = field.NewConfiguration( Fields: []field.SchemaField{ AccountIdField, DatabricksClientIdField, DatabricksClientSecretField, HostnameField, AccountHostnameField, WorkspacesField, ExcludeWorkspacesField, + SyncUnityCatalogField, SyncUnityCatalogTablesField, }, Default: true, }, diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 10996986..af989701 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -16,19 +16,25 @@ import ( ) type Databricks struct { - client *databricks.Client - workspaces []string + client *databricks.Client + workspaces []string + syncUnityCatalog bool + syncUnityCatalogTables bool } // ResourceSyncers returns a ResourceSyncerV2 for each resource type that should be synced from the upstream service. func (d *Databricks) ResourceSyncers(ctx context.Context) []connectorbuilder.ResourceSyncerV2 { + resolver := newUCResolver(d.client) syncers := []connectorbuilder.ResourceSyncerV2{ newAccountBuilder(d.client), newGroupBuilder(d.client), newServicePrincipalBuilder(d.client), newUserBuilder(d.client), - newWorkspaceBuilder(d.client, d.workspaces), + newWorkspaceBuilder(d.client, d.workspaces, d.syncUnityCatalog), newRoleBuilder(d.client), + newCatalogBuilder(resolver), + newSchemaBuilder(resolver, d.syncUnityCatalogTables), + newTableBuilder(resolver), } return syncers @@ -174,6 +180,8 @@ func New( auth databricks.Auth, excludeWorkspaces []string, workspaces []string, + syncUnityCatalog bool, + syncUnityCatalogTables bool, ) (*Databricks, error) { httpClient, err := auth.GetClient(ctx) if err != nil { @@ -186,8 +194,10 @@ func New( } return &Databricks{ - client: client, - workspaces: workspaces, + client: client, + workspaces: workspaces, + syncUnityCatalog: syncUnityCatalog, + syncUnityCatalogTables: syncUnityCatalogTables, }, nil } @@ -216,6 +226,8 @@ func NewConnector(ctx context.Context, cfg *config.Databricks, opts *cli.Connect auth, cfg.DatabricksExcludeWorkspaces, cfg.Workspaces, + cfg.SyncUnityCatalog, + cfg.SyncUnityCatalogTables, ) if err != nil { return nil, nil, err diff --git a/pkg/connector/pagination.go b/pkg/connector/pagination.go index ca0c111d..c302bc33 100644 --- a/pkg/connector/pagination.go +++ b/pkg/connector/pagination.go @@ -58,3 +58,23 @@ func prepareNextToken(page uint, pageTotal int, total uint) string { return token } + +// parseCursorToken initializes a pagination Bag for cursor-based (opaque +// string token) APIs such as Unity Catalog. Unlike parsePageToken it returns +// the current cursor string directly instead of a numeric offset. Advance with +// bag.NextToken(next) where next is the API's next_page_token ("" ends). +func parseCursorToken(i string, resourceID *v2.ResourceId) (*pagination.Bag, string, error) { + b := &pagination.Bag{} + if err := b.Unmarshal(i); err != nil { + return nil, "", err + } + + if b.Current() == nil { + b.Push(pagination.PageState{ + ResourceTypeID: resourceID.ResourceType, + ResourceID: resourceID.Resource, + }) + } + + return b, b.PageToken(), nil +} diff --git a/pkg/connector/resource_types.go b/pkg/connector/resource_types.go index 748bdbef..b351547e 100644 --- a/pkg/connector/resource_types.go +++ b/pkg/connector/resource_types.go @@ -46,4 +46,22 @@ var ( Id: "account", DisplayName: "Account", } + + // The catalog resource type is a Unity Catalog top-level data container. + catalogResourceType = &v2.ResourceType{ + Id: "catalog", + DisplayName: "Catalog", + } + + // The schema resource type is a Unity Catalog schema within a catalog. + schemaResourceType = &v2.ResourceType{ + Id: "schema", + DisplayName: "Schema", + } + + // The table resource type is a Unity Catalog table or view within a schema. + tableResourceType = &v2.ResourceType{ + Id: "table", + DisplayName: "Table", + } ) diff --git a/pkg/connector/unity_catalog.go b/pkg/connector/unity_catalog.go new file mode 100644 index 00000000..307a78c7 --- /dev/null +++ b/pkg/connector/unity_catalog.go @@ -0,0 +1,668 @@ +package connector + +import ( + "context" + "fmt" + "strings" + "sync" + + "github.com/conductorone/baton-databricks/pkg/databricks" + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "github.com/conductorone/baton-sdk/pkg/annotations" + ent "github.com/conductorone/baton-sdk/pkg/types/entitlement" + "github.com/conductorone/baton-sdk/pkg/types/grant" + rs "github.com/conductorone/baton-sdk/pkg/types/resource" + "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" + "go.uber.org/zap" +) + +// securableIDSep separates the workspace deployment name from the Unity Catalog +// dotted full name inside a resource ID. UC full names use dots, so a distinct +// separator keeps the two parts unambiguous. Child List() calls only receive the +// parent's ResourceId, so everything needed to enumerate children (workspace + +// full name) must be encoded here. +const securableIDSep = "::" + +// ownerEntitlement is a read-only entitlement marking the Unity Catalog owner of +// a securable. Ownership carries implicit full control that the permissions API +// does not report, so it is surfaced explicitly for access reviews. It is not +// provisionable through the permissions API. +const ownerEntitlement = "owner" + +// Unity Catalog privileges per securable level. These are the entitlements +// exposed for each level; a grant referencing a privilege outside its level's +// set is skipped (with a warning) so no grant points at an entitlement that +// Entitlements() never produced. +var ( + catalogPrivileges = []string{ + "ALL_PRIVILEGES", "USE_CATALOG", "USE_SCHEMA", "CREATE_SCHEMA", "CREATE_TABLE", + "CREATE_FUNCTION", "CREATE_MODEL", "CREATE_VOLUME", "CREATE_MATERIALIZED_VIEW", + "MODIFY", "SELECT", "EXECUTE", "READ_VOLUME", "WRITE_VOLUME", "REFRESH", + "APPLY_TAG", "BROWSE", "EXTERNAL_USE_SCHEMA", "MANAGE", + } + schemaPrivileges = []string{ + "ALL_PRIVILEGES", "USE_SCHEMA", "CREATE_TABLE", "CREATE_FUNCTION", "CREATE_MODEL", + "CREATE_VOLUME", "CREATE_MATERIALIZED_VIEW", "MODIFY", "SELECT", "EXECUTE", + "READ_VOLUME", "WRITE_VOLUME", "REFRESH", "APPLY_TAG", "BROWSE", + "EXTERNAL_USE_SCHEMA", "MANAGE", + } + tablePrivileges = []string{ + "ALL_PRIVILEGES", "SELECT", "MODIFY", "APPLY_TAG", "BROWSE", "MANAGE", + } +) + +func makeSecurableID(workspaceId, fullName string) string { + return workspaceId + securableIDSep + fullName +} + +// splitSecurableID reverses makeSecurableID into (workspaceId, fullName). +func splitSecurableID(id string) (string, string) { + workspaceId, fullName, found := strings.Cut(id, securableIDSep) + if !found { + return "", id + } + return workspaceId, fullName +} + +func securableProfile(workspaceId, fullName, owner, securableType string) map[string]interface{} { + return map[string]interface{}{ + "workspace_id": workspaceId, + "full_name": fullName, + "owner": owner, + "securable_type": securableType, + } +} + +// securablePrivilegeEntitlements builds one permission entitlement per privilege +// plus the read-only owner entitlement. +func securablePrivilegeEntitlements(resource *v2.Resource, privileges []string) []*v2.Entitlement { + rv := make([]*v2.Entitlement, 0, len(privileges)+1) + for _, priv := range privileges { + rv = append(rv, ent.NewPermissionEntitlement( + resource, + priv, + ent.WithGrantableTo(userResourceType, groupResourceType, servicePrincipalResourceType), + ent.WithDisplayName(fmt.Sprintf("%s %s", resource.DisplayName, priv)), + ent.WithDescription(fmt.Sprintf("%s privilege on %s", priv, resource.DisplayName)), + )) + } + rv = append(rv, ent.NewPermissionEntitlement( + resource, + ownerEntitlement, + ent.WithGrantableTo(userResourceType, groupResourceType, servicePrincipalResourceType), + ent.WithDisplayName(fmt.Sprintf("%s owner", resource.DisplayName)), + ent.WithDescription(fmt.Sprintf("Owner of %s", resource.DisplayName)), + )) + return rv +} + +// ucResolver resolves Unity Catalog principal strings to/from connector resource +// IDs, caching results per (scope, principal) to avoid a SCIM lookup storm: +// Grants() runs for every catalog, schema and table, and each grant would +// otherwise cost up to three uncached SCIM list calls per principal. +type ucResolver struct { + client *databricks.Client + mu sync.Mutex + cache map[string]*v2.ResourceId +} + +func newUCResolver(client *databricks.Client) *ucResolver { + return &ucResolver{client: client, cache: make(map[string]*v2.ResourceId)} +} + +// scimScopes returns the SCIM scopes to search, in order. Unity Catalog grants +// are typically held by account-level identities, which are synced from account +// SCIM (scope "") whenever the account API is available; the workspace scope is +// a fallback for workspace-local identities (and the only scope under token +// auth, where identities are synced per workspace). +func (r *ucResolver) scimScopes(workspaceId string) []string { + if r.client.IsAccountAPIAvailable() { + return []string{"", workspaceId} + } + return []string{workspaceId} +} + +// resolve maps a Unity Catalog principal string (user name/email, group display +// name, or service principal application ID) to a connector resource ID. Returns +// nil when no identity matches. Results (including negatives) are cached. +func (r *ucResolver) resolve(ctx context.Context, workspaceId, principal string) (*v2.ResourceId, error) { + key := workspaceId + "\x00" + principal + + r.mu.Lock() + if id, ok := r.cache[key]; ok { + r.mu.Unlock() + return id, nil + } + r.mu.Unlock() + + var result *v2.ResourceId + for _, scope := range r.scimScopes(workspaceId) { + id, err := resolvePrincipalResourceID(ctx, r.client, scope, principal) + if err != nil { + return nil, err + } + if id != nil { + result = id + break + } + } + + r.mu.Lock() + r.cache[key] = result + r.mu.Unlock() + return result, nil +} + +// resolveName converts a connector principal resource ID back into the Unity +// Catalog principal string expected by the permissions API, searching the same +// scopes as resolve. +func (r *ucResolver) resolveName(ctx context.Context, workspaceId string, principal *v2.ResourceId) (string, error) { + for _, scope := range r.scimScopes(workspaceId) { + name, err := resolvePrincipalName(ctx, r.client, scope, principal) + if err != nil { + return "", err + } + if name != "" { + return name, nil + } + } + return "", nil +} + +// resolvePrincipalResourceID resolves a Unity Catalog principal string against a +// single SCIM scope (workspaceId "" targets the account API). Returns nil when +// the principal cannot be matched to any known identity in that scope. +func resolvePrincipalResourceID(ctx context.Context, c *databricks.Client, workspaceId, principal string) (*v2.ResourceId, error) { + if userID, _, err := c.FindUserID(ctx, workspaceId, principal); err == nil && userID != "" { + return &v2.ResourceId{ResourceType: userResourceType.Id, Resource: userID}, nil + } else if err != nil { + return nil, err + } + + if groupID, _, err := c.FindGroupID(ctx, workspaceId, principal); err == nil && groupID != "" { + return &v2.ResourceId{ResourceType: groupResourceType.Id, Resource: groupID}, nil + } else if err != nil { + return nil, err + } + + if spID, _, err := c.FindServicePrincipalID(ctx, workspaceId, principal); err == nil && spID != "" { + return &v2.ResourceId{ResourceType: servicePrincipalResourceType.Id, Resource: spID}, nil + } else if err != nil { + return nil, err + } + + return nil, nil +} + +// resolvePrincipalName converts a connector principal resource ID into the Unity +// Catalog principal string within a single SCIM scope. Returns "" when not found. +func resolvePrincipalName(ctx context.Context, c *databricks.Client, workspaceId string, principal *v2.ResourceId) (string, error) { + switch principal.ResourceType { + case userResourceType.Id: + name, _, err := c.FindUsername(ctx, workspaceId, principal.Resource) + return name, err + case groupResourceType.Id: + name, _, err := c.FindGroupDisplayName(ctx, workspaceId, principal.Resource) + return name, err + case servicePrincipalResourceType.Id: + name, _, err := c.FindServicePrincipalAppID(ctx, workspaceId, principal.Resource) + return name, err + default: + return "", fmt.Errorf("databricks-connector: unsupported principal type: %s", principal.ResourceType) + } +} + +// securableGrants lists the direct privilege assignments on a securable and +// converts them into grants, expanding group principals and adding the owner. +// Failures degrade to an empty result (logged) rather than aborting the sync; +// context cancellation still propagates. +func securableGrants(ctx context.Context, r *ucResolver, resource *v2.Resource, securableType string, privileges []string) ([]*v2.Grant, annotations.Annotations, error) { + l := ctxzap.Extract(ctx) + workspaceId, fullName := splitSecurableID(resource.Id.Resource) + + assignments, ratelimit, err := r.client.ListPermissions(ctx, workspaceId, securableType, fullName) + if err != nil { + if ctx.Err() != nil { + return nil, nil, ctx.Err() + } + l.Warn("databricks-connector: unable to list unity catalog permissions, skipping securable grants", + zap.String("securable_type", securableType), + zap.String("securable", fullName), + zap.Error(err), + ) + return nil, nil, nil + } + + annos := annotations.Annotations{} + if ratelimit != nil { + annos.WithRateLimiting(ratelimit) + } + + known := make(map[string]struct{}, len(privileges)) + for _, p := range privileges { + known[p] = struct{}{} + } + + var rv []*v2.Grant + for _, a := range assignments { + principalID, err := r.resolve(ctx, workspaceId, a.Principal) + if err != nil { + if ctx.Err() != nil { + return nil, nil, ctx.Err() + } + l.Warn("databricks-connector: failed to resolve unity catalog principal, skipping", + zap.String("principal", a.Principal), + zap.String("securable", fullName), + zap.Error(err), + ) + continue + } + if principalID == nil { + l.Warn("databricks-connector: skipping unresolved unity catalog principal", + zap.String("principal", a.Principal), + zap.String("securable", fullName), + ) + continue + } + + for _, priv := range a.Privileges { + if _, ok := known[priv]; !ok { + l.Warn("databricks-connector: skipping unmodeled unity catalog privilege", + zap.String("privilege", priv), + zap.String("securable_type", securableType), + zap.String("securable", fullName), + ) + continue + } + g, err := securableGrant(ctx, r.client, resource, workspaceId, priv, principalID) + if err != nil { + return nil, nil, err + } + rv = append(rv, g) + } + } + + // Owners hold implicit full control the permissions API omits. + if owner, ok := rs.GetProfileStringValue(rs.GetProfile(resource), "owner"); ok && owner != "" { + principalID, err := r.resolve(ctx, workspaceId, owner) + if err != nil { + if ctx.Err() != nil { + return nil, nil, ctx.Err() + } + l.Warn("databricks-connector: failed to resolve unity catalog owner, skipping owner grant", + zap.String("owner", owner), + zap.String("securable", fullName), + zap.Error(err), + ) + return rv, annos, nil + } + if principalID == nil { + l.Warn("databricks-connector: skipping unresolved unity catalog owner", + zap.String("owner", owner), + zap.String("securable", fullName), + ) + return rv, annos, nil + } + + g, err := securableGrant(ctx, r.client, resource, workspaceId, ownerEntitlement, principalID) + if err != nil { + return nil, nil, err + } + rv = append(rv, g) + } + + return rv, annos, nil +} + +// securableGrant builds a single grant, adding group-membership expansion when +// the principal is a group so reviews see the effective members. +func securableGrant(ctx context.Context, c *databricks.Client, resource *v2.Resource, workspaceId, entitlement string, principalID *v2.ResourceId) (*v2.Grant, error) { + if principalID.ResourceType != groupResourceType.Id { + return grant.NewGrant(resource, entitlement, principalID), nil + } + + groupParent, err := groupGrantParent(c.IsAccountAPIAvailable(), c.GetAccountId(), workspaceId) + if err != nil { + return nil, err + } + expandedID, expandAnnotation, err := groupGrantExpansion(ctx, principalID.Resource, groupParent) + if err != nil { + return nil, err + } + return grant.NewGrant(resource, entitlement, expandedID, grant.WithAnnotation(expandAnnotation)), nil +} + +// securableGrantChange applies a single add/remove privilege change on a securable. +func securableGrantChange(ctx context.Context, r *ucResolver, entitlement *v2.Entitlement, principal *v2.Resource, add bool) (annotations.Annotations, error) { + l := ctxzap.Extract(ctx) + + if !isValidPrincipal(principal.Id) { + l.Warn("databricks-connector: only users, groups and service principals can hold unity catalog privileges", + zap.String("principal_id", principal.Id.String()), + zap.String("principal_type", principal.Id.ResourceType), + ) + return nil, fmt.Errorf("databricks-connector: only users, groups and service principals can hold unity catalog privileges") + } + + privilege := entitlement.Slug + if privilege == ownerEntitlement { + return nil, fmt.Errorf("databricks-connector: unity catalog ownership cannot be provisioned via the permissions API") + } + + workspaceId, fullName := splitSecurableID(entitlement.Resource.Id.Resource) + securableType, ok := rs.GetProfileStringValue(rs.GetProfile(entitlement.Resource), "securable_type") + if !ok { + return nil, fmt.Errorf("databricks-connector: securable type not found on entitlement resource") + } + + principalName, err := r.resolveName(ctx, workspaceId, principal.Id) + if err != nil { + return nil, fmt.Errorf("databricks-connector: failed to resolve principal name: %w", err) + } + if principalName == "" { + return nil, fmt.Errorf("databricks-connector: could not resolve principal %s to a unity catalog principal", principal.Id.String()) + } + + change := databricks.PermissionsChange{Principal: principalName} + if add { + change.Add = []string{privilege} + } else { + change.Remove = []string{privilege} + } + + if _, err := r.client.UpdatePermissions(ctx, workspaceId, securableType, fullName, []databricks.PermissionsChange{change}); err != nil { + return nil, fmt.Errorf("databricks-connector: failed to update %s permissions for %q: %w", securableType, fullName, err) + } + + return nil, nil +} + +// syncOpResults builds a SyncOpResults carrying the next page token and any +// rate-limit feedback so the SDK can back off on the highest-volume calls. +func syncOpResults(nextPage string, ratelimit *v2.RateLimitDescription) *rs.SyncOpResults { + res := &rs.SyncOpResults{NextPageToken: nextPage} + if ratelimit != nil { + res.Annotations = annotations.Annotations{} + res.Annotations.WithRateLimiting(ratelimit) + } + return res +} + +// ---- Catalog builder ---- + +type catalogBuilder struct { + resolver *ucResolver + resourceType *v2.ResourceType +} + +func (b *catalogBuilder) ResourceType(ctx context.Context) *v2.ResourceType { + return catalogResourceType +} + +func catalogResource(workspaceId string, catalog *databricks.Catalog, parent *v2.ResourceId) (*v2.Resource, error) { + return rs.NewResource( + catalog.Name, + catalogResourceType, + makeSecurableID(workspaceId, catalog.Name), + rs.WithResourceProfile(securableProfile(workspaceId, catalog.Name, catalog.Owner, databricks.SecurableCatalog)), + rs.WithParentResourceID(parent), + rs.WithAnnotation(&v2.ChildResourceType{ResourceTypeId: schemaResourceType.Id}), + ) +} + +func (b *catalogBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId, attr rs.SyncOpAttrs) ([]*v2.Resource, *rs.SyncOpResults, error) { + if parentResourceID == nil || parentResourceID.ResourceType != workspaceResourceType.Id { + return nil, nil, nil + } + workspaceId := parentResourceID.Resource + + bag, pageToken, err := parseCursorToken(attr.PageToken.Token, &v2.ResourceId{ResourceType: catalogResourceType.Id}) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to parse page token: %w", err) + } + + catalogs, nextToken, ratelimit, err := b.resolver.client.ListCatalogs(ctx, workspaceId, pageToken, ResourcesPageSize) + if err != nil { + if ctx.Err() != nil { + return nil, nil, ctx.Err() + } + // Degrade gracefully: workspaces without Unity Catalog (or where the + // service principal lacks UC access) must not fail the sync. + ctxzap.Extract(ctx).Warn("databricks-connector: unable to list catalogs, skipping unity catalog for workspace", + zap.String("workspace", workspaceId), + zap.Error(err), + ) + return nil, nil, nil + } + + var rv []*v2.Resource + for _, catalog := range catalogs { + cCopy := catalog + cr, err := catalogResource(workspaceId, &cCopy, parentResourceID) + if err != nil { + return nil, nil, err + } + rv = append(rv, cr) + } + + nextPage, err := bag.NextToken(nextToken) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to create next page token: %w", err) + } + + return rv, syncOpResults(nextPage, ratelimit), nil +} + +func (b *catalogBuilder) Entitlements(_ context.Context, resource *v2.Resource, _ rs.SyncOpAttrs) ([]*v2.Entitlement, *rs.SyncOpResults, error) { + return securablePrivilegeEntitlements(resource, catalogPrivileges), nil, nil +} + +func (b *catalogBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.SyncOpAttrs) ([]*v2.Grant, *rs.SyncOpResults, error) { + grants, annos, err := securableGrants(ctx, b.resolver, resource, databricks.SecurableCatalog, catalogPrivileges) + if err != nil { + return nil, nil, err + } + return grants, &rs.SyncOpResults{Annotations: annos}, nil +} + +func (b *catalogBuilder) Grant(ctx context.Context, principal *v2.Resource, entitlement *v2.Entitlement) (annotations.Annotations, error) { + return securableGrantChange(ctx, b.resolver, entitlement, principal, true) +} + +func (b *catalogBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotations.Annotations, error) { + return securableGrantChange(ctx, b.resolver, grant.Entitlement, grant.Principal, false) +} + +func newCatalogBuilder(resolver *ucResolver) *catalogBuilder { + return &catalogBuilder{resolver: resolver, resourceType: catalogResourceType} +} + +// ---- Schema builder ---- + +type schemaBuilder struct { + resolver *ucResolver + resourceType *v2.ResourceType + syncTables bool +} + +func (b *schemaBuilder) ResourceType(ctx context.Context) *v2.ResourceType { + return schemaResourceType +} + +func schemaResource(workspaceId string, schema *databricks.Schema, parent *v2.ResourceId, syncTables bool) (*v2.Resource, error) { + opts := []rs.ResourceOption{ + rs.WithResourceProfile(securableProfile(workspaceId, schema.FullName, schema.Owner, databricks.SecurableSchema)), + rs.WithParentResourceID(parent), + } + if syncTables { + opts = append(opts, rs.WithAnnotation(&v2.ChildResourceType{ResourceTypeId: tableResourceType.Id})) + } + + return rs.NewResource( + schema.Name, + schemaResourceType, + makeSecurableID(workspaceId, schema.FullName), + opts..., + ) +} + +func (b *schemaBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId, attr rs.SyncOpAttrs) ([]*v2.Resource, *rs.SyncOpResults, error) { + if parentResourceID == nil || parentResourceID.ResourceType != catalogResourceType.Id { + return nil, nil, nil + } + workspaceId, catalogName := splitSecurableID(parentResourceID.Resource) + + bag, pageToken, err := parseCursorToken(attr.PageToken.Token, &v2.ResourceId{ResourceType: schemaResourceType.Id}) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to parse page token: %w", err) + } + + schemas, nextToken, ratelimit, err := b.resolver.client.ListSchemas(ctx, workspaceId, catalogName, pageToken, ResourcesPageSize) + if err != nil { + if ctx.Err() != nil { + return nil, nil, ctx.Err() + } + ctxzap.Extract(ctx).Warn("databricks-connector: unable to list schemas, skipping catalog", + zap.String("workspace", workspaceId), + zap.String("catalog", catalogName), + zap.Error(err), + ) + return nil, nil, nil + } + + var rv []*v2.Resource + for _, schema := range schemas { + sCopy := schema + sr, err := schemaResource(workspaceId, &sCopy, parentResourceID, b.syncTables) + if err != nil { + return nil, nil, err + } + rv = append(rv, sr) + } + + nextPage, err := bag.NextToken(nextToken) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to create next page token: %w", err) + } + + return rv, syncOpResults(nextPage, ratelimit), nil +} + +func (b *schemaBuilder) Entitlements(_ context.Context, resource *v2.Resource, _ rs.SyncOpAttrs) ([]*v2.Entitlement, *rs.SyncOpResults, error) { + return securablePrivilegeEntitlements(resource, schemaPrivileges), nil, nil +} + +func (b *schemaBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.SyncOpAttrs) ([]*v2.Grant, *rs.SyncOpResults, error) { + grants, annos, err := securableGrants(ctx, b.resolver, resource, databricks.SecurableSchema, schemaPrivileges) + if err != nil { + return nil, nil, err + } + return grants, &rs.SyncOpResults{Annotations: annos}, nil +} + +func (b *schemaBuilder) Grant(ctx context.Context, principal *v2.Resource, entitlement *v2.Entitlement) (annotations.Annotations, error) { + return securableGrantChange(ctx, b.resolver, entitlement, principal, true) +} + +func (b *schemaBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotations.Annotations, error) { + return securableGrantChange(ctx, b.resolver, grant.Entitlement, grant.Principal, false) +} + +func newSchemaBuilder(resolver *ucResolver, syncTables bool) *schemaBuilder { + return &schemaBuilder{resolver: resolver, resourceType: schemaResourceType, syncTables: syncTables} +} + +// ---- Table builder ---- + +type tableBuilder struct { + resolver *ucResolver + resourceType *v2.ResourceType +} + +func (b *tableBuilder) ResourceType(ctx context.Context) *v2.ResourceType { + return tableResourceType +} + +func tableResource(workspaceId string, table *databricks.Table, parent *v2.ResourceId) (*v2.Resource, error) { + profile := securableProfile(workspaceId, table.FullName, table.Owner, databricks.SecurableTable) + profile["table_type"] = table.TableType + + return rs.NewResource( + table.Name, + tableResourceType, + makeSecurableID(workspaceId, table.FullName), + rs.WithResourceProfile(profile), + rs.WithParentResourceID(parent), + ) +} + +func (b *tableBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId, attr rs.SyncOpAttrs) ([]*v2.Resource, *rs.SyncOpResults, error) { + if parentResourceID == nil || parentResourceID.ResourceType != schemaResourceType.Id { + return nil, nil, nil + } + workspaceId, schemaFullName := splitSecurableID(parentResourceID.Resource) + catalogName, schemaName, found := strings.Cut(schemaFullName, ".") + if !found { + return nil, nil, fmt.Errorf("databricks-connector: malformed schema id: %q", parentResourceID.Resource) + } + + bag, pageToken, err := parseCursorToken(attr.PageToken.Token, &v2.ResourceId{ResourceType: tableResourceType.Id}) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to parse page token: %w", err) + } + + tables, nextToken, ratelimit, err := b.resolver.client.ListTables(ctx, workspaceId, catalogName, schemaName, pageToken, ResourcesPageSize) + if err != nil { + if ctx.Err() != nil { + return nil, nil, ctx.Err() + } + ctxzap.Extract(ctx).Warn("databricks-connector: unable to list tables, skipping schema", + zap.String("workspace", workspaceId), + zap.String("catalog", catalogName), + zap.String("schema", schemaName), + zap.Error(err), + ) + return nil, nil, nil + } + + var rv []*v2.Resource + for _, table := range tables { + tCopy := table + tr, err := tableResource(workspaceId, &tCopy, parentResourceID) + if err != nil { + return nil, nil, err + } + rv = append(rv, tr) + } + + nextPage, err := bag.NextToken(nextToken) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to create next page token: %w", err) + } + + return rv, syncOpResults(nextPage, ratelimit), nil +} + +func (b *tableBuilder) Entitlements(_ context.Context, resource *v2.Resource, _ rs.SyncOpAttrs) ([]*v2.Entitlement, *rs.SyncOpResults, error) { + return securablePrivilegeEntitlements(resource, tablePrivileges), nil, nil +} + +func (b *tableBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.SyncOpAttrs) ([]*v2.Grant, *rs.SyncOpResults, error) { + grants, annos, err := securableGrants(ctx, b.resolver, resource, databricks.SecurableTable, tablePrivileges) + if err != nil { + return nil, nil, err + } + return grants, &rs.SyncOpResults{Annotations: annos}, nil +} + +func (b *tableBuilder) Grant(ctx context.Context, principal *v2.Resource, entitlement *v2.Entitlement) (annotations.Annotations, error) { + return securableGrantChange(ctx, b.resolver, entitlement, principal, true) +} + +func (b *tableBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotations.Annotations, error) { + return securableGrantChange(ctx, b.resolver, grant.Entitlement, grant.Principal, false) +} + +func newTableBuilder(resolver *ucResolver) *tableBuilder { + return &tableBuilder{resolver: resolver, resourceType: tableResourceType} +} diff --git a/pkg/connector/unity_catalog_test.go b/pkg/connector/unity_catalog_test.go new file mode 100644 index 00000000..9ebb8cfb --- /dev/null +++ b/pkg/connector/unity_catalog_test.go @@ -0,0 +1,136 @@ +package connector + +import ( + "bytes" + "context" + "io" + "net/http" + "slices" + "testing" + + "github.com/conductorone/baton-databricks/pkg/databricks" + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + rs "github.com/conductorone/baton-sdk/pkg/types/resource" +) + +type statusRT struct { + status int + body string +} + +func (rt *statusRT) RoundTrip(r *http.Request) (*http.Response, error) { + body := rt.body + if body == "" { + body = "{}" + } + return &http.Response{ + StatusCode: rt.status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewBufferString(body)), + Request: r, + }, nil +} + +func newErrClient(t *testing.T, status int) *databricks.Client { + t.Helper() + c, err := databricks.NewClient( + context.Background(), + &http.Client{Transport: &statusRT{status: status, body: `{"message":"denied"}`}}, + "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acc-1", "", + &databricks.NoAuth{}, nil, + ) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + return c +} + +// A workspace without Unity Catalog, or a service principal lacking UC access, +// must not fail the sync — the catalog List degrades to empty instead of +// returning an error that aborts every other resource type. +func TestCatalogListDegradesOnForbidden(t *testing.T) { + b := newCatalogBuilder(newUCResolver(newErrClient(t, http.StatusForbidden))) + parent := &v2.ResourceId{ResourceType: workspaceResourceType.Id, Resource: "dbc-abc"} + + resources, results, err := b.List(context.Background(), parent, rs.SyncOpAttrs{}) + if err != nil { + t.Fatalf("List returned error on 403, want graceful skip: %v", err) + } + if len(resources) != 0 { + t.Fatalf("List returned %d resources on 403, want 0", len(resources)) + } + if results != nil && results.NextPageToken != "" { + t.Fatalf("List returned next page token on 403, want none") + } +} + +// Context cancellation must propagate rather than be swallowed as degradation. +func TestCatalogListPropagatesContextCancellation(t *testing.T) { + b := newCatalogBuilder(newUCResolver(newErrClient(t, http.StatusForbidden))) + parent := &v2.ResourceId{ResourceType: workspaceResourceType.Id, Resource: "dbc-abc"} + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + if _, _, err := b.List(ctx, parent, rs.SyncOpAttrs{}); err == nil { + t.Fatalf("List swallowed cancelled context, want error") + } +} + +func childResourceTypeIDs(t *testing.T, r *v2.Resource) []string { + t.Helper() + var out []string + for _, a := range r.GetAnnotations() { + crt := &v2.ChildResourceType{} + if a.MessageIs(crt) { + if err := a.UnmarshalTo(crt); err != nil { + t.Fatalf("unmarshal child resource type: %v", err) + } + out = append(out, crt.ResourceTypeId) + } + } + return out +} + +// Unity Catalog is opt-in: the catalog child (and thus UC sync) must be absent +// unless enabled, and table sync is separately gated on the schema child. +func TestUnityCatalogSyncIsGated(t *testing.T) { + ctx := context.Background() + ws := &databricks.Workspace{Name: "prod", DeploymentName: "dbc-abc", ID: 1} + acct := &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: "acc-1"} + + off, err := workspaceResource(ctx, ws, acct, false) + if err != nil { + t.Fatalf("workspaceResource: %v", err) + } + if slices.Contains(childResourceTypeIDs(t, off), catalogResourceType.Id) { + t.Fatalf("catalog child present when UC sync disabled") + } + + on, err := workspaceResource(ctx, ws, acct, true) + if err != nil { + t.Fatalf("workspaceResource: %v", err) + } + if !slices.Contains(childResourceTypeIDs(t, on), catalogResourceType.Id) { + t.Fatalf("catalog child missing when UC sync enabled") + } + + catParent := &v2.ResourceId{ResourceType: catalogResourceType.Id, Resource: makeSecurableID("dbc-abc", "prod_evaluation")} + schema := &databricks.Schema{Name: "reports", FullName: "prod_evaluation.reports"} + + noTables, err := schemaResource("dbc-abc", schema, catParent, false) + if err != nil { + t.Fatalf("schemaResource: %v", err) + } + if slices.Contains(childResourceTypeIDs(t, noTables), tableResourceType.Id) { + t.Fatalf("table child present when table sync disabled") + } + + withTables, err := schemaResource("dbc-abc", schema, catParent, true) + if err != nil { + t.Fatalf("schemaResource: %v", err) + } + if !slices.Contains(childResourceTypeIDs(t, withTables), tableResourceType.Id) { + t.Fatalf("table child missing when table sync enabled") + } +} diff --git a/pkg/connector/workspaces.go b/pkg/connector/workspaces.go index 8a594192..7367d68f 100644 --- a/pkg/connector/workspaces.go +++ b/pkg/connector/workspaces.go @@ -22,9 +22,10 @@ import ( const workspaceMemberEntitlement = "member" type workspaceBuilder struct { - client *databricks.Client - resourceType *v2.ResourceType - workspaces map[string]struct{} + client *databricks.Client + resourceType *v2.ResourceType + workspaces map[string]struct{} + syncUnityCatalog bool } func (w *workspaceBuilder) ResourceType(ctx context.Context) *v2.ResourceType { @@ -36,27 +37,39 @@ func (w *workspaceBuilder) ResourceType(ctx context.Context) *v2.ResourceType { // Deployment names are unique per Databricks cloud (they form the workspace's // canonical hostname), so they're safe as the resource ID here. // Users, groups and service principals hang off the workspace here instead of the account. -func minimalWorkspaceResource(_ context.Context, workspace *databricks.Workspace, parent *v2.ResourceId) (*v2.Resource, error) { +func minimalWorkspaceResource(_ context.Context, workspace *databricks.Workspace, parent *v2.ResourceId, syncUnityCatalog bool) (*v2.Resource, error) { + children := []protoreflect.ProtoMessage{ + &v2.ChildResourceType{ResourceTypeId: userResourceType.Id}, + &v2.ChildResourceType{ResourceTypeId: groupResourceType.Id}, + &v2.ChildResourceType{ResourceTypeId: servicePrincipalResourceType.Id}, + &v2.ChildResourceType{ResourceTypeId: roleResourceType.Id}, + } + if syncUnityCatalog { + children = append(children, &v2.ChildResourceType{ResourceTypeId: catalogResourceType.Id}) + } + return rs.NewGroupResource( workspace.DeploymentName, workspaceResourceType, workspace.DeploymentName, nil, rs.WithParentResourceID(parent), - rs.WithAnnotation( - &v2.ChildResourceType{ResourceTypeId: userResourceType.Id}, - &v2.ChildResourceType{ResourceTypeId: groupResourceType.Id}, - &v2.ChildResourceType{ResourceTypeId: servicePrincipalResourceType.Id}, - &v2.ChildResourceType{ResourceTypeId: roleResourceType.Id}, - ), + rs.WithAnnotation(children...), ) } -func workspaceResource(_ context.Context, workspace *databricks.Workspace, parent *v2.ResourceId) (*v2.Resource, error) { +func workspaceResource(_ context.Context, workspace *databricks.Workspace, parent *v2.ResourceId, syncUnityCatalog bool) (*v2.Resource, error) { profile := map[string]interface{}{ "workspace_id": workspace.ID, } + children := []protoreflect.ProtoMessage{ + &v2.ChildResourceType{ResourceTypeId: roleResourceType.Id}, + } + if syncUnityCatalog { + children = append(children, &v2.ChildResourceType{ResourceTypeId: catalogResourceType.Id}) + } + resource, err := rs.NewGroupResource( workspace.Name, workspaceResourceType, @@ -64,9 +77,7 @@ func workspaceResource(_ context.Context, workspace *databricks.Workspace, paren nil, rs.WithResourceProfile(profile), rs.WithParentResourceID(parent), - rs.WithAnnotation( - &v2.ChildResourceType{ResourceTypeId: roleResourceType.Id}, - ), + rs.WithAnnotation(children...), ) if err != nil { @@ -92,7 +103,7 @@ func (w *workspaceBuilder) List(ctx context.Context, parentResourceID *v2.Resour ws := &databricks.Workspace{DeploymentName: workspace} - wr, err := minimalWorkspaceResource(ctx, ws, parentResourceID) + wr, err := minimalWorkspaceResource(ctx, ws, parentResourceID, w.syncUnityCatalog) if err != nil { return nil, nil, err } @@ -127,7 +138,7 @@ func (w *workspaceBuilder) List(ctx context.Context, parentResourceID *v2.Resour wCopy := workspace - wr, err := workspaceResource(ctx, &wCopy, parentResourceID) + wr, err := workspaceResource(ctx, &wCopy, parentResourceID, w.syncUnityCatalog) if err != nil { return nil, nil, err } @@ -345,15 +356,16 @@ func (w *workspaceBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotat return nil, nil } -func newWorkspaceBuilder(client *databricks.Client, workspaces []string) *workspaceBuilder { +func newWorkspaceBuilder(client *databricks.Client, workspaces []string, syncUnityCatalog bool) *workspaceBuilder { wMap := make(map[string]struct{}, len(workspaces)) for _, w := range workspaces { wMap[w] = struct{}{} } return &workspaceBuilder{ - client: client, - resourceType: workspaceResourceType, - workspaces: wMap, + client: client, + resourceType: workspaceResourceType, + workspaces: wMap, + syncUnityCatalog: syncUnityCatalog, } } diff --git a/pkg/databricks/unity_catalog.go b/pkg/databricks/unity_catalog.go new file mode 100644 index 00000000..e9f85d61 --- /dev/null +++ b/pkg/databricks/unity_catalog.go @@ -0,0 +1,228 @@ +package databricks + +import ( + "context" + "net/http" + "net/url" + "strconv" + + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" +) + +const ( + unityCatalogCatalogsEndpoint = "/api/2.1/unity-catalog/catalogs" + unityCatalogSchemasEndpoint = "/api/2.1/unity-catalog/schemas" + unityCatalogTablesEndpoint = "/api/2.1/unity-catalog/tables" + unityCatalogPermissionsEndpoint = "/api/2.1/unity-catalog/permissions" + + // Securable types as they appear in the Unity Catalog permissions path. + SecurableCatalog = "catalog" + SecurableSchema = "schema" + SecurableTable = "table" +) + +// Catalog is a Unity Catalog top-level data container. +type Catalog struct { + Name string `json:"name"` + Owner string `json:"owner"` + Comment string `json:"comment"` + MetastoreID string `json:"metastore_id"` +} + +// Schema is a grouping of tables/views within a catalog. +type Schema struct { + Name string `json:"name"` + CatalogName string `json:"catalog_name"` + FullName string `json:"full_name"` + Owner string `json:"owner"` + Comment string `json:"comment"` +} + +// Table is a data asset (table or view) within a schema. +type Table struct { + Name string `json:"name"` + CatalogName string `json:"catalog_name"` + SchemaName string `json:"schema_name"` + FullName string `json:"full_name"` + TableType string `json:"table_type"` + Owner string `json:"owner"` +} + +// PrivilegeAssignment maps a principal to the privileges it holds on a securable. +type PrivilegeAssignment struct { + Principal string `json:"principal"` + Privileges []string `json:"privileges"` +} + +// PermissionsChange is a single add/remove delta applied to a securable's grants. +type PermissionsChange struct { + Principal string `json:"principal"` + Add []string `json:"add,omitempty"` + Remove []string `json:"remove,omitempty"` +} + +type catalogsResponse struct { + Catalogs []Catalog `json:"catalogs"` + NextPageToken string `json:"next_page_token"` +} + +type schemasResponse struct { + Schemas []Schema `json:"schemas"` + NextPageToken string `json:"next_page_token"` +} + +type tablesResponse struct { + Tables []Table `json:"tables"` + NextPageToken string `json:"next_page_token"` +} + +type permissionsResponse struct { + PrivilegeAssignments []PrivilegeAssignment `json:"privilege_assignments"` +} + +// ucListVars applies Unity Catalog list/query parameters. +type ucListVars struct { + catalogName string + schemaName string + principal string + pageToken string + maxResults uint +} + +func (v *ucListVars) Apply(params *url.Values) { + if v.catalogName != "" { + params.Add("catalog_name", v.catalogName) + } + if v.schemaName != "" { + params.Add("schema_name", v.schemaName) + } + if v.principal != "" { + params.Add("principal", v.principal) + } + if v.maxResults > 0 { + params.Add("max_results", strconv.FormatUint(uint64(v.maxResults), 10)) + } + if v.pageToken != "" { + params.Add("page_token", v.pageToken) + } +} + +// ListCatalogs returns a page of catalogs in the workspace's metastore. +func (c *Client) ListCatalogs( + ctx context.Context, + workspaceId string, + pageToken string, + maxResults uint, +) ( + []Catalog, + string, + *v2.RateLimitDescription, + error, +) { + u := c.workspaceUrl(workspaceId).JoinPath(unityCatalogCatalogsEndpoint) + + var res catalogsResponse + ratelimitData, err := c.Get(ctx, u, &res, &ucListVars{pageToken: pageToken, maxResults: maxResults}) + if err != nil { + return nil, "", ratelimitData, nameWorkspace403Remedy(workspaceId, err) + } + + return res.Catalogs, res.NextPageToken, ratelimitData, nil +} + +// ListSchemas returns a page of schemas within a catalog. +func (c *Client) ListSchemas( + ctx context.Context, + workspaceId string, + catalogName string, + pageToken string, + maxResults uint, +) ( + []Schema, + string, + *v2.RateLimitDescription, + error, +) { + u := c.workspaceUrl(workspaceId).JoinPath(unityCatalogSchemasEndpoint) + + var res schemasResponse + ratelimitData, err := c.Get(ctx, u, &res, &ucListVars{catalogName: catalogName, pageToken: pageToken, maxResults: maxResults}) + if err != nil { + return nil, "", ratelimitData, nameWorkspace403Remedy(workspaceId, err) + } + + return res.Schemas, res.NextPageToken, ratelimitData, nil +} + +// ListTables returns a page of tables within a schema. +func (c *Client) ListTables( + ctx context.Context, + workspaceId string, + catalogName string, + schemaName string, + pageToken string, + maxResults uint, +) ( + []Table, + string, + *v2.RateLimitDescription, + error, +) { + u := c.workspaceUrl(workspaceId).JoinPath(unityCatalogTablesEndpoint) + + var res tablesResponse + ratelimitData, err := c.Get(ctx, u, &res, &ucListVars{catalogName: catalogName, schemaName: schemaName, pageToken: pageToken, maxResults: maxResults}) + if err != nil { + return nil, "", ratelimitData, nameWorkspace403Remedy(workspaceId, err) + } + + return res.Tables, res.NextPageToken, ratelimitData, nil +} + +// ListPermissions returns the direct privilege assignments on a securable. +// securableType is one of Securable{Catalog,Schema,Table}; fullName is the +// dotted Unity Catalog name (catalog, catalog.schema, or catalog.schema.table). +func (c *Client) ListPermissions( + ctx context.Context, + workspaceId string, + securableType string, + fullName string, +) ( + []PrivilegeAssignment, + *v2.RateLimitDescription, + error, +) { + u := c.workspaceUrl(workspaceId).JoinPath(unityCatalogPermissionsEndpoint, securableType, fullName) + + var res permissionsResponse + ratelimitData, err := c.Get(ctx, u, &res) + if err != nil { + return nil, ratelimitData, nameWorkspace403Remedy(workspaceId, err) + } + + return res.PrivilegeAssignments, ratelimitData, nil +} + +// UpdatePermissions applies add/remove privilege changes to a securable. +func (c *Client) UpdatePermissions( + ctx context.Context, + workspaceId string, + securableType string, + fullName string, + changes []PermissionsChange, +) ( + *v2.RateLimitDescription, + error, +) { + u := c.workspaceUrl(workspaceId).JoinPath(unityCatalogPermissionsEndpoint, securableType, fullName) + + body := map[string]interface{}{"changes": changes} + + var res permissionsResponse + ratelimitData, err := c.doRequest(ctx, u, http.MethodPatch, body, &res) + if err != nil { + return ratelimitData, nameWorkspace403Remedy(workspaceId, err) + } + + return ratelimitData, nil +} diff --git a/pkg/databricks/unity_catalog_test.go b/pkg/databricks/unity_catalog_test.go new file mode 100644 index 00000000..1acb1837 --- /dev/null +++ b/pkg/databricks/unity_catalog_test.go @@ -0,0 +1,131 @@ +package databricks + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "testing" +) + +type capturedReq struct { + method string + url string + body string +} + +type captureRT struct { + reqs []capturedReq + body string +} + +func (rt *captureRT) RoundTrip(r *http.Request) (*http.Response, error) { + var b string + if r.Body != nil { + data, _ := io.ReadAll(r.Body) + b = string(data) + } + rt.reqs = append(rt.reqs, capturedReq{method: r.Method, url: r.URL.String(), body: b}) + resp := rt.body + if resp == "" { + resp = "{}" + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewBufferString(resp)), + Request: r, + }, nil +} + +func TestUnityCatalogRequests(t *testing.T) { + ctx := context.Background() + + run := func(body string, fn func(c *Client)) []capturedReq { + rt := &captureRT{body: body} + c, err := NewClient(ctx, &http.Client{Transport: rt}, "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acc-1", "", &NoAuth{}, nil) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + fn(c) + return rt.reqs + } + + // ListCatalogs: workspace subdomain + UC path + query params, decode payload. + reqs := run(`{"catalogs":[{"name":"prod_evaluation","owner":"SCRUM-AIFUN"}],"next_page_token":"tok2"}`, func(c *Client) { + cats, next, _, err := c.ListCatalogs(ctx, "dbc-abc", "", 50) + if err != nil { + t.Fatalf("ListCatalogs: %v", err) + } + if len(cats) != 1 || cats[0].Name != "prod_evaluation" || cats[0].Owner != "SCRUM-AIFUN" { + t.Fatalf("ListCatalogs decode = %+v", cats) + } + if next != "tok2" { + t.Fatalf("ListCatalogs next = %q, want tok2", next) + } + }) + got := reqs[0] + want := "https://dbc-abc.example.cloud.databricks.com/api/2.1/unity-catalog/catalogs?max_results=50" + if got.method != http.MethodGet || got.url != want { + t.Fatalf("ListCatalogs req = %s %s, want GET %s", got.method, got.url, want) + } + + // ListSchemas: catalog_name query param. + reqs = run(`{"schemas":[]}`, func(c *Client) { + if _, _, _, err := c.ListSchemas(ctx, "dbc-abc", "prod_evaluation", "", 50); err != nil { + t.Fatalf("ListSchemas: %v", err) + } + }) + if u := reqs[0].url; u != "https://dbc-abc.example.cloud.databricks.com/api/2.1/unity-catalog/schemas?catalog_name=prod_evaluation&max_results=50" { + t.Fatalf("ListSchemas url = %s", u) + } + + // ListTables: catalog_name + schema_name query params. + reqs = run(`{"tables":[]}`, func(c *Client) { + if _, _, _, err := c.ListTables(ctx, "dbc-abc", "prod_evaluation", "reports", "", 50); err != nil { + t.Fatalf("ListTables: %v", err) + } + }) + if u := reqs[0].url; u != "https://dbc-abc.example.cloud.databricks.com/api/2.1/unity-catalog/tables?catalog_name=prod_evaluation&max_results=50&schema_name=reports" { + t.Fatalf("ListTables url = %s", u) + } + + // ListPermissions on a table: dotted full_name must survive as one path segment. + reqs = run(`{"privilege_assignments":[{"principal":"alice@corp.com","privileges":["SELECT","MODIFY"]}]}`, func(c *Client) { + pas, _, err := c.ListPermissions(ctx, "dbc-abc", SecurableTable, "prod_evaluation.reports.performance_report") + if err != nil { + t.Fatalf("ListPermissions: %v", err) + } + if len(pas) != 1 || pas[0].Principal != "alice@corp.com" || len(pas[0].Privileges) != 2 { + t.Fatalf("ListPermissions decode = %+v", pas) + } + }) + if u := reqs[0].url; u != "https://dbc-abc.example.cloud.databricks.com/api/2.1/unity-catalog/permissions/table/prod_evaluation.reports.performance_report" { + t.Fatalf("ListPermissions url = %s", u) + } + + // UpdatePermissions: PATCH with changes body. + reqs = run(`{"privilege_assignments":[]}`, func(c *Client) { + if _, err := c.UpdatePermissions(ctx, "dbc-abc", SecurableSchema, "prod_evaluation.reports", + []PermissionsChange{{Principal: "SCRUM-AIFUN", Add: []string{"SELECT"}}}); err != nil { + t.Fatalf("UpdatePermissions: %v", err) + } + }) + got = reqs[0] + if got.method != http.MethodPatch { + t.Fatalf("UpdatePermissions method = %s, want PATCH", got.method) + } + if u := got.url; u != "https://dbc-abc.example.cloud.databricks.com/api/2.1/unity-catalog/permissions/schema/prod_evaluation.reports" { + t.Fatalf("UpdatePermissions url = %s", u) + } + var payload struct { + Changes []PermissionsChange `json:"changes"` + } + if err := json.Unmarshal([]byte(got.body), &payload); err != nil { + t.Fatalf("UpdatePermissions body not json: %v (%s)", err, got.body) + } + if len(payload.Changes) != 1 || payload.Changes[0].Principal != "SCRUM-AIFUN" || payload.Changes[0].Add[0] != "SELECT" { + t.Fatalf("UpdatePermissions body = %s", got.body) + } +}