diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 5a0408d0..eec514cd 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -8,6 +8,12 @@ on: jobs: test: runs-on: ubuntu-latest + # This job grants and revokes one hardcoded group membership on the shared + # Databricks tenant, so two PRs running it at once read each other's writes. + # The group is deliberately not per-ref: it serializes across every branch. + concurrency: + group: databricks-sync-test + cancel-in-progress: false env: BATON_LOG_LEVEL: debug # Add any environment variables needed to run baton-bitbucket-datacenter diff --git a/README.md b/README.md index 65930d04..585e49d5 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ `baton-databricks` is a connector for Databricks built using the [Baton SDK](https://github.com/conductorone/baton-sdk). It communicates with the Databricks API, to sync data about Databricks identities (users, groups and -service principals), roles and workspaces. +service principals), roles, workspaces and Unity Catalog metastores. Check out [Baton](https://github.com/conductorone/baton) to learn more about the project in general. @@ -30,6 +30,11 @@ the Generate secret button. You can use this secret to authenticate across all workspaces that service principal has access to. This requires admin access to the Databricks account and each workspace you want to sync. +Syncing Unity Catalog needs more than account admin on its own: the service +principal also has to be able to read the privileges recorded on each metastore. +See [Unity Catalog permissions](#unity-catalog-permissions) for the exact +requirement. + # Using Azure Databricks To work with Azure Databricks, you need to provide the hostname flag. @@ -76,12 +81,21 @@ baton resources - Service Principals - Users - Roles +- Metastores The connector fetches all resources from the account and every workspace the service principal can access. To limit the scope, pass a comma-separated list of workspaces to the `--workspaces` flag. Each entry can be a workspace name, deployment name, or numeric workspace ID, matched case-insensitively. +`--workspaces` also scopes which metastores sync: only a metastore attached to a +listed workspace is synced, so narrowing the list removes the other metastores +and their grants from ConductorOne. A list that names no workspace on the +account is rejected during credential validation, before any sync runs. + +Metastores are Unity Catalog securables and are opt-in: the type stays off until +it is enabled for the connector in ConductorOne. + ## Authentication OAuth is the only authentication method: an account-level service principal's @@ -97,7 +111,9 @@ To instead exclude specific workspaces from the sync, pass them to the `--databricks-exclude-workspaces` flag (or the `BATON_DATABRICKS_EXCLUDE_WORKSPACES` environment variable) as a comma-separated list. Each entry can be a workspace name, deployment name, or numeric workspace -ID. Excluded workspaces and their roles are skipped entirely. +ID. Excluded workspaces and their roles are skipped entirely. A metastore +attached only to excluded workspaces is not synced, and its grants are removed +from ConductorOne. ## Incremental sync @@ -188,6 +204,55 @@ membership from a group removes the group's entire workspace assignment, including admin access if the group is a workspace admin, so its members lose it. +## Unity Catalog + +A metastore is the top of the Databricks Unity Catalog hierarchy. The resource +tree the connector syncs is: + +``` +account +├── workspace +│ └── role +├── user +├── group +├── service principal +└── metastore +``` + +A metastore exposes its privileges as entitlements (`CREATE_CATALOG`, +`READ_METADATA`, `USE_SHARE` and so on) plus a non-grantable `owner`. A +privilege held by a group is expandable, so ConductorOne resolves the group's +members as holders. + +`MANAGE` and `ALL_PRIVILEGES` are not offered: Databricks rejects both at the +metastore level. + +A metastore lives on the account plane, but its privileges are read through a +workspace attached to it. The connector picks that workspace itself: the first +running one in scope where it can read every privilege, which it proves by +owning the metastore or by holding Admin directly on that workspace. Grant and +Revoke resolve the same route from the metastore id. + +A metastore syncs through a running workspace in scope. The sync fails, naming +the metastore, while an attached workspace in scope is still provisioning, or +while the metastore's workspace list cannot be read. + +## Unity Catalog provisioning + +A metastore privilege can be granted to and revoked from a user, group or +service principal. The `owner` entitlement is read-only: granting or revoking it +fails with an `InvalidArgument` error. + +## Unity Catalog permissions + +The service principal must own the metastore, or be assigned **Admin** directly +on at least one running workspace attached to it. Admin inherited only through a +group is not enough. Otherwise the sync fails naming the metastore. + +The connector requests the `all-apis` OAuth scope. Each resource type declares +the scopes, roles and privileges it needs, so `./baton-databricks capabilities` +reports them per type. + # Contributing, Support and Issues We started Baton because we were tired of taking screenshots and manually diff --git a/baton_capabilities.json b/baton_capabilities.json index 57cfd894..3f948d7d 100644 --- a/baton_capabilities.json +++ b/baton_capabilities.json @@ -4,14 +4,36 @@ { "resourceType": { "id": "account", - "displayName": "Account" + "displayName": "Account", + "annotations": [ + { + "@type": "type.googleapis.com/c1.connector.v2.CapabilityPermissions", + "permissions": [ + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } + ] }, "capabilities": [ "CAPABILITY_SYNC", "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], - "permissions": {} + "permissions": { + "permissions": [ + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } }, { "resourceType": { @@ -19,6 +41,22 @@ "displayName": "Group", "traits": [ "TRAIT_GROUP" + ], + "annotations": [ + { + "@type": "type.googleapis.com/c1.connector.v2.CapabilityPermissions", + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } ] }, "capabilities": [ @@ -26,7 +64,86 @@ "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], - "permissions": {} + "permissions": { + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } + }, + { + "resourceType": { + "id": "metastore", + "displayName": "Metastore", + "traits": [ + "TRAIT_APP" + ], + "annotations": [ + { + "@type": "type.googleapis.com/c1.connector.v2.OptInRequired" + }, + { + "@type": "type.googleapis.com/c1.connector.v2.CapabilityPermissions", + "permissions": [ + { + "permission": "unity-catalog" + }, + { + "permission": "scim" + }, + { + "permission": "access-management" + }, + { + "permission": "provisioning" + }, + { + "permission": "Account admin" + }, + { + "permission": "Workspace admin" + } + ] + }, + { + "@type": "type.googleapis.com/c1.connector.v2.SkipEntitlements" + } + ] + }, + "capabilities": [ + "CAPABILITY_SYNC", + "CAPABILITY_PROVISION" + ], + "permissions": { + "permissions": [ + { + "permission": "unity-catalog" + }, + { + "permission": "scim" + }, + { + "permission": "access-management" + }, + { + "permission": "provisioning" + }, + { + "permission": "Account admin" + }, + { + "permission": "Workspace admin" + } + ] + }, + "optInRequired": true }, { "resourceType": { @@ -34,6 +151,19 @@ "displayName": "Role", "traits": [ "TRAIT_ROLE" + ], + "annotations": [ + { + "@type": "type.googleapis.com/c1.connector.v2.CapabilityPermissions", + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "Account admin" + } + ] + } ] }, "capabilities": [ @@ -41,7 +171,16 @@ "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], - "permissions": {} + "permissions": { + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "Account admin" + } + ] + } }, { "resourceType": { @@ -49,6 +188,22 @@ "displayName": "Service Principal", "traits": [ "TRAIT_GROUP" + ], + "annotations": [ + { + "@type": "type.googleapis.com/c1.connector.v2.CapabilityPermissions", + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } ] }, "capabilities": [ @@ -56,7 +211,19 @@ "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], - "permissions": {} + "permissions": { + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } }, { "resourceType": { @@ -68,6 +235,17 @@ "annotations": [ { "@type": "type.googleapis.com/c1.connector.v2.SkipEntitlementsAndGrants" + }, + { + "@type": "type.googleapis.com/c1.connector.v2.CapabilityPermissions", + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "Account admin" + } + ] } ] }, @@ -77,7 +255,16 @@ "CAPABILITY_ACCOUNT_PROVISIONING", "CAPABILITY_RESOURCE_DELETE" ], - "permissions": {} + "permissions": { + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "Account admin" + } + ] + } }, { "resourceType": { @@ -85,6 +272,22 @@ "displayName": "Workspace", "traits": [ "TRAIT_GROUP" + ], + "annotations": [ + { + "@type": "type.googleapis.com/c1.connector.v2.CapabilityPermissions", + "permissions": [ + { + "permission": "provisioning" + }, + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } ] }, "capabilities": [ @@ -92,7 +295,19 @@ "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], - "permissions": {} + "permissions": { + "permissions": [ + { + "permission": "provisioning" + }, + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } } ], "connectorCapabilities": [ diff --git a/docs/connector.mdx b/docs/connector.mdx index 3c649d1a..5c2edbfc 100644 --- a/docs/connector.mdx +++ b/docs/connector.mdx @@ -12,6 +12,7 @@ sidebarTitle: "Databricks" | :--- | :--- | :--- | | Accounts | | | | Groups | | | +| Metastores | | | | Roles | | | | Service principals | | | | Users | | | @@ -27,6 +28,28 @@ Revoking the workspace **member** entitlement removes the principal's entire wor [This connector syncs non-human identities](/product/admin/nhi) and displays them on the **Identities overview** dashboard. +## Metastore access in C1 + +The connector syncs the Databricks [Unity Catalog](https://docs.databricks.com/aws/en/data-governance/unity-catalog/manage-privileges/privileges) +metastore as a C1 resource. Each privilege Databricks records on it +(`CREATE_CATALOG`, `READ_METADATA`, `USE_SHARE` and so on) becomes an +entitlement in C1, so you can review who holds it, request it, and revoke it +for users, groups and service principals. A grant held by a group is expandable +so C1 can resolve the group's members as holders. Ownership is shown as a +read-only `owner` entitlement — it cannot be granted or revoked through C1. + + +The Metastores resource type is turned off by default. Turn it on from the connector's configuration page in C1. + + + +**Notes on Unity Catalog access in C1:** + +- Ownership is shown but cannot be provisioned. An owner that is not a Databricks account identity, such as a system account or a workspace admins group, is not shown at all. +- **MANAGE** and **ALL_PRIVILEGES** are not offered as entitlements, because Databricks rejects both at the metastore level. +- Unity Catalog tags are not synced. + + ## Authentication methods The connector authenticates with **OAuth** — an account-level service principal's client ID and secret. This is the only authentication method the connector supports. @@ -97,6 +120,8 @@ Configuring the connector requires you to pass in credentials generated in Datab A user with the **Account admin** role in your Databricks account must perform this task. Account admin is an account-level role, managed in the Databricks account console. +Every step below happens in the Databricks **account console**. On AWS that is [https://accounts.cloud.databricks.com](https://accounts.cloud.databricks.com). Azure Databricks and Google Cloud Platform customers use their own account console host. When they configure the connector they only need their Databricks domain in **Hostname**; Account Hostname is optional. + ### Look up your Databricks account ID @@ -111,22 +136,47 @@ A user with the **Account admin** role in your Databricks account must perform t -### Generate Databricks credentials +### Create a service principal and an OAuth secret The Databricks connector authenticates with OAuth, which syncs info from all Databricks workspaces the service principal can access. - Follow the [Databricks OAuth authentication documentation](https://docs.databricks.com/en/dev-tools/auth/oauth-m2m.html) to create a service principal and create an OAuth secret. + In the account console, go to **User management** > **Service principals** and click **Add service principal**. Give it a name you will recognize, such as `ConductorOne`. - Carefully copy and save the OAuth client ID and secret. + Open the new service principal, go to its **Secrets** tab, and generate an OAuth secret. + + Carefully copy and save the **Client ID** and **Secret**. The secret is shown only once. + + For the full walkthrough, see the [Databricks OAuth machine-to-machine documentation](https://docs.databricks.com/en/dev-tools/auth/oauth-m2m.html). Give the service principal the **Account admin** role: in the account console, go to **User management** > **Service principals**, click the service principal, and on the **Roles** tab select **Account admin**. Databricks requires this role for any service principal that calls account-level APIs, and the connector fails validation without it. + +There is no scope to pick in the console. The connector requests the `all-apis` scope itself when it exchanges the client ID and secret for a token, which covers the four API surfaces it reads: Unity Catalog, SCIM, access management and account provisioning. + + +### Grant the service principal the roles it needs + + + + Make the service principal an account admin: in the account console, go to **User management** > **Service principals**, open the service principal, select its **Roles** tab, and enable **Account admin**. + + + Assign it to each workspace you want C1 to sync, as an admin: in the account console, go to **Workspaces**, open the workspace, select **Permissions**, click **Add permissions**, select the service principal, and choose **Admin**. Assign **Admin** directly to the service principal — the connector does not treat Admin inherited through a group as enough. + + Repeat this for every workspace you want synced. Workspace **Admin** is enough to read Unity Catalog privileges on a metastore reached through that workspace. If the service principal can't be assigned Admin directly on a synced workspace attached to the metastore, make it the **Metastore Admin** instead: in the account console, go to **Catalog**, open the metastore, and set **Metastore Admin** to the service principal's application ID. A metastore with no synced workspace attached is left out either way. + + + + +Without this, the metastore sync fails with a permission error naming the metastore. Databricks returns a partial privilege list rather than an error, and C1 would treat the missing grants as revoked. + + **Done.** Here's the set of credentials you'll need when setting up the connector: - Account ID @@ -180,10 +230,19 @@ To complete this task, you'll need: Enter your OAuth client ID and client secret into the **OAuth2 Client ID** and **OAuth2 Client Secret** fields. - **Google Cloud Platform and Azure Databricks customers only:** Enter your Databricks account hostname and hostname in the relevant fields. + Optionally, narrow what the connector syncs: + + - **Workspaces**: limit the sync to these workspaces, by workspace name, deployment name, or numeric workspace ID. This also scopes which metastores sync: only a metastore attached to one of these workspaces is synced, so narrowing the list removes the other metastores and their grants from C1. + - **Exclude Workspaces**: workspaces to leave out of the sync, by workspace name, deployment name, or numeric workspace ID. A metastore attached only to excluded workspaces is not synced either, and C1 removes it and its grants. + + Fill in **Workspaces** or **Exclude Workspaces**, not both. + + + **Google Cloud Platform and Azure Databricks customers only:** Enter your Databricks domain in the **Hostname** field, for example `gcp.databricks.com` or `azuredatabricks.net`. This is the domain your workspace URLs end with, not the URL of one workspace. Account Hostname is optional: the connector derives it from Hostname when unset. - For more on how to look up these values for GCP, see the GCP Databricks [REST API reference](https://docs.databricks.com/api/gcp/account/introduction) and [identifiers for workspace objects](https://docs.gcp.databricks.com/en/workspace/workspace-details.html) documentation. - - For more on how to look up these values for Azure Databricks, see the the Azure Databricks [REST API reference](https://docs.databricks.com/api/azure/account/introduction) and [identifiers for workspace objects](https://learn.microsoft.com/en-us/azure/databricks/workspace/workspace-details) documentation. + - For more on how to look up these values for Azure Databricks, see the Azure Databricks [REST API reference](https://docs.databricks.com/api/azure/account/introduction) and [identifiers for workspace objects](https://learn.microsoft.com/en-us/azure/databricks/workspace/workspace-details) documentation. + Click **Save**. @@ -267,6 +326,7 @@ stringData: # Optional: limit the sync to specific workspaces # (workspace name, deployment name, or numeric ID). + # Also scopes which metastores sync: narrowing it removes metastores that no listed workspace is attached to. # Mutually exclusive with BATON_DATABRICKS_EXCLUDE_WORKSPACES — set one or the other, never both. # BATON_WORKSPACES: @@ -340,4 +400,8 @@ spec: **Done.** Your Databricks connector is now pulling access data into C1. - \ No newline at end of file + + + +Stable binaries and container images for the Databricks connector are available from the [official download center](https://dist.conductorone.com/ConductorOne/baton-databricks). + \ No newline at end of file diff --git a/docs/docs-info.md b/docs/docs-info.md index fcebe05c..de5f520d 100644 --- a/docs/docs-info.md +++ b/docs/docs-info.md @@ -5,6 +5,8 @@ While developing the connector, please fill out this form. This information is n 1. What resources does the connector sync? > The connector syncs the Databricks account, workspaces, users, groups, service principals, and roles. + > + > It also syncs Unity Catalog metastores. That resource type is opt-in — it stays off until it is enabled per resource type in ConductorOne. 2. Can the connector provision any resources? If so, which ones? @@ -12,6 +14,7 @@ While developing the connector, please fill out this form. This information is n > > - **User accounts**: create and delete account users. > - **Entitlements**: grant and revoke role and membership assignments on accounts, workspaces, groups, service principals, and roles. + > - **Unity Catalog privileges**: grant and revoke a metastore privilege for a user, group or service principal. The `owner` entitlement is read-only and cannot be granted or revoked. ## Connector credentials @@ -21,7 +24,7 @@ While developing the connector, please fill out this form. This information is n > > - **OAuth**: a service principal OAuth client ID and client secret. Syncs the account and every workspace the service principal can access. This is the connector's only authentication method. > - > Google Cloud Platform and Azure Databricks customers also provide the account hostname and hostname. + > Google Cloud Platform and Azure Databricks customers also provide their Databricks domain, such as `gcp.databricks.com` or `azuredatabricks.net`. Account Hostname is optional: the connector derives it from Hostname when unset. 2. For each item in the list above: @@ -33,7 +36,11 @@ While developing the connector, please fill out this form. This information is n * Does the credential need any specific scopes or permissions? If so, list them here. - > The credential must have admin access to each resource it reads or writes: account-admin on the Databricks account (for account-level sync and provisioning) and workspace-admin on each workspace being synced. + > The credential must have admin access to each resource it reads or writes: account-admin on the Databricks account (for account-level sync and provisioning) and workspace-admin assigned directly to the service principal on each workspace being synced (Admin inherited only through a group is not enough). + > + > The OAuth token carries scopes on top of that. The connector requests `all-apis`, which covers the four it uses: `unity-catalog` (the securables and the permissions endpoint), `scim` (account principals), `access-management` (roles and workspace permission assignments) and `provisioning` (the account workspaces listing). + > + > Unity Catalog adds a third layer on top of scopes and roles: standing on the securable itself. For each metastore being synced, the service principal must either own the metastore or be assigned Admin directly on at least one running workspace attached to it. Without that standing the sync fails, and ConductorOne keeps the grants it already has. * If applicable: Is the list of scopes or permissions different to sync (read) versus provision (read-write)? If so, list the difference here. @@ -42,3 +49,13 @@ While developing the connector, please fill out this form. This information is n * What level of access or permissions does the user need in order to create the credentials? (For example, must be a super administrator, must have access to the admin console, etc.) > Account admin access to the Databricks account console (to create the service principal and OAuth secret), and workspace admin on each workspace being synced. + +## Limitations + +- **Direct grants only.** Databricks inherits privileges downward. The connector reports the assignments stored on the object. The inheritance endpoint does not return principal IDs. +- **Direct Admin, or the sync fails.** Admin inherited through a group does not count. Databricks then returns HTTP 200 with only the privileges the caller can see. The sync fails instead of reporting that list, so ConductorOne keeps the grants it already has. +- **Ownership is read-only.** It is shown and cannot be granted or revoked. An owner that is not an account identity (`System user`, `_workspace_admins_…`) is not shown. +- **No `MANAGE` or `ALL_PRIVILEGES` on a metastore.** Databricks rejects both, so they are not offered. +- **Read through one running workspace.** Privileges are not served on the account plane. The connector reads each metastore through the first running workspace in scope where it can see every privilege. With no running workspace in scope, the metastore is not synced and ConductorOne removes it. A workspace that is still provisioning fails the sync until it is `RUNNING`. An unreadable assignment list fails the sync and ConductorOne keeps the metastore. +- **`workspaces` scopes metastores.** An allowlist that names no workspace on the account fails at credential validation, before any sync. +- **Not synced yet:** catalogs, schemas, tables, volumes, functions, registered models, external locations, storage credentials, connections, shares, recipients, and tags. diff --git a/pkg/connector/account.go b/pkg/connector/account.go index c2726501..b6581e58 100644 --- a/pkg/connector/account.go +++ b/pkg/connector/account.go @@ -55,6 +55,7 @@ func (a *accountBuilder) accountResource(_ context.Context) (*v2.Resource, error &v2.ChildResourceType{ResourceTypeId: groupResourceType.Id}, &v2.ChildResourceType{ResourceTypeId: servicePrincipalResourceType.Id}, &v2.ChildResourceType{ResourceTypeId: roleResourceType.Id}, + &v2.ChildResourceType{ResourceTypeId: metastoreResourceType.Id}, ) } diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 082ad834..09b739d2 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -4,6 +4,8 @@ import ( "context" "fmt" "io" + "slices" + "strconv" "strings" "github.com/conductorone/baton-databricks/pkg/config" @@ -18,23 +20,36 @@ import ( type Databricks struct { client *databricks.Client + clientID string workspaces []string + opts *cli.ConnectorOpts // sqlWarehouseID enables incremental sync when non-empty. sqlWarehouseID string } // ResourceSyncers returns a ResourceSyncerV2 for each resource type that should be synced from the upstream service. func (d *Databricks) ResourceSyncers(_ context.Context) []connectorbuilder.ResourceSyncerV2 { - syncers := []connectorbuilder.ResourceSyncerV2{ + return []connectorbuilder.ResourceSyncerV2{ newAccountBuilder(d.client), newGroupBuilder(d.client), newServicePrincipalBuilder(d.client), newUserBuilder(d.client), newWorkspaceBuilder(d.client, d.workspaces), newRoleBuilder(d.client), + newMetastoreBuilder(d.client, d.clientID, configuredWorkspaceSet(d.workspaces), syncPrincipals{ + groups: d.willSyncResourceType(groupResourceType.Id), + servicePrincipals: d.willSyncResourceType(servicePrincipalResourceType.Id), + }), } +} - return syncers +// willSyncResourceType is true when the type is in this sync. Nil opts (the +// capabilities prototype) means no filter, so every advertised type is in scope. +func (d *Databricks) willSyncResourceType(resourceTypeID string) bool { + if d.opts == nil { + return true + } + return d.opts.WillSyncResourceType(resourceTypeID) } // EventFeeds registers the audit-log event feed unconditionally; whether sql-warehouse-id @@ -55,7 +70,8 @@ func (d *Databricks) Asset(_ context.Context, _ *v2.AssetRef) (string, io.ReadCl func (d *Databricks) Metadata(_ context.Context) (*v2.ConnectorMetadata, error) { return &v2.ConnectorMetadata{ DisplayName: "Databricks", - Description: "Connector syncing Databricks workspaces, users, groups, service principals and roles to Baton", + Description: "Connector syncing Databricks workspaces, users, groups, service principals and roles to Baton, " + + "along with Unity Catalog metastores and the privileges granted on them", AccountCreationSchema: &v2.ConnectorAccountCreationSchema{ FieldMap: map[string]*v2.ConnectorAccountCreationSchema_Field{ "email": { @@ -126,7 +142,7 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err // Per-workspace probing is deliberately not done here: a workspace the service // principal can't reach is skipped during sync rather than failing validation. annos := annotations.Annotations{} - _, rateLimit, err := d.client.ListWorkspaces(ctx) + workspaces, rateLimit, err := d.client.ListWorkspaces(ctx) if rateLimit != nil { annos.WithRateLimiting(rateLimit) } @@ -134,6 +150,15 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err return annos, fmt.Errorf("databricks-connector: failed to list workspaces: %w", err) } + if len(d.workspaces) > 0 && !slices.ContainsFunc(workspaces, func(w databricks.Workspace) bool { + _, ok := matchConfiguredWorkspace(configuredWorkspaceSet(d.workspaces), w.DeploymentName, w.Name, strconv.Itoa(w.ID)) + return ok + }) { + return annos, fmt.Errorf( + "databricks-connector: none of the configured workspaces (%s) is on the account", + strings.Join(d.workspaces, ", ")) + } + l := ctxzap.Extract(ctx) if d.sqlWarehouseID == "" { l.Info("databricks-connector: incremental sync disabled (sql-warehouse-id not set)") @@ -141,20 +166,24 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err } l.Info("databricks-connector: incremental sync enabled", zap.String("sql_warehouse_id", d.sqlWarehouseID)) + return annos, nil } -// New returns a new instance of the connector. +// New returns a new instance of the connector. clientID is the application id of +// the service principal the connector authenticates as. func New( ctx context.Context, hostname, accountHostname, accountID, - baseURL string, + baseURL, + clientID string, auth databricks.Auth, excludeWorkspaces []string, workspaces []string, sqlWarehouseID string, + opts *cli.ConnectorOpts, ) (*Databricks, error) { httpClient, err := auth.GetClient(ctx) if err != nil { @@ -168,14 +197,16 @@ func New( return &Databricks{ client: client, + clientID: clientID, workspaces: workspaces, + opts: opts, // Trimmed so a whitespace-only value counts as unset (incremental sync off). sqlWarehouseID: strings.TrimSpace(sqlWarehouseID), }, nil } // NewConnector returns a new connector builder from a configuration struct. -func NewConnector(ctx context.Context, cfg *config.Databricks, _ *cli.ConnectorOpts) (connectorbuilder.ConnectorBuilderV2, []connectorbuilder.Opt, error) { +func NewConnector(ctx context.Context, cfg *config.Databricks, opts *cli.ConnectorOpts) (connectorbuilder.ConnectorBuilderV2, []connectorbuilder.Opt, error) { accountHostname := getAccountHostname(cfg, cfg.Hostname) auth := prepareClientAuth(cfg) @@ -185,10 +216,12 @@ func NewConnector(ctx context.Context, cfg *config.Databricks, _ *cli.ConnectorO accountHostname, cfg.AccountId, cfg.BaseUrl, + cfg.DatabricksClientId, auth, cfg.DatabricksExcludeWorkspaces, cfg.Workspaces, cfg.SqlWarehouseId, + opts, ) if err != nil { return nil, nil, err diff --git a/pkg/connector/connector_test.go b/pkg/connector/connector_test.go index a26b87ff..e64c4558 100644 --- a/pkg/connector/connector_test.go +++ b/pkg/connector/connector_test.go @@ -1,9 +1,12 @@ package connector import ( + "context" + "slices" "testing" "github.com/conductorone/baton-databricks/pkg/config" + "github.com/conductorone/baton-sdk/pkg/cli" ) // Regression for CXH-2165: with account-hostname unset the Azure/GCP hostname @@ -35,3 +38,56 @@ func TestAccountHostnameFieldHasNoDefault(t *testing.T) { t.Errorf("account-hostname field must have no default, got %q", v) } } + +func TestWillSyncResourceType(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts *cli.ConnectorOpts + want map[string]bool + }{ + {"nil opts is the capabilities prototype", nil, map[string]bool{ + metastoreResourceType.Id: true, + groupResourceType.Id: true, + servicePrincipalResourceType.Id: true, + }}, + {"a filter leaving out group", &cli.ConnectorOpts{ + SyncResourceTypeIDs: []string{accountResourceType.Id, servicePrincipalResourceType.Id, metastoreResourceType.Id}, + }, map[string]bool{ + groupResourceType.Id: false, + servicePrincipalResourceType.Id: true, + }}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + d := &Databricks{opts: tc.opts} + for resourceTypeID, want := range tc.want { + if got := d.willSyncResourceType(resourceTypeID); got != want { + t.Errorf("willSyncResourceType(%q) = %v, want %v", resourceTypeID, got, want) + } + } + }) + } +} + +// The SDK derives capabilities and Grant/Revoke dispatch from the registered +// syncers, so leaving metastore out of a sync must not unregister it. +func TestResourceSyncersAlwaysRegisterTheMetastoreBuilder(t *testing.T) { + t.Parallel() + + d := &Databricks{opts: &cli.ConnectorOpts{ + SyncResourceTypeIDs: []string{accountResourceType.Id, groupResourceType.Id, userResourceType.Id}, + }} + + ctx := context.Background() + var registered []string + for _, syncer := range d.ResourceSyncers(ctx) { + registered = append(registered, syncer.ResourceType(ctx).GetId()) + } + if !slices.Contains(registered, metastoreResourceType.Id) { + t.Fatalf("registered syncers = %v, want them to include %s", registered, metastoreResourceType.Id) + } +} diff --git a/pkg/connector/groups.go b/pkg/connector/groups.go index 5053fb28..38b7ba37 100644 --- a/pkg/connector/groups.go +++ b/pkg/connector/groups.go @@ -41,7 +41,7 @@ func groupResourceId(_ context.Context, groupId string, parentResourceId *v2.Res } func groupResource(ctx context.Context, group *databricks.Group, parent *v2.ResourceId) (*v2.Resource, error) { - profile := map[string]interface{}{ + profile := map[string]any{ "display_name": group.DisplayName, "group_id": group.ID, "parent_type": parent.GetResourceType(), diff --git a/pkg/connector/groups_test.go b/pkg/connector/groups_test.go index 9882f7d6..6a2fbdf4 100644 --- a/pkg/connector/groups_test.go +++ b/pkg/connector/groups_test.go @@ -10,6 +10,7 @@ import ( "net/http/httptest" "net/url" "slices" + "strconv" "strings" "testing" @@ -134,7 +135,7 @@ func testUserPrincipal(t *testing.T, id string) *v2.Resource { // TestGroupRuleSetCompoundGroupPrincipal covers the CXH-2347 rule-set path: granting or // revoking a group's manager role to a group principal must look the principal up by its -// native id (the filter "id eq 'account/...'" matches nothing) and write it to the rule set +// native id (the filter id eq "account/..." matches nothing) and write it to the rule set // as "groups/". func TestGroupRuleSetCompoundGroupPrincipal(t *testing.T) { const wantPrincipal = "groups/Member Group" @@ -176,7 +177,7 @@ func TestGroupRuleSetCompoundGroupPrincipal(t *testing.T) { filter := req.URL.Query().Get("filter") filters = append(filters, filter) w.Header().Set("Content-Type", "application/json") - if filter == "id eq 'g-member-native'" { + if filter == "id eq "+strconv.Quote("g-member-native") { _, _ = w.Write([]byte(`{"Resources":[{"id":"g-member-native","displayName":"Member Group"}],"totalResults":1}`)) return } @@ -215,7 +216,7 @@ func TestGroupRuleSetCompoundGroupPrincipal(t *testing.T) { t.Fatalf("%s returned error: %v", tc.name, err) } - if !slices.Equal(filters, []string{"id eq 'g-member-native'"}) { + if !slices.Equal(filters, []string{"id eq " + strconv.Quote("g-member-native")}) { t.Errorf("group lookup filters = %q, want a single lookup by the native id (never %q)", filters, compoundPrincipalID) } @@ -314,7 +315,7 @@ func TestGroupRuleSetPrincipalLookup(t *testing.T) { switch { case req.Method == http.MethodGet && req.URL.Path == tc.scimPath: w.Header().Set("Content-Type", "application/json") - if exists && req.URL.Query().Get("filter") == "id eq '"+tc.nativeID+"'" { + if exists && req.URL.Query().Get("filter") == "id eq "+strconv.Quote(tc.nativeID) { _, _ = w.Write([]byte(`{"Resources":[` + tc.found + `],"totalResults":1}`)) return } diff --git a/pkg/connector/helpers.go b/pkg/connector/helpers.go index 3bb9a946..baf378a9 100644 --- a/pkg/connector/helpers.go +++ b/pkg/connector/helpers.go @@ -51,9 +51,42 @@ func groupGrantExpansion(ctx context.Context, groupId string, parentResource *v2 }, nil } +func capabilityPermissions(perms ...string) *v2.CapabilityPermissions { + cp := &v2.CapabilityPermissions{Permissions: make([]*v2.CapabilityPermission, 0, len(perms))} + for _, perm := range perms { + cp.Permissions = append(cp.Permissions, &v2.CapabilityPermission{Permission: perm}) + } + + return cp +} + +// noteRateLimit puts the descriptor doRequest parsed on the sync annotations, +// which is the only place the SDK reads it from. +func noteRateLimit(annos *annotations.Annotations, rateLimit *v2.RateLimitDescription) { + if rateLimit != nil { + annos.WithRateLimiting(rateLimit) + } +} + func annotationsForUserResourceType() annotations.Annotations { annos := annotations.Annotations{} annos.Update(&v2.SkipEntitlementsAndGrants{}) + annos.Update(capabilityPermissions(scopeSCIM, permissionAccountAdmin)) + return annos +} + +func annotationsForSecurableResourceType(perms ...string) annotations.Annotations { + annos := annotations.Annotations{} + annos.Update(&v2.OptInRequired{}) + annos.Update(capabilityPermissions(perms...)) + return annos +} + +// StaticEntitlements is a separate SDK pass, so Entitlements() must not also run. +// Grants() still runs. +func annotationsForStaticSecurableResourceType(perms ...string) annotations.Annotations { + annos := annotationsForSecurableResourceType(perms...) + annos.Update(&v2.SkipEntitlements{}) return annos } diff --git a/pkg/connector/metastore.go b/pkg/connector/metastore.go new file mode 100644 index 00000000..726f8208 --- /dev/null +++ b/pkg/connector/metastore.go @@ -0,0 +1,400 @@ +package connector + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "slices" + "strconv" + "strings" + + "github.com/conductorone/baton-databricks/pkg/databricks" + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "github.com/conductorone/baton-sdk/pkg/annotations" + "github.com/conductorone/baton-sdk/pkg/connectorbuilder" + 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/conductorone/baton-sdk/pkg/uhttp" + "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" + "go.uber.org/zap" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +var _ connectorbuilder.StaticEntitlementSyncerV2 = (*metastoreBuilder)(nil) + +type metastoreBuilder struct { + client *databricks.Client + clientID string + // workspaces is the configured allowlist, empty when none was given. + workspaces map[string]struct{} + sync syncPrincipals +} + +func newMetastoreBuilder(client *databricks.Client, clientID string, workspaces map[string]struct{}, sync syncPrincipals) *metastoreBuilder { + return &metastoreBuilder{client: client, clientID: clientID, workspaces: workspaces, sync: sync} +} + +func (m *metastoreBuilder) ResourceType(_ context.Context) *v2.ResourceType { + return metastoreResourceType +} + +func metastoreResource(_ context.Context, metastore *databricks.Metastore, parent *v2.ResourceId, workspace, workspaceID, workspaceName string) (*v2.Resource, error) { + profile := map[string]any{ + profileKeyMetastoreID: metastore.MetastoreID, + profileKeyRegion: metastore.Region, + profileKeyOwner: metastore.Owner, + profileKeySecurableType: securableTypeMetastore, + profileKeyWorkspace: workspace, + profileKeyWorkspaceID: workspaceID, + profileKeyWorkspaceName: workspaceName, + } + + name := metastore.Name + if name == "" { + name = metastore.MetastoreID + } + + // The permissions endpoint addresses a metastore by UUID, so that is the id. + return rs.NewAppResource( + name, + metastoreResourceType, + metastore.MetastoreID, + nil, + rs.WithResourceProfile(profile), + rs.WithParentResourceID(parent), + rs.WithAnnotation(&v2.RawId{Id: metastore.MetastoreID}), + ) +} + +// List emits each metastore with the workspace its permissions are read through. +// A metastore is left out only when no attached workspace is in scope or usable. +// A pending workspace fails the sync instead, so ConductorOne keeps the metastore. +func (m *metastoreBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId, _ rs.SyncOpAttrs) ([]*v2.Resource, *rs.SyncOpResults, error) { + if parentResourceID.GetResourceType() != accountResourceType.Id { + return nil, nil, nil + } + + annos := annotations.Annotations{} + metastores, rateLimit, err := m.client.ListMetastores(ctx) + noteRateLimit(&annos, rateLimit) + if err != nil { + return nil, &rs.SyncOpResults{Annotations: annos}, refuseUnreadable("databricks-connector: failed to list metastores", err) + } + + scoped, running, err := m.scopedWorkspaces(ctx, &annos) + if err != nil { + return nil, &rs.SyncOpResults{Annotations: annos}, err + } + + l := ctxzap.Extract(ctx) + var rv []*v2.Resource + for _, metastore := range metastores { + if metastore.MetastoreID == "" { + continue + } + + // A 404 here does not mean the metastore is gone. Skipping it would make ConductorOne delete it. + workspace, pending, err := m.routeMetastore(ctx, metastore.MetastoreID, metastore.Owner, scoped, running, &annos) + if err != nil { + return nil, &rs.SyncOpResults{Annotations: annos}, err + } + if workspace == nil { + if pending { + return nil, &rs.SyncOpResults{Annotations: annos}, status.Errorf(codes.FailedPrecondition, + "databricks-connector: metastore %s has no running workspace in sync scope", metastore.MetastoreID) + } + l.Debug("databricks-connector: metastore has no workspace in sync scope; skipping it", + zap.String("metastore", metastore.MetastoreID)) + continue + } + + mr, err := metastoreResource(ctx, &metastore, parentResourceID, workspace.DeploymentName, strconv.Itoa(workspace.ID), workspace.Name) + if err != nil { + return nil, &rs.SyncOpResults{Annotations: annos}, err + } + rv = append(rv, mr) + } + + return rv, &rs.SyncOpResults{Annotations: annos}, nil +} + +func (m *metastoreBuilder) StaticEntitlements(_ context.Context, _ rs.SyncOpAttrs) ([]*v2.Entitlement, *rs.SyncOpResults, error) { + grantable := ent.WithGrantableTo(userResourceType, groupResourceType, servicePrincipalResourceType) + rv := make([]*v2.Entitlement, 0, len(metastorePrivileges)+1) + for _, privilege := range metastorePrivileges { + rv = append(rv, ent.NewPermissionEntitlement(nil, privilege, + grantable, + ent.WithDescription(fmt.Sprintf("%s privilege on a metastore in Databricks", privilege)), + )) + } + + rv = append(rv, ent.NewOwnershipEntitlement(nil, ownerEntitlement, + ent.WithDescription("Owns a metastore in Databricks"), + )) + + return rv, nil, nil +} + +func (m *metastoreBuilder) Entitlements(context.Context, *v2.Resource, rs.SyncOpAttrs) ([]*v2.Entitlement, *rs.SyncOpResults, error) { + return nil, nil, nil +} + +func (m *metastoreBuilder) Grants(ctx context.Context, resource *v2.Resource, attr rs.SyncOpAttrs) ([]*v2.Grant, *rs.SyncOpResults, error) { + l := ctxzap.Extract(ctx) + metastoreID := resource.GetId().GetResource() + profile := rs.GetProfile(resource) + // List writes the route. An empty one is a resource this sync did not emit. + workspace, _ := rs.GetProfileStringValue(profile, profileKeyWorkspace) + if workspace == "" { + return nil, nil, status.Errorf(codes.FailedPrecondition, + "databricks-connector: metastore %s has no workspace on its profile", metastoreID) + } + + annos := annotations.Annotations{} + pageToken := attr.PageToken.Token + assignments, nextPageToken, rateLimit, err := m.client.ListPermissions(ctx, workspace, databricks.SecurableMetastore, metastoreID, "", pageToken) + noteRateLimit(&annos, rateLimit) + if err != nil { + // Only a NotFound is FailedPrecondition. The SDK reads NotFound as "no + // grants", and whether the metastore is gone is settled by the next List. + // A throttle or a 5xx keeps the status the client returned, so the SDK retries it. + return nil, &rs.SyncOpResults{Annotations: annos}, refuseUnreadable(fmt.Sprintf( + "databricks-connector: failed to list permissions on metastore %s through workspace %s", metastoreID, workspace), err) + } + + resolver := newPrincipalResolver(m.client, &annos) + var rv []*v2.Grant + for _, assignment := range assignments { + principalId, err := resolver.resolve(ctx, assignment.PrincipalID.String(), assignment.Principal) + if err != nil { + return nil, &rs.SyncOpResults{Annotations: annos}, fmt.Errorf("databricks-connector: %w", err) + } + if principalId == nil { + l.Debug("databricks-connector: skipping a grant whose principal does not resolve", + zap.String("metastore", metastoreID), + zap.String("principal", assignment.Principal), + zap.String("principal_id", assignment.PrincipalID.String())) + continue + } + if !emitPrincipal(m.sync, principalId.GetResourceType()) { + continue + } + + options := groupSecurableGrantOptions(principalId) + for _, privilege := range assignment.Privileges { + if slices.Contains(metastorePrivileges, privilege) { + rv = append(rv, grant.NewGrant(resource, privilege, principalId, options...)) + } + } + } + + // The owner is emitted once, with the first page. + if owner, _ := rs.GetProfileStringValue(profile, profileKeyOwner); pageToken == "" && owner != "" { + principalId, err := resolver.resolve(ctx, "", owner) + if err != nil { + return nil, &rs.SyncOpResults{Annotations: annos}, fmt.Errorf("databricks-connector: %w", err) + } + if principalId != nil && emitPrincipal(m.sync, principalId.GetResourceType()) { + options := append(groupSecurableGrantOptions(principalId), grant.WithAnnotation(&v2.GrantImmutable{})) + rv = append(rv, grant.NewGrant(resource, ownerEntitlement, principalId, options...)) + } + } + + if nextPageToken != "" && nextPageToken == pageToken { + return nil, &rs.SyncOpResults{Annotations: annos}, status.Errorf(codes.Internal, + "databricks-connector: listing privileges on metastore %s repeated page token %q", metastoreID, nextPageToken) + } + + return rv, &rs.SyncOpResults{Annotations: annos, NextPageToken: nextPageToken}, nil +} + +type metastoreChange struct { + metastoreID string + workspace string + privilege string + scimID string + principalName string +} + +// resolveChange validates a Grant/Revoke and resolves what the PATCH needs. +func (m *metastoreBuilder) resolveChange( + ctx context.Context, + principalID *v2.ResourceId, + entitlement *v2.Entitlement, + isGrant bool, + annos *annotations.Annotations, +) (metastoreChange, error) { + if !isValidPrincipal(principalID) { + return metastoreChange{}, status.Errorf(codes.InvalidArgument, + "databricks-connector: only users, groups and service principals can hold a metastore privilege, got %s", principalID.GetResourceType()) + } + + resource := entitlement.GetResource() + metastoreID := resource.GetId().GetResource() + if metastoreID == "" { + return metastoreChange{}, status.Error(codes.InvalidArgument, "databricks-connector: missing metastore resource id") + } + + privilege, err := parseEntitlementPrivilege(entitlement.GetId()) + if err != nil { + return metastoreChange{}, err + } + switch { + case privilege == ownerEntitlement: + return metastoreChange{}, status.Errorf(codes.InvalidArgument, + "databricks-connector: metastore ownership cannot be provisioned; change the owner in Databricks") + case slices.Contains(legacyPrivileges, strings.ToUpper(strings.TrimSpace(privilege))): + return metastoreChange{}, status.Errorf(codes.InvalidArgument, + "databricks-connector: %s is a legacy alias Databricks rewrites; use the explicit privilege", privilege) + case !slices.Contains(metastorePrivileges, privilege): + return metastoreChange{}, status.Errorf(codes.InvalidArgument, + "databricks-connector: %s is not a grantable metastore privilege", privilege) + } + + // Provisioning carries ids, not the profile a sync wrote. + routed, err := m.resolveMetastoreRoute(ctx, metastoreID, annos) + if err != nil { + if errors.Is(err, errMetastoreGone) && isGrant { + return metastoreChange{}, status.Errorf(codes.NotFound, + "databricks-connector: metastore %s is no longer on the account", metastoreID) + } + + return metastoreChange{}, err + } + workspace := routed.DeploymentName + + scimID, err := unityPrincipalScimID(principalID) + if err != nil { + return metastoreChange{}, fmt.Errorf("databricks-connector: %w", err) + } + + principalName, err := unityPrincipalName(ctx, m.client, principalID) + if err != nil { + if !errors.Is(err, errPrincipalNotFound) { + return metastoreChange{}, fmt.Errorf("databricks-connector: failed to resolve principal %s: %w", principalID.GetResource(), err) + } + // Revoke addresses a vanished principal's rows by id. + if isGrant { + return metastoreChange{}, status.Errorf(codes.NotFound, + "databricks-connector: principal %s no longer exists in Databricks", scimID) + } + } + + return metastoreChange{ + metastoreID: metastoreID, + workspace: workspace, + privilege: privilege, + scimID: scimID, + principalName: principalName, + }, nil +} + +func (c metastoreChange) matchHeld(assignments []databricks.PrivilegeAssignment) (databricks.PrivilegeAssignment, bool) { + return matchHeldPrivilege(assignments, c.scimID, c.principalName, c.privilege) +} + +func (c metastoreChange) updateError(verb string, err error) error { + if isForbiddenError(err) { + return uhttp.WrapErrors(codes.PermissionDenied, + fmt.Sprintf("databricks-connector: cannot %s %s on metastore %s: the connector must own the metastore or be Admin on an attached workspace", + verb, c.privilege, c.metastoreID), + err) + } + + return fmt.Errorf("databricks-connector: failed to %s %s on metastore %s: %w", verb, c.privilege, c.metastoreID, err) +} + +func (m *metastoreBuilder) Grant(ctx context.Context, principal *v2.Resource, entitlement *v2.Entitlement) (annotations.Annotations, error) { + annos := annotations.Annotations{} + change, err := m.resolveChange(ctx, principal.GetId(), entitlement, true, &annos) + if err != nil { + return annos, err + } + + // Read past the cache: a stale page would skip a needed PATCH. + assignments, rateLimit, err := m.client.DrainPermissions(ctx, change.workspace, databricks.SecurableMetastore, change.metastoreID) + noteRateLimit(&annos, rateLimit) + if err != nil { + return annos, fmt.Errorf("databricks-connector: failed to read privileges on metastore %s: %w", change.metastoreID, err) + } + if _, held := change.matchHeld(assignments); held { + annos.Update(&v2.GrantAlreadyExists{}) + return annos, nil + } + + var echoed []databricks.PrivilegeAssignment + rateLimit, err = m.client.UpdatePermissionsUntil(ctx, change.workspace, databricks.SecurableMetastore, change.metastoreID, + []databricks.PermissionsChange{{Principal: change.principalName, Add: []string{change.privilege}}}, + func(page []databricks.PrivilegeAssignment) (bool, error) { + echoed = append(echoed, page...) + _, held := change.matchHeld(echoed) + return held, nil + }) + noteRateLimit(&annos, rateLimit) + if err != nil { + return annos, change.updateError("grant", err) + } + + // Databricks no-ops a PATCH it cannot apply, so only the echo confirms it. + if _, held := change.matchHeld(echoed); !held { + return annos, status.Errorf(codes.Internal, + "databricks-connector: granting %s on metastore %s was accepted but not applied", change.privilege, change.metastoreID) + } + + return annos, nil +} + +func (m *metastoreBuilder) Revoke(ctx context.Context, revoked *v2.Grant) (annotations.Annotations, error) { + annos := annotations.Annotations{} + change, err := m.resolveChange(ctx, revoked.GetPrincipal().GetId(), revoked.GetEntitlement(), false, &annos) + if err != nil { + if errors.Is(err, errMetastoreGone) { + annos.Update(&v2.GrantAlreadyRevoked{}) + return annos, nil + } + + return annos, err + } + + var seen []databricks.PrivilegeAssignment + rateLimit, err := m.client.ForEachUncachedPermissionsPage(ctx, change.workspace, databricks.SecurableMetastore, change.metastoreID, "", + func(page []databricks.PrivilegeAssignment) (bool, error) { + seen = append(seen, page...) + matched, ok := change.matchHeld(seen) + // Only an id match is final; a name match may be superseded on a later page. + return ok && change.scimID != "" && matched.PrincipalID.String() == change.scimID && change.scimID != "0", nil + }) + noteRateLimit(&annos, rateLimit) + if err != nil { + return annos, fmt.Errorf("databricks-connector: failed to read privileges on metastore %s: %w", change.metastoreID, err) + } + + // A missed row can be a partial read, and removing a privilege nobody holds is a no-op. + held, ok := change.matchHeld(seen) + removal := databricks.PermissionsChange{Principal: change.principalName, Remove: []string{change.privilege}} + switch heldID := held.PrincipalID.String(); { + case ok && heldID != "" && heldID != "0": + removal = databricks.PermissionsChange{PrincipalID: held.PrincipalID, Remove: []string{change.privilege}} + case change.principalName == "": + removal = databricks.PermissionsChange{PrincipalID: json.Number(change.scimID), Remove: []string{change.privilege}} + } + + result, rateLimit, err := m.client.UpdatePermissions(ctx, change.workspace, databricks.SecurableMetastore, change.metastoreID, + []databricks.PermissionsChange{removal}) + noteRateLimit(&annos, rateLimit) + if err != nil { + return annos, change.updateError("revoke", err) + } + if _, stillHeld := change.matchHeld(result); stillHeld { + return annos, status.Errorf(codes.Internal, + "databricks-connector: revoking %s on metastore %s was accepted but not applied", change.privilege, change.metastoreID) + } + if !ok { + annos.Update(&v2.GrantAlreadyRevoked{}) + } + + return annos, nil +} diff --git a/pkg/connector/metastore_helpers.go b/pkg/connector/metastore_helpers.go new file mode 100644 index 00000000..4687597f --- /dev/null +++ b/pkg/connector/metastore_helpers.go @@ -0,0 +1,579 @@ +package connector + +import ( + "context" + "errors" + "fmt" + "net/http" + "slices" + "strconv" + "strings" + + "github.com/conductorone/baton-databricks/pkg/databricks" + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "github.com/conductorone/baton-sdk/pkg/annotations" + "github.com/conductorone/baton-sdk/pkg/types/grant" + rs "github.com/conductorone/baton-sdk/pkg/types/resource" + "github.com/conductorone/baton-sdk/pkg/uhttp" + "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" + "go.uber.org/zap" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +const ( + securableTypeMetastore = "METASTORE" + + // workspacePermissionAdmin is the workspace-level assignment that makes its + // holder a metastore admin for the workspace's Unity Catalog reads. + workspacePermissionAdmin = "ADMIN" + + // Statuses the account API documents for a workspace. PROVISIONING and + // NOT_PROVISIONED can still become RUNNING. CANCELLING is a workspace being + // deleted: it will not become RUNNING. FAILED and BANNED cannot either. + workspaceStatusRunning = "RUNNING" + workspaceStatusProvisioning = "PROVISIONING" + workspaceStatusNotProvisioned = "NOT_PROVISIONED" + workspaceStatusCancelling = "CANCELLING" + + profileKeyMetastoreID = "metastore_id" + profileKeyOwner = "owner" + profileKeySecurableType = "securable_type" + profileKeyRegion = "region" + profileKeyWorkspace = "workspace" + profileKeyWorkspaceID = "workspace_id" + profileKeyWorkspaceName = "workspace_name" + + // ownerEntitlement is single-valued and Databricks cannot remove an owner, so it is never grantable. + ownerEntitlement = "owner" +) + +var ( + // Canonical UPPER_SNAKE names, the form the permissions API answers with. + // ALL_PRIVILEGES and MANAGE are rejected on a metastore. + metastorePrivileges = []string{ + "CREATE_CATALOG", + "CREATE_CLEAN_ROOM", + "CREATE_CONNECTION", + "CREATE_EXTERNAL_LOCATION", + "CREATE_EXTERNAL_METADATA", + "CREATE_PROVIDER", + "CREATE_RECIPIENT", + "CREATE_SHARE", + "CREATE_SERVICE_CREDENTIAL", + "CREATE_STORAGE_CREDENTIAL", + "MANAGE_ALLOWLIST", + "READ_METADATA", + "SET_SHARE_PERMISSION", + "USE_MARKETPLACE_ASSETS", + "USE_PROVIDER", + "USE_RECIPIENT", + "USE_SHARE", + } + + // Hive-era aliases Databricks accepts and rewrites server-side, so they never read back. + legacyPrivileges = []string{"CREATE", "CREATE_VIEW", "USAGE"} +) + +// syncPrincipals says which principal types this sync includes. +type syncPrincipals struct { + groups bool + servicePrincipals bool +} + +func emitPrincipal(sync syncPrincipals, resourceType string) bool { + switch resourceType { + case userResourceType.Id: + return true + case groupResourceType.Id: + return sync.groups + case servicePrincipalResourceType.Id: + return sync.servicePrincipals + default: + return false + } +} + +func parseEntitlementPrivilege(entitlementID string) (string, error) { + separator := strings.LastIndex(entitlementID, ":") + if separator < 0 || separator == len(entitlementID)-1 { + return "", status.Errorf(codes.InvalidArgument, "databricks-connector: invalid entitlement id %q", entitlementID) + } + + return entitlementID[separator+1:], nil +} + +// matchHeldPrivilege finds the assignment carrying privilege for this principal. +// The SCIM id wins; the name is used only when the assignment has no id. +func matchHeldPrivilege(assignments []databricks.PrivilegeAssignment, scimID, principalName, privilege string) (databricks.PrivilegeAssignment, bool) { + var nameOnly databricks.PrivilegeAssignment + var sawNameOnly bool + for _, assignment := range assignments { + if !slices.Contains(assignment.Privileges, privilege) { + continue + } + + assignmentID := assignment.PrincipalID.String() + if scimID != "" && assignmentID == scimID && assignmentID != "0" { + return assignment, true + } + if principalName == "" || assignment.Principal != principalName { + continue + } + if scimID != "" && assignmentID != "" && assignmentID != "0" && assignmentID != scimID { + continue + } + if !sawNameOnly { + nameOnly, sawNameOnly = assignment, true + } + } + + return nameOnly, sawNameOnly +} + +// groupSecurableGrantOptions lets C1 expand a group-held privilege through the +// group's members. Shallow stays unset so nested groups propagate. +func groupSecurableGrantOptions(principalId *v2.ResourceId) []grant.GrantOption { + if principalId.GetResourceType() != groupResourceType.Id { + return nil + } + + return []grant.GrantOption{grant.WithAnnotation(&v2.GrantExpandable{ + EntitlementIds: []string{fmt.Sprintf("group:%s:%s", principalId.GetResource(), groupMemberEntitlement)}, + ResourceTypeIds: []string{userResourceType.Id, groupResourceType.Id, servicePrincipalResourceType.Id}, + })} +} + +func isForbiddenError(err error) bool { + return apiStatus(err) == http.StatusForbidden +} + +func apiStatus(err error) int { + var apiErr *databricks.APIError + if errors.As(err, &apiErr) { + return apiErr.StatusCode + } + + return 0 +} + +// refuseUnreadable keeps a NotFound from reading as an empty result: the SDK +// skips a child List that answers NotFound, and ConductorOne then deletes what +// that parent contained. Every other status, a throttle included, keeps the +// code the client returned so the SDK can still back off and retry. +func refuseUnreadable(msg string, err error) error { + if err == nil { + return nil + } + if apiStatus(err) == http.StatusNotFound || status.Code(err) == codes.NotFound { + return uhttp.WrapErrors(codes.FailedPrecondition, msg, err) + } + + return fmt.Errorf("%s: %w", msg, err) +} + +// unityPrincipalScimID is the numeric SCIM id a privilege assignment carries in principal_id. +func unityPrincipalScimID(principal *v2.ResourceId) (string, error) { + if principal.GetResourceType() != groupResourceType.Id { + return principal.GetResource(), nil + } + + _, parsed, err := parseResourceId(principal.GetResource()) + if err != nil { + return "", fmt.Errorf("failed to parse group resource id: %w", err) + } + + return parsed.GetResource(), nil +} + +// unityPrincipalName is the user's userName, the group's displayName or the +// service principal's applicationId. Empty when the principal no longer exists. +func unityPrincipalName(ctx context.Context, c *databricks.Client, principal *v2.ResourceId) (string, error) { + principalId, err := unityPrincipalScimID(principal) + if err != nil { + return "", err + } + + qualified, err := preparePrincipalId(ctx, c, "", principal.GetResourceType(), principalId) + if err != nil { + return "", err + } + + _, name, _ := strings.Cut(qualified, "/") + + return name, nil +} + +// scopedWorkspaces indexes, by numeric id, the workspaces this sync may read +// through, and marks which of them are RUNNING. ListWorkspaces already drops the +// excluded ones. +func (m *metastoreBuilder) scopedWorkspaces(ctx context.Context, annos *annotations.Annotations) (map[string]databricks.Workspace, map[string]bool, error) { + workspaces, rateLimit, err := m.client.ListWorkspaces(ctx) + noteRateLimit(annos, rateLimit) + if err != nil { + return nil, nil, refuseUnreadable("databricks-connector: failed to list workspaces", err) + } + + scoped := make(map[string]databricks.Workspace, len(workspaces)) + running := make(map[string]bool, len(workspaces)) + for _, workspace := range workspaces { + id := strconv.Itoa(workspace.ID) + if workspace.DeploymentName == "" { + continue + } + if len(m.workspaces) > 0 { + if _, ok := matchConfiguredWorkspace(m.workspaces, workspace.DeploymentName, workspace.Name, id); !ok { + continue + } + } + // A workspace that will not become RUNNING is left out of scope entirely, + // so a metastore attached only to one is skipped rather than failing every + // sync. Only a status that can still change keeps the metastore waiting. + if !workspaceIsUsable(workspace.Status) { + continue + } + scoped[id] = workspace + running[id] = workspace.Status == "" || strings.EqualFold(workspace.Status, workspaceStatusRunning) + } + + return scoped, running, nil +} + +// workspaceIsUsable reports a status that is RUNNING or can still become it. +// FAILED and BANNED are terminal, and a status Databricks has not documented +// takes that path too: treating one as pending would fail the sync on every run. +func workspaceIsUsable(workspaceStatus string) bool { + switch strings.ToUpper(workspaceStatus) { + case "", workspaceStatusRunning, workspaceStatusProvisioning, workspaceStatusNotProvisioned, workspaceStatusCancelling: + return true + default: + return false + } +} + +// errMetastoreGone marks a metastore the account no longer lists. +var errMetastoreGone = errors.New("metastore is no longer on the account") + +// routeMetastore picks a workspace the connector can read every privilege through: +// it owns the metastore, or it is a direct Admin there. A nil workspace with no +// error means none is running and in scope; pending says one can still become RUNNING. +func (m *metastoreBuilder) routeMetastore( + ctx context.Context, + metastoreID, owner string, + scoped map[string]databricks.Workspace, + running map[string]bool, + annos *annotations.Annotations, +) (*databricks.Workspace, bool, error) { + ids, rateLimit, err := m.client.ListMetastoreWorkspaces(ctx, metastoreID) + noteRateLimit(annos, rateLimit) + if err != nil { + // A 404 here is not an empty assignment list. NotFound on this List makes + // the SDK skip the metastore, and ConductorOne deletes it. + return nil, false, refuseUnreadable(fmt.Sprintf( + "databricks-connector: failed to list the workspaces assigned to metastore %s, so the sync cannot tell which serve it", + metastoreID), err) + } + + pending := false + var candidates []databricks.Workspace + for _, id := range ids { + workspace, ok := scoped[id] + if !ok { + continue + } + if running[id] { + candidates = append(candidates, workspace) + continue + } + pending = true + } + if len(candidates) == 0 { + return nil, pending, nil + } + + if ownsPrincipal(m.clientID, owner) { + return &candidates[0], false, nil + } + + for i := range candidates { + admin, err := isWorkspaceAdmin(ctx, m.client, m.clientID, strconv.Itoa(candidates[i].ID), annos) + if err != nil { + return nil, false, err + } + if admin { + return &candidates[i], false, nil + } + } + + return nil, false, untrustedMetastoreError(metastoreID, candidates) +} + +// resolveMetastoreRoute routes a metastore from its id, reading owner and standing live. +func (m *metastoreBuilder) resolveMetastoreRoute( + ctx context.Context, + metastoreID string, + annos *annotations.Annotations, +) (*databricks.Workspace, error) { + metastores, rateLimit, err := m.client.ListMetastores(ctx) + noteRateLimit(annos, rateLimit) + if err != nil { + return nil, fmt.Errorf("databricks-connector: failed to list metastores: %w", err) + } + + found := slices.IndexFunc(metastores, func(metastore databricks.Metastore) bool { + return metastore.MetastoreID == metastoreID + }) + if found < 0 { + return nil, errMetastoreGone + } + + scoped, running, err := m.scopedWorkspaces(ctx, annos) + if err != nil { + return nil, err + } + + workspace, pending, err := m.routeMetastore(ctx, metastoreID, metastores[found].Owner, scoped, running, annos) + if err != nil { + return nil, err + } + if workspace == nil { + reason := "no workspace in sync scope is attached to it" + if pending { + reason = "every workspace in sync scope attached to it is still on its way to RUNNING" + } + + return nil, status.Errorf(codes.FailedPrecondition, + "databricks-connector: cannot route metastore %s: %s", metastoreID, reason) + } + + return workspace, nil +} + +// ownsPrincipal reports whether name is the service principal the connector +// authenticates as. Databricks names a service principal by its applicationId +// in an owner field, and does not normalise its case. +func ownsPrincipal(ownPrincipal, name string) bool { + return ownPrincipal != "" && strings.EqualFold(name, ownPrincipal) +} + +// isWorkspaceAdmin lists the workspace's permission assignments on every call. +// An assignment list the connector may not read leaves Admin unproven, which is +// the same answer as not holding it. +func isWorkspaceAdmin( + ctx context.Context, + client *databricks.Client, + ownPrincipal, workspaceID string, + annos *annotations.Annotations, +) (bool, error) { + if ownPrincipal == "" || workspaceID == "" { + return false, nil + } + + assignments, rateLimit, err := client.ListWorkspaceMembers(ctx, workspaceID) + noteRateLimit(annos, rateLimit) + if err != nil { + if !deniesTrust(err) { + return false, fmt.Errorf("databricks-connector: failed to list the permission assignments of workspace %s: %w", workspaceID, err) + } + ctxzap.Extract(ctx).Debug("databricks-connector: the permission assignments of the workspace cannot be read, so Admin stays unproven", + zap.String("workspace_id", workspaceID), + zap.Error(err)) + + return false, nil + } + + return holdsWorkspaceAdmin(ctx, client, ownPrincipal, assignments, annos) +} + +// holdsWorkspaceAdmin is true only when the connector principal itself holds ADMIN. Admin inherited through a group does not count. +func holdsWorkspaceAdmin( + ctx context.Context, + client *databricks.Client, + ownPrincipal string, + assignments []databricks.WorkspaceAssignment, + annos *annotations.Annotations, +) (bool, error) { + var adminIDs []string + for _, assignment := range assignments { + principal := assignment.Principal + if principal == nil || !slices.Contains(assignment.Permissions, workspacePermissionAdmin) { + continue + } + if ownsPrincipal(ownPrincipal, principal.ServicePrincipalAppID) { + return true, nil + } + if principal.ID != 0 { + adminIDs = append(adminIDs, strconv.Itoa(principal.ID)) + } + } + + if len(adminIDs) == 0 { + return false, nil + } + + ownID, err := ownPrincipalScimID(ctx, client, ownPrincipal, annos) + if err != nil { + return false, err + } + + return ownID != "" && slices.Contains(adminIDs, ownID), nil +} + +// ownPrincipalScimID looks up the connector's service principal by application id. +func ownPrincipalScimID( + ctx context.Context, + client *databricks.Client, + ownPrincipal string, + annos *annotations.Annotations, +) (string, error) { + id, rateLimit, err := client.FindServicePrincipalID(ctx, "", ownPrincipal) + noteRateLimit(annos, rateLimit) + if err != nil { + return "", fmt.Errorf("databricks-connector: failed to resolve the connector's own service principal: %w", err) + } + + return id, nil +} + +// deniesTrust is an answer that withholds what the connector asked for, as +// opposed to one that says the account plane is unavailable. +func deniesTrust(err error) bool { + switch apiStatus(err) { + case http.StatusUnauthorized, http.StatusForbidden, http.StatusNotFound: + return true + default: + return false + } +} + +func untrustedMetastoreError(metastoreID string, tried []databricks.Workspace) error { + names := make([]string, 0, len(tried)) + for _, workspace := range tried { + names = append(names, workspace.DeploymentName) + } + + return status.Errorf(codes.PermissionDenied, + "databricks-connector: cannot read every privilege on metastore %s through any workspace attached to it (%s), "+ + "and Databricks answers a partial read with 200, "+ + "so reporting it would revoke the access ConductorOne cannot see; "+ + "the connector's service principal must own the metastore or be assigned Admin directly on one of those workspaces", + metastoreID, strings.Join(names, ", ")) +} + +// principalResolver maps a privilege assignment to a C1 resource id. The id is used when present; the name is only the fallback. +type principalResolver struct { + client *databricks.Client + annos *annotations.Annotations + byID map[string]*v2.ResourceId + byName map[string]*v2.ResourceId +} + +func newPrincipalResolver(client *databricks.Client, annos *annotations.Annotations) *principalResolver { + return &principalResolver{client: client, annos: annos, byID: map[string]*v2.ResourceId{}, byName: map[string]*v2.ResourceId{}} +} + +// resolve prefers principal_id and never falls back to the name when the id is +// present: a deleted principal keeps only its id, and the name is ambiguous. +func (p *principalResolver) resolve(ctx context.Context, id, name string) (*v2.ResourceId, error) { + if id != "" && id != "0" { + if resourceId, ok := p.byID[id]; ok { + return resourceId, nil + } + resourceId, err := p.lookupID(ctx, id) + if err != nil { + return nil, err + } + p.byID[id] = resourceId + if _, ok := p.byName[name]; !ok && name != "" && resourceId != nil { + p.byName[name] = resourceId + } + + return resourceId, nil + } + if name == "" { + return nil, nil + } + if resourceId, ok := p.byName[name]; ok { + return resourceId, nil + } + resourceId, err := p.lookupName(ctx, name) + if err != nil { + return nil, err + } + p.byName[name] = resourceId + + return resourceId, nil +} + +func (p *principalResolver) lookupID(ctx context.Context, id string) (*v2.ResourceId, error) { + user, rateLimit, err := p.client.GetUser(ctx, "", id) + noteRateLimit(p.annos, rateLimit) + if err != nil && !isNotFoundError(err) { + return nil, fmt.Errorf("failed to resolve principal %s as a user: %w", id, err) + } + if err == nil && user != nil && user.ID != "" { + return &v2.ResourceId{ResourceType: userResourceType.Id, Resource: user.ID}, nil + } + + // A filtered list answers this without pulling the group's members. + displayName, rateLimit, err := p.client.FindGroupDisplayName(ctx, "", id) + noteRateLimit(p.annos, rateLimit) + if err != nil { + return nil, fmt.Errorf("failed to resolve principal %s as a group: %w", id, err) + } + if displayName != "" { + return p.groupID(ctx, id) + } + + servicePrincipal, rateLimit, err := p.client.GetServicePrincipal(ctx, "", id) + noteRateLimit(p.annos, rateLimit) + if err != nil && !isNotFoundError(err) { + return nil, fmt.Errorf("failed to resolve principal %s as a service principal: %w", id, err) + } + if err == nil && servicePrincipal != nil && servicePrincipal.ID != "" { + return &v2.ResourceId{ResourceType: servicePrincipalResourceType.Id, Resource: servicePrincipal.ID}, nil + } + + return nil, nil +} + +func (p *principalResolver) lookupName(ctx context.Context, name string) (*v2.ResourceId, error) { + userID, rateLimit, err := p.client.FindUserID(ctx, "", name) + noteRateLimit(p.annos, rateLimit) + if err != nil { + return nil, fmt.Errorf("failed to resolve %q as a user: %w", name, err) + } + if userID != "" { + return &v2.ResourceId{ResourceType: userResourceType.Id, Resource: userID}, nil + } + + groupID, rateLimit, err := p.client.FindGroupID(ctx, "", name) + noteRateLimit(p.annos, rateLimit) + if err != nil { + return nil, fmt.Errorf("failed to resolve %q as a group: %w", name, err) + } + if groupID != "" { + return p.groupID(ctx, groupID) + } + + servicePrincipalID, rateLimit, err := p.client.FindServicePrincipalID(ctx, "", name) + noteRateLimit(p.annos, rateLimit) + if err != nil { + return nil, fmt.Errorf("failed to resolve %q as a service principal: %w", name, err) + } + if servicePrincipalID != "" { + return &v2.ResourceId{ResourceType: servicePrincipalResourceType.Id, Resource: servicePrincipalID}, nil + } + + return nil, nil +} + +// groupID matches the id the group syncer emits for an account-level group. +func (p *principalResolver) groupID(ctx context.Context, groupID string) (*v2.ResourceId, error) { + account, err := rs.NewResourceID(accountResourceType, p.client.GetAccountId()) + if err != nil { + return nil, err + } + + return rs.NewResourceID(groupResourceType, groupResourceId(ctx, groupID, account)) +} diff --git a/pkg/connector/metastore_test.go b/pkg/connector/metastore_test.go new file mode 100644 index 00000000..0db0c170 --- /dev/null +++ b/pkg/connector/metastore_test.go @@ -0,0 +1,744 @@ +package connector + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "slices" + "strconv" + "strings" + "testing" + + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "github.com/conductorone/baton-sdk/pkg/types/grant" + rs "github.com/conductorone/baton-sdk/pkg/types/resource" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/conductorone/baton-databricks/pkg/databricks" +) + +// workspaceRouter dispatches on the deployment-name label of the request host, +// the way Databricks routes an account whose workspaces sit on separate hosts. +type workspaceRouter struct { + target string + hits map[string]int +} + +func (r *workspaceRouter) RoundTrip(req *http.Request) (*http.Response, error) { + workspace, _, _ := strings.Cut(req.URL.Host, ".") + r.hits[workspace]++ + + routed := req.Clone(req.Context()) + routed.URL.Scheme = "http" + routed.URL.Host = r.target + routed.Header.Set("X-Test-Workspace", workspace) + + return http.DefaultTransport.RoundTrip(routed) +} + +func newWorkspaceRoutedClient(t *testing.T, handler http.HandlerFunc) (*databricks.Client, *workspaceRouter) { + t.Helper() + + srv := httptest.NewServer(handler) + t.Cleanup(srv.Close) + + router := &workspaceRouter{target: strings.TrimPrefix(srv.URL, "http://"), hits: map[string]int{}} + client, err := databricks.NewClient(context.Background(), &http.Client{Transport: router}, + testHostname, testAccountHost, testAccountId, "", &databricks.NoAuth{}, nil) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + + return client, router +} + +// fakeAccount serves the account, Unity Catalog and SCIM endpoints the metastore syncer calls. +type fakeAccount struct { + metastores []map[string]any + assignments map[string]any // metastore id -> []int workspace ids, or an HTTP status + workspaces []map[string]any + permissions []map[string]any + permStatus int + patchStatus int + users map[string]string // id -> userName + groups map[string]string // id -> displayName + sps map[string]string // id -> applicationId + + // workspaceMembers maps a workspace id to its permission assignments, or to + // the HTTP status the account API answers the list with. + workspaceMembers map[string]any + + permHosts []string + patches []map[string]any + scimGets map[string]int + memberLists []string +} + +func defaultFakeAccount() *fakeAccount { + return &fakeAccount{ + metastores: []map[string]any{{"metastore_id": "ms-a", "name": "main", "owner": "eng", "region": "us-east-1"}}, + assignments: map[string]any{"ms-a": []int{1}}, + workspaces: []map[string]any{ + {"workspace_id": 1, "workspace_name": "prod", "deployment_name": "wsa", "workspace_status": "RUNNING"}, + }, + users: map[string]string{"1": "alice@example.com"}, + groups: map[string]string{"2": "eng"}, + sps: map[string]string{"3": "app-9"}, + scimGets: map[string]int{}, + } +} + +func scimList(field string, records map[string]string, filter string) []map[string]any { + key, value, _ := strings.Cut(filter, " eq ") + value, _ = strconv.Unquote(value) + var out []map[string]any + for id, name := range records { + if (key == "id" && id == value) || (key != "id" && name == value) { + out = append(out, map[string]any{"id": id, field: name}) + } + } + + return out +} + +func (f *fakeAccount) handler(t *testing.T) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + path := r.URL.Path + encode := func(v any) { _ = json.NewEncoder(w).Encode(v) } + switch { + case strings.HasSuffix(path, "/metastores"): + encode(map[string]any{"metastores": f.metastores}) + case strings.Contains(path, "/metastores/") && strings.HasSuffix(path, "/workspaces"): + parts := strings.Split(path, "/") + switch v := f.assignments[parts[len(parts)-2]].(type) { + case int: + http.Error(w, `{"message":"denied"}`, v) + default: + encode(map[string]any{"workspace_ids": v}) + } + case strings.HasSuffix(path, "/permissionassignments"): + parts := strings.Split(path, "/") + workspaceID := parts[len(parts)-2] + f.memberLists = append(f.memberLists, workspaceID) + switch v := f.workspaceMembers[workspaceID].(type) { + case int: + http.Error(w, `{"message":"denied"}`, v) + default: + encode(map[string]any{"permission_assignments": v}) + } + case strings.HasSuffix(path, "/workspaces"): + encode(f.workspaces) + case strings.Contains(path, "/unity-catalog/permissions/metastore/"): + f.permHosts = append(f.permHosts, r.Header.Get("X-Test-Workspace")) + if f.permStatus != 0 { + http.Error(w, `{"message":"nope"}`, f.permStatus) + return + } + if r.Method == http.MethodPatch { + if f.patchStatus != 0 { + http.Error(w, `{"message":"denied"}`, f.patchStatus) + return + } + var body struct { + Changes []map[string]any `json:"changes"` + } + decoder := json.NewDecoder(r.Body) + decoder.UseNumber() + _ = decoder.Decode(&body) + for _, change := range body.Changes { + f.patches = append(f.patches, change) + f.apply(change) + } + } + encode(map[string]any{"privilege_assignments": f.permissions}) + case strings.Contains(path, "/scim/v2/"): + f.scim(w, r) + default: + t.Errorf("unexpected request %s %s", r.Method, path) + http.Error(w, `{"message":"unexpected"}`, http.StatusInternalServerError) + } + } +} + +func (f *fakeAccount) scim(w http.ResponseWriter, r *http.Request) { + _, rest, _ := strings.Cut(r.URL.Path, "/scim/v2/") + kind, id, _ := strings.Cut(rest, "/") + field, records := map[string]string{"Users": "userName", "Groups": "displayName", "ServicePrincipals": "applicationId"}[kind], + map[string]map[string]string{"Users": f.users, "Groups": f.groups, "ServicePrincipals": f.sps}[kind] + if id != "" { + f.scimGets[kind+"/"+id]++ + if name, ok := records[id]; ok { + _ = json.NewEncoder(w).Encode(map[string]any{"id": id, field: name}) + return + } + http.Error(w, `{"message":"missing"}`, http.StatusNotFound) + return + } + found := scimList(field, records, r.URL.Query().Get("filter")) + _ = json.NewEncoder(w).Encode(map[string]any{"totalResults": len(found), "Resources": found}) +} + +// principalID is the SCIM id the account holds for a principal name. +func (f *fakeAccount) principalID(name string) int { + for _, records := range []map[string]string{f.users, f.groups, f.sps} { + for id, recorded := range records { + if recorded == name { + numeric, _ := strconv.Atoi(id) + return numeric + } + } + } + + return 0 +} + +func (f *fakeAccount) apply(change map[string]any) { + added, _ := change["add"].([]any) + for _, privilege := range added { + f.permissions = append(f.permissions, map[string]any{ + "principal": change["principal"], + "principal_id": f.principalID(fmt.Sprint(change["principal"])), + "privileges": []string{privilege.(string)}, + }) + } + if remove, ok := change["remove"].([]any); ok { + kept := f.permissions[:0] + for _, p := range f.permissions { + if fmt.Sprint(p["principal_id"]) != fmt.Sprint(change["principal_id"]) || !slices.Contains(p["privileges"].([]string), remove[0].(string)) { + kept = append(kept, p) + } + } + f.permissions = kept + } +} + +// testOwnPrincipal is the applicationId the connector authenticates as. It is +// the owner of the metastore in defaultFakeAccount, which is the live shape: +// ownership settles grant visibility without any probe call. +const testOwnPrincipal = "eng" + +func newTestMetastoreBuilder(t *testing.T, f *fakeAccount, allowlist []string, sync syncPrincipals) *metastoreBuilder { + return newTestMetastoreBuilderAs(t, f, allowlist, sync, testOwnPrincipal) +} + +func newTestMetastoreBuilderAs(t *testing.T, f *fakeAccount, allowlist []string, sync syncPrincipals, clientID string) *metastoreBuilder { + t.Setenv("BATON_HTTP_CACHE_TTL", "0") + client, _ := newWorkspaceRoutedClient(t, f.handler(t)) + + return newMetastoreBuilder(client, clientID, configuredWorkspaceSet(allowlist), sync) +} + +var testAccountResourceID = &v2.ResourceId{ResourceType: "account", Resource: testAccountId} + +// testSyncedMetastore is the resource List emits, route included. +func testSyncedMetastore(t *testing.T, workspace string) *v2.Resource { + t.Helper() + r, err := metastoreResource(context.Background(), &databricks.Metastore{MetastoreID: "ms-a", Name: "main", Owner: "eng"}, + testAccountResourceID, workspace, "1", "prod") + if err != nil { + t.Fatalf("metastoreResource: %v", err) + } + + return r +} + +// testMetastore is the resource a provisioning call carries: an id, no profile. +func testMetastore(t *testing.T) *v2.Resource { + t.Helper() + r, err := rs.NewAppResource("main", metastoreResourceType, "ms-a", nil, rs.WithParentResourceID(testAccountResourceID)) + if err != nil { + t.Fatalf("NewAppResource: %v", err) + } + + return r +} + +// Which workspaces serve a metastore decides whether it syncs at all, so an +// assignment list the connector cannot read is unknown, not absent. Skipping it +// would drop the metastore and make ConductorOne remove its grants. Only the +// NotFound has to be re-coded: the SDK runs the metastore List with +// allowWarning, reads NotFound as "the account has no metastores" and lets the +// sync succeed with none. A 401 or 403 already fails it. +func TestMetastoreListFailsWhenAnAssignmentListCannotBeRead(t *testing.T) { + for code, want := range map[int]codes.Code{ + http.StatusNotFound: codes.FailedPrecondition, + http.StatusForbidden: codes.PermissionDenied, + http.StatusUnauthorized: codes.Unauthenticated, + } { + t.Run(http.StatusText(code), func(t *testing.T) { + f := defaultFakeAccount() + f.assignments = map[string]any{"ms-a": code} + m := newTestMetastoreBuilder(t, f, nil, syncPrincipals{}) + + resources, _, err := m.List(context.Background(), testAccountResourceID, rs.SyncOpAttrs{}) + if status.Code(err) != want { + t.Fatalf("List = %v, want %s", err, want) + } + if !strings.Contains(err.Error(), "ms-a") { + t.Errorf("error %q does not name the metastore", err) + } + if len(resources) != 0 { + t.Errorf("listed %d metastores, want none", len(resources)) + } + }) + } +} + +// A workspace that will not become RUNNING leaves the metastore out of the sync +// instead of failing it on every run; one that can still change keeps it waiting. +func TestMetastoreListSeparatesTerminalFromPendingWorkspaces(t *testing.T) { + tests := []struct { + workspaceStatus string + wantErr bool + wantSynced bool + }{ + {"RUNNING", false, true}, + {"PROVISIONING", true, false}, + {"NOT_PROVISIONED", true, false}, + {"CANCELLING", true, false}, + {"FAILED", false, false}, + {"BANNED", false, false}, + {"SOMETHING_NEW", false, false}, + } + for _, test := range tests { + t.Run(test.workspaceStatus, func(t *testing.T) { + f := defaultFakeAccount() + f.assignments = map[string]any{"ms-a": []int{1}} + f.workspaces = []map[string]any{ + {"workspace_id": 1, "workspace_name": "prod", "deployment_name": "wsa", "workspace_status": test.workspaceStatus}, + } + m := newTestMetastoreBuilder(t, f, nil, syncPrincipals{}) + + resources, _, err := m.List(context.Background(), testAccountResourceID, rs.SyncOpAttrs{}) + if test.wantErr != (err != nil) { + t.Fatalf("List error = %v, want error: %v", err, test.wantErr) + } + if got := len(resources) > 0; got != test.wantSynced { + t.Errorf("metastore synced = %v, want %v", got, test.wantSynced) + } + }) + } +} + +// An owner sees every privilege, so the route is the first running workspace in scope. +func TestMetastoreListRoutesTheOwnerThroughTheFirstScopedWorkspace(t *testing.T) { + tests := []struct { + name string + allowlist []string + want map[string]string // metastore id -> profile workspace + }{ + {"no allowlist", nil, map[string]string{"ms-a": "wsa"}}, + {"allowlist by name", []string{"third"}, map[string]string{"ms-a": "wsc"}}, + {"allowlist matching nothing", []string{"nope"}, map[string]string{}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + f := defaultFakeAccount() + f.metastores = append(f.metastores, map[string]any{"metastore_id": "ms-elsewhere"}) + f.assignments = map[string]any{ + "ms-a": []int{2, 1, 3}, + "ms-elsewhere": []int{4}, + } + f.workspaces = []map[string]any{ + {"workspace_id": 1, "workspace_name": "prod", "deployment_name": "wsa", "workspace_status": "RUNNING"}, + {"workspace_id": 2, "workspace_name": "dev", "deployment_name": "wsb", "workspace_status": "PROVISIONING"}, + {"workspace_id": 3, "workspace_name": "third", "deployment_name": "wsc", "workspace_status": "RUNNING"}, + } + m := newTestMetastoreBuilder(t, f, test.allowlist, syncPrincipals{}) + + resources, _, err := m.List(context.Background(), testAccountResourceID, rs.SyncOpAttrs{}) + if err != nil { + t.Fatalf("List: %v", err) + } + got := map[string]string{} + for _, r := range resources { + got[r.GetId().GetResource()], _ = rs.GetProfileStringValue(rs.GetProfile(r), profileKeyWorkspace) + } + if len(got) != len(test.want) { + t.Fatalf("List = %v, want %v", got, test.want) + } + for id, workspace := range test.want { + if got[id] != workspace { + t.Errorf("metastore %s routed through %q, want %q", id, got[id], workspace) + } + } + }) + } +} + +func grantKeys(grants []*v2.Grant) []string { + keys := make([]string, 0, len(grants)) + for _, g := range grants { + keys = append(keys, g.GetPrincipal().GetId().GetResource()+" "+g.GetEntitlement().GetId()) + } + slices.Sort(keys) + + return keys +} + +var ( + adminByAppID = []map[string]any{ + {"principal": map[string]any{"service_principal_name": "app-9"}, "permissions": []string{"ADMIN"}}, + } + adminByScimID = []map[string]any{ + {"principal": map[string]any{"principal_id": 3}, "permissions": []string{"ADMIN"}}, + } + notAdmin = []map[string]any{ + {"principal": map[string]any{"service_principal_name": "app-9"}, "permissions": []string{"USER"}}, + {"principal": map[string]any{"group_name": "admins"}, "permissions": []string{"ADMIN"}}, + } +) + +// A read without standing returns 200 with a partial list, so List routes only +// through a workspace it can fully read. +func TestMetastoreListRoutesOnlyThroughAWorkspaceItCanFullyRead(t *testing.T) { + tests := []struct { + name string + clientID string + members any + wantRouted bool + wantProbes int + }{ + {"owner of the metastore", testOwnPrincipal, nil, true, 0}, + {"direct workspace admin", "app-9", adminByAppID, true, 1}, + {"direct workspace admin matched by scim id", "app-9", adminByScimID, true, 1}, + {"admin held only through a group", "app-9", notAdmin, false, 1}, + {"assignment list denied", "app-9", http.StatusForbidden, false, 1}, + {"no principal to compare", "", adminByAppID, false, 0}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + f := defaultFakeAccount() + f.workspaceMembers = map[string]any{"1": test.members} + m := newTestMetastoreBuilderAs(t, f, nil, syncPrincipals{}, test.clientID) + + resources, _, err := m.List(context.Background(), testAccountResourceID, rs.SyncOpAttrs{}) + switch { + case test.wantRouted && err != nil: + t.Fatalf("List: %v", err) + case test.wantRouted && len(resources) != 1: + t.Fatalf("List returned %d resources, want the metastore", len(resources)) + case !test.wantRouted: + if status.Code(err) != codes.PermissionDenied { + t.Fatalf("List = %v, want PermissionDenied", err) + } + if len(resources) != 0 { + t.Errorf("resources = %v, want none", resources) + } + if !strings.Contains(err.Error(), "ms-a") { + t.Errorf("error %q does not name the metastore", err) + } + } + if len(f.memberLists) != test.wantProbes { + t.Errorf("workspace assignments listed %d times, want %d", len(f.memberLists), test.wantProbes) + } + }) + } +} + +// The route is where standing holds, not the first workspace the assignment list names. +func TestMetastoreListRoutesThroughTheWorkspaceWhereStandingHolds(t *testing.T) { + f := defaultFakeAccount() + f.metastores = []map[string]any{{"metastore_id": "ms-a", "name": "main", "owner": "someone-else"}} + f.assignments = map[string]any{"ms-a": []int{1, 2}} + f.workspaces = []map[string]any{ + {"workspace_id": 1, "workspace_name": "prod", "deployment_name": "wsa", "workspace_status": "RUNNING"}, + {"workspace_id": 2, "workspace_name": "dev", "deployment_name": "wsb", "workspace_status": "RUNNING"}, + } + f.workspaceMembers = map[string]any{"1": nil, "2": adminByAppID} + m := newTestMetastoreBuilderAs(t, f, nil, syncPrincipals{}, "app-9") + + resources, _, err := m.List(context.Background(), testAccountResourceID, rs.SyncOpAttrs{}) + if err != nil { + t.Fatalf("List: %v", err) + } + if len(resources) != 1 { + t.Fatalf("List returned %d resources, want 1", len(resources)) + } + workspace, _ := rs.GetProfileStringValue(rs.GetProfile(resources[0]), profileKeyWorkspace) + if workspace != "wsb" { + t.Errorf("profile workspace = %q, want wsb, the workspace where the connector is Admin", workspace) + } + if !slices.Equal(f.memberLists, []string{"1", "2"}) { + t.Errorf("probed %v, want both workspaces in order until one holds Admin", f.memberLists) + } +} + +// Grants uses the route List already chose and does not probe again. +func TestMetastoreGrantsReadThroughTheRoutedWorkspace(t *testing.T) { + f := defaultFakeAccount() + f.permissions = []map[string]any{ + {"principal": "alice@example.com", "principal_id": 1, "privileges": []string{"CREATE_CATALOG"}}, + } + m := newTestMetastoreBuilderAs(t, f, nil, syncPrincipals{}, "app-9") + + grants, _, err := m.Grants(context.Background(), testSyncedMetastore(t, "wsb"), rs.SyncOpAttrs{}) + if err != nil { + t.Fatalf("Grants: %v", err) + } + if want := []string{"1 metastore:ms-a:CREATE_CATALOG"}; !slices.Equal(grantKeys(grants), want) { + t.Errorf("grants = %v, want %v", grantKeys(grants), want) + } + if !slices.Equal(f.permHosts, []string{"wsb"}) { + t.Errorf("permissions read through %v, want the routed workspace wsb", f.permHosts) + } + if len(f.memberLists) != 0 { + t.Errorf("workspace assignments listed %d times, want none: List already settled the route", len(f.memberLists)) + } +} + +// A NotFound must not look like an empty grant page, and a throttle must stay +// retryable. Wrapping every permissions error as FailedPrecondition does both wrong. +func TestMetastoreGrantsKeepThePermissionsStatus(t *testing.T) { + tests := []struct { + name string + code int + want codes.Code + }{ + {name: "not found", code: http.StatusNotFound, want: codes.FailedPrecondition}, + {name: "throttle", code: http.StatusTooManyRequests, want: codes.Unavailable}, + {name: "unavailable", code: http.StatusServiceUnavailable, want: codes.Unavailable}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + f := defaultFakeAccount() + f.permStatus = test.code + m := newTestMetastoreBuilder(t, f, nil, syncPrincipals{}) + + grants, _, err := m.Grants(context.Background(), testSyncedMetastore(t, "wsa"), rs.SyncOpAttrs{}) + if status.Code(err) != test.want { + t.Fatalf("Grants = %v, want %s", err, test.want) + } + if len(grants) != 0 { + t.Errorf("grants = %v, want none", grantKeys(grants)) + } + }) + } +} + +// No grants and no error would revoke every privilege, so an unrouted resource fails. +func TestMetastoreGrantsFailWithoutARouteOnTheProfile(t *testing.T) { + f := defaultFakeAccount() + m := newTestMetastoreBuilder(t, f, nil, syncPrincipals{}) + + grants, _, err := m.Grants(context.Background(), testMetastore(t), rs.SyncOpAttrs{}) + if status.Code(err) != codes.FailedPrecondition { + t.Fatalf("Grants = %v, want FailedPrecondition", err) + } + if len(grants) != 0 { + t.Errorf("grants = %v, want none", grantKeys(grants)) + } +} + +var alice = &v2.Resource{Id: &v2.ResourceId{ResourceType: "user", Resource: "1"}} + +func TestMetastoreGrantRejections(t *testing.T) { + tests := []struct { + name string + privilege string + patchStatus int + code codes.Code + }{ + {"owner", ownerEntitlement, 0, codes.InvalidArgument}, + {"legacy alias", "USAGE", 0, codes.InvalidArgument}, + {"not offered", "MANAGE", 0, codes.InvalidArgument}, + {"forbidden", "CREATE_CATALOG", http.StatusForbidden, codes.PermissionDenied}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + f := defaultFakeAccount() + m := newTestMetastoreBuilder(t, f, nil, syncPrincipals{}) + g := grant.NewGrant(testMetastore(t), test.privilege, alice.GetId()) + f.patchStatus = test.patchStatus + + _, err := m.Grant(context.Background(), alice, g.GetEntitlement()) + if status.Code(err) != test.code { + t.Fatalf("Grant = %v, want %s", err, test.code) + } + if len(f.patches) != 0 { + t.Errorf("patches = %v, want none", f.patches) + } + }) + } +} + +// Grant and Revoke resolve the route from the metastore id. +func TestMetastoreProvisioningRoutesFromTheMetastoreID(t *testing.T) { + f := defaultFakeAccount() + f.metastores = []map[string]any{{"metastore_id": "ms-a", "name": "main", "owner": "someone-else"}} + f.assignments = map[string]any{"ms-a": []int{1, 2}} + f.workspaces = []map[string]any{ + {"workspace_id": 1, "workspace_name": "prod", "deployment_name": "wsa", "workspace_status": "RUNNING"}, + {"workspace_id": 2, "workspace_name": "dev", "deployment_name": "wsb", "workspace_status": "RUNNING"}, + } + f.workspaceMembers = map[string]any{"1": nil, "2": adminByAppID} + m := newTestMetastoreBuilderAs(t, f, nil, syncPrincipals{}, "app-9") + g := grant.NewGrant(testMetastore(t), "CREATE_CATALOG", alice.GetId()) + + if _, err := m.Grant(context.Background(), alice, g.GetEntitlement()); err != nil { + t.Fatalf("Grant: %v", err) + } + if len(f.patches) != 1 { + t.Fatalf("patches = %v, want one", f.patches) + } + for _, host := range f.permHosts { + if host != "wsb" { + t.Fatalf("permissions addressed %v, want only wsb, the workspace where the connector is Admin", f.permHosts) + } + } +} + +// A deleted metastore has nothing to grant and nothing left to revoke. +func TestMetastoreProvisioningOnADeletedMetastore(t *testing.T) { + newBuilder := func(t *testing.T) (*metastoreBuilder, *fakeAccount) { + f := defaultFakeAccount() + f.metastores = []map[string]any{{"metastore_id": "ms-other", "name": "other", "owner": "eng"}} + + return newTestMetastoreBuilder(t, f, nil, syncPrincipals{}), f + } + + t.Run("grant", func(t *testing.T) { + m, f := newBuilder(t) + g := grant.NewGrant(testMetastore(t), "CREATE_CATALOG", alice.GetId()) + + _, err := m.Grant(context.Background(), alice, g.GetEntitlement()) + if status.Code(err) != codes.NotFound { + t.Fatalf("Grant = %v, want NotFound", err) + } + if len(f.patches) != 0 { + t.Errorf("patches = %v, want none", f.patches) + } + }) + + t.Run("revoke", func(t *testing.T) { + m, f := newBuilder(t) + g := grant.NewGrant(testMetastore(t), "CREATE_CATALOG", alice.GetId()) + + annos, err := m.Revoke(context.Background(), g) + if err != nil || !annos.Contains(&v2.GrantAlreadyRevoked{}) { + t.Fatalf("Revoke = %v, %v; want GrantAlreadyRevoked", annos, err) + } + if len(f.patches) != 0 { + t.Errorf("patches = %v, want none", f.patches) + } + }) +} + +func TestMetastoreRevoke(t *testing.T) { + f := defaultFakeAccount() + f.permissions = []map[string]any{{"principal": "alice@example.com", "principal_id": 1, "privileges": []string{"CREATE_CATALOG"}}} + m := newTestMetastoreBuilder(t, f, nil, syncPrincipals{}) + g := grant.NewGrant(testMetastore(t), "CREATE_CATALOG", alice.GetId()) + + if _, err := m.Revoke(context.Background(), g); err != nil { + t.Fatalf("Revoke: %v", err) + } + if len(f.patches) != 1 || fmt.Sprint(f.patches[0]["principal_id"]) != "1" { + t.Fatalf("patches = %v, want one removal by principal_id 1", f.patches) + } + + annos, err := m.Revoke(context.Background(), g) + if err != nil || !annos.Contains(&v2.GrantAlreadyRevoked{}) { + t.Fatalf("second Revoke = %v, %v; want GrantAlreadyRevoked", annos, err) + } +} + +// A row the pre-read missed may be a partial read, so the removal still goes out. +func TestMetastoreRevokeWritesWhenThePreReadMissesTheRow(t *testing.T) { + f := defaultFakeAccount() + f.permissions = nil + m := newTestMetastoreBuilder(t, f, nil, syncPrincipals{}) + g := grant.NewGrant(testMetastore(t), "CREATE_CATALOG", alice.GetId()) + + annos, err := m.Revoke(context.Background(), g) + if err != nil { + t.Fatalf("Revoke: %v", err) + } + if len(f.patches) != 1 { + t.Fatalf("patches = %v, want the removal to go out anyway", f.patches) + } + if got := f.patches[0]["remove"]; fmt.Sprint(got) != "[CREATE_CATALOG]" { + t.Errorf("removal = %v, want CREATE_CATALOG", got) + } + if got := f.patches[0]["principal"]; fmt.Sprint(got) != "alice@example.com" { + t.Errorf("removal addressed %v, want the principal name alice@example.com", got) + } + if _, byID := f.patches[0]["principal_id"]; byID { + t.Errorf("removal = %v, want no principal_id: the pre-read saw no row to take one from", f.patches[0]) + } + if !annos.Contains(&v2.GrantAlreadyRevoked{}) { + t.Error("want GrantAlreadyRevoked: the echo proves the privilege is not held") + } +} + +// A deleted principal's rows are only addressable by id. +func TestMetastoreRevokeRemovesADeletedPrincipalByID(t *testing.T) { + f := defaultFakeAccount() + f.users = map[string]string{} + f.permissions = []map[string]any{ + {"principal": "alice@example.com", "principal_id": 1, "privileges": []string{"CREATE_CATALOG"}}, + } + m := newTestMetastoreBuilder(t, f, nil, syncPrincipals{}) + g := grant.NewGrant(testMetastore(t), "CREATE_CATALOG", alice.GetId()) + + if _, err := m.Revoke(context.Background(), g); err != nil { + t.Fatalf("Revoke: %v", err) + } + if len(f.patches) != 1 { + t.Fatalf("patches = %v, want one removal", f.patches) + } + if got := f.patches[0]["principal_id"]; fmt.Sprint(got) != "1" { + t.Errorf("removal addressed %v, want principal_id 1", got) + } + if _, byName := f.patches[0]["principal"]; byName { + t.Errorf("removal = %v, want no principal name: Databricks no longer has one", f.patches[0]) + } +} + +// A grant to a deleted principal would write an empty name. +func TestMetastoreGrantToADeletedPrincipalFails(t *testing.T) { + f := defaultFakeAccount() + f.users = map[string]string{} + m := newTestMetastoreBuilder(t, f, nil, syncPrincipals{}) + g := grant.NewGrant(testMetastore(t), "CREATE_CATALOG", alice.GetId()) + + _, err := m.Grant(context.Background(), alice, g.GetEntitlement()) + if status.Code(err) != codes.NotFound { + t.Fatalf("Grant = %v, want NotFound", err) + } + if len(f.patches) != 0 { + t.Errorf("patches = %v, want none", f.patches) + } +} + +// A privilege list the connector cannot fully read cannot be written to either: +// the change would be decided from the subset Databricks chose to show. +func TestMetastoreProvisioningRefusesAPartialPrivilegeList(t *testing.T) { + for _, verb := range []string{"grant", "revoke"} { + t.Run(verb, func(t *testing.T) { + f := defaultFakeAccount() + f.permissions = []map[string]any{ + {"principal": "alice@example.com", "principal_id": 1, "privileges": []string{"CREATE_CATALOG"}}, + } + m := newTestMetastoreBuilderAs(t, f, nil, syncPrincipals{}, "app-9") + g := grant.NewGrant(testMetastore(t), "CREATE_CATALOG", alice.GetId()) + + var err error + if verb == "grant" { + _, err = m.Grant(context.Background(), alice, g.GetEntitlement()) + } else { + _, err = m.Revoke(context.Background(), g) + } + if status.Code(err) != codes.PermissionDenied { + t.Fatalf("%s = %v, want PermissionDenied", verb, err) + } + if len(f.patches) != 0 { + t.Errorf("patches = %v, want none", f.patches) + } + }) + } +} diff --git a/pkg/connector/raw_id_test.go b/pkg/connector/raw_id_test.go index d12a99ba..2809a80b 100644 --- a/pkg/connector/raw_id_test.go +++ b/pkg/connector/raw_id_test.go @@ -65,6 +65,11 @@ func TestResourcesHaveRawId(t *testing.T) { build((&servicePrincipalBuilder{}).servicePrincipalResource(ctx, &databricks.ServicePrincipal{BaseResponse: databricks.BaseResponse{ID: "sp1"}}, account)), "sp1", }, + { + "metastore", + build(metastoreResource(ctx, &databricks.Metastore{MetastoreID: "ms1", Name: "ms"}, account, "", "", "")), + "ms1", + }, } for _, tc := range cases { diff --git a/pkg/connector/resource_types.go b/pkg/connector/resource_types.go index 748bdbef..2e07ce95 100644 --- a/pkg/connector/resource_types.go +++ b/pkg/connector/resource_types.go @@ -2,8 +2,36 @@ package connector import ( v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "github.com/conductorone/baton-sdk/pkg/annotations" ) +// The Databricks permissions this connector's reads depend on. +const ( + scopeUnityCatalog = "unity-catalog" + scopeSCIM = "scim" + scopeAccessManagement = "access-management" + scopeProvisioning = "provisioning" +) + +const ( + permissionAccountAdmin = "Account admin" + permissionWorkspaceAdmin = "Workspace admin" +) + +var securableScopes = []string{ + scopeUnityCatalog, + scopeSCIM, + scopeAccessManagement, + scopeProvisioning, +} + +func securablePermissions(privileges ...string) []string { + perms := append([]string{}, securableScopes...) + perms = append(perms, permissionAccountAdmin, permissionWorkspaceAdmin) + + return append(perms, privileges...) +} + var ( // The user resource type is for all user objects from the database. userResourceType = &v2.ResourceType{ @@ -18,6 +46,7 @@ var ( Id: "group", DisplayName: "Group", Traits: []v2.ResourceType_Trait{v2.ResourceType_TRAIT_GROUP}, + Annotations: annotations.New(capabilityPermissions(scopeSCIM, scopeAccessManagement, permissionAccountAdmin)), } // The service principal resource type is for all service principal objects from the database. @@ -25,6 +54,7 @@ var ( Id: "service_principal", DisplayName: "Service Principal", Traits: []v2.ResourceType_Trait{v2.ResourceType_TRAIT_GROUP}, + Annotations: annotations.New(capabilityPermissions(scopeSCIM, scopeAccessManagement, permissionAccountAdmin)), } // The role resource type is for all static roles and entitlements available in API. @@ -32,6 +62,7 @@ var ( Id: "role", DisplayName: "Role", Traits: []v2.ResourceType_Trait{v2.ResourceType_TRAIT_ROLE}, + Annotations: annotations.New(capabilityPermissions(scopeSCIM, permissionAccountAdmin)), } // The workspace resource type is for all workspace objects from the database. @@ -39,11 +70,19 @@ var ( Id: "workspace", DisplayName: "Workspace", Traits: []v2.ResourceType_Trait{v2.ResourceType_TRAIT_GROUP}, + Annotations: annotations.New(capabilityPermissions(scopeProvisioning, scopeAccessManagement, permissionAccountAdmin)), } // The account resource type is for top level resource type. accountResourceType = &v2.ResourceType{ Id: "account", DisplayName: "Account", + Annotations: annotations.New(capabilityPermissions(scopeAccessManagement, permissionAccountAdmin)), + } + metastoreResourceType = &v2.ResourceType{ + Id: "metastore", + DisplayName: "Metastore", + Traits: []v2.ResourceType_Trait{v2.ResourceType_TRAIT_APP}, + Annotations: annotationsForStaticSecurableResourceType(securablePermissions()...), } ) diff --git a/pkg/connector/roles.go b/pkg/connector/roles.go index 127589c8..066d10c9 100644 --- a/pkg/connector/roles.go +++ b/pkg/connector/roles.go @@ -53,7 +53,7 @@ func roleResourceId(role string, parent *v2.ResourceId) string { } func roleResource(ctx context.Context, role string, parent *v2.ResourceId) (*v2.Resource, error) { - profile := map[string]interface{}{ + profile := map[string]any{ "role_name": role, "parent_type": parent.GetResourceType(), "parent_id": parent.GetResource(), diff --git a/pkg/connector/service-principals.go b/pkg/connector/service-principals.go index 51bc131d..a868f054 100644 --- a/pkg/connector/service-principals.go +++ b/pkg/connector/service-principals.go @@ -27,7 +27,7 @@ func (s *servicePrincipalBuilder) ResourceType(ctx context.Context) *v2.Resource } func (s *servicePrincipalBuilder) servicePrincipalResource(ctx context.Context, servicePrincipal *databricks.ServicePrincipal, parent *v2.ResourceId) (*v2.Resource, error) { - profile := map[string]interface{}{ + profile := map[string]any{ "application_id": servicePrincipal.ApplicationID, "display_name": servicePrincipal.DisplayName, "parent_type": parent.GetResourceType(), diff --git a/pkg/connector/users.go b/pkg/connector/users.go index 028e3e29..f18bc213 100644 --- a/pkg/connector/users.go +++ b/pkg/connector/users.go @@ -48,7 +48,7 @@ func (u *userBuilder) userResource(ctx context.Context, user *databricks.User, p } firstName, lastName := rs.SplitFullName(user.DisplayName) - profile := map[string]interface{}{ + profile := map[string]any{ "first_name": firstName, "last_name": lastName, "email": primaryEmail, diff --git a/pkg/connector/validate_test.go b/pkg/connector/validate_test.go index 04678e99..28799667 100644 --- a/pkg/connector/validate_test.go +++ b/pkg/connector/validate_test.go @@ -110,15 +110,16 @@ func TestValidateWithoutWarehouseIDDisablesIncrementalSync(t *testing.T) { } } -// Incremental sync misconfiguration (unknown warehouse, no workspace in scope) must not -// block full syncs, so Validate succeeds without probing the warehouse. +// Incremental sync misconfiguration (an unknown warehouse) must not block full syncs, +// so Validate succeeds without probing the warehouse. A workspace allowlist that names +// no workspace on the account is a different failure, covered by its own test. func TestValidateWithWarehouseIDDoesNotProbeIncrementalSync(t *testing.T) { cases := []struct { name string workspaces []string }{ {name: "unknown warehouse"}, - {name: "no workspace in scope", workspaces: []string{"not-a-workspace"}}, + {name: "one workspace in scope", workspaces: []string{"ws1"}}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { @@ -139,7 +140,8 @@ func TestValidateWithWarehouseIDDoesNotProbeIncrementalSync(t *testing.T) { // A whitespace-only sql-warehouse-id counts as unset, so incremental sync stays off. func TestNewTreatsBlankWarehouseIDAsDisabled(t *testing.T) { - d, err := New(context.Background(), "cloud.databricks.com", "accounts.cloud.databricks.com", "acct-1", "", &databricks.NoAuth{}, nil, nil, " ") + d, err := New(context.Background(), "cloud.databricks.com", "accounts.cloud.databricks.com", "acct-1", "", "", + &databricks.NoAuth{}, nil, nil, " ", nil) if err != nil { t.Fatalf("New: %v", err) } diff --git a/pkg/connector/workspaces.go b/pkg/connector/workspaces.go index 1dee3c4a..4ed74764 100644 --- a/pkg/connector/workspaces.go +++ b/pkg/connector/workspaces.go @@ -41,7 +41,7 @@ func (w *workspaceBuilder) ResourceType(ctx context.Context) *v2.ResourceType { } func workspaceResource(_ context.Context, workspace *databricks.Workspace, parent *v2.ResourceId) (*v2.Resource, error) { - profile := map[string]interface{}{ + profile := map[string]any{ "workspace_id": workspace.ID, } @@ -386,15 +386,21 @@ func (w *workspaceBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotat return nil, nil } -func newWorkspaceBuilder(client *databricks.Client, workspaces []string) *workspaceBuilder { +// configuredWorkspaceSet indexes the allowlist. Both the workspace builder and +// Unity Catalog routing filter on it, so they read it from one place. +func configuredWorkspaceSet(workspaces []string) map[string]struct{} { wMap := make(map[string]struct{}, len(workspaces)) for _, w := range workspaces { wMap[w] = struct{}{} } + return wMap +} + +func newWorkspaceBuilder(client *databricks.Client, workspaces []string) *workspaceBuilder { return &workspaceBuilder{ client: client, resourceType: workspaceResourceType, - workspaces: wMap, + workspaces: configuredWorkspaceSet(workspaces), } } diff --git a/pkg/databricks/client.go b/pkg/databricks/client.go index c01e3899..2ff8722c 100644 --- a/pkg/databricks/client.go +++ b/pkg/databricks/client.go @@ -105,8 +105,10 @@ func NewClient(ctx context.Context, httpClient *http.Client, hostname, accountHo Host: accountHostname, } - // Key the cache on Cache-Control so read-before-write no-cache reads don't overwrite the shared sync entry. - cli, err := uhttp.NewBaseHttpClientWithContext(ctx, httpClient, uhttp.WithCacheKeyHeaders("Cache-Control")) + // Key the cache on the scope header so two workspaces answering the same path keep + // separate entries, and on Cache-Control so a read-before-write no-cache read does + // not overwrite the shared sync entry. + cli, err := uhttp.NewBaseHttpClientWithContext(ctx, httpClient, uhttp.WithCacheKeyHeaders(cacheScopeHeader, "Cache-Control")) return &Client{ httpClient: cli, auth: auth, @@ -201,6 +203,8 @@ func (c *Client) ListUsers( return res.Resources, res.Total, ratelimitData, nil } +// GetUser reads one SCIM user by id. A 404 means no user has that id +// (RFC 7644 §3.4.1), which is how a group or service principal id answers here. func (c *Client) GetUser( ctx context.Context, workspaceId string, @@ -250,7 +254,7 @@ func (c *Client) FindUserID( ctx, workspaceId, &PaginationVars{Count: 1}, - NewFilterVars(fmt.Sprintf("userName eq '%s'", username)), + NewFilterVars("userName eq "+strconv.Quote(username)), ) if err != nil { @@ -305,7 +309,7 @@ func (c *Client) FindUsername( ctx, workspaceId, &PaginationVars{Count: 1}, - NewFilterVars(fmt.Sprintf("id eq '%s'", userID)), + NewFilterVars("id eq "+strconv.Quote(userID)), ) if err != nil { @@ -345,6 +349,8 @@ func (c *Client) ListGroups( return res.Resources, res.Total, ratelimitData, nil } +// GetGroup reads one SCIM group by id. A 404 means no group has that id +// (RFC 7644 §3.4.1), which is how a user or service principal id answers here. func (c *Client) GetGroup(ctx context.Context, workspaceId, groupId string, vars ...Vars) ( *Group, *v2.RateLimitDescription, @@ -360,7 +366,7 @@ func (c *Client) GetGroupNoCache(ctx context.Context, workspaceId, groupId strin *v2.RateLimitDescription, error, ) { - return c.getGroup(ctx, c.getNoCache, workspaceId, groupId, vars...) + return c.getGroup(ctx, c.GetUncached, workspaceId, groupId, vars...) } func (c *Client) getGroup( @@ -420,7 +426,7 @@ func (c *Client) FindGroupID( ctx, workspaceId, &PaginationVars{Count: 1}, - NewFilterVars(fmt.Sprintf("displayName eq '%s'", displayName)), + NewFilterVars("displayName eq "+strconv.Quote(displayName)), ) if err != nil { @@ -447,7 +453,7 @@ func (c *Client) FindGroupDisplayName( ctx, workspaceId, &PaginationVars{Count: 1}, - NewFilterVars(fmt.Sprintf("id eq '%s'", groupID)), + NewFilterVars("id eq "+strconv.Quote(groupID)), ) if err != nil { @@ -487,6 +493,8 @@ func (c *Client) ListServicePrincipals( return res.Resources, res.Total, ratelimitData, nil } +// GetServicePrincipal reads one SCIM service principal by id. A 404 means no +// service principal has that id (RFC 7644 §3.4.1). func (c *Client) GetServicePrincipal( ctx context.Context, workspaceId string, @@ -538,7 +546,7 @@ func (c *Client) FindServicePrincipalID( ctx, workspaceId, &PaginationVars{Count: 1}, - NewFilterVars(fmt.Sprintf("applicationId eq '%s'", appID)), + NewFilterVars("applicationId eq "+strconv.Quote(appID)), ) if err != nil { @@ -565,7 +573,7 @@ func (c *Client) FindServicePrincipalAppID( ctx, workspaceId, &PaginationVars{Count: 1}, - NewFilterVars(fmt.Sprintf("id eq '%s'", servicePrincipalID)), + NewFilterVars("id eq "+strconv.Quote(servicePrincipalID)), ) if err != nil { @@ -713,7 +721,7 @@ func (c *Client) listWorkspaceAssignments( u := c.accountBaseUrl.JoinPath(fmt.Sprintf(accountWorkspaceAssignmentsEndpoint, c.accountId, workspaceId)) get := c.Get if noCache { - get = c.getNoCache + get = c.GetUncached } ratelimitData, err := get(ctx, u, &res) if err != nil { diff --git a/pkg/databricks/client_test.go b/pkg/databricks/client_test.go index 2425c83a..6d01e384 100644 --- a/pkg/databricks/client_test.go +++ b/pkg/databricks/client_test.go @@ -2,6 +2,7 @@ package databricks import ( "context" + "encoding/json" "net/http" "net/http/httptest" "net/url" @@ -72,6 +73,61 @@ func TestGetAccountHostname(t *testing.T) { } } +// A single-quoted value is not RFC 7644 filter syntax, and an apostrophe in a +// userName would end it early, so the value is sent as a quoted JSON string. +func TestFindFiltersQuoteTheValue(t *testing.T) { + t.Setenv("BATON_HTTP_CACHE_TTL", "0") + + tests := []struct { + name string + find func(c *Client) error + want string + }{ + {"user by userName", func(c *Client) error { + _, _, err := c.FindUserID(context.Background(), "", "o'brien@example.com") + return err + }, `userName eq "o'brien@example.com"`}, + {"group by displayName", func(c *Client) error { + _, _, err := c.FindGroupID(context.Background(), "", `data "eng"`) + return err + }, `displayName eq "data \"eng\""`}, + {"service principal by applicationId", func(c *Client) error { + _, _, err := c.FindServicePrincipalID(context.Background(), "", "app-1") + return err + }, `applicationId eq "app-1"`}, + {"user by id", func(c *Client) error { + _, _, err := c.FindUsername(context.Background(), "", "42") + return err + }, `id eq "42"`}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = r.URL.Query().Get("filter") + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"Resources": []any{}, "totalResults": 0}) + })) + t.Cleanup(srv.Close) + + client := srv.Client() + client.Transport = rewriteHost(srv.Listener.Addr().String(), client.Transport) + c, err := NewClient(context.Background(), client, "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acc-1", "", &NoAuth{}, nil) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + + if err := tt.find(c); err != nil { + t.Fatalf("find: %v", err) + } + if got != tt.want { + t.Errorf("filter = %q, want %q", got, tt.want) + } + }) + } +} + func TestIsWorkspaceExcludedReturnsEveryMatchedKey(t *testing.T) { ws := Workspace{ID: 12345, Name: "prod", DeploymentName: "dbc-abc"} c := newTestClient(t, []string{"prod", "dbc-abc", "staging"}) diff --git a/pkg/databricks/models.go b/pkg/databricks/models.go index 97a3db82..1d811724 100644 --- a/pkg/databricks/models.go +++ b/pkg/databricks/models.go @@ -1,5 +1,7 @@ package databricks +import "encoding/json" + type BaseResponse struct { ID string `json:"id"` } @@ -164,3 +166,51 @@ type RuleSet struct { Principals []string `json:"principals"` Role string `json:"role"` } + +type Metastore struct { + MetastoreID string `json:"metastore_id"` + Name string `json:"name"` + Region string `json:"region"` + Owner string `json:"owner"` +} + +type PrivilegeAssignment struct { + Principal string `json:"principal"` + PrincipalID json.Number `json:"principal_id"` + Privileges []string `json:"privileges"` +} + +type PermissionsChange struct { + Principal string `json:"principal,omitempty"` + PrincipalID json.Number `json:"principal_id,omitempty"` + Add []string `json:"add,omitempty"` + Remove []string `json:"remove,omitempty"` +} + +func (p PermissionsChange) validate() error { + switch { + case p.Principal == "" && p.PrincipalID == "": + return invalidArgument("a principal or a principal id is required") + case p.Principal != "" && p.PrincipalID != "": + return invalidArgument("principal %q and principal id %s cannot both be set", p.Principal, p.PrincipalID) + case p.PrincipalID != "" && len(p.Add) > 0: + return invalidArgument("principal id %s addresses removals only, so it cannot add %v", p.PrincipalID, p.Add) + case len(p.Add) == 0 && len(p.Remove) == 0: + return invalidArgument("no privileges to add or remove") + } + + return nil +} + +type metastoresResponse struct { + Metastores []Metastore `json:"metastores"` +} + +type metastoreWorkspacesResponse struct { + WorkspaceIDs []json.Number `json:"workspace_ids"` +} + +type permissionsResponse struct { + PrivilegeAssignments []PrivilegeAssignment `json:"privilege_assignments"` + NextPageToken string `json:"next_page_token"` +} diff --git a/pkg/databricks/pagination.go b/pkg/databricks/pagination.go new file mode 100644 index 00000000..219486ad --- /dev/null +++ b/pkg/databricks/pagination.go @@ -0,0 +1,80 @@ +package databricks + +import ( + "context" + "fmt" + + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" +) + +const ( + // The SDK caps the walks it drives; this one is driven here. + maxCursorPages = 1000 +) + +// forEachPageFrom walks listPage from start. A page can come back empty and still +// carry a next token, so only an absent token ends the walk, and a repeated token +// is an error: Databricks restarts at page one on a token it does not recognize. +func forEachPageFrom[T any]( + ctx context.Context, + start string, + listPage func(pageToken string) ([]T, string, *v2.RateLimitDescription, error), + visit func(page []T) (stop bool, err error), +) (*v2.RateLimitDescription, error) { + var rateLimit *v2.RateLimitDescription + seen := map[string]struct{}{} + + for pageToken, pages := start, 0; ; pages++ { + if pageToken != "" { + if _, ok := seen[pageToken]; ok { + return rateLimit, fmt.Errorf("cursor walk repeated a page token after %d pages", pages) + } + seen[pageToken] = struct{}{} + } + if pages >= maxCursorPages { + return rateLimit, fmt.Errorf("cursor walk exceeded %d pages", maxCursorPages) + } + if err := ctx.Err(); err != nil { + return rateLimit, err + } + + page, next, pageRateLimit, err := listPage(pageToken) + if pageRateLimit != nil { + rateLimit = pageRateLimit + } + if err != nil { + return rateLimit, err + } + + stop, err := visit(page) + if err != nil { + return rateLimit, err + } + if stop || next == "" { + return rateLimit, nil + } + pageToken = next + } +} + +// drainPagesFrom walks listPage to exhaustion. +func drainPagesFrom[T any]( + ctx context.Context, + start string, + listPage func(pageToken string) ([]T, string, *v2.RateLimitDescription, error), +) ( + []T, + *v2.RateLimitDescription, + error, +) { + var items []T + rateLimit, err := forEachPageFrom(ctx, start, listPage, func(page []T) (bool, error) { + items = append(items, page...) + return false, nil + }) + if err != nil { + return nil, rateLimit, err + } + + return items, rateLimit, nil +} diff --git a/pkg/databricks/permissions_cache_test.go b/pkg/databricks/permissions_cache_test.go new file mode 100644 index 00000000..a63302eb --- /dev/null +++ b/pkg/databricks/permissions_cache_test.go @@ -0,0 +1,128 @@ +package databricks + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" +) + +// This runs at the cache's default settings on purpose: BATON_HTTP_CACHE_TTL is +// deliberately not set, because turning the cache off is what hides this. +// +// Grant and Revoke both decide whether to PATCH by reading the securable's +// current assignments first, and both reads produce the same URL and query, so +// they share one cache entry. uhttp holds a GET 200 for an hour and a PATCH does +// not invalidate it, so a cached pre-read makes Revoke see the state from before +// the Grant: it reports GrantAlreadyRevoked and never sends the PATCH, leaving +// the principal with access that C1 believes is gone. +func TestProvisioningPreReadsAreNotServedFromTheCache(t *testing.T) { + var gets int + granted := false + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.Method { + case http.MethodGet: + gets++ + var assignments []PrivilegeAssignment + if granted { + assignments = []PrivilegeAssignment{{Principal: "alice", Privileges: []string{"SELECT"}}} + } + _ = json.NewEncoder(w).Encode(permissionsResponse{PrivilegeAssignments: assignments}) + case http.MethodPatch: + granted = true + _ = json.NewEncoder(w).Encode(permissionsResponse{ + PrivilegeAssignments: []PrivilegeAssignment{{Principal: "alice", Privileges: []string{"SELECT"}}}, + }) + default: + http.Error(w, "unexpected method", http.StatusMethodNotAllowed) + } + })) + t.Cleanup(srv.Close) + + client := srv.Client() + client.Transport = rewriteHost(srv.Listener.Addr().String(), client.Transport) + c, err := NewClient(context.Background(), client, "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acc-1", srv.URL, &NoAuth{}, nil) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + ctx := context.Background() + + // Grant's pre-read: nothing is held yet, so the Grant proceeds to the PATCH. + held, _, err := c.DrainPermissions(ctx, "ws", SecurableMetastore, "ms-1") + if err != nil { + t.Fatalf("Grant pre-read: %v", err) + } + if len(held) != 0 { + t.Fatalf("Grant pre-read saw %d assignment(s), want none", len(held)) + } + + if _, _, err := c.UpdatePermissions(ctx, "ws", SecurableMetastore, "ms-1", + []PermissionsChange{{Principal: "alice", Add: []string{"SELECT"}}}); err != nil { + t.Fatalf("UpdatePermissions: %v", err) + } + + // Revoke's pre-read, same process and same hour, against the same cache key. + var seen []PrivilegeAssignment + if _, err := c.ForEachUncachedPermissionsPage(ctx, "ws", SecurableMetastore, "ms-1", "", + func(page []PrivilegeAssignment) (bool, error) { + seen = append(seen, page...) + return false, nil + }); err != nil { + t.Fatalf("Revoke pre-read: %v", err) + } + + if len(seen) != 1 { + t.Fatalf("Revoke pre-read saw %d assignment(s), want the privilege the Grant just added: "+ + "a cached page here reports GrantAlreadyRevoked and skips the PATCH", len(seen)) + } + if gets != 3 { + t.Errorf("%d GET(s) reached the API, want 3: both pre-reads and the post-PATCH refresh must bypass the cache", gets) + } +} + +// The pre-read stores page one under the key the sync's cached ListPermissions +// uses. Without a refresh after the PATCH, a sync in the same process within the +// cache TTL reads the state from before the grant. +func TestASyncReadAfterAGrantSeesThePostChangeState(t *testing.T) { + granted := false + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodPatch { + granted = true + } + var assignments []PrivilegeAssignment + if granted { + assignments = []PrivilegeAssignment{{Principal: "alice", Privileges: []string{"SELECT"}}} + } + _ = json.NewEncoder(w).Encode(permissionsResponse{PrivilegeAssignments: assignments}) + })) + t.Cleanup(srv.Close) + + client := srv.Client() + client.Transport = rewriteHost(srv.Listener.Addr().String(), client.Transport) + c, err := NewClient(context.Background(), client, "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acc-1", srv.URL, &NoAuth{}, nil) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + ctx := context.Background() + + if _, _, err := c.DrainPermissions(ctx, "ws", SecurableMetastore, "ms-1"); err != nil { + t.Fatalf("Grant pre-read: %v", err) + } + if _, _, err := c.UpdatePermissions(ctx, "ws", SecurableMetastore, "ms-1", + []PermissionsChange{{Principal: "alice", Add: []string{"SELECT"}}}); err != nil { + t.Fatalf("UpdatePermissions: %v", err) + } + + assignments, _, _, err := c.ListPermissions(ctx, "ws", SecurableMetastore, "ms-1", "", "") + if err != nil { + t.Fatalf("ListPermissions: %v", err) + } + if len(assignments) != 1 { + t.Fatalf("the cached sync read saw %d assignment(s), want the privilege the Grant just added", len(assignments)) + } +} diff --git a/pkg/databricks/request.go b/pkg/databricks/request.go index 468c8a3e..77c6d073 100644 --- a/pkg/databricks/request.go +++ b/pkg/databricks/request.go @@ -51,6 +51,7 @@ func wrapTransportAuthError(err error) error { // APIError represents an error response from the Databricks API. type APIError struct { StatusCode int + ErrorCode string Detail string Message string Err error @@ -58,8 +59,9 @@ type APIError struct { func (e *APIError) Error() string { return fmt.Sprintf( - "unexpected status code %d: %s %s %v", + "unexpected status code %d: %s %s %s %v", e.StatusCode, + e.ErrorCode, e.Detail, e.Message, e.Err, @@ -88,6 +90,11 @@ func nameWorkspace403Remedy(workspaceId string, err error) error { return err } +// invalidArgument marks a malformed request so it does not reach C1 as Unknown and get retried. +func invalidArgument(format string, args ...any) error { + return uhttp.WrapErrors(codes.InvalidArgument, fmt.Sprintf(format, args...)) +} + type noCacheKey struct{} // WithoutCache marks ctx so GETs made with it bypass the uhttp response cache, for @@ -104,7 +111,7 @@ func isNoCache(ctx context.Context) bool { func (c *Client) Get( ctx context.Context, urlAddress *url.URL, - response interface{}, + response any, params ...Vars, ) (*v2.RateLimitDescription, error) { return c.doRequest( @@ -113,6 +120,30 @@ func (c *Client) Get( http.MethodGet, nil, response, + nil, + params..., + ) +} + +// GetUncached is Get with the uhttp response cache bypassed on the read side. +// The cache holds a GET 200 for an hour and a PATCH does not invalidate it, so a +// read that decides whether a write still has to happen must not be served from +// it: the write would be skipped against a page that predates the last change. +// uhttp still records the fresh response, so the entry a later reader sees is the +// one this call just observed. +func (c *Client) GetUncached( + ctx context.Context, + urlAddress *url.URL, + response any, + params ...Vars, +) (*v2.RateLimitDescription, error) { + return c.doRequest( + ctx, + urlAddress, + http.MethodGet, + nil, + response, + []uhttp.RequestOption{uhttp.WithNoCache()}, params..., ) } @@ -120,8 +151,8 @@ func (c *Client) Get( func (c *Client) Put( ctx context.Context, urlAddress *url.URL, - body interface{}, - response interface{}, + body any, + response any, params ...Vars, ) (*v2.RateLimitDescription, error) { return c.doRequest( @@ -130,40 +161,43 @@ func (c *Client) Put( http.MethodPut, body, response, + nil, params..., ) } -func (c *Client) Patch( +func (c *Client) Post( ctx context.Context, urlAddress *url.URL, - body interface{}, - response interface{}, + body any, + response any, params ...Vars, ) (*v2.RateLimitDescription, error) { return c.doRequest( ctx, urlAddress, - http.MethodPatch, + http.MethodPost, body, response, + nil, params..., ) } -func (c *Client) Post( +func (c *Client) Patch( ctx context.Context, urlAddress *url.URL, - body interface{}, - response interface{}, + body any, + response any, params ...Vars, ) (*v2.RateLimitDescription, error) { return c.doRequest( ctx, urlAddress, - http.MethodPost, + http.MethodPatch, body, response, + nil, params..., ) } @@ -182,7 +216,7 @@ func (c *Client) Delete( ) } -func parseJSON(body io.Reader, res interface{}) error { +func parseJSON(body io.Reader, res any) error { // Databricks seems to return content-type text/plain even though it's json, // so don't check content type. if err := json.NewDecoder(body).Decode(res); err != nil { @@ -192,52 +226,30 @@ func parseJSON(body io.Reader, res interface{}) error { return nil } -// getNoCache is a GET that bypasses the uhttp response cache. -func (c *Client) getNoCache( - ctx context.Context, - urlAddress *url.URL, - response interface{}, - params ...Vars, -) (*v2.RateLimitDescription, error) { - return c.doRequestWithOptions( - ctx, - urlAddress, - http.MethodGet, - nil, - response, - []uhttp.RequestOption{uhttp.WithNoCache()}, - params..., - ) -} - -func (c *Client) doRequest( - ctx context.Context, - urlAddress *url.URL, - method string, - body interface{}, - response interface{}, - params ...Vars, -) (*v2.RateLimitDescription, error) { - return c.doRequestWithOptions(ctx, urlAddress, method, body, response, nil, params...) -} +// cacheScopeHeader pins a cached response to the host it was read from: uhttp +// keys its cache on path, query and headers with no host component, so two +// workspace deployments answering the same path share one entry. net/http builds +// the Host line from the URL and never sends this header on the wire. +const cacheScopeHeader = "Host" -// doRequestWithOptions is doRequest with extra request options appended. -func (c *Client) doRequestWithOptions( +// prepareRequest builds every request this client sends, so none can reach the +// cache without the scope header. uhttp writes an entry even for a caller that +// asked not to read one. +func (c *Client) prepareRequest( ctx context.Context, urlAddress *url.URL, method string, - body interface{}, - response interface{}, - extraOptions []uhttp.RequestOption, + body any, + requestOptions []uhttp.RequestOption, params ...Vars, -) (*v2.RateLimitDescription, error) { +) (*http.Request, error) { // TODO(marcos): Refactor URLs so that we don't have to unescape. - u, err := url.PathUnescape(urlAddress.String()) + unescaped, err := url.PathUnescape(urlAddress.String()) if err != nil { return nil, err } - uri, err := url.Parse(u) + requestURL, err := url.Parse(unescaped) if err != nil { return nil, err } @@ -251,9 +263,9 @@ func (c *Client) doRequestWithOptions( if method == http.MethodGet && isNoCache(ctx) { options = append(options, uhttp.WithNoCache()) } - options = append(options, extraOptions...) + options = append(options, requestOptions...) - req, err := c.httpClient.NewRequest(ctx, method, uri, options...) + req, err := c.httpClient.NewRequest(ctx, method, requestURL, options...) if err != nil { return nil, err } @@ -267,8 +279,27 @@ func (c *Client) doRequestWithOptions( req.URL.RawQuery = query.Encode() } + req.Header.Set(cacheScopeHeader, req.URL.Host) + c.auth.Apply(req) + return req, nil +} + +func (c *Client) doRequest( + ctx context.Context, + urlAddress *url.URL, + method string, + body any, + response any, + requestOptions []uhttp.RequestOption, + params ...Vars, +) (*v2.RateLimitDescription, error) { + req, err := c.prepareRequest(ctx, urlAddress, method, body, requestOptions, params...) + if err != nil { + return nil, err + } + ratelimitData := &v2.RateLimitDescription{} doOptions := []uhttp.DoOption{uhttp.WithRatelimitData(ratelimitData)} // A nil response means the caller ignores the body, which may be empty (SCIM PATCH answers 204). @@ -289,8 +320,9 @@ func (c *Client) doRequestWithOptions( } var errorResponse struct { - Detail string `json:"detail"` - Message string `json:"message"` + ErrorCode string `json:"error_code"` + Detail string `json:"detail"` + Message string `json:"message"` } if err := parseJSON(resp.Body, &errorResponse); err != nil { return nil, err @@ -298,6 +330,7 @@ func (c *Client) doRequestWithOptions( return ratelimitData, &APIError{ StatusCode: resp.StatusCode, + ErrorCode: errorResponse.ErrorCode, Detail: errorResponse.Detail, Message: errorResponse.Message, Err: err, @@ -308,43 +341,15 @@ func (c *Client) doRequestNoResponse( ctx context.Context, urlAddress *url.URL, method string, - body interface{}, - response interface{}, + body any, + response any, params ...Vars, ) (*v2.RateLimitDescription, error) { - u, err := url.PathUnescape(urlAddress.String()) - if err != nil { - return nil, err - } - - uri, err := url.Parse(u) + req, err := c.prepareRequest(ctx, urlAddress, method, body, nil, params...) if err != nil { return nil, err } - options := []uhttp.RequestOption{ - uhttp.WithAcceptJSONHeader(), - } - if body != nil { - options = append(options, uhttp.WithJSONBody(body)) - } - - req, err := c.httpClient.NewRequest(ctx, method, uri, options...) - if err != nil { - return nil, err - } - - if len(params) > 0 { - query := url.Values{} - for _, param := range params { - param.Apply(&query) - } - - req.URL.RawQuery = query.Encode() - } - - c.auth.Apply(req) - ratelimitData := &v2.RateLimitDescription{} resp, err := c.httpClient.Do( req, diff --git a/pkg/databricks/unity_catalog.go b/pkg/databricks/unity_catalog.go new file mode 100644 index 00000000..45a2acee --- /dev/null +++ b/pkg/databricks/unity_catalog.go @@ -0,0 +1,308 @@ +package databricks + +import ( + "context" + "fmt" + "net/url" + "strings" + + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" + "go.uber.org/zap" +) + +// A metastore lives on the account plane; every securable below it is only +// reachable through a workspace host the metastore is assigned to. +const ( + // https://docs.databricks.com/api/account/accountmetastores/list + accountMetastoresEndpoint = "/api/2.0/accounts/%s/metastores" + + // https://docs.databricks.com/api/account/accountmetastoreassignments/list + accountMetastoreWorkspacesEndpoint = "/api/2.0/accounts/%s/metastores/%s/workspaces" + + // GET https://docs.databricks.com/api/workspace/grants/get + // PATCH https://docs.databricks.com/api/workspace/grants/update + unityCatalogPermissionsEndpoint = "/api/2.1/unity-catalog/permissions" + + // Path segments for the permissions endpoint, not interchangeable with the + // uppercase securable_type a listed securable carries. A metastore is + // addressed by its UUID: its name is rejected. + SecurableMetastore = "metastore" + + // The permissions endpoint rejects any max_results from 1 to 149 with a 400. + // Zero selects the paginated response that replaces the unpaginated form. + permissionsMaxResults uint = 0 +) + +// ListMetastores (GET /api/2.0/accounts/{account_id}/metastores). Not paginated. +func (c *Client) ListMetastores( + ctx context.Context, +) ( + []Metastore, + *v2.RateLimitDescription, + error, +) { + var res metastoresResponse + + u := c.accountBaseUrl.JoinPath(fmt.Sprintf(accountMetastoresEndpoint, c.accountId)) + ratelimitData, err := c.Get(ctx, u, &res) + if err != nil { + return nil, ratelimitData, fmt.Errorf("failed to list metastores: %w", err) + } + + return res.Metastores, ratelimitData, nil +} + +// ListMetastoreWorkspaces (GET /api/2.0/accounts/{account_id}/metastores/{metastore_id}/workspaces) +// returns the numeric workspace ids assigned to that metastore. +func (c *Client) ListMetastoreWorkspaces( + ctx context.Context, + metastoreID string, +) ( + []string, + *v2.RateLimitDescription, + error, +) { + var res metastoreWorkspacesResponse + + u := c.accountBaseUrl.JoinPath(fmt.Sprintf(accountMetastoreWorkspacesEndpoint, c.accountId, metastoreID)) + ratelimitData, err := c.Get(ctx, u, &res) + if err != nil { + return nil, ratelimitData, fmt.Errorf("failed to list workspaces for metastore %s: %w", metastoreID, err) + } + + ids := make([]string, 0, len(res.WorkspaceIDs)) + for _, id := range res.WorkspaceIDs { + if id == "" { + continue + } + ids = append(ids, id.String()) + } + + return ids, ratelimitData, nil +} + +// ListPermissions (GET /api/2.1/unity-catalog/permissions/{securable_type}/{full_name}). +// The principal query filters effectively, not literally: an inherited assignment +// comes back naming the group it is inherited through, so a caller must read every +// row's privileges and never match on the principal field. Empty lists every row. +func (c *Client) ListPermissions( + ctx context.Context, + workspaceId, securableType, fullName, principal, pageToken string, +) ( + []PrivilegeAssignment, + string, + *v2.RateLimitDescription, + error, +) { + return c.listPermissions(ctx, workspaceId, securableType, fullName, principal, pageToken, false) +} + +func (c *Client) listPermissions( + ctx context.Context, + workspaceId, securableType, fullName, principal, pageToken string, + uncached bool, +) ( + []PrivilegeAssignment, + string, + *v2.RateLimitDescription, + error, +) { + u, err := c.permissionsUrl(workspaceId, securableType, fullName) + if err != nil { + return nil, "", nil, err + } + + var res permissionsResponse + + maxResults := permissionsMaxResults + read := c.Get + if uncached { + read = c.GetUncached + } + ratelimitData, err := read(ctx, u, &res, &unityCatalogVars{ + principal: principal, + maxResults: &maxResults, + pageToken: pageToken, + }) + if err != nil { + if principal != "" { + return nil, "", ratelimitData, fmt.Errorf("failed to list permissions for principal %s on %s %s: %w", principal, securableType, fullName, err) + } + + return nil, "", ratelimitData, fmt.Errorf("failed to list permissions on %s %s: %w", securableType, fullName, err) + } + + return res.PrivilegeAssignments, res.NextPageToken, ratelimitData, nil +} + +// DrainPermissions walks ListPermissions to exhaustion for every principal, +// bypassing the response cache. Its only caller is the pre-read that decides +// whether a Grant still has to PATCH, and that decision is wrong when it is made +// against a page the cache has been holding since before the last change. +func (c *Client) DrainPermissions( + ctx context.Context, + workspaceId, securableType, fullName string, +) ( + []PrivilegeAssignment, + *v2.RateLimitDescription, + error, +) { + return drainPagesFrom(ctx, "", func(pageToken string) ([]PrivilegeAssignment, string, *v2.RateLimitDescription, error) { + return c.listPermissions(ctx, workspaceId, securableType, fullName, "", pageToken, true) + }) +} + +// ForEachUncachedPermissionsPage walks ListPermissions past the cache, without +// retaining the pages, for the pre-read that decides whether a Revoke still has +// to PATCH. See DrainPermissions. +func (c *Client) ForEachUncachedPermissionsPage( + ctx context.Context, + workspaceId, securableType, fullName, principal string, + visit func([]PrivilegeAssignment) (bool, error), +) (*v2.RateLimitDescription, error) { + return forEachPageFrom(ctx, "", func(pageToken string) ([]PrivilegeAssignment, string, *v2.RateLimitDescription, error) { + return c.listPermissions(ctx, workspaceId, securableType, fullName, principal, pageToken, true) + }, visit) +} + +// UpdatePermissions (PATCH /api/2.1/unity-catalog/permissions/{securable_type}/{full_name}). +// Databricks no-ops a privilege already held or already absent, so the echoed +// assignments are the only confirmation the change landed — and that echo is +// itself paginated, so a principal on a later page must not read as success. +func (c *Client) UpdatePermissions( + ctx context.Context, + workspaceId, securableType, fullName string, + changes []PermissionsChange, +) ( + []PrivilegeAssignment, + *v2.RateLimitDescription, + error, +) { + var assignments []PrivilegeAssignment + ratelimitData, err := c.updatePermissions(ctx, workspaceId, securableType, fullName, changes, + func(page []PrivilegeAssignment) (bool, error) { + assignments = append(assignments, page...) + return false, nil + }) + if err != nil { + return nil, ratelimitData, err + } + + return assignments, ratelimitData, nil +} + +// UpdatePermissionsUntil is UpdatePermissions for a caller that can stop early. +func (c *Client) UpdatePermissionsUntil( + ctx context.Context, + workspaceId, securableType, fullName string, + changes []PermissionsChange, + visit func([]PrivilegeAssignment) (bool, error), +) (*v2.RateLimitDescription, error) { + return c.updatePermissions(ctx, workspaceId, securableType, fullName, changes, visit) +} + +func (c *Client) updatePermissions( + ctx context.Context, + workspaceId, securableType, fullName string, + changes []PermissionsChange, + visit func([]PrivilegeAssignment) (bool, error), +) (*v2.RateLimitDescription, error) { + if len(changes) == 0 { + return nil, invalidArgument("no permission changes given for %s %s", securableType, fullName) + } + + for i, change := range changes { + if err := change.validate(); err != nil { + return nil, fmt.Errorf("invalid permission change %d for %s %s: %w", i, securableType, fullName, err) + } + } + + u, err := c.permissionsUrl(workspaceId, securableType, fullName) + if err != nil { + return nil, err + } + + payload := struct { + Changes []PermissionsChange `json:"changes"` + }{ + Changes: changes, + } + + var res permissionsResponse + + ratelimitData, err := c.Patch(ctx, u, payload, &res) + if err != nil { + return ratelimitData, fmt.Errorf("failed to update permissions on %s %s: %w", securableType, fullName, err) + } + + // The provisioning pre-read stored page one under the key the sync reads, + // and a PATCH does not invalidate it. Re-reading past the cache overwrites + // that entry with the post-change state. The write already landed, so a + // failed refresh must not fail it. A read that stops before the last page + // leaves those later pages until the next cycle; max_results 0 makes that rare. + _, _, refreshRateLimit, refreshErr := c.listPermissions(ctx, workspaceId, securableType, fullName, "", "", true) + if refreshErr != nil { + ctxzap.Extract(ctx).Debug("could not refresh the permissions cache after the write", + zap.String("securable_type", securableType), + zap.String("full_name", fullName), + zap.Error(refreshErr), + ) + } else if refreshRateLimit != nil { + ratelimitData = refreshRateLimit + } + + stop, err := visit(res.PrivilegeAssignments) + if err != nil || stop || res.NextPageToken == "" { + return ratelimitData, err + } + + // The PATCH body is already page one. The remaining pages confirm the change + // landed, so they are read past the cache for the same reason the pre-read is. + drainRateLimit, drainErr := forEachPageFrom(ctx, res.NextPageToken, + func(pageToken string) ([]PrivilegeAssignment, string, *v2.RateLimitDescription, error) { + return c.listPermissions(ctx, workspaceId, securableType, fullName, "", pageToken, true) + }, visit) + if drainRateLimit != nil { + ratelimitData = drainRateLimit + } + if drainErr != nil { + return ratelimitData, fmt.Errorf( + "failed to read the remaining permissions pages after updating %s %s: %w", + securableType, fullName, drainErr) + } + + return ratelimitData, nil +} + +// A name that does not survive URL assembly addresses a different securable, +// which on a PATCH writes someone else's permissions. +func (c *Client) permissionsUrl(workspaceId, securableType, fullName string) (*url.URL, error) { + if securableType == "" || fullName == "" { + return nil, invalidArgument("securable type and full name are required, got %q and %q", securableType, fullName) + } + + u := c.workspaceUrl(workspaceId).JoinPath(unityCatalogPermissionsEndpoint, securableType, fullName) + if !addressesSecurable(u, "/"+securableType+"/"+fullName) { + return nil, invalidArgument("cannot address %s %q: the name does not survive URL assembly", securableType, fullName) + } + + return u, nil +} + +// JoinPath reads each element as pre-escaped, so a bare "%" in a name drops the +// whole path, and the request layer's unescape round trip then turns a "?" into +// a query string and a "#" into a fragment. +func addressesSecurable(u *url.URL, suffix string) bool { + unescaped, err := url.PathUnescape(u.String()) + if err != nil { + return false + } + + sent, err := url.Parse(unescaped) + if err != nil { + return false + } + + return sent.RawQuery == "" && sent.Fragment == "" && strings.HasSuffix(sent.Path, suffix) +} diff --git a/pkg/databricks/unity_catalog_test.go b/pkg/databricks/unity_catalog_test.go new file mode 100644 index 00000000..3c5e66d8 --- /dev/null +++ b/pkg/databricks/unity_catalog_test.go @@ -0,0 +1,113 @@ +package databricks + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" +) + +// Databricks ignores a page token it cannot deserialize and silently restarts at +// page one, so a walk that accepts a repeated token re-reads the same page forever. +func TestDrainPagesFromStartsAtSeedAndRejectsARepeat(t *testing.T) { + t.Parallel() + + var tokens []string + _, _, err := drainPagesFrom(context.Background(), "page-2", func(pageToken string) ([]string, string, *v2.RateLimitDescription, error) { + tokens = append(tokens, pageToken) + if pageToken == "page-2" { + return []string{"from-seed"}, "page-2", nil, nil + } + return []string{"should-not-run"}, "", nil, nil + }) + if err == nil { + t.Fatal("drainPagesFrom returned nil error for a repeated token") + } + if len(tokens) != 1 || tokens[0] != "page-2" { + t.Fatalf("fetched tokens %v, want only the seed page-2", tokens) + } +} + +// The PATCH echo is itself paginated, so reading just the PATCH body makes a Revoke +// report success while the principal still holds the privilege on a later page. The +// walk continues from the PATCH's own token rather than re-fetching page one. +func TestUpdatePermissionsDrainsFromPatchToken(t *testing.T) { + t.Setenv("BATON_HTTP_CACHE_TTL", "0") + + var gets []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.Method { + case http.MethodPatch: + body, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read patch body: %v", err) + } + if len(body) == 0 { + t.Error("patch body is empty") + } + _ = json.NewEncoder(w).Encode(permissionsResponse{ + PrivilegeAssignments: []PrivilegeAssignment{{Principal: "alice", Privileges: []string{"SELECT"}}}, + NextPageToken: "page-2", + }) + case http.MethodGet: + token := r.URL.Query().Get("page_token") + gets = append(gets, token) + principal := "bob" + if token != "page-2" { + principal = "unexpected" + } + _ = json.NewEncoder(w).Encode(permissionsResponse{ + PrivilegeAssignments: []PrivilegeAssignment{{Principal: principal, Privileges: []string{"MODIFY"}}}, + }) + default: + http.Error(w, "unexpected method", http.StatusMethodNotAllowed) + } + })) + t.Cleanup(srv.Close) + + client := srv.Client() + client.Transport = rewriteHost(srv.Listener.Addr().String(), client.Transport) + c, err := NewClient(context.Background(), client, "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acc-1", srv.URL, &NoAuth{}, nil) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + + got, _, err := c.UpdatePermissions(context.Background(), "ws", SecurableMetastore, "ms-1", []PermissionsChange{{ + Principal: "alice", + Add: []string{"CREATE_CATALOG"}, + }}) + if err != nil { + t.Fatalf("UpdatePermissions: %v", err) + } + // Page one is re-read only to refresh the cache; confirmation comes from the PATCH body. + if len(gets) != 2 || gets[0] != "" || gets[1] != "page-2" { + t.Fatalf("GET page tokens %v, want the cache refresh then page-2", gets) + } + if len(got) != 2 || got[0].Principal != "alice" || got[1].Principal != "bob" { + t.Fatalf("assignments = %+v, want alice then bob", got) + } +} + +type hostRewriter struct { + host http.RoundTripper + addr string +} + +func rewriteHost(addr string, next http.RoundTripper) http.RoundTripper { + if next == nil { + next = http.DefaultTransport + } + return hostRewriter{host: next, addr: addr} +} + +func (h hostRewriter) RoundTrip(req *http.Request) (*http.Response, error) { + clone := req.Clone(req.Context()) + clone.URL.Scheme = "http" + clone.URL.Host = h.addr + return h.host.RoundTrip(clone) +} diff --git a/pkg/databricks/vars.go b/pkg/databricks/vars.go index cbf06e05..9e9c7cfe 100644 --- a/pkg/databricks/vars.go +++ b/pkg/databricks/vars.go @@ -3,6 +3,7 @@ package databricks import ( "fmt" "net/url" + "strconv" ) type Vars interface { @@ -167,3 +168,26 @@ func NewNameVars(name string, etag string) *NameVars { Etag: etag, } } + +// Databricks reads max_results=0 as "page at the server-configured length" +// while an absent max_results asks for the deprecated unpaginated response, so +// maxResults is a pointer: its presence is the signal. +type unityCatalogVars struct { + principal string + maxResults *uint + pageToken string +} + +func (v *unityCatalogVars) Apply(params *url.Values) { + if v.principal != "" { + params.Add("principal", v.principal) + } + + if v.maxResults != nil { + params.Add("max_results", strconv.FormatUint(uint64(*v.maxResults), 10)) + } + + if v.pageToken != "" { + params.Add("page_token", v.pageToken) + } +}