diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 6030d7f1..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 @@ -34,3 +40,61 @@ jobs: baton-entitlement: 'group:account/8c6f99ec-78ce-4654-8f92-e716b3dd67a7/group/66271399072731:member' baton-principal: '5346803201281760' baton-principal-type: 'user' + # Runs the incremental sync event feed against the local test server (test-server/), so it + # needs no Databricks account or secrets and also runs on pull requests from forks. + test-server: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + - name: Install Go + uses: actions/setup-go@v5 + with: + go-version-file: go.mod + - name: Build connector and test server + run: | + go build -o baton-databricks ./cmd/baton-databricks + go build -o /tmp/databricks-test-server ./test-server + - name: Start test server + run: | + /tmp/databricks-test-server --ca-out /tmp/test-server-ca.pem > test-server.log 2>&1 & + for i in $(seq 1 30); do + if curl -sf --cacert /tmp/test-server-ca.pem https://localhost:8765/health > /dev/null; then + echo "test server is up"; exit 0 + fi + sleep 1 + done + echo "test server did not come up in time"; cat test-server.log; exit 1 + - name: Run incremental sync against the test server + env: + SSL_CERT_FILE: /tmp/test-server-ca.pem + run: | + ./baton-databricks \ + --hostname localhost:8765 --account-hostname localhost:8765 \ + --account-id test-account \ + --databricks-client-id test-client --databricks-client-secret test-secret \ + --sql-warehouse-id wh-serverless \ + --event-feed true --event-feed-id databricks_audit_log \ + --event-feed-start-at "$(date -u -d '2 days ago' +%Y-%m-%dT%H:%M:%SZ)" \ + > events.jsonl + - name: Compare events with the expected output + run: | + normalize() { jq -c -S 'del(.occurredAt)' "$1" | sort; } + diff <(normalize test-server/testdata/expected-events.jsonl) <(normalize events.jsonl) + - name: Check the cold-start timeout message + env: + SSL_CERT_FILE: /tmp/test-server-ca.pem + run: | + if ./baton-databricks \ + --hostname localhost:8765 --account-hostname localhost:8765 \ + --account-id test-account \ + --databricks-client-id test-client --databricks-client-secret test-secret \ + --sql-warehouse-id wh-broken \ + --event-feed true --event-feed-id databricks_audit_log \ + --event-feed-start-at "$(date -u -d '2 days ago' +%Y-%m-%dT%H:%M:%SZ)" 2> err.log; then + echo "expected the event feed to fail with a broken warehouse"; exit 1 + fi + grep -q "has not started" err.log + - name: Show test server log on failure + if: failure() + run: cat test-server.log diff --git a/.gitignore b/.gitignore index b2118062..e9ac02eb 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,6 @@ # Dependency directories (remove the comment below to include it) # vendor/ dist/ + +# CA certificate written by the local test server (test-server/) +test-server-ca.pem diff --git a/README.md b/README.md index 371bbe5e..ec789267 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,90 @@ 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 + +By default, `baton-databricks` does a full resync of every resource on every run. +You can opt into an additional, cheap pathway that polls a Databricks audit log +between full syncs to pick up access changes early. Full syncs still run as the +correctness backstop. Incremental sync picks up added and removed group +memberships, but deleted users, groups and service principals are only caught +by the next full sync. + +Role assignments that Databricks stores in access-control rule sets (group roles +such as `roles/group.manager`, service principal roles, and the account +`marketplace.admin` role) are also only caught by the next full sync: the +connector does not read the audit events Databricks writes for rule-set +changes. A `roles/group.manager` change, for example, is logged as +`updateRuleSet` under the `accountsAccessControl` audit service, which the +connector does not query. + +**Incremental sync is enabled by setting `--sql-warehouse-id` (or +`BATON_SQL_WAREHOUSE_ID`) and disabled by leaving it unset or empty.** There is +no separate on/off flag. The incremental-sync setup (warehouse found, +`system.access.audit` queryable, at least one workspace in sync scope) is not +checked at startup, so a misconfiguration never blocks full syncs or +provisioning. Instead, each event-feed poll fails with an error naming the +problem until it is fixed. + +Incremental sync requires a reachable Account API to resolve audit events back +to synced resources, same as the rest of the connector. Audit events are only +resolved for workspaces in sync scope, so `--workspaces` and +`--databricks-exclude-workspaces` apply to incremental sync the same way they +apply to full syncs. + +### Requirements + +Incremental sync works only when all of the following are true: + +- **Unity Catalog.** The SQL warehouse is in a workspace enabled for Unity + Catalog, and the metastore is on Unity Catalog Privilege Model Version 1.0 + ([system tables + requirements](https://docs.databricks.com/aws/en/admin/system-tables/)). +- **System tables.** The connector queries the `system.access.audit` system + table, which Databricks currently offers as a Public Preview. +- **A SQL warehouse, preferably serverless.** Pass its ID with + `--sql-warehouse-id`. Any warehouse works, including the **Starter + Warehouse** Databricks creates in each workspace, so you don't need to create + one. +- **Grants for the service principal** the connector authenticates as, on top of + the Account admin role it already needs: + - `USE CATALOG` on the `system` catalog + - `USE SCHEMA` and `SELECT` on the `system.access` schema + - `CAN USE` on the SQL warehouse + + For example, run as a metastore admin, replacing `` with the + service principal's application (client) ID: + + ```sql + GRANT USE CATALOG ON CATALOG system TO ``; + GRANT USE SCHEMA, SELECT ON SCHEMA system.access TO ``; + ``` + +**We recommend a serverless SQL warehouse.** It starts in seconds, while a pro +or classic warehouse can take several minutes. Each statement gets up to 30 +seconds on the initial request, including warehouse startup, then up to 5 more +minutes of polling. A statement still unfinished after that is cancelled, so a +poll can take close to 6 minutes before it fails with a `DeadlineExceeded` +error and is retried on the next poll. Serverless SQL warehouses require the +Premium plan or above. + +**Incremental sync covers one cloud region.** `system.access.audit` holds +workspace-level events only for workspaces in the same cloud region as the +workspace that queries it; account-level events are global ([system tables +reference](https://docs.databricks.com/aws/en/admin/system-tables/)). The +connector finds the warehouse's workspace automatically by checking each +workspace the service principal can access, except those excluded with +`--databricks-exclude-workspaces`; the `--workspaces` allowlist does not limit +this search. It receives +incremental updates only for workspaces in that region. Workspaces in other +regions are still covered by full syncs. + +Audit log entries are delivered with a delay, so the connector reads them about +four hours after they happen. ## Group provisioning @@ -122,6 +219,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 @@ -175,6 +321,7 @@ Flags: -p, --provisioning This must be set in order for provisioning actions to be enabled ($BATON_PROVISIONING) --skip-entitlements-and-grants This must be set to skip syncing of entitlements and grants ($BATON_SKIP_ENTITLEMENTS_AND_GRANTS) --skip-full-sync This must be set to skip a full sync ($BATON_SKIP_FULL_SYNC) + --sql-warehouse-id string Setting this enables incremental sync; leaving it empty disables it. ID of the Databricks SQL warehouse used to query the system.access.audit log between full syncs, so access changes show up before the next full sync (deleted users, groups and service principals are still only caught by full syncs). The warehouse can be in any workspace not excluded by databricks-exclude-workspaces, in the same cloud region as the workspaces to update incrementally; Databricks only exposes workspace-level audit events for that region. The connector discovers the warehouse's workspace automatically. ($BATON_SQL_WAREHOUSE_ID) --storage-engine string The storage engine to use when opening the sync c1z file: sqlite or pebble. Defaults to pebble when unset. ($BATON_STORAGE_ENGINE) --sync-resource-types strings The resource type IDs to sync ($BATON_SYNC_RESOURCE_TYPES) --sync-resources strings The resource IDs to sync ($BATON_SYNC_RESOURCES) diff --git a/baton_capabilities.json b/baton_capabilities.json index 91985cec..3f948d7d 100644 --- a/baton_capabilities.json +++ b/baton_capabilities.json @@ -4,13 +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": { @@ -18,13 +41,109 @@ "displayName": "Group", "traits": [ "TRAIT_GROUP" + ], + "annotations": [ + { + "@type": "type.googleapis.com/c1.connector.v2.CapabilityPermissions", + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } ] }, "capabilities": [ "CAPABILITY_SYNC", + "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": { @@ -32,13 +151,36 @@ "displayName": "Role", "traits": [ "TRAIT_ROLE" + ], + "annotations": [ + { + "@type": "type.googleapis.com/c1.connector.v2.CapabilityPermissions", + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "Account admin" + } + ] + } ] }, "capabilities": [ "CAPABILITY_SYNC", + "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], - "permissions": {} + "permissions": { + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "Account admin" + } + ] + } }, { "resourceType": { @@ -46,13 +188,42 @@ "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": [ "CAPABILITY_SYNC", + "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], - "permissions": {} + "permissions": { + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } }, { "resourceType": { @@ -64,15 +235,36 @@ "annotations": [ { "@type": "type.googleapis.com/c1.connector.v2.SkipEntitlementsAndGrants" + }, + { + "@type": "type.googleapis.com/c1.connector.v2.CapabilityPermissions", + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "Account admin" + } + ] } ] }, "capabilities": [ "CAPABILITY_SYNC", + "CAPABILITY_TARGETED_SYNC", "CAPABILITY_ACCOUNT_PROVISIONING", "CAPABILITY_RESOURCE_DELETE" ], - "permissions": {} + "permissions": { + "permissions": [ + { + "permission": "scim" + }, + { + "permission": "Account admin" + } + ] + } }, { "resourceType": { @@ -80,20 +272,52 @@ "displayName": "Workspace", "traits": [ "TRAIT_GROUP" + ], + "annotations": [ + { + "@type": "type.googleapis.com/c1.connector.v2.CapabilityPermissions", + "permissions": [ + { + "permission": "provisioning" + }, + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } ] }, "capabilities": [ "CAPABILITY_SYNC", + "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], - "permissions": {} + "permissions": { + "permissions": [ + { + "permission": "provisioning" + }, + { + "permission": "access-management" + }, + { + "permission": "Account admin" + } + ] + } } ], "connectorCapabilities": [ "CAPABILITY_PROVISION", "CAPABILITY_SYNC", "CAPABILITY_ACCOUNT_PROVISIONING", - "CAPABILITY_RESOURCE_DELETE" + "CAPABILITY_RESOURCE_DELETE", + "CAPABILITY_TARGETED_SYNC", + "CAPABILITY_EVENT_FEED_V2", + "CAPABILITY_SERVICE_MODE_TARGETED_SYNC" ], "credentialDetails": { "capabilityAccountProvisioning": { diff --git a/config_schema.json b/config_schema.json index e387c1ed..65d90702 100644 --- a/config_schema.json +++ b/config_schema.json @@ -151,6 +151,12 @@ "displayName": "Exclude Workspaces", "description": "Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces.", "stringSliceField": {} + }, + { + "name": "sql-warehouse-id", + "displayName": "SQL Warehouse ID", + "description": "Setting this enables incremental sync; leaving it empty disables it. ID of the Databricks SQL warehouse used to query the system.access.audit log between full syncs, so access changes show up before the next full sync (deleted users, groups and service principals are still only caught by full syncs). The warehouse can be in any workspace not excluded by databricks-exclude-workspaces, in the same cloud region as the workspaces to update incrementally; Databricks only exposes workspace-level audit events for that region. The connector discovers the warehouse's workspace automatically.", + "stringField": {} } ], "constraints": [ diff --git a/docs/connector.mdx b/docs/connector.mdx index f57273e3..db250bf5 100644 --- a/docs/connector.mdx +++ b/docs/connector.mdx @@ -12,6 +12,7 @@ sidebarTitle: "Databricks" | :--- | :--- | :--- | | Accounts | | | | Groups | | | +| Metastores | | | | Roles | | | | Service principals | | | | Users | | | @@ -27,45 +28,157 @@ 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. +### Optional: faster updates between syncs + +By default, the Databricks connector picks up access changes on its regular sync schedule. You can optionally turn on incremental sync, which checks Databricks' activity log between full syncs so that changes like new group members show up in C1 sooner. Full syncs still run as usual and remain the source of truth. Removed group memberships are picked up incrementally too, but deleted users, groups, and service principals are only reflected after the next full sync. + +Role assignments that Databricks stores in access-control rule sets are also only reflected after the next full sync. These are group roles such as group manager, service principal roles, and the account `marketplace.admin` role. + +Incremental sync is turned on by providing the ID of a Databricks SQL warehouse (the **SQL Warehouse ID** setting, `BATON_SQL_WAREHOUSE_ID`) and turned off by leaving it empty — there is no separate on/off switch. Ask your connector operator or C1 support contact to enable it for you. + +#### Incremental sync requirements + +Your Databricks account must meet all of the following: + +- **Unity Catalog.** The SQL warehouse must be in a workspace enabled for Unity Catalog, and the metastore must be on Unity Catalog Privilege Model Version 1.0. Databricks enables most workspaces created after November 2023 for Unity Catalog automatically. +- **System tables.** The connector reads Databricks' audit log system table, `system.access.audit`. Databricks currently offers this table as a Public Preview. +- **A SQL warehouse, preferably serverless.** You don't need to create a new one: Databricks creates a **Starter Warehouse** in each workspace automatically, and any SQL warehouse works. See the recommendation below about the warehouse type. +- **Permissions for the connector's service principal**, in addition to the **Account admin** role it already needs: + - `USE CATALOG` on the `system` catalog + - `USE SCHEMA` and `SELECT` on the `system.access` schema + - **Can use** on the SQL warehouse + + +**We recommend a serverless SQL warehouse.** It starts in seconds, while a pro or classic warehouse can take several minutes to start. The connector waits up to about 6 minutes for the warehouse to start and each query to finish, so a slower start can delay updates until the next one. Serverless SQL warehouses require the Premium plan or above. + + + +**Incremental sync covers one cloud region.** Databricks keeps workspace activity in the audit log of the region where each workspace is deployed. The connector only receives incremental updates for workspaces in the same region as the SQL warehouse. Workspaces in other regions are still updated by full syncs. + + +#### Set up incremental sync + + +A Databricks user with the **Account admin** role, and the **Metastore admin** role or the right to grant access to the `system` catalog, must perform this task. + + + + + Choose a SQL warehouse in a Unity Catalog-enabled workspace, in the same region as the workspaces you want to update incrementally. In that workspace, go to **SQL Warehouses**, open the warehouse, and copy its ID. Check its type: if it isn't **Serverless**, switch it to serverless if your plan allows it (see the recommendation above). + + + On the **SQL Warehouses** page, click the kebab menu (three-dot icon) next to the warehouse, click **Permissions**, and give the connector's service principal the **Can use** permission. + + + In the same workspace, run the following in the SQL editor, replacing `` with the service principal's application (client) ID: + + ```sql + GRANT USE CATALOG ON CATALOG system TO ``; + GRANT USE SCHEMA, SELECT ON SCHEMA system.access TO ``; + ``` + + + Enter the warehouse ID in the connector's **SQL Warehouse ID** setting (`BATON_SQL_WAREHOUSE_ID`). + + + +The SQL warehouse can live in any workspace in that region that is not listed in **Exclude Workspaces**, including a workspace outside **Workspaces** — the connector finds it automatically. **Workspaces** and **Exclude Workspaces** also limit which workspaces incremental sync updates, as they do for full syncs. Changes reach C1 roughly four hours after they happen: Databricks delivers audit log entries with a delay, so the connector waits before reading them. + + +The incremental sync setup isn't checked when the connector starts, so a problem with it never blocks full syncs or provisioning. If something is missing, such as a permission or a warehouse that can't be found, every incremental update fails and the connector logs an error that names the cause until you fix it. + + ## Gather Databricks credentials Configuring the connector requires you to pass in credentials generated in Databricks. Gather these credentials before you move on. -A user with the **Account admin** role in each Databricks workspace you want to sync must perform this task. +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 - Ensure that your Databricks user account has the **Account admin** role and is assigned to each Databricks workspace you want to sync to C1. + Log in to the Databricks account console with a user that has the **Account admin** role. In the Databricks account console, open the menu that appears next to your username in the upper right corner. - Carefully and copy and save account ID. + Carefully copy and save the account ID. -### 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`. + + + 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**. - Carefully copy and save the OAuth client ID and secret. + 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 @@ -119,10 +232,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**. @@ -197,7 +319,7 @@ stringData: BATON_CLIENT_ID: BATON_CLIENT_SECRET: - # Databricks account ID (required for both credential options below) + # Databricks account ID (required) BATON_ACCOUNT_ID: # Databricks credentials: OAuth @@ -206,13 +328,18 @@ 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: # Optional: exclude specific workspaces from the sync # (workspace name, deployment name, or numeric ID). # Mutually exclusive with BATON_WORKSPACES — set one or the other, never both. - BATON_DATABRICKS_EXCLUDE_WORKSPACES: + # BATON_DATABRICKS_EXCLUDE_WORKSPACES: + + # Optional: setting this enables incremental sync; leave it unset to disable it. + # ID of the SQL warehouse used to query the system.access.audit log between full syncs. + # BATON_SQL_WAREHOUSE_ID: # Optional: include if you want C1 to provision access using this connector BATON_PROVISIONING: true @@ -275,4 +402,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/config/conf.gen.go b/pkg/config/conf.gen.go index dd2a3bba..7094ed29 100644 --- a/pkg/config/conf.gen.go +++ b/pkg/config/conf.gen.go @@ -12,6 +12,7 @@ type Databricks struct { Workspaces []string `mapstructure:"workspaces"` BaseUrl string `mapstructure:"base-url"` DatabricksExcludeWorkspaces []string `mapstructure:"databricks-exclude-workspaces"` + SqlWarehouseId string `mapstructure:"sql-warehouse-id"` } func (c *Databricks) findFieldByTag(tagValue string) (any, bool) { diff --git a/pkg/config/config.go b/pkg/config/config.go index 35afe969..900f0e8d 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -54,6 +54,20 @@ var ( field.WithDescription("Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces."), field.WithDisplayName("Exclude Workspaces"), ) + // SQLWarehouseIDField doubles as the incremental-sync switch: set enables it, empty disables it. + // system.access.audit only has workspace-level events from the warehouse's region, so the + // warehouse must share a region with the workspaces incremental sync should cover. + SQLWarehouseIDField = field.StringField( + "sql-warehouse-id", + field.WithDescription( + "Setting this enables incremental sync; leaving it empty disables it. "+ + "ID of the Databricks SQL warehouse used to query the system.access.audit log between full syncs, "+ + "so access changes show up before the next full sync (deleted users, groups and service principals are still only caught by full syncs). "+ + "The warehouse can be in any workspace not excluded by databricks-exclude-workspaces, in the same cloud region as the workspaces to update incrementally; "+ + "Databricks only exposes workspace-level audit events for that region. The connector discovers the warehouse's workspace automatically.", + ), + field.WithDisplayName("SQL Warehouse ID"), + ) configFields = []field.SchemaField{ AccountHostnameField, AccountIdField, @@ -63,6 +77,7 @@ var ( WorkspacesField, BaseURLField, ExcludeWorkspacesField, + SQLWarehouseIDField, } ) diff --git a/pkg/connector/account.go b/pkg/connector/account.go index 3684ed45..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}, ) } @@ -219,6 +220,20 @@ func (a *accountBuilder) Grant(ctx context.Context, principal *v2.Resource, enti return nil, nil } +// Get returns the singleton account resource, used to re-sync it after a RESOURCE_CHANGE event. +func (a *accountBuilder) Get(ctx context.Context, resourceId *v2.ResourceId, parentResourceId *v2.ResourceId) (*v2.Resource, annotations.Annotations, error) { + resource, err := a.accountResource(ctx) + if err != nil { + return nil, nil, err + } + + ctxzap.Extract(ctx).Debug("databricks-connector: Get: refreshed account", + zap.String("resource_id", resourceId.Resource), + zap.String("parent_resource_id", parentResourceId.GetResource()), + ) + return resource, nil, nil +} + func (a *accountBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotations.Annotations, error) { l := ctxzap.Extract(ctx) diff --git a/pkg/connector/audit_event_feed.go b/pkg/connector/audit_event_feed.go new file mode 100644 index 00000000..9298246e --- /dev/null +++ b/pkg/connector/audit_event_feed.go @@ -0,0 +1,963 @@ +package connector + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "net/http" + "sort" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + "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/pagination" + ent "github.com/conductorone/baton-sdk/pkg/types/entitlement" + "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" + "google.golang.org/protobuf/types/known/timestamppb" +) + +const ( + auditEventFeedId = "databricks_audit_log" + + // Trail the watermark by this much instead of the newest event seen, since slower-indexing + // areas of the audited system could otherwise have events skipped permanently. + auditLogTrailingLag = 4 * time.Hour + + // The first poll looks back this far since there's no prior watermark yet. Must be >= + // auditLogTrailingLag, or the lagCutoff guard in ListEvents skips every poll until it is. + auditLogLookback = auditLogTrailingLag + + auditLogPageLimit = 1000 + + // auditLogRetention caps how far back the feed reads; C1 doesn't request event-feed data + // older than ~30 days, so an older cursor is treated as stale. + auditLogRetention = 30 * 24 * time.Hour + + auditServiceAccounts = "accounts" + auditServiceWorkspace = "workspace" +) + +// auditActionMapping describes what an audit log action_name affects: an optional primary +// resource (resourceType + the request_params key holding its native ID), an optional +// account-scoped role, and/or optional workspace-scoped roles/entitlements. grant, when set, +// means the row is an atomic entitlement assignment/revocation rather than a generic change. +type auditActionMapping struct { + resourceType *v2.ResourceType + idParam string + // workspaceIdParam names a request_params key holding the workspace id, for account-level + // actions on a workspace, whose workspace_id column is 0. + workspaceIdParam string + accountRole string + roleNames []string + grant *auditGrantMapping +} + +// auditGrantMapping names the principal and direction of a grant-mapped action. The +// entitlement's resource is resourceType/idParam on the enclosing auditActionMapping. +type auditGrantMapping struct { + entitlement string + principalParam string + revoke bool +} + +// auditActionKey identifies an audit log action by (service_name, action_name), since +// action_name alone is ambiguous across services (e.g. "delete" also means "cluster terminated"). +type auditActionKey struct { + Service string + Action string +} + +// auditLogActions maps (service_name, action_name) pairs to the resources they affect, per +// https://docs.databricks.com/aws/en/admin/account-settings/audit-logs. +var auditLogActions = map[auditActionKey]auditActionMapping{ + {auditServiceAccounts, "createGroup"}: {resourceType: groupResourceType, idParam: "targetGroupId"}, + // Databricks documents targetUserId as the added/removed member's id on all four of these + // actions, singular and plural alike: https://docs.databricks.com/aws/en/admin/account-settings/audit-logs. + {auditServiceAccounts, "addPrincipalToGroup"}: { + resourceType: groupResourceType, idParam: "targetGroupId", + grant: &auditGrantMapping{entitlement: groupMemberEntitlement, principalParam: "targetUserId"}, + }, + {auditServiceAccounts, "removePrincipalFromGroup"}: { + resourceType: groupResourceType, idParam: "targetGroupId", + grant: &auditGrantMapping{entitlement: groupMemberEntitlement, principalParam: "targetUserId", revoke: true}, + }, + {auditServiceAccounts, "addPrincipalsToGroup"}: { + resourceType: groupResourceType, idParam: "targetGroupId", + grant: &auditGrantMapping{entitlement: groupMemberEntitlement, principalParam: "targetUserId"}, + }, + {auditServiceAccounts, "removePrincipalsFromGroup"}: { + resourceType: groupResourceType, idParam: "targetGroupId", + grant: &auditGrantMapping{entitlement: groupMemberEntitlement, principalParam: "targetUserId", revoke: true}, + }, + {auditServiceAccounts, "removeGroup"}: {resourceType: groupResourceType, idParam: "targetGroupId"}, + {auditServiceAccounts, "updateGroup"}: { + resourceType: groupResourceType, idParam: "targetGroupId", + roleNames: []string{ClusterCreateRole, InstancePoolCreateRole}, + }, + // "add"/"delete" are the real user-lifecycle events; deleteUser is a parameterless PII purge. + {auditServiceAccounts, "add"}: {resourceType: userResourceType, idParam: "targetUserId"}, + {auditServiceAccounts, "updateUser"}: { + resourceType: userResourceType, idParam: "targetUserId", + roleNames: []string{ClusterCreateRole, InstancePoolCreateRole}, + }, + {auditServiceAccounts, "delete"}: {resourceType: userResourceType, idParam: "targetUserId"}, + {auditServiceAccounts, "activateUser"}: {resourceType: userResourceType, idParam: "targetUserId"}, + {auditServiceAccounts, "deactivateUser"}: {resourceType: userResourceType, idParam: "targetUserId"}, + // createServicePrincipal/updateServicePrincipal/deleteServicePrincipal: commented out, not + // real action_name values. Not found anywhere in the official audit log reference (AWS/GCP/ + // Azure) — the only documented service-principal action is changeServicePrincipalAcls. + // {auditServiceAccounts, "createServicePrincipal"}: {resourceType: servicePrincipalResourceType, idParam: "targetServicePrincipalId"}, + // {auditServiceAccounts, "updateServicePrincipal"}: { + // resourceType: servicePrincipalResourceType, idParam: "targetServicePrincipalId", + // roleNames: []string{ClusterCreateRole, InstancePoolCreateRole}, + // }, + // {auditServiceAccounts, "deleteServicePrincipal"}: {resourceType: servicePrincipalResourceType, idParam: "targetServicePrincipalId"}, + {auditServiceAccounts, "changeDatabricksWorkspaceAcl"}: {resourceType: workspaceResourceType, roleNames: []string{WorkspaceAccessRole}}, + {auditServiceAccounts, "changeDatabricksSqlAcl"}: {roleNames: []string{SQLAccessRole}}, + // The reference describes setAdmin as account admin but lists it with the workspace-level + // events, so the account-admin role is refreshed too in case it is. + {auditServiceAccounts, "setAdmin"}: {resourceType: userResourceType, idParam: "targetUserId", accountRole: AccountAdminRole}, + {auditServiceAccounts, "removeAdmin"}: {resourceType: userResourceType, idParam: "targetUserId"}, + {auditServiceAccounts, "setAccountAdmin"}: {resourceType: userResourceType, idParam: "targetUserId", accountRole: AccountAdminRole}, + {auditServiceAccounts, "removeAccountAdmin"}: {resourceType: userResourceType, idParam: "targetUserId", accountRole: AccountAdminRole}, + // Workspace membership (permission assignments). The GCP reference only lists + // updatePermissionAssignment; AWS and Azure list all three. + {auditServiceWorkspace, "addPermissionAssignment"}: {resourceType: workspaceResourceType, workspaceIdParam: "workspace_id"}, + {auditServiceWorkspace, "updatePermissionAssignment"}: {resourceType: workspaceResourceType, workspaceIdParam: "workspace_id"}, + {auditServiceWorkspace, "deletePermissionAssignment"}: {resourceType: workspaceResourceType, workspaceIdParam: "workspace_id"}, +} + +func auditLogActionNames() []string { + seen := make(map[string]struct{}, len(auditLogActions)) + for key := range auditLogActions { + seen[key.Action] = struct{}{} + } + names := make([]string, 0, len(seen)) + for name := range seen { + names = append(names, name) + } + sort.Strings(names) + return names +} + +func auditLogServiceNames() []string { + seen := make(map[string]struct{}, len(auditLogActions)) + for key := range auditLogActions { + seen[key.Service] = struct{}{} + } + names := make([]string, 0, len(seen)) + for name := range seen { + names = append(names, name) + } + sort.Strings(names) + return names +} + +// eventPageCursor is the opaque state persisted between ListEvents calls. (StartAt, +// StartAfterEventID) form a composite boundary: unprocessed rows are those with +// event_time > StartAt, or event_time == StartAt AND event_id > StartAfterEventID. This +// keeps the boundary well-ordered even when many rows share the same event_time. +type eventPageCursor struct { + StartAt time.Time `json:"start_at"` + StartAfterEventID string `json:"start_after_event_id"` +} + +func encodeEventCursor(c eventPageCursor) (string, error) { + b, err := json.Marshal(c) + if err != nil { + return "", fmt.Errorf("failed to marshal event cursor: %w", err) + } + return base64.StdEncoding.EncodeToString(b), nil +} + +// decodeEventCursor returns a zero-value cursor (self-healing to the lookback default) when +// missing, corrupt, or stale. All three are skip-and-continue, not connector bugs, so all three +// log at Debug, not Warn (which C1 surfaces as an incident). +func decodeEventCursor(ctx context.Context, s string, now time.Time) eventPageCursor { + l := ctxzap.Extract(ctx) + + if s == "" { + return eventPageCursor{} + } + + raw, err := base64.StdEncoding.DecodeString(s) + if err != nil { + l.Debug("databricks-connector: corrupt event cursor, resetting to lookback default", zap.Error(err)) + return eventPageCursor{} + } + + var c eventPageCursor + if err := json.Unmarshal(raw, &c); err != nil { + l.Debug("databricks-connector: corrupt event cursor, resetting to lookback default", zap.Error(err)) + return eventPageCursor{} + } + + if !c.StartAt.IsZero() && now.Sub(c.StartAt) > auditLogRetention { + l.Debug("databricks-connector: event cursor is older than the audit log retention window, resetting to lookback default", + zap.Time("cursor_start_at", c.StartAt), + ) + return eventPageCursor{} + } + + return c +} + +type auditLogRow struct { + EventID string + EventTime time.Time + WorkspaceID int64 + ActionName string + ServiceName string + RequestParams map[string]string +} + +type auditEventFeed struct { + client *databricks.Client + workspaces []string + // sqlWarehouseID enables the feed when non-empty; empty means incremental sync is off. + sqlWarehouseID string + + // queryWorkspaceID caches which workspace hosts sqlWarehouseID across polls. It's a discovered value, not provided by configs. + queryWorkspaceMu sync.Mutex + queryWorkspaceID string + + skippedLookupFailures atomic.Int64 +} + +func newAuditEventFeed( + client *databricks.Client, + workspaces []string, + sqlWarehouseID string, +) *auditEventFeed { + return &auditEventFeed{ + client: client, + workspaces: workspaces, + sqlWarehouseID: sqlWarehouseID, + } +} + +// resolveQueryWorkspaceID returns which workspace hosts f.sqlWarehouseID, resolving it +// once via resolveWarehouseWorkspace rather than probing on every poll. allWorkspaces +// must be every workspace in the account, not filtered by the --workspaces allowlist: +// the warehouse can live in a workspace outside that scope (in the same region). +func (f *auditEventFeed) resolveQueryWorkspaceID(ctx context.Context, allWorkspaces []databricks.Workspace) (string, *v2.RateLimitDescription, error) { + f.queryWorkspaceMu.Lock() + defer f.queryWorkspaceMu.Unlock() + + l := ctxzap.Extract(ctx) + if f.queryWorkspaceID != "" { + l.Debug("databricks-connector: resolveQueryWorkspaceID: using cached sql warehouse workspace", + zap.String("warehouse_id", f.sqlWarehouseID), + zap.String("workspace_id", f.queryWorkspaceID), + ) + return f.queryWorkspaceID, nil, nil + } + + id, rateLimit, err := resolveWarehouseWorkspace(ctx, f.client, allWorkspaces, f.sqlWarehouseID) + if err != nil { + return "", rateLimit, err + } + + l.Debug("databricks-connector: resolveQueryWorkspaceID: discovered sql warehouse workspace", + zap.String("warehouse_id", f.sqlWarehouseID), + zap.String("workspace_id", id), + zap.Int("candidate_workspace_count", len(allWorkspaces)), + ) + f.queryWorkspaceID = id + return id, rateLimit, nil +} + +// EventFeedMetadata is registered unconditionally; whether sql-warehouse-id is set gates +// behavior inside ListEvents instead, to avoid confusing "feed not found" errors when it's off. +func (f *auditEventFeed) EventFeedMetadata(_ context.Context) *v2.EventFeedMetadata { + return &v2.EventFeedMetadata{ + Id: auditEventFeedId, + SupportedEventTypes: []v2.EventType{ + v2.EventType_EVENT_TYPE_RESOURCE_CHANGE, + v2.EventType_EVENT_TYPE_CREATE_GRANT, + v2.EventType_EVENT_TYPE_CREATE_REVOKE, + }, + } +} + +// bootstrapStartAt picks the watermark for a feed with no prior cursor. A zero-valued +// earliestEvent (unset) is treated like nil; a real one is clamped to auditLogRetention so a +// very old value can't trigger a full-table scan against system.access.audit. +func bootstrapStartAt(now time.Time, earliestEvent *timestamppb.Timestamp) time.Time { + if earliestEvent == nil || earliestEvent.AsTime().Unix() == 0 { + return now.Add(-auditLogLookback) + } + + start := earliestEvent.AsTime() + if floor := now.Add(-auditLogRetention); start.Before(floor) { + return floor + } + return start +} + +func (f *auditEventFeed) ListEvents( + ctx context.Context, + earliestEvent *timestamppb.Timestamp, + pToken *pagination.StreamToken, +) ([]*v2.Event, *pagination.StreamState, annotations.Annotations, error) { + l := ctxzap.Extract(ctx) + annos := annotations.Annotations{} + + if f.sqlWarehouseID == "" { + return nil, &pagination.StreamState{}, nil, nil + } + + now := time.Now() + cursor := decodeEventCursor(ctx, pToken.Cursor, now) + + if cursor.StartAt.IsZero() { + cursor = eventPageCursor{StartAt: bootstrapStartAt(now, earliestEvent)} + } + + // Rows are only ever fetched up to lagCutoff (see queryAuditLog), so slow-indexing rows + // have auditLogTrailingLag to appear before we consider that window done. If the cursor + // is already past it (e.g. right after bootstrap, where the lookback is narrower than the + // lag), there's nothing to query yet. + lagCutoff := now.Add(-auditLogTrailingLag) + l.Debug("databricks-connector: ListEvents: poll started", + zap.Bool("resumed_from_cursor", pToken.Cursor != ""), + zap.Time("window_start", cursor.StartAt), + zap.String("window_start_after_event_id", cursor.StartAfterEventID), + zap.Time("window_end", lagCutoff), + ) + if !cursor.StartAt.Before(lagCutoff) { + l.Debug("databricks-connector: ListEvents: nothing to query yet, window start is within the trailing lag", + zap.Time("window_start", cursor.StartAt), + zap.Time("window_end", lagCutoff), + ) + encoded, err := encodeEventCursor(cursor) + if err != nil { + return nil, nil, nil, fmt.Errorf("databricks-connector: failed to encode event cursor: %w", err) + } + return nil, &pagination.StreamState{Cursor: encoded, HasMore: false}, annos, nil + } + + // ListWorkspaces already applies --databricks-exclude-workspaces. allWorkspaces (not + // narrowed by --workspaces) locates the query warehouse; scopedWorkspaces decides which + // workspaces' audit rows get resolved, so events never surface a workspace a full sync skips. + allWorkspaces, rateLimit, err := f.client.ListWorkspaces(ctx) + if rateLimit != nil { + annos.WithRateLimiting(rateLimit) + } + if err != nil { + return nil, nil, annos, fmt.Errorf("databricks-connector: failed to list workspaces: %w", err) + } + scopedWorkspaces := filterConfiguredWorkspaces(allWorkspaces, f.workspaces) + l.Debug("databricks-connector: ListEvents: listed workspaces", + zap.Int("workspace_count", len(allWorkspaces)), + zap.Int("in_scope_workspace_count", len(scopedWorkspaces)), + ) + if len(scopedWorkspaces) == 0 { + return nil, nil, annos, uhttp.WrapErrors(codes.FailedPrecondition, "databricks-connector: no workspace available to query system.access.audit") + } + + workspaceLookup := make(map[int64]string, len(scopedWorkspaces)) + for _, w := range scopedWorkspaces { + workspaceLookup[int64(w.ID)] = w.DeploymentName + } + + queryWorkspaceId, rateLimit, err := f.resolveQueryWorkspaceID(ctx, allWorkspaces) + if rateLimit != nil { + annos.WithRateLimiting(rateLimit) + } + if err != nil { + return nil, nil, annos, err + } + + rows, rawRowCount, lastRawBoundary, rateLimit, err := f.queryAuditLog(ctx, queryWorkspaceId, cursor, lagCutoff) + if err != nil { + if rateLimit != nil { + annos.WithRateLimiting(rateLimit) + } + return nil, nil, annos, fmt.Errorf("databricks-connector: failed to query audit log: %w", err) + } + + if rateLimit != nil { + annos.WithRateLimiting(rateLimit) + } + + var events []*v2.Event + var unmappedRows, skippedLookupRows, resourceChangeEvents, grantEvents, revokeEvents int + principalTypes := principalTypeCache{} + for _, row := range rows { + affected, rowRateLimit, err := mapAuditRowToResource(ctx, f.client, principalTypes, row, f.client.GetAccountId(), workspaceLookup) + if rowRateLimit != nil { + annos.WithRateLimiting(rowRateLimit) + } + // Unmappable rows come back as no error, so this is a real lookup failure. Transient + // ones are returned so the next poll retries the row; others would block the cursor + // on every poll, so they're skipped and left to the next full sync. + if err != nil { + if isRetryableLookupError(err) { + return nil, nil, annos, fmt.Errorf("databricks-connector: failed to map audit row %s (%s): %w", row.EventID, row.ActionName, err) + } + if count := f.skippedLookupFailures.Add(1); count == 1 || count == 10 || count == 100 || count%1000 == 0 { + l.Warn("databricks-connector: skipping audit row, principal lookup failed", + zap.String("action_name", row.ActionName), + zap.String("event_id", row.EventID), + zap.Int64("total_occurrences", count), + zap.Error(err), + ) + } + skippedLookupRows++ + continue + } + if len(affected) == 0 { + unmappedRows++ + l.Debug("databricks-connector: skipping audit row with no resource mapping", + zap.String("action_name", row.ActionName), + zap.String("event_id", row.EventID), + ) + continue + } + + l.Debug("databricks-connector: ListEvents: mapped audit row to events", + zap.String("event_id", row.EventID), + zap.String("service_name", row.ServiceName), + zap.String("action_name", row.ActionName), + zap.Int("event_count", len(affected)), + ) + for i, a := range affected { + event := &v2.Event{ + Id: fmt.Sprintf("%s/%d", row.EventID, i), + OccurredAt: timestamppb.New(row.EventTime), + } + switch { + case a.grant != nil && a.grant.revoke: + revokeEvents++ + event.Event = &v2.Event_CreateRevokeEvent{ + CreateRevokeEvent: &v2.CreateRevokeEvent{Entitlement: a.grant.entitlement, Principal: a.grant.principal}, + } + case a.grant != nil: + grantEvents++ + event.Event = &v2.Event_CreateGrantEvent{ + CreateGrantEvent: &v2.CreateGrantEvent{Entitlement: a.grant.entitlement, Principal: a.grant.principal}, + } + default: + resourceChangeEvents++ + event.Event = &v2.Event_ResourceChangeEvent{ + ResourceChangeEvent: &v2.ResourceChangeEvent{ResourceId: a.resourceId, ParentResourceId: a.parentResourceId}, + } + } + events = append(events, event) + } + } + + hasMore := rawRowCount >= auditLogPageLimit + nextCursor, advanced := advanceEventCursor(cursor, rows, lastRawBoundary, hasMore, lagCutoff) + if !advanced { + hasMore = false + } + + encoded, err := encodeEventCursor(nextCursor) + if err != nil { + return nil, nil, annos, fmt.Errorf("databricks-connector: failed to encode event cursor: %w", err) + } + + l.Debug("databricks-connector: ListEvents: poll finished", + zap.Int("audit_rows_fetched", rawRowCount), + zap.Int("audit_rows_parsed", len(rows)), + zap.Int("audit_rows_unmapped", unmappedRows), + zap.Int("audit_rows_skipped_lookup_failure", skippedLookupRows), + zap.Int("events_total", len(events)), + zap.Int("events_resource_change", resourceChangeEvents), + zap.Int("events_create_grant", grantEvents), + zap.Int("events_create_revoke", revokeEvents), + zap.Bool("has_more", hasMore), + zap.Time("next_window_start", nextCursor.StartAt), + zap.String("next_window_start_after_event_id", nextCursor.StartAfterEventID), + ) + return events, &pagination.StreamState{Cursor: encoded, HasMore: hasMore}, annos, nil +} + +func isRetryableLookupError(err error) bool { + switch status.Code(err) { + case codes.Unavailable, codes.ResourceExhausted, codes.DeadlineExceeded, codes.Internal, codes.Unknown, codes.Canceled: + return true + default: + return false + } +} + +// advanceEventCursor advances to the last row processed while more of the page remains, or +// to lagCutoff once drained: rows are already bounded by lagCutoff (see queryAuditLog), so a +// drained page proves nothing else exists up to that point. Prefers lastRawBoundary (from the +// raw page, so a page where every row fails to parse still advances), falling back to the last +// successfully parsed row. advanced is false only if neither is available, so the caller can +// avoid reporting HasMore with an unchanged cursor. +func advanceEventCursor(cursor eventPageCursor, rows []auditLogRow, lastRawBoundary eventPageCursor, hasMore bool, lagCutoff time.Time) (eventPageCursor, bool) { + if hasMore { + if !lastRawBoundary.StartAt.IsZero() { + return lastRawBoundary, true + } + if len(rows) > 0 { + last := rows[len(rows)-1] + return eventPageCursor{StartAt: last.EventTime, StartAfterEventID: last.EventID}, true + } + return cursor, false + } + return eventPageCursor{StartAt: lagCutoff}, true +} + +// affectedResource is either one resource a mapped audit row's action changed (RESOURCE_CHANGE), +// or, when grant is set, the entitlement/principal it assigned or revoked (CREATE_GRANT/CREATE_REVOKE). +type affectedResource struct { + resourceId *v2.ResourceId + parentResourceId *v2.ResourceId + grant *affectedGrant +} + +type affectedGrant struct { + entitlement *v2.Entitlement + principal *v2.Resource + revoke bool +} + +// mapAuditRowToResource maps an audit row to every Baton resource its action affects, skipping +// anything unresolvable. Principals are always account-parented, mirroring full sync, regardless +// of the workspace the audit row occurred in. +func mapAuditRowToResource( + ctx context.Context, + client *databricks.Client, + principalTypes principalTypeCache, + row auditLogRow, + accountId string, + workspaceLookup map[int64]string, +) ([]affectedResource, *v2.RateLimitDescription, error) { + mapping, ok := auditLogActions[auditActionKey{Service: row.ServiceName, Action: row.ActionName}] + if !ok { + return nil, nil, nil + } + + accountParent := &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: accountId} + + workspaceId := row.WorkspaceID + if workspaceId == 0 && mapping.workspaceIdParam != "" { + // An unparseable value leaves workspaceId at 0, so the row is skipped below. + if parsed, err := strconv.ParseInt(row.RequestParams[mapping.workspaceIdParam], 10, 64); err == nil { + workspaceId = parsed + } + } + + var workspaceParent *v2.ResourceId + if workspaceId != 0 { + deploymentName, found := workspaceLookup[workspaceId] + if !found { + return nil, nil, nil + } + workspaceParent = &v2.ResourceId{ResourceType: workspaceResourceType.Id, Resource: deploymentName} + } + + var affected []affectedResource + + switch { + case mapping.resourceType == workspaceResourceType: + if workspaceParent == nil { + return nil, nil, nil + } + affected = append(affected, affectedResource{resourceId: workspaceParent, parentResourceId: accountParent}) + case mapping.resourceType != nil: + nativeId, ok := row.RequestParams[mapping.idParam] + if !ok || nativeId == "" { + return nil, nil, nil + } + + resourceId := &v2.ResourceId{ResourceType: mapping.resourceType.Id, Resource: nativeId} + if mapping.resourceType == groupResourceType { + resourceId.Resource = groupResourceId(ctx, nativeId, accountParent) + } + + if mapping.grant != nil { + g, rateLimit, err := buildAffectedGrant(ctx, client, principalTypes, row, mapping.grant, resourceId, accountParent) + if err != nil { + return nil, rateLimit, fmt.Errorf("databricks-connector: failed to resolve grant principal for event %s: %w", row.EventID, err) + } + if g == nil { + // No resolvable principal: still refresh the group's grants. + return []affectedResource{{resourceId: resourceId, parentResourceId: accountParent}}, rateLimit, nil + } + return []affectedResource{{grant: g}}, rateLimit, nil + } + + affected = append(affected, affectedResource{resourceId: resourceId, parentResourceId: accountParent}) + } + + if mapping.accountRole != "" { + affected = append(affected, affectedResource{ + resourceId: &v2.ResourceId{ResourceType: roleResourceType.Id, Resource: roleResourceId(mapping.accountRole, accountParent)}, + parentResourceId: accountParent, + }) + } + + if workspaceParent != nil { + for _, roleName := range mapping.roleNames { + affected = append(affected, affectedResource{ + resourceId: &v2.ResourceId{ResourceType: roleResourceType.Id, Resource: roleResourceId(roleName, workspaceParent)}, + parentResourceId: workspaceParent, + }) + } + } + + return affected, nil, nil +} + +// buildAffectedGrant builds the entitlement (on entitlementResourceId/parent) and principal a +// grant-mapped row names, skipping rows missing the principal's native id. +func buildAffectedGrant( + ctx context.Context, + client *databricks.Client, + principalTypes principalTypeCache, + row auditLogRow, + gm *auditGrantMapping, + entitlementResourceId, parent *v2.ResourceId, +) (*affectedGrant, *v2.RateLimitDescription, error) { + principalNativeId, ok := row.RequestParams[gm.principalParam] + if !ok || principalNativeId == "" { + return nil, nil, nil + } + + principalType, rateLimit, err := resolvePrincipalType(ctx, client, principalTypes, principalNativeId) + if err != nil { + return nil, rateLimit, err + } + if principalType == nil { + return nil, rateLimit, nil + } + + principalResourceId := principalNativeId + if principalType == groupResourceType { + principalResourceId = groupResourceId(ctx, principalNativeId, parent) + } + + entitlementResource := &v2.Resource{Id: entitlementResourceId, ParentResourceId: parent} + return &affectedGrant{ + entitlement: ent.NewAssignmentEntitlement(entitlementResource, gm.entitlement), + principal: &v2.Resource{Id: &v2.ResourceId{ResourceType: principalType.Id, Resource: principalResourceId}}, + revoke: gm.revoke, + }, rateLimit, nil +} + +// principalTypeCache memoizes resolvePrincipalType for one ListEvents call, so a principal +// named on many rows is looked up once. A nil value means the id matched no principal. +type principalTypeCache map[string]*v2.ResourceType + +// resolvePrincipalType finds which resource type nativeId belongs to, trying user, group, then +// service principal in that order via the Account API and stopping at the first match: the audit log's +// targetUserId doesn't record which kind of principal it actually is. Returns a nil type (no +// error) if nativeId doesn't match any of the three. Results, including no match, are cached +// in principalTypes; errors aren't. +func resolvePrincipalType(ctx context.Context, client *databricks.Client, principalTypes principalTypeCache, nativeId string) (*v2.ResourceType, *v2.RateLimitDescription, error) { + if resourceType, ok := principalTypes[nativeId]; ok { + return resourceType, nil, nil + } + + candidates := []struct { + resourceType *v2.ResourceType + found func() (bool, *v2.RateLimitDescription, error) + }{ + {userResourceType, func() (bool, *v2.RateLimitDescription, error) { + return principalExists(client.GetUser(ctx, "", nativeId)) + }}, + {groupResourceType, func() (bool, *v2.RateLimitDescription, error) { + return principalExists(client.GetGroup(ctx, "", nativeId)) + }}, + {servicePrincipalResourceType, func() (bool, *v2.RateLimitDescription, error) { + return principalExists(client.GetServicePrincipal(ctx, "", nativeId)) + }}, + } + + var rateLimit *v2.RateLimitDescription + for _, candidate := range candidates { + found, rl, err := candidate.found() + if rl != nil { + rateLimit = rl + } + if err != nil { + return nil, rateLimit, fmt.Errorf("failed to check principal %s: %w", nativeId, err) + } + if found { + ctxzap.Extract(ctx).Debug("databricks-connector: resolvePrincipalType: resolved audit principal", + zap.String("principal_id", nativeId), + zap.String("resource_type", candidate.resourceType.Id), + ) + principalTypes[nativeId] = candidate.resourceType + return candidate.resourceType, rateLimit, nil + } + } + + ctxzap.Extract(ctx).Debug("databricks-connector: resolvePrincipalType: audit principal matched no user, group or service principal", + zap.String("principal_id", nativeId), + ) + principalTypes[nativeId] = nil + return nil, rateLimit, nil +} + +// principalExists adapts a client Get* call into a found/not-found result: a 404 means not +// found (not an error), anything else is a real failure. +func principalExists[T any](_ T, rateLimit *v2.RateLimitDescription, err error) (bool, *v2.RateLimitDescription, error) { + if err == nil { + return true, rateLimit, nil + } + var apiErr *databricks.APIError + if errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusNotFound { + return false, rateLimit, nil + } + return false, rateLimit, err +} + +// filterConfiguredWorkspaces narrows workspaces to configuredWorkspaces (the --workspaces +// allowlist), or returns workspaces unchanged when the allowlist is empty. This scopes +// which workspaces' audit rows get resolved to resources — it must NOT be applied before +// locating the query warehouse (resolveWarehouseWorkspace), which can live in any workspace +// in the account regardless of this allowlist. +func filterConfiguredWorkspaces(workspaces []databricks.Workspace, configuredWorkspaces []string) []databricks.Workspace { + if len(configuredWorkspaces) == 0 { + return workspaces + } + + configured := make(map[string]struct{}, len(configuredWorkspaces)) + for _, name := range configuredWorkspaces { + configured[name] = struct{}{} + } + + filtered := make([]databricks.Workspace, 0, len(workspaces)) + for _, w := range workspaces { + if _, ok := matchConfiguredWorkspace(configured, w.DeploymentName, w.Name, strconv.Itoa(w.ID)); ok { + filtered = append(filtered, w) + } + } + + return filtered +} + +// resolveWarehouseWorkspace finds which workspace hosts warehouseId by probing each +// candidate workspace, since Databricks has no account-level lookup for this. +func resolveWarehouseWorkspace(ctx context.Context, client *databricks.Client, workspaces []databricks.Workspace, warehouseId string) (string, *v2.RateLimitDescription, error) { + var rateLimit *v2.RateLimitDescription + var inaccessible []string + for _, w := range workspaces { + found, rl, err := client.WarehouseExists(ctx, w.DeploymentName, warehouseId) + if rl != nil { + rateLimit = rl + } + if err != nil { + // The service principal needn't be assigned to every workspace; keep looking. + if isAccessDenied(err) { + inaccessible = append(inaccessible, w.DeploymentName) + continue + } + return "", rateLimit, fmt.Errorf( + "databricks-connector: failed to check workspace %s for sql-warehouse-id %s: %w", + w.DeploymentName, warehouseId, err, + ) + } + if found { + return w.DeploymentName, rateLimit, nil + } + } + + skipped := "" + if len(inaccessible) > 0 { + skipped = fmt.Sprintf(" (skipped %d inaccessible: %s)", len(inaccessible), strings.Join(inaccessible, ", ")) + } + if len(workspaces) == 1 { + return "", rateLimit, uhttp.WrapErrors(codes.NotFound, fmt.Sprintf( + "databricks-connector: sql-warehouse-id %q was not found in workspace %s%s", + warehouseId, workspaces[0].DeploymentName, skipped, + )) + } + return "", rateLimit, uhttp.WrapErrors(codes.NotFound, fmt.Sprintf( + "databricks-connector: sql-warehouse-id %q was not found in any of the %d available workspaces%s", + warehouseId, len(workspaces), skipped, + )) +} + +// isAccessDenied reports whether err is a 401/403 from the Databricks API. +func isAccessDenied(err error) bool { + var apiErr *databricks.APIError + return errors.As(err, &apiErr) && + (apiErr.StatusCode == http.StatusUnauthorized || apiErr.StatusCode == http.StatusForbidden) +} + +// queryAuditLog returns rows in (cursor, lagCutoff], plus the raw row count (before +// parseAuditLogRows may drop malformed ones) so callers can tell if the page was full. +func (f *auditEventFeed) queryAuditLog(ctx context.Context, workspaceId string, cursor eventPageCursor, lagCutoff time.Time) ([]auditLogRow, int, eventPageCursor, *v2.RateLimitDescription, error) { + // The (event_time, event_id) tiebreaker keeps ordering deterministic and lets us page + // with a composite > predicate, so progress never stalls even if many rows share one + // event_time (see advanceEventCursor); this holds as long as event_id compares consistently + // under Databricks SQL's ">"/ORDER BY, which is true for this table's opaque IDs. + statement := fmt.Sprintf(` + SELECT event_id, event_time, workspace_id, action_name, service_name, request_params + FROM system.access.audit + WHERE event_date >= :start_date + AND (event_time > :start_time OR (event_time = :start_time AND event_id > :start_after_event_id)) + AND event_time <= :lag_cutoff + AND service_name IN (%s) + AND action_name IN (%s) + ORDER BY event_time ASC, event_id ASC + LIMIT %d + `, quotedInClause(auditLogServiceNames()), quotedInClause(auditLogActionNames()), auditLogPageLimit) + + result, rateLimit, err := f.client.ExecuteStatement( + ctx, + workspaceId, + f.sqlWarehouseID, + statement, + databricks.StatementParameter{Name: "start_date", Value: cursor.StartAt.UTC().Format("2006-01-02"), Type: "DATE"}, + // RFC3339Nano, not RFC3339: the (event_time, event_id) tiebreaker needs start_time + // to round-trip at the same sub-second precision parseAuditLogRows parses. + databricks.StatementParameter{Name: "start_time", Value: cursor.StartAt.UTC().Format(time.RFC3339Nano), Type: "TIMESTAMP"}, + databricks.StatementParameter{Name: "start_after_event_id", Value: cursor.StartAfterEventID, Type: "STRING"}, + databricks.StatementParameter{Name: "lag_cutoff", Value: lagCutoff.UTC().Format(time.RFC3339Nano), Type: "TIMESTAMP"}, + ) + if err != nil { + return nil, 0, eventPageCursor{}, rateLimit, err + } + + rows, lastRawBoundary, err := parseAuditLogRows(ctx, result) + if err != nil { + return nil, len(result.Rows), lastRawBoundary, rateLimit, err + } + + ctxzap.Extract(ctx).Debug("databricks-connector: queryAuditLog: audit log query succeeded", + zap.String("workspace_id", workspaceId), + zap.String("warehouse_id", f.sqlWarehouseID), + zap.Int("audit_rows_fetched", len(result.Rows)), + zap.Int("audit_rows_parsed", len(rows)), + zap.Int("audit_rows_malformed", len(result.Rows)-len(rows)), + ) + return rows, len(result.Rows), lastRawBoundary, rateLimit, nil +} + +func quotedInClause(values []string) string { + quoted := make([]string, len(values)) + for i, v := range values { + quoted[i] = "'" + v + "'" + } + + return strings.Join(quoted, ", ") +} + +const ( + colEventID = "event_id" + colEventTime = "event_time" + colWorkspaceID = "workspace_id" + colActionName = "action_name" + colServiceName = "service_name" + colRequestParams = "request_params" +) + +// parseAuditLogRows skips (and logs) any individual row that fails to parse instead of +// failing the whole page; a missing expected column is a schema problem, so that still fails. +// lastRawBoundary is the last row's (event_time, event_id) regardless of whether the rest of +// that row parsed, so a page where every row fails can still page forward. +func parseAuditLogRows(ctx context.Context, result *databricks.StatementResult) ([]auditLogRow, eventPageCursor, error) { + colIndex := make(map[string]int, len(result.Columns)) + for i, name := range result.Columns { + colIndex[name] = i + } + + maxColIndex := 0 + for _, name := range []string{colEventID, colEventTime, colWorkspaceID, colActionName, colServiceName, colRequestParams} { + idx, ok := colIndex[name] + if !ok { + return nil, eventPageCursor{}, fmt.Errorf("audit log query result missing column %q", name) + } + if idx > maxColIndex { + maxColIndex = idx + } + } + + l := ctxzap.Extract(ctx) + + rows := make([]auditLogRow, 0, len(result.Rows)) + for _, r := range result.Rows { + row, err := parseAuditLogRow(r, colIndex, maxColIndex) + if err != nil { + l.Debug("databricks-connector: skipping malformed audit log row", zap.Error(err)) + continue + } + rows = append(rows, row) + } + + var lastRawBoundary eventPageCursor + if len(result.Rows) > 0 { + var err error + lastRawBoundary, err = parseEventBoundary(result.Rows[len(result.Rows)-1], colIndex) + if err != nil { + l.Debug("databricks-connector: failed to parse last raw row's event boundary", zap.Error(err)) + } + } + + return rows, lastRawBoundary, nil +} + +// parseEventBoundary extracts just event_id/event_time from a raw row, tolerating failures +// elsewhere in the row (e.g. malformed request_params) that would fail parseAuditLogRow. +func parseEventBoundary(r []string, colIndex map[string]int) (eventPageCursor, error) { + idIdx, timeIdx := colIndex[colEventID], colIndex[colEventTime] + if len(r) <= idIdx || len(r) <= timeIdx { + return eventPageCursor{}, fmt.Errorf("audit log query result row has %d columns, expected event_id/event_time", len(r)) + } + + eventTime, err := time.Parse("2006-01-02 15:04:05.999", r[timeIdx]) + if err != nil { + eventTime, err = time.Parse(time.RFC3339, r[timeIdx]) + if err != nil { + return eventPageCursor{}, fmt.Errorf("failed to parse event_time %q: %w", r[timeIdx], err) + } + } + + return eventPageCursor{StartAt: eventTime, StartAfterEventID: r[idIdx]}, nil +} + +func parseAuditLogRow(r []string, colIndex map[string]int, maxColIndex int) (auditLogRow, error) { + if len(r) <= maxColIndex { + return auditLogRow{}, fmt.Errorf("audit log query result row has %d columns, expected at least %d", len(r), maxColIndex+1) + } + + boundary, err := parseEventBoundary(r, colIndex) + if err != nil { + return auditLogRow{}, err + } + + var workspaceId int64 + if v := r[colIndex[colWorkspaceID]]; v != "" { + workspaceId, err = strconv.ParseInt(v, 10, 64) + if err != nil { + return auditLogRow{}, fmt.Errorf("failed to parse workspace_id %q: %w", v, err) + } + } + + requestParams := map[string]string{} + if v := r[colIndex[colRequestParams]]; v != "" { + if err := json.Unmarshal([]byte(v), &requestParams); err != nil { + return auditLogRow{}, fmt.Errorf("failed to parse request_params %q: %w", v, err) + } + } + + return auditLogRow{ + EventID: boundary.StartAfterEventID, + EventTime: boundary.StartAt, + WorkspaceID: workspaceId, + ActionName: r[colIndex[colActionName]], + ServiceName: r[colIndex[colServiceName]], + RequestParams: requestParams, + }, nil +} diff --git a/pkg/connector/audit_event_feed_test.go b/pkg/connector/audit_event_feed_test.go new file mode 100644 index 00000000..60ad8572 --- /dev/null +++ b/pkg/connector/audit_event_feed_test.go @@ -0,0 +1,1467 @@ +package connector + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/conductorone/baton-databricks/pkg/databricks" + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "github.com/conductorone/baton-sdk/pkg/pagination" + ent "github.com/conductorone/baton-sdk/pkg/types/entitlement" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/types/known/timestamppb" +) + +// testEarliestEvent starts a fresh cursor before lagCutoff, so ListEvents' bootstrap +// doesn't hit the guard for cursor.StartAt >= lagCutoff (see TestListEventsSkipsQueryPastLagCutoff). +func testEarliestEvent() *timestamppb.Timestamp { + return timestamppb.New(time.Now().Add(-2 * auditLogTrailingLag)) +} + +func TestBootstrapStartAt(t *testing.T) { + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + + cases := []struct { + name string + earliestEvent *timestamppb.Timestamp + want time.Time + }{ + {"nil earliestEvent uses the lookback default", nil, now.Add(-auditLogLookback)}, + {"zero-valued earliestEvent is treated like nil", timestamppb.New(time.Unix(0, 0)), now.Add(-auditLogLookback)}, + {"recent earliestEvent is used as-is", timestamppb.New(now.Add(-time.Hour)), now.Add(-time.Hour)}, + {"earliestEvent older than retention is clamped", timestamppb.New(now.AddDate(-2, 0, 0)), now.Add(-auditLogRetention)}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := bootstrapStartAt(now, tc.earliestEvent); !got.Equal(tc.want) { + t.Errorf("bootstrapStartAt() = %v, want %v", got, tc.want) + } + }) + } +} + +// writeJSONNotFound writes a 404 with a JSON body; a plain-text one (e.g. http.NotFound) +// breaks the client's JSON decoder. +func writeJSONNotFound(w http.ResponseWriter) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"error_code":"RESOURCE_DOES_NOT_EXIST","message":"not found"}`)) +} + +// newProbeTestClient builds a databricks.Client whose requests are redirected to a local +// httptest.Server running handler (see redirectTransport). +func newProbeTestClient(t *testing.T, handler http.HandlerFunc) *databricks.Client { + t.Helper() + + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + + target, err := url.Parse(server.URL) + if err != nil { + t.Fatalf("failed to parse test server URL: %v", err) + } + + httpClient := &http.Client{Transport: &redirectTransport{target: target}} + auth := &databricks.NoAuth{} + client, err := databricks.NewClient(context.Background(), httpClient, "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acct-1", "", auth, nil) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + return client +} + +func TestResolveWarehouseWorkspace(t *testing.T) { + const warehouseId = "wh-123" + + t.Run("single workspace is still probed", func(t *testing.T) { + probed := false + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + probed = true + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"id":%q}`, warehouseId) + }) + workspaces := []databricks.Workspace{{ID: 1, DeploymentName: "dbc-only"}} + + got, _, err := resolveWarehouseWorkspace(context.Background(), client, workspaces, warehouseId) + if err != nil { + t.Fatalf("resolveWarehouseWorkspace() error = %v", err) + } + if got != "dbc-only" { + t.Errorf("got %q, want %q", got, "dbc-only") + } + if !probed { + t.Error("WarehouseExists was never called; a single workspace must still be probed so a bad sql-warehouse-id is caught with a clear error") + } + }) + + t.Run("single workspace missing the warehouse gets a specific not-found error", func(t *testing.T) { + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + writeJSONNotFound(w) + }) + workspaces := []databricks.Workspace{{ID: 1, DeploymentName: "dbc-only"}} + + _, _, err := resolveWarehouseWorkspace(context.Background(), client, workspaces, warehouseId) + if err == nil { + t.Fatal("resolveWarehouseWorkspace() error = nil, want error when the single workspace doesn't have the warehouse") + } + wantMsg := `sql-warehouse-id "wh-123" was not found in workspace dbc-only` + if !strings.Contains(err.Error(), wantMsg) { + t.Errorf("error = %q, want it to contain %q", err.Error(), wantMsg) + } + }) + + t.Run("probes each workspace until the warehouse is found", func(t *testing.T) { + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + if strings.HasPrefix(r.Header.Get("X-Test-Original-Host"), "dbc-bbb.") { + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"id":%q}`, warehouseId) + return + } + writeJSONNotFound(w) + }) + workspaces := []databricks.Workspace{ + {ID: 1, DeploymentName: "dbc-aaa"}, + {ID: 2, DeploymentName: "dbc-bbb"}, + } + + got, _, err := resolveWarehouseWorkspace(context.Background(), client, workspaces, warehouseId) + if err != nil { + t.Fatalf("resolveWarehouseWorkspace() error = %v", err) + } + if got != "dbc-bbb" { + t.Errorf("got %q, want %q", got, "dbc-bbb") + } + }) + + t.Run("warehouse not found anywhere is a clear error", func(t *testing.T) { + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + writeJSONNotFound(w) + }) + workspaces := []databricks.Workspace{ + {ID: 1, DeploymentName: "dbc-aaa"}, + {ID: 2, DeploymentName: "dbc-bbb"}, + } + + _, _, err := resolveWarehouseWorkspace(context.Background(), client, workspaces, warehouseId) + if err == nil { + t.Fatal("resolveWarehouseWorkspace() error = nil, want error when no workspace has the warehouse") + } + }) + + t.Run("an inaccessible workspace is skipped, not fatal", func(t *testing.T) { + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + if strings.HasPrefix(r.Header.Get("X-Test-Original-Host"), "dbc-aaa.") { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"message":"not assigned to workspace"}`)) + return + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"id":%q}`, warehouseId) + }) + workspaces := []databricks.Workspace{ + {ID: 1, DeploymentName: "dbc-aaa"}, + {ID: 2, DeploymentName: "dbc-bbb"}, + } + + got, _, err := resolveWarehouseWorkspace(context.Background(), client, workspaces, warehouseId) + if err != nil { + t.Fatalf("resolveWarehouseWorkspace() error = %v", err) + } + if got != "dbc-bbb" { + t.Errorf("got %q, want %q", got, "dbc-bbb") + } + }) + + t.Run("all workspaces inaccessible is a not-found naming them", func(t *testing.T) { + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"message":"unauthorized"}`)) + }) + workspaces := []databricks.Workspace{ + {ID: 1, DeploymentName: "dbc-aaa"}, + {ID: 2, DeploymentName: "dbc-bbb"}, + } + + _, _, err := resolveWarehouseWorkspace(context.Background(), client, workspaces, warehouseId) + if got := status.Code(err); got != codes.NotFound { + t.Fatalf("error code = %s (%v), want %s", got, err, codes.NotFound) + } + if !strings.Contains(err.Error(), "skipped 2 inaccessible: dbc-aaa, dbc-bbb") { + t.Errorf("error = %q, want it to name the skipped workspaces", err.Error()) + } + }) + + t.Run("a real error from a probe is returned, not swallowed as not-found", func(t *testing.T) { + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"message":"boom"}`)) + }) + workspaces := []databricks.Workspace{ + {ID: 1, DeploymentName: "dbc-aaa"}, + {ID: 2, DeploymentName: "dbc-bbb"}, + } + + _, _, err := resolveWarehouseWorkspace(context.Background(), client, workspaces, warehouseId) + if err == nil { + t.Fatal("resolveWarehouseWorkspace() error = nil, want a propagated probe error") + } + }) +} + +func TestEventCursorRoundTrip(t *testing.T) { + want := eventPageCursor{ + StartAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), + StartAfterEventID: "b", + } + + encoded, err := encodeEventCursor(want) + if err != nil { + t.Fatalf("encodeEventCursor() error = %v", err) + } + + got := decodeEventCursor(context.Background(), encoded, want.StartAt) + if !got.StartAt.Equal(want.StartAt) || got.StartAfterEventID != want.StartAfterEventID { + t.Errorf("decodeEventCursor() = %+v, want %+v", got, want) + } +} + +func TestDecodeEventCursorSelfHeals(t *testing.T) { + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + cases := []string{"", "not-base64!!!", "aW52YWxpZC1qc29u"} // last one is base64("invalid-json") + for _, c := range cases { + got := decodeEventCursor(context.Background(), c, now) + if !got.StartAt.IsZero() { + t.Errorf("decodeEventCursor(%q) = %+v, want zero-value cursor", c, got) + } + } +} + +// TestDecodeEventCursorResetsStaleCursor covers a valid cursor whose StartAt has aged past +// the audit log retention window, which must self-heal like a corrupt cursor. +func TestDecodeEventCursorResetsStaleCursor(t *testing.T) { + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + + stale := eventPageCursor{StartAt: now.Add(-auditLogRetention - time.Hour), StartAfterEventID: "x"} + encoded, err := encodeEventCursor(stale) + if err != nil { + t.Fatalf("encodeEventCursor() error = %v", err) + } + + got := decodeEventCursor(context.Background(), encoded, now) + if !got.StartAt.IsZero() || got.StartAfterEventID != "" { + t.Errorf("decodeEventCursor() = %+v, want zero-value cursor for a stale StartAt", got) + } + + fresh := eventPageCursor{StartAt: now.Add(-auditLogRetention + time.Hour), StartAfterEventID: "y"} + encoded, err = encodeEventCursor(fresh) + if err != nil { + t.Fatalf("encodeEventCursor() error = %v", err) + } + + got = decodeEventCursor(context.Background(), encoded, now) + if !got.StartAt.Equal(fresh.StartAt) || got.StartAfterEventID != fresh.StartAfterEventID { + t.Errorf("decodeEventCursor() = %+v, want unchanged %+v (within retention)", got, fresh) + } +} + +func TestAdvanceEventCursorAdvancesToLastRawBoundary(t *testing.T) { + cursor := eventPageCursor{StartAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)} + lastRawBoundary := eventPageCursor{StartAt: cursor.StartAt.Add(time.Minute), StartAfterEventID: "evt-999"} + + next, advanced := advanceEventCursor(cursor, nil, lastRawBoundary, true, cursor.StartAt.Add(auditLogTrailingLag)) + + if !advanced || next != lastRawBoundary { + t.Errorf("next = %+v, advanced = %v, want %+v, true", next, advanced, lastRawBoundary) + } +} + +// TestAdvanceEventCursorFallsBackToLastParsedRow covers a full page where only the last raw +// row fails to parse: lastRawBoundary is zero-value, but the other rows parsed fine and must +// still advance the cursor instead of stalling on a page that otherwise made progress. +func TestAdvanceEventCursorFallsBackToLastParsedRow(t *testing.T) { + cursor := eventPageCursor{StartAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)} + rows := []auditLogRow{ + {EventID: "evt-1", EventTime: cursor.StartAt.Add(time.Minute)}, + {EventID: "evt-2", EventTime: cursor.StartAt.Add(2 * time.Minute)}, + } + + next, advanced := advanceEventCursor(cursor, rows, eventPageCursor{}, true, cursor.StartAt.Add(auditLogTrailingLag)) + + want := eventPageCursor{StartAt: rows[1].EventTime, StartAfterEventID: rows[1].EventID} + if !advanced || next != want { + t.Errorf("next = %+v, advanced = %v, want %+v, true", next, advanced, want) + } +} + +// TestAdvanceEventCursorFullyMalformedPageDoesNotStall covers a full page where every row +// failed to parse: neither lastRawBoundary nor rows can advance the cursor, so this must +// signal advanced=false (letting the caller stop claiming HasMore) instead of looping forever +// on an unchanged cursor. It also must not panic indexing an empty rows slice. +func TestAdvanceEventCursorFullyMalformedPageDoesNotStall(t *testing.T) { + cursor := eventPageCursor{StartAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), StartAfterEventID: "evt-1"} + + next, advanced := advanceEventCursor(cursor, nil, eventPageCursor{}, true, cursor.StartAt.Add(auditLogTrailingLag)) + + if advanced || next != cursor { + t.Errorf("next = %+v, advanced = %v, want unchanged cursor %+v, false", next, advanced, cursor) + } +} + +// TestAdvanceEventCursorDrainedJumpsToLagCutoff verifies a drained page advances straight to +// lagCutoff: rows are already bounded by it (see queryAuditLog), so draining proves nothing +// else exists up to that point. +func TestAdvanceEventCursorDrainedJumpsToLagCutoff(t *testing.T) { + cursor := eventPageCursor{StartAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)} + lagCutoff := cursor.StartAt.Add(4 * time.Hour) + + next, advanced := advanceEventCursor(cursor, nil, eventPageCursor{}, false, lagCutoff) + + if !advanced || !next.StartAt.Equal(lagCutoff) || next.StartAfterEventID != "" { + t.Errorf("next = %+v, advanced = %v, want StartAt=%v with no tiebreaker, true", next, advanced, lagCutoff) + } +} + +func TestMapAuditRowToResource(t *testing.T) { + // None of these cases are grant-mapped, so mapAuditRowToResource should never reach the + // principal-type lookup; a client that fails any request proves that. + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + t.Fatalf("unexpected request %s %s: no case here is grant-mapped", r.Method, r.URL.Path) + }) + workspaceLookup := map[int64]string{123: "my-workspace"} + accountId := "acct-1" + + accountParent := &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: accountId} + workspaceParent := &v2.ResourceId{ResourceType: workspaceResourceType.Id, Resource: "my-workspace"} + + type wantResource struct { + resourceType string + resource string + parentType string + parentID string + } + + cases := []struct { + name string + row auditLogRow + want []wantResource + }{ + { + name: "account-level group create", + row: auditLogRow{ + ActionName: "createGroup", + ServiceName: "accounts", + WorkspaceID: 0, + RequestParams: map[string]string{"targetGroupId": "g-1"}, + }, + want: []wantResource{ + {groupResourceType.Id, groupResourceId(context.Background(), "g-1", accountParent), accountResourceType.Id, accountId}, + }, + }, + { + name: "workspace-scoped group change stays account-parented", + row: auditLogRow{ + ActionName: "removeGroup", + ServiceName: "accounts", + WorkspaceID: 123, + RequestParams: map[string]string{"targetGroupId": "g-1"}, + }, + want: []wantResource{ + // Groups are only ever synced as children of the account when the Account + // API is reachable, regardless of which workspace the change occurred in. + {groupResourceType.Id, groupResourceId(context.Background(), "g-1", accountParent), accountResourceType.Id, accountId}, + }, + }, + { + name: "workspace-scoped acl change also refreshes the workspace-access role", + row: auditLogRow{ + ActionName: "changeDatabricksWorkspaceAcl", + ServiceName: "accounts", + WorkspaceID: 123, + }, + want: []wantResource{ + {workspaceResourceType.Id, "my-workspace", accountResourceType.Id, accountId}, + {roleResourceType.Id, roleResourceId(WorkspaceAccessRole, workspaceParent), workspaceResourceType.Id, "my-workspace"}, + }, + }, + { + name: "setAdmin refreshes the user and the account-admin role", + row: auditLogRow{ + ActionName: "setAdmin", + ServiceName: "accounts", + RequestParams: map[string]string{"targetUserId": "u-1"}, + }, + want: []wantResource{ + {userResourceType.Id, "u-1", accountResourceType.Id, accountId}, + {roleResourceType.Id, roleResourceId(AccountAdminRole, accountParent), accountResourceType.Id, accountId}, + }, + }, + { + name: "setAccountAdmin refreshes the user and the account-admin role", + row: auditLogRow{ + ActionName: "setAccountAdmin", + ServiceName: "accounts", + RequestParams: map[string]string{"targetUserId": "u-1"}, + }, + want: []wantResource{ + {userResourceType.Id, "u-1", accountResourceType.Id, accountId}, + {roleResourceType.Id, roleResourceId(AccountAdminRole, accountParent), accountResourceType.Id, accountId}, + }, + }, + { + name: "removeAccountAdmin refreshes the user and the account-admin role", + row: auditLogRow{ + ActionName: "removeAccountAdmin", + ServiceName: "accounts", + RequestParams: map[string]string{"targetUserId": "u-1"}, + }, + want: []wantResource{ + {userResourceType.Id, "u-1", accountResourceType.Id, accountId}, + {roleResourceType.Id, roleResourceId(AccountAdminRole, accountParent), accountResourceType.Id, accountId}, + }, + }, + { + name: "deactivateUser refreshes the user", + row: auditLogRow{ + ActionName: "deactivateUser", + ServiceName: "accounts", + RequestParams: map[string]string{"targetUserId": "u-1"}, + }, + want: []wantResource{ + {userResourceType.Id, "u-1", accountResourceType.Id, accountId}, + }, + }, + { + name: "account-level addPermissionAssignment resolves the workspace from request_params", + row: auditLogRow{ + ActionName: "addPermissionAssignment", + ServiceName: "workspace", + WorkspaceID: 0, + RequestParams: map[string]string{"principal_id": "u-1", "account_id": accountId, "workspace_id": "123"}, + }, + want: []wantResource{ + {workspaceResourceType.Id, "my-workspace", accountResourceType.Id, accountId}, + }, + }, + { + name: "account-level updatePermissionAssignment resolves the workspace from request_params", + row: auditLogRow{ + ActionName: "updatePermissionAssignment", + ServiceName: "workspace", + WorkspaceID: 0, + RequestParams: map[string]string{"principal_id": "u-1", "workspace_id": "123"}, + }, + want: []wantResource{ + {workspaceResourceType.Id, "my-workspace", accountResourceType.Id, accountId}, + }, + }, + { + name: "updatePermissionAssignment uses the row's workspace_id", + row: auditLogRow{ + ActionName: "updatePermissionAssignment", + ServiceName: "workspace", + WorkspaceID: 123, + RequestParams: map[string]string{"principal_id": "u-1"}, + }, + want: []wantResource{ + {workspaceResourceType.Id, "my-workspace", accountResourceType.Id, accountId}, + }, + }, + { + name: "deletePermissionAssignment for a workspace outside sync scope is skipped", + row: auditLogRow{ + ActionName: "deletePermissionAssignment", + ServiceName: "workspace", + RequestParams: map[string]string{"principal_id": "u-1", "workspace_id": "999"}, + }, + }, + { + name: "permission assignment with an unparseable workspace_id is skipped", + row: auditLogRow{ + ActionName: "deletePermissionAssignment", + ServiceName: "workspace", + RequestParams: map[string]string{"principal_id": "u-1", "workspace_id": "not-a-number"}, + }, + }, + { + name: "updateUser stays account-parented, but workspace roles still refresh", + row: auditLogRow{ + ActionName: "updateUser", + ServiceName: "accounts", + WorkspaceID: 123, + RequestParams: map[string]string{"targetUserId": "u-1"}, + }, + want: []wantResource{ + {userResourceType.Id, "u-1", accountResourceType.Id, accountId}, + {roleResourceType.Id, roleResourceId(ClusterCreateRole, workspaceParent), workspaceResourceType.Id, "my-workspace"}, + {roleResourceType.Id, roleResourceId(InstancePoolCreateRole, workspaceParent), workspaceResourceType.Id, "my-workspace"}, + }, + }, + { + name: "unknown action is skipped", + row: auditLogRow{ActionName: "someUnityCatalogAction"}, + }, + { + name: "known action_name under an unmapped service_name is skipped", + row: auditLogRow{ActionName: "delete", ServiceName: "clusters"}, + }, + { + name: "unresolvable workspace is skipped", + row: auditLogRow{ + ActionName: "add", + ServiceName: "accounts", + WorkspaceID: 999, + RequestParams: map[string]string{"targetUserId": "u-1"}, + }, + }, + { + name: "missing id param is skipped", + row: auditLogRow{ActionName: "add", ServiceName: "accounts", WorkspaceID: 0}, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, _, err := mapAuditRowToResource(context.Background(), client, principalTypeCache{}, tc.row, accountId, workspaceLookup) + if err != nil { + t.Fatalf("mapAuditRowToResource() error = %v", err) + } + if len(got) != len(tc.want) { + t.Fatalf("got %d affected resources, want %d: %+v", len(got), len(tc.want), got) + } + for i, w := range tc.want { + if got[i].resourceId.ResourceType != w.resourceType || got[i].resourceId.Resource != w.resource { + t.Errorf("[%d] resourceId = %+v, want type=%s id=%s", i, got[i].resourceId, w.resourceType, w.resource) + } + if got[i].parentResourceId.ResourceType != w.parentType || got[i].parentResourceId.Resource != w.parentID { + t.Errorf("[%d] parentResourceId = %+v, want type=%s id=%s", i, got[i].parentResourceId, w.parentType, w.parentID) + } + } + }) + } +} + +func assertGroupRefresh(t *testing.T, got []affectedResource, wantGroupId string) { + t.Helper() + if len(got) != 1 || got[0].grant != nil { + t.Fatalf("got %+v, want exactly one group refresh", got) + } + if id := got[0].resourceId; id.GetResourceType() != groupResourceType.Id || id.GetResource() != wantGroupId { + t.Errorf("resourceId = %+v, want type=%s id=%s", id, groupResourceType.Id, wantGroupId) + } +} + +// principalKindTestClient serves GetUser/GetGroup/GetServicePrincipal lookups for a single +// native id, returning 200 for kind ("Users", "Groups", or "ServicePrincipals") and 404 for +// the other two - or 500 for all three if kind is "error". +func principalKindTestClient(t *testing.T, nativeId, kind string) *databricks.Client { + t.Helper() + return newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + if !strings.HasSuffix(r.URL.Path, "/"+nativeId) { + writeJSONNotFound(w) + return + } + if kind == "error" { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"message":"boom"}`)) + return + } + if strings.Contains(r.URL.Path, "/"+kind+"/") { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{}`)) + return + } + writeJSONNotFound(w) + }) +} + +// TestMapAuditRowToResourceGrantMapping covers the group-membership actions that map to an +// atomic CREATE_GRANT/CREATE_REVOKE instead of a RESOURCE_CHANGE. +func TestMapAuditRowToResourceGrantMapping(t *testing.T) { + accountId := "acct-1" + accountParent := &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: accountId} + wantGroupId := groupResourceId(context.Background(), "g-1", accountParent) + wantEntitlementId := ent.NewEntitlementID(&v2.Resource{Id: &v2.ResourceId{ResourceType: groupResourceType.Id, Resource: wantGroupId}}, groupMemberEntitlement) + + cases := []struct { + name string + actionName string + wantRevoke bool + }{ + {"addPrincipalToGroup grants", "addPrincipalToGroup", false}, + {"removePrincipalFromGroup revokes", "removePrincipalFromGroup", true}, + {"addPrincipalsToGroup grants", "addPrincipalsToGroup", false}, + {"removePrincipalsFromGroup revokes", "removePrincipalsFromGroup", true}, + } + + client := principalKindTestClient(t, "u-1", "Users") + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + row := auditLogRow{ + ActionName: tc.actionName, + ServiceName: auditServiceAccounts, + RequestParams: map[string]string{"targetGroupId": "g-1", "targetUserId": "u-1"}, + } + + got, _, err := mapAuditRowToResource(context.Background(), client, principalTypeCache{}, row, accountId, nil) + if err != nil { + t.Fatalf("mapAuditRowToResource() error = %v", err) + } + if len(got) != 1 || got[0].grant == nil { + t.Fatalf("got %+v, want exactly one grant-mapped affected resource", got) + } + + g := got[0].grant + if g.revoke != tc.wantRevoke { + t.Errorf("revoke = %v, want %v", g.revoke, tc.wantRevoke) + } + if g.entitlement.GetId() != wantEntitlementId { + t.Errorf("entitlement id = %q, want %q", g.entitlement.GetId(), wantEntitlementId) + } + if g.principal.GetId().GetResourceType() != userResourceType.Id || g.principal.GetId().GetResource() != "u-1" { + t.Errorf("principal = %+v, want type=%s id=u-1", g.principal.GetId(), userResourceType.Id) + } + }) + } + + t.Run("missing principal id still refreshes the group", func(t *testing.T) { + row := auditLogRow{ + ActionName: "addPrincipalToGroup", + ServiceName: auditServiceAccounts, + RequestParams: map[string]string{"targetGroupId": "g-1"}, + } + got, _, err := mapAuditRowToResource(context.Background(), client, principalTypeCache{}, row, accountId, nil) + if err != nil { + t.Fatalf("mapAuditRowToResource() error = %v", err) + } + assertGroupRefresh(t, got, wantGroupId) + }) + + t.Run("principal resolves to a nested group", func(t *testing.T) { + groupClient := principalKindTestClient(t, "g-2", "Groups") + row := auditLogRow{ + ActionName: "addPrincipalToGroup", + ServiceName: auditServiceAccounts, + RequestParams: map[string]string{"targetGroupId": "g-1", "targetUserId": "g-2"}, + } + got, _, err := mapAuditRowToResource(context.Background(), groupClient, principalTypeCache{}, row, accountId, nil) + if err != nil { + t.Fatalf("mapAuditRowToResource() error = %v", err) + } + if len(got) != 1 || got[0].grant == nil { + t.Fatalf("got %+v, want exactly one grant-mapped affected resource", got) + } + wantId := groupResourceId(context.Background(), "g-2", accountParent) + if p := got[0].grant.principal.GetId(); p.GetResourceType() != groupResourceType.Id || p.GetResource() != wantId { + t.Errorf("principal = %+v, want type=%s id=%s", p, groupResourceType.Id, wantId) + } + }) + + t.Run("principal resolves to a service principal", func(t *testing.T) { + spClient := principalKindTestClient(t, "sp-1", "ServicePrincipals") + row := auditLogRow{ + ActionName: "addPrincipalToGroup", + ServiceName: auditServiceAccounts, + RequestParams: map[string]string{"targetGroupId": "g-1", "targetUserId": "sp-1"}, + } + got, _, err := mapAuditRowToResource(context.Background(), spClient, principalTypeCache{}, row, accountId, nil) + if err != nil { + t.Fatalf("mapAuditRowToResource() error = %v", err) + } + if len(got) != 1 || got[0].grant == nil { + t.Fatalf("got %+v, want exactly one grant-mapped affected resource", got) + } + if p := got[0].grant.principal.GetId(); p.GetResourceType() != servicePrincipalResourceType.Id || p.GetResource() != "sp-1" { + t.Errorf("principal = %+v, want type=%s id=sp-1", p, servicePrincipalResourceType.Id) + } + }) + + t.Run("principal not found anywhere still refreshes the group", func(t *testing.T) { + notFoundClient := principalKindTestClient(t, "ghost", "none") + row := auditLogRow{ + ActionName: "addPrincipalToGroup", + ServiceName: auditServiceAccounts, + RequestParams: map[string]string{"targetGroupId": "g-1", "targetUserId": "ghost"}, + } + got, _, err := mapAuditRowToResource(context.Background(), notFoundClient, principalTypeCache{}, row, accountId, nil) + if err != nil { + t.Fatalf("mapAuditRowToResource() error = %v", err) + } + assertGroupRefresh(t, got, wantGroupId) + }) + + t.Run("a real lookup error propagates instead of being swallowed", func(t *testing.T) { + errClient := principalKindTestClient(t, "u-1", "error") + row := auditLogRow{ + ActionName: "addPrincipalToGroup", + ServiceName: auditServiceAccounts, + RequestParams: map[string]string{"targetGroupId": "g-1", "targetUserId": "u-1"}, + } + if _, _, err := mapAuditRowToResource(context.Background(), errClient, principalTypeCache{}, row, accountId, nil); err == nil { + t.Error("mapAuditRowToResource() error = nil, want a propagated lookup error") + } + }) +} + +func TestParseAuditLogRowsDedupesNothingAndParsesFields(t *testing.T) { + result := &databricks.StatementResult{ + Columns: []string{"event_id", "event_time", "workspace_id", "action_name", "service_name", "request_params"}, + Rows: [][]string{ + {"evt-1", "2026-01-01 00:00:00.000", "123", "createGroup", "accounts", `{"targetGroupId":"g-1"}`}, + {"evt-2", "2026-01-01T00:01:00Z", "0", "add", "accounts", `{"targetUserId":"u-1"}`}, + }, + } + + rows, _, err := parseAuditLogRows(context.Background(), result) + if err != nil { + t.Fatalf("parseAuditLogRows() error = %v", err) + } + if len(rows) != 2 { + t.Fatalf("len(rows) = %d, want 2", len(rows)) + } + if rows[0].WorkspaceID != 123 || rows[0].ServiceName != "accounts" || rows[0].RequestParams["targetGroupId"] != "g-1" { + t.Errorf("row[0] = %+v", rows[0]) + } + if rows[1].WorkspaceID != 0 || rows[1].RequestParams["targetUserId"] != "u-1" { + t.Errorf("row[1] = %+v", rows[1]) + } +} + +func TestParseAuditLogRowsMissingColumnErrors(t *testing.T) { + result := &databricks.StatementResult{ + Columns: []string{"event_id", "event_time"}, + Rows: [][]string{{"evt-1", "2026-01-01 00:00:00.000"}}, + } + + if _, _, err := parseAuditLogRows(context.Background(), result); err == nil { + t.Error("parseAuditLogRows() error = nil, want error for missing required column") + } +} + +// TestParseAuditLogRowsSkipsMalformedRow verifies a single poisoned row (unparseable +// event_time here) is skipped rather than failing the whole page. +func TestParseAuditLogRowsSkipsMalformedRow(t *testing.T) { + result := &databricks.StatementResult{ + Columns: []string{"event_id", "event_time", "workspace_id", "action_name", "service_name", "request_params"}, + Rows: [][]string{ + {"evt-1", "not-a-timestamp", "123", "createGroup", "accounts", `{"targetGroupId":"g-1"}`}, + {"evt-2", "2026-01-01T00:01:00Z", "0", "add", "accounts", `{"targetUserId":"u-1"}`}, + }, + } + + rows, _, err := parseAuditLogRows(context.Background(), result) + if err != nil { + t.Fatalf("parseAuditLogRows() error = %v, want the malformed row skipped instead", err) + } + if len(rows) != 1 { + t.Fatalf("len(rows) = %d, want 1 (malformed row skipped)", len(rows)) + } + if rows[0].EventID != "evt-2" { + t.Errorf("rows[0].EventID = %q, want %q", rows[0].EventID, "evt-2") + } +} + +// redirectTransport rewrites every outgoing request to target, since workspaceUrl always +// builds a "." subdomain a local httptest.Server can't listen on directly. +type redirectTransport struct { + target *url.URL +} + +func (t *redirectTransport) RoundTrip(req *http.Request) (*http.Response, error) { + req = req.Clone(req.Context()) + originalHost := req.URL.Host + req.URL.Scheme = t.target.Scheme + req.URL.Host = t.target.Host + req.Host = t.target.Host + // Preserve the workspace-specific host the client intended, since it's otherwise lost + // once every request is rewritten to the same local test server. + req.Header.Set("X-Test-Original-Host", originalHost) + return http.DefaultTransport.RoundTrip(req) +} + +// TestListEventsSkipsQueryPastLagCutoff covers the bootstrap default landing exactly at +// lagCutoff (auditLogLookback == auditLogTrailingLag): ListEvents must not query an +// inverted (start > end) range. +func TestListEventsSkipsQueryPastLagCutoff(t *testing.T) { + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + t.Fatalf("unexpected request %s %s: should have been skipped by the lagCutoff guard", r.Method, r.URL.Path) + }) + + feed := newAuditEventFeed(client, nil, "wh-1") + events, streamState, _, err := feed.ListEvents(context.Background(), nil, &pagination.StreamToken{Cursor: ""}) + if err != nil { + t.Fatalf("ListEvents() error = %v", err) + } + if len(events) != 0 { + t.Errorf("len(events) = %d, want 0", len(events)) + } + if streamState.HasMore { + t.Error("HasMore = true, want false") + } + if streamState.Cursor == "" { + t.Error("Cursor is empty, want the bootstrap cursor to still be encoded") + } +} + +// TestListEventsEndToEnd exercises ListEvents against a mocked Statement Execution API, +// verifying the service_name query filter, resource mapping, and rate-limit propagation. +func TestListEventsEndToEnd(t *testing.T) { + const wantLimit = 100 + const wantRemaining = 42 + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // ListEvents always lists workspaces via the Account API before querying the audit log. + if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/accounts/acct-1/workspaces") { + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode([]map[string]any{ + {"workspace_id": 1, "workspace_name": "ws1", "deployment_name": "ws1"}, + }); err != nil { + t.Fatalf("failed to encode mock workspaces response: %v", err) + } + return + } + + // resolveWarehouseWorkspace probes every candidate workspace (including a lone + // one) via WarehouseExists before running the audit-log query. + if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/sql/warehouses/wh-1") { + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"id":"wh-1"}`) + return + } + + if r.Method != http.MethodPost || !strings.HasSuffix(r.URL.Path, "/api/2.0/sql/statements") { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + http.NotFound(w, r) + return + } + + var body struct { + Statement string `json:"statement"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatalf("failed to decode statement request body: %v", err) + } + if !strings.Contains(body.Statement, "service_name") { + t.Errorf("statement does not select/filter service_name: %s", body.Statement) + } + if !strings.Contains(body.Statement, "'accounts'") { + t.Errorf("statement does not filter service_name IN ('accounts', ...): %s", body.Statement) + } + + w.Header().Set("X-Ratelimit-Limit", strconv.Itoa(wantLimit)) + w.Header().Set("X-Ratelimit-Remaining", strconv.Itoa(wantRemaining)) + w.Header().Set("Content-Type", "application/json") + + resp := map[string]any{ + "statement_id": "stmt-1", + "status": map[string]any{"state": "SUCCEEDED"}, + "manifest": map[string]any{ + "schema": map[string]any{ + "columns": []map[string]any{ + {"name": "event_id"}, + {"name": "event_time"}, + {"name": "workspace_id"}, + {"name": "action_name"}, + {"name": "service_name"}, + {"name": "request_params"}, + }, + }, + }, + "result": map[string]any{ + "data_array": [][]string{ + {"evt-1", "2026-01-01T00:00:00Z", "0", "add", "accounts", `{"targetUserId":"u-1"}`}, + }, + }, + } + if err := json.NewEncoder(w).Encode(resp); err != nil { + t.Fatalf("failed to encode mock statement response: %v", err) + } + })) + defer server.Close() + + target, err := url.Parse(server.URL) + if err != nil { + t.Fatalf("failed to parse test server URL: %v", err) + } + + httpClient := &http.Client{Transport: &redirectTransport{target: target}} + auth := &databricks.NoAuth{} + client, err := databricks.NewClient(context.Background(), httpClient, "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acct-1", "", auth, nil) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + // Mirrors what Validate() sets before any sync/event-feed call runs in production. + + feed := newAuditEventFeed(client, nil, "wh-1") + + events, streamState, annos, err := feed.ListEvents(context.Background(), testEarliestEvent(), &pagination.StreamToken{Cursor: ""}) + if err != nil { + t.Fatalf("ListEvents() error = %v", err) + } + if len(events) != 1 { + t.Fatalf("len(events) = %d, want 1: %+v", len(events), events) + } + + rc := events[0].GetResourceChangeEvent() + if rc.GetResourceId().GetResourceType() != userResourceType.Id || rc.GetResourceId().GetResource() != "u-1" { + t.Errorf("event resource = %+v, want type=%s id=u-1", rc.GetResourceId(), userResourceType.Id) + } + if streamState.Cursor == "" { + t.Error("StreamState.Cursor is empty, want an encoded cursor") + } + + rld := &v2.RateLimitDescription{} + ok, err := annos.Pick(rld) + if err != nil { + t.Fatalf("annos.Pick() error = %v", err) + } + if !ok { + t.Fatal("annotations do not carry a RateLimitDescription, want the rate-limit headers propagated") + } + if rld.GetLimit() != wantLimit || rld.GetRemaining() != wantRemaining { + t.Errorf("RateLimitDescription = %+v, want limit=%d remaining=%d", rld, wantLimit, wantRemaining) + } +} + +// TestListEventsEmitsCreateGrantEvent verifies ListEvents wires a grant-mapped audit row all +// the way to a CreateGrantEvent, not a ResourceChangeEvent. +func TestListEventsEmitsCreateGrantEvent(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/accounts/acct-1/workspaces") { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode([]map[string]any{{"workspace_id": 1, "workspace_name": "ws1", "deployment_name": "ws1"}}) + return + } + if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/sql/warehouses/wh-1") { + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"id":"wh-1"}`) + return + } + if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/accounts/acct-1/scim/v2/Users/u-1") { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{}`) + return + } + if r.Method != http.MethodPost || !strings.HasSuffix(r.URL.Path, "/api/2.0/sql/statements") { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + http.NotFound(w, r) + return + } + + w.Header().Set("Content-Type", "application/json") + resp := map[string]any{ + "statement_id": "stmt-1", + "status": map[string]any{"state": "SUCCEEDED"}, + "manifest": map[string]any{ + "schema": map[string]any{ + "columns": []map[string]any{ + {"name": "event_id"}, {"name": "event_time"}, {"name": "workspace_id"}, + {"name": "action_name"}, {"name": "service_name"}, {"name": "request_params"}, + }, + }, + }, + "result": map[string]any{ + "data_array": [][]string{ + {"evt-1", "2026-01-01T00:00:00Z", "0", "addPrincipalToGroup", "accounts", `{"targetGroupId":"g-1","targetUserId":"u-1"}`}, + }, + }, + } + if err := json.NewEncoder(w).Encode(resp); err != nil { + t.Fatalf("failed to encode mock statement response: %v", err) + } + })) + defer server.Close() + + target, err := url.Parse(server.URL) + if err != nil { + t.Fatalf("failed to parse test server URL: %v", err) + } + + httpClient := &http.Client{Transport: &redirectTransport{target: target}} + auth := &databricks.NoAuth{} + client, err := databricks.NewClient(context.Background(), httpClient, "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acct-1", "", auth, nil) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + + feed := newAuditEventFeed(client, nil, "wh-1") + events, _, _, err := feed.ListEvents(context.Background(), testEarliestEvent(), &pagination.StreamToken{Cursor: ""}) + if err != nil { + t.Fatalf("ListEvents() error = %v", err) + } + if len(events) != 1 { + t.Fatalf("len(events) = %d, want 1: %+v", len(events), events) + } + + cg := events[0].GetCreateGrantEvent() + if cg == nil { + t.Fatalf("event = %+v, want a CreateGrantEvent", events[0]) + } + accountParent := &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: "acct-1"} + wantGroupId := groupResourceId(context.Background(), "g-1", accountParent) + if cg.GetEntitlement().GetResource().GetId().GetResource() != wantGroupId { + t.Errorf("entitlement resource = %q, want %q", cg.GetEntitlement().GetResource().GetId().GetResource(), wantGroupId) + } + if cg.GetPrincipal().GetId().GetResourceType() != userResourceType.Id || cg.GetPrincipal().GetId().GetResource() != "u-1" { + t.Errorf("principal = %+v, want type=%s id=u-1", cg.GetPrincipal().GetId(), userResourceType.Id) + } +} + +// newLookupFailureFeed returns a feed whose single grant-mapped audit row's principal lookup +// fails with lookupStatus. +func newLookupFailureFeed(t *testing.T, lookupStatus int) *auditEventFeed { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/accounts/acct-1/workspaces") { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode([]map[string]any{{"workspace_id": 1, "workspace_name": "ws1", "deployment_name": "ws1"}}) + return + } + if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/sql/warehouses/wh-1") { + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"id":"wh-1"}`) + return + } + if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/accounts/acct-1/scim/v2/Users/u-1") { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(lookupStatus) + fmt.Fprint(w, `{"detail":"lookup failed"}`) + return + } + if r.Method != http.MethodPost || !strings.HasSuffix(r.URL.Path, "/api/2.0/sql/statements") { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + http.NotFound(w, r) + return + } + + w.Header().Set("Content-Type", "application/json") + resp := map[string]any{ + "statement_id": "stmt-1", + "status": map[string]any{"state": "SUCCEEDED"}, + "manifest": map[string]any{ + "schema": map[string]any{ + "columns": []map[string]any{ + {"name": "event_id"}, {"name": "event_time"}, {"name": "workspace_id"}, + {"name": "action_name"}, {"name": "service_name"}, {"name": "request_params"}, + }, + }, + }, + "result": map[string]any{ + "data_array": [][]string{ + {"evt-1", "2026-01-01T00:00:00Z", "0", "addPrincipalToGroup", "accounts", `{"targetGroupId":"g-1","targetUserId":"u-1"}`}, + }, + }, + } + if err := json.NewEncoder(w).Encode(resp); err != nil { + t.Fatalf("failed to encode mock statement response: %v", err) + } + })) + t.Cleanup(server.Close) + + target, err := url.Parse(server.URL) + if err != nil { + t.Fatalf("failed to parse test server URL: %v", err) + } + + httpClient := &http.Client{Transport: &redirectTransport{target: target}} + auth := &databricks.NoAuth{} + client, err := databricks.NewClient(context.Background(), httpClient, "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acct-1", "", auth, nil) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + + return newAuditEventFeed(client, nil, "wh-1") +} + +// TestListEventsReturnsErrorOnRetryableLookupFailure verifies that a transient principal lookup +// failure fails the poll, so the cursor isn't advanced past the row. +func TestListEventsReturnsErrorOnRetryableLookupFailure(t *testing.T) { + feed := newLookupFailureFeed(t, http.StatusServiceUnavailable) + events, streamState, _, err := feed.ListEvents(context.Background(), testEarliestEvent(), &pagination.StreamToken{Cursor: ""}) + if err == nil { + t.Fatalf("ListEvents() error = nil, want the principal lookup failure; events=%+v", events) + } + if events != nil || streamState != nil { + t.Errorf("ListEvents() = (%+v, %+v), want no events and no stream state so the cursor is not advanced", events, streamState) + } +} + +// TestListEventsSkipsRowOnPermanentLookupFailure verifies that a non-retryable principal lookup +// failure skips the row instead of blocking the cursor on every poll. +func TestListEventsSkipsRowOnPermanentLookupFailure(t *testing.T) { + feed := newLookupFailureFeed(t, http.StatusForbidden) + events, streamState, _, err := feed.ListEvents(context.Background(), testEarliestEvent(), &pagination.StreamToken{Cursor: ""}) + if err != nil { + t.Fatalf("ListEvents() error = %v, want the row skipped", err) + } + if len(events) != 0 { + t.Errorf("len(events) = %d, want 0", len(events)) + } + if streamState == nil || streamState.Cursor == "" { + t.Errorf("streamState = %+v, want an advanced cursor", streamState) + } +} + +// TestListEventsFullPageCursorAdvance verifies the HasMore wiring on a full page +// (auditLogPageLimit raw rows): it pages forward from the last raw row's boundary, but when +// no boundary can be parsed it reports HasMore=false with the cursor unchanged rather than +// claiming more data behind a cursor that would repeat the same query. +func TestListEventsFullPageCursorAdvance(t *testing.T) { + start := eventPageCursor{StartAt: time.Now().Add(-2 * auditLogTrailingLag).UTC().Truncate(time.Second), StartAfterEventID: "evt-0"} + lastTime := start.StartAt.Add(time.Minute) + + cases := []struct { + name string + lastRowTime string + wantHasMore bool + wantCursor eventPageCursor + }{ + { + name: "parseable last row pages forward", + lastRowTime: lastTime.Format(time.RFC3339), + wantHasMore: true, + wantCursor: eventPageCursor{StartAt: lastTime, StartAfterEventID: fmt.Sprintf("evt-%d", auditLogPageLimit)}, + }, + { + name: "no parseable boundary stops without advancing", + lastRowTime: "not-a-time", + wantHasMore: false, + wantCursor: start, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/accounts/acct-1/workspaces"): + _ = json.NewEncoder(w).Encode([]map[string]any{{"workspace_id": 1, "workspace_name": "ws1", "deployment_name": "ws1"}}) + case r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/sql/warehouses/wh-1"): + fmt.Fprint(w, `{"id":"wh-1"}`) + case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/api/2.0/sql/statements"): + // Every row but the last has an unparseable event_time, so no row parses + // and only the last raw row can supply the boundary. The action is + // unmapped, so no principal lookups are made either way. + data := make([][]string, auditLogPageLimit) + for i := range data { + data[i] = []string{fmt.Sprintf("evt-%d", i+1), "not-a-time", "0", "someUnmappedAction", "accounts", `{}`} + } + data[len(data)-1][1] = tc.lastRowTime + _ = json.NewEncoder(w).Encode(map[string]any{ + "statement_id": "stmt-1", + "status": map[string]any{"state": "SUCCEEDED"}, + "manifest": map[string]any{"schema": map[string]any{"columns": []map[string]any{ + {"name": "event_id"}, {"name": "event_time"}, {"name": "workspace_id"}, + {"name": "action_name"}, {"name": "service_name"}, {"name": "request_params"}, + }}}, + "result": map[string]any{"data_array": data}, + }) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + http.NotFound(w, r) + } + }) + + encoded, err := encodeEventCursor(start) + if err != nil { + t.Fatalf("encodeEventCursor() error = %v", err) + } + + feed := newAuditEventFeed(client, nil, "wh-1") + events, streamState, _, err := feed.ListEvents(context.Background(), nil, &pagination.StreamToken{Cursor: encoded}) + if err != nil { + t.Fatalf("ListEvents() error = %v", err) + } + if len(events) != 0 { + t.Errorf("len(events) = %d, want 0", len(events)) + } + if streamState.HasMore != tc.wantHasMore { + t.Errorf("HasMore = %v, want %v", streamState.HasMore, tc.wantHasMore) + } + got := decodeEventCursor(context.Background(), streamState.Cursor, time.Now()) + if !got.StartAt.Equal(tc.wantCursor.StartAt) || got.StartAfterEventID != tc.wantCursor.StartAfterEventID { + t.Errorf("cursor = %+v, want %+v", got, tc.wantCursor) + } + }) + } +} + +// TestListEventsFindsWarehouseOutsideWorkspacesAllowlist verifies that resolving the SQL +// warehouse's workspace is NOT scoped by --workspaces: the warehouse can live in any +// workspace in the account, so the allowlist must only narrow which workspaces' audit +// rows get resolved to resources, not which workspaces are searched for the warehouse. +func TestListEventsFindsWarehouseOutsideWorkspacesAllowlist(t *testing.T) { + queriedWorkspace := "" + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/accounts/acct-1/workspaces") { + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode([]map[string]any{ + {"workspace_id": 1, "workspace_name": "ws1", "deployment_name": "ws1"}, + {"workspace_id": 2, "workspace_name": "ws2", "deployment_name": "ws2"}, + }); err != nil { + t.Fatalf("failed to encode mock workspaces response: %v", err) + } + return + } + + // The warehouse only exists in ws1, which is NOT in the --workspaces allowlist + // below (only ws2 is configured). Resolution must still find it in ws1. + if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/sql/warehouses/wh-1") { + isWs1 := strings.HasPrefix(r.Header.Get("X-Test-Original-Host"), "ws1.") + w.Header().Set("Content-Type", "application/json") + if isWs1 { + fmt.Fprintf(w, `{"id":"wh-1"}`) + return + } + writeJSONNotFound(w) + return + } + + if r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/api/2.0/sql/statements") { + queriedWorkspace = r.Header.Get("X-Test-Original-Host") + w.Header().Set("Content-Type", "application/json") + resp := map[string]any{ + "statement_id": "stmt-1", + "status": map[string]any{"state": "SUCCEEDED"}, + "manifest": map[string]any{ + "schema": map[string]any{ + "columns": []map[string]any{ + {"name": "event_id"}, {"name": "event_time"}, {"name": "workspace_id"}, + {"name": "action_name"}, {"name": "service_name"}, {"name": "request_params"}, + }, + }, + }, + "result": map[string]any{"data_array": [][]string{}}, + } + if err := json.NewEncoder(w).Encode(resp); err != nil { + t.Fatalf("failed to encode mock statement response: %v", err) + } + return + } + + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + http.NotFound(w, r) + })) + defer server.Close() + + target, err := url.Parse(server.URL) + if err != nil { + t.Fatalf("failed to parse test server URL: %v", err) + } + + httpClient := &http.Client{Transport: &redirectTransport{target: target}} + auth := &databricks.NoAuth{} + client, err := databricks.NewClient(context.Background(), httpClient, "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acct-1", "", auth, nil) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + + // --workspaces is scoped to ws2 only; the warehouse lives in ws1. + feed := newAuditEventFeed(client, []string{"ws2"}, "wh-1") + + _, _, _, err = feed.ListEvents(context.Background(), testEarliestEvent(), &pagination.StreamToken{Cursor: ""}) + if err != nil { + t.Fatalf("ListEvents() error = %v, want warehouse resolution to succeed despite living outside --workspaces", err) + } + if !strings.HasPrefix(queriedWorkspace, "ws1.") { + t.Errorf("audit query ran against %q, want it to run against ws1 (where the warehouse actually lives)", queriedWorkspace) + } +} + +func TestFilterConfiguredWorkspaces(t *testing.T) { + all := []databricks.Workspace{ + {ID: 1, Name: "Prod", DeploymentName: "dbc-prod"}, + {ID: 2, Name: "Dev", DeploymentName: "dbc-dev"}, + } + + tests := []struct { + name string + configured []string + want []int + }{ + {"empty allowlist keeps all", nil, []int{1, 2}}, + {"deployment name", []string{"dbc-dev"}, []int{2}}, + {"workspace name, case-insensitive", []string{"PROD"}, []int{1}}, + {"numeric ID", []string{"2"}, []int{2}}, + {"no match", []string{"other"}, nil}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := filterConfiguredWorkspaces(all, tt.configured) + var ids []int + for _, w := range got { + ids = append(ids, w.ID) + } + if fmt.Sprint(ids) != fmt.Sprint(tt.want) { + t.Errorf("filterConfiguredWorkspaces(%v) = %v, want %v", tt.configured, ids, tt.want) + } + }) + } +} + +// TestExecuteStatementPollsPastCache guards against the uhttp GET cache freezing +// pollStatement: the poll URL never changes, so a cached RUNNING response would be +// returned until statementPollMaxWait. +func TestExecuteStatementPollsPastCache(t *testing.T) { + var polls atomic.Int32 + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + state := "PENDING" + if r.Method == http.MethodGet { + state = "RUNNING" + if polls.Add(1) >= 2 { + state = "SUCCEEDED" + } + } + fmt.Fprintf(w, `{"statement_id":"stmt-1","status":{"state":%q},"manifest":{"schema":{"columns":[{"name":"x"}]}},"result":{"data_array":[]}}`, state) + }) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + if _, _, err := client.ExecuteStatement(ctx, "dbc-aaa", "wh-1", "SELECT 1"); err != nil { + t.Fatalf("ExecuteStatement() error = %v", err) + } + if got := polls.Load(); got < 2 { + t.Errorf("server received %d poll GETs, want >= 2 (poll response was served from cache)", got) + } +} + +// TestGetBypassesCacheButSyncCallersDont verifies incremental-sync refreshes (builder Get) +// always hit the API, while the same client call from the regular sync path stays cached. +func TestGetBypassesCacheButSyncCallersDont(t *testing.T) { + var hits atomic.Int32 + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"id":"u-1","userName":"alice@example.com","active":true}`) + }) + ctx := context.Background() + accountParent := &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: "acct-1"} + userId := &v2.ResourceId{ResourceType: userResourceType.Id, Resource: "u-1"} + + for range 2 { + if _, _, err := client.GetUser(ctx, "", "u-1"); err != nil { + t.Fatalf("GetUser() error = %v", err) + } + } + if got := hits.Load(); got != 1 { + t.Fatalf("regular GetUser calls made %d API calls, want 1 (second must be cached)", got) + } + + builder := newUserBuilder(client) + for range 2 { + if _, _, err := builder.Get(ctx, userId, accountParent); err != nil { + t.Fatalf("Get() error = %v", err) + } + } + if got := hits.Load() - 1; got != 2 { + t.Errorf("builder Get made %d API calls, want 2 (must bypass the cache)", got) + } +} + +// TestListEventsCachesPrincipalTypes verifies a principal named on several rows of one page is +// looked up once, including an id that matches no principal at all. +func TestListEventsCachesPrincipalTypes(t *testing.T) { + var mu sync.Mutex + hits := map[string]int{} + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/accounts/acct-1/workspaces"): + _ = json.NewEncoder(w).Encode([]map[string]any{{"workspace_id": 1, "workspace_name": "ws1", "deployment_name": "ws1"}}) + case r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/api/2.0/sql/warehouses/wh-1"): + fmt.Fprint(w, `{"id":"wh-1"}`) + case r.Method == http.MethodGet && strings.Contains(r.URL.Path, "/scim/v2/"): + mu.Lock() + hits[r.URL.Path]++ + mu.Unlock() + if strings.HasSuffix(r.URL.Path, "/ServicePrincipals/sp-1") { + fmt.Fprint(w, `{}`) + return + } + w.WriteHeader(http.StatusNotFound) + fmt.Fprint(w, `{"detail":"not found"}`) + case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/api/2.0/sql/statements"): + row := func(id, principal string) []string { + return []string{id, "2026-01-01T00:00:00Z", "0", "addPrincipalToGroup", "accounts", fmt.Sprintf(`{"targetGroupId":"g-1","targetUserId":%q}`, principal)} + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "statement_id": "stmt-1", + "status": map[string]any{"state": "SUCCEEDED"}, + "manifest": map[string]any{"schema": map[string]any{"columns": []map[string]any{ + {"name": "event_id"}, {"name": "event_time"}, {"name": "workspace_id"}, + {"name": "action_name"}, {"name": "service_name"}, {"name": "request_params"}, + }}}, + "result": map[string]any{"data_array": [][]string{ + row("evt-1", "sp-1"), row("evt-2", "sp-1"), row("evt-3", "sp-1"), + row("evt-4", "gone-1"), row("evt-5", "gone-1"), + }}, + }) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + http.NotFound(w, r) + } + }) + + feed := newAuditEventFeed(client, nil, "wh-1") + events, _, _, err := feed.ListEvents(context.Background(), testEarliestEvent(), &pagination.StreamToken{Cursor: ""}) + if err != nil { + t.Fatalf("ListEvents() error = %v", err) + } + if len(events) != 5 { + t.Fatalf("len(events) = %d, want 5", len(events)) + } + for _, id := range []string{"sp-1", "gone-1"} { + for _, kind := range []string{"Users", "Groups", "ServicePrincipals"} { + path := fmt.Sprintf("/api/2.0/accounts/acct-1/scim/v2/%s/%s", kind, id) + if got := hits[path]; got != 1 { + t.Errorf("%s hit %d times, want 1", path, got) + } + } + } +} + +// TestResolvePrincipalTypeDoesNotCacheErrors verifies a failed lookup is retried on the next +// row instead of being remembered. +func TestResolvePrincipalTypeDoesNotCacheErrors(t *testing.T) { + var calls atomic.Int32 + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if calls.Add(1) == 1 { + w.WriteHeader(http.StatusForbidden) + fmt.Fprint(w, `{"detail":"denied"}`) + return + } + fmt.Fprint(w, `{}`) + }) + + principalTypes := principalTypeCache{} + if _, _, err := resolvePrincipalType(context.Background(), client, principalTypes, "u-1"); err == nil { + t.Fatal("first resolvePrincipalType() error = nil, want the lookup failure") + } + got, _, err := resolvePrincipalType(context.Background(), client, principalTypes, "u-1") + if err != nil { + t.Fatalf("second resolvePrincipalType() error = %v", err) + } + if got != userResourceType { + t.Errorf("second resolvePrincipalType() = %v, want %s", got, userResourceType.Id) + } + if n := calls.Load(); n != 2 { + t.Errorf("lookup made %d calls, want 2 (the error must not be cached)", n) + } +} diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index f97021e6..09b739d2 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -4,6 +4,9 @@ import ( "context" "fmt" "io" + "slices" + "strconv" + "strings" "github.com/conductorone/baton-databricks/pkg/config" "github.com/conductorone/baton-databricks/pkg/databricks" @@ -11,38 +14,64 @@ import ( "github.com/conductorone/baton-sdk/pkg/annotations" "github.com/conductorone/baton-sdk/pkg/cli" "github.com/conductorone/baton-sdk/pkg/connectorbuilder" + "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" + "go.uber.org/zap" ) 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(ctx context.Context) []connectorbuilder.ResourceSyncerV2 { - syncers := []connectorbuilder.ResourceSyncerV2{ +func (d *Databricks) ResourceSyncers(_ context.Context) []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), + }), } +} + +// 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) +} - return syncers +// EventFeeds registers the audit-log event feed unconditionally; whether sql-warehouse-id +// is set gates its behavior inside ListEvents instead. +func (d *Databricks) EventFeeds(_ context.Context) []connectorbuilder.EventFeed { + return []connectorbuilder.EventFeed{ + newAuditEventFeed(d.client, d.workspaces, d.sqlWarehouseID), + } } // Asset takes an input AssetRef and attempts to fetch it using the connector's authenticated http client // It streams a response, always starting with a metadata object, following by chunked payloads for the asset. -func (d *Databricks) Asset(ctx context.Context, asset *v2.AssetRef) (string, io.ReadCloser, error) { +func (d *Databricks) Asset(_ context.Context, _ *v2.AssetRef) (string, io.ReadCloser, error) { return "", nil, nil } // Metadata returns metadata about the connector. -func (d *Databricks) Metadata(ctx context.Context) (*v2.ConnectorMetadata, error) { +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": { @@ -112,23 +141,49 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err // Workspace enumeration is the other account-plane call every sync depends on. // Per-workspace probing is deliberately not done here: a workspace the service // principal can't reach is skipped during sync rather than failing validation. - if _, _, err := d.client.ListWorkspaces(ctx); err != nil { - return nil, fmt.Errorf("databricks-connector: failed to list workspaces: %w", err) + annos := annotations.Annotations{} + workspaces, rateLimit, err := d.client.ListWorkspaces(ctx) + if rateLimit != nil { + annos.WithRateLimiting(rateLimit) + } + if err != nil { + 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, ", ")) } - return nil, nil + l := ctxzap.Extract(ctx) + if d.sqlWarehouseID == "" { + l.Info("databricks-connector: incremental sync disabled (sql-warehouse-id not set)") + return annos, nil + } + + 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 { @@ -142,12 +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) @@ -157,9 +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 b5979b73..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(), @@ -288,6 +288,40 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, attr r return rv, &rs.SyncOpResults{NextPageToken: nextPage, Annotations: annos}, nil } +// Get re-fetches a single group, used to re-sync it after a RESOURCE_CHANGE event; fetched +// without members since groupResource doesn't need them (Grants fetches them separately). +func (g *groupBuilder) Get(ctx context.Context, resourceId *v2.ResourceId, parentResourceId *v2.ResourceId) (*v2.Resource, annotations.Annotations, error) { + parentId, groupId, err := parseResourceId(resourceId.Resource) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to parse group resource id: %w", err) + } + + var workspaceId string + if parentId != nil && parentId.ResourceType == workspaceResourceType.Id { + workspaceId = parentId.Resource + } + + group, rateLimitData, err := g.client.GetGroup(databricks.WithoutCache(ctx), workspaceId, groupId.Resource, databricks.NewGroupAttrVars()) + annos := annotations.Annotations{} + if rateLimitData != nil { + annos.WithRateLimiting(rateLimitData) + } + if err != nil { + return nil, annos, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err) + } + + resource, err := groupResource(ctx, group, parentResourceId) + if err != nil { + return nil, nil, err + } + + ctxzap.Extract(ctx).Debug("databricks-connector: Get: refreshed group", + zap.String("resource_id", resourceId.Resource), + zap.String("parent_resource_id", parentResourceId.GetResource()), + ) + return resource, annos, nil +} + func (g *groupBuilder) Grant(ctx context.Context, principal *v2.Resource, entitlement *v2.Entitlement) (annotations.Annotations, error) { l := ctxzap.Extract(ctx) 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 175f6ae5..066d10c9 100644 --- a/pkg/connector/roles.go +++ b/pkg/connector/roles.go @@ -3,6 +3,7 @@ package connector import ( "context" "fmt" + "strings" "github.com/conductorone/baton-databricks/pkg/databricks" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" @@ -43,21 +44,22 @@ func (r *roleBuilder) ResourceType(ctx context.Context) *v2.ResourceType { return roleResourceType } +// roleResourceId builds a role's resource ID, namespaced by workspace for workspace roles. +func roleResourceId(role string, parent *v2.ResourceId) string { + if parent.GetResourceType() == workspaceResourceType.Id { + return fmt.Sprintf("%s:%s", parent.Resource, role) + } + return role +} + func roleResource(ctx context.Context, role string, parent *v2.ResourceId) (*v2.Resource, error) { - var roleID string - profile := map[string]interface{}{ + profile := map[string]any{ "role_name": role, - "parent_type": parent.ResourceType, - "parent_id": parent.Resource, + "parent_type": parent.GetResourceType(), + "parent_id": parent.GetResource(), } - // To differentiate between what type of role does the resource represent. - switch parent.ResourceType { - case workspaceResourceType.Id: - roleID = fmt.Sprintf("%s:%s", parent.Resource, role) - case accountResourceType.Id: - roleID = role - } + roleID := roleResourceId(role, parent) resource, err := rs.NewRoleResource( role, @@ -297,6 +299,31 @@ func (r *roleBuilder) Grants(ctx context.Context, resource *v2.Resource, attr rs return rv, &rs.SyncOpResults{NextPageToken: nextPage}, nil } +// Get rebuilds a single role resource, used to re-sync it after a RESOURCE_CHANGE event. +// Roles are synthetic (not fetched from an API), so this just reconstructs the resource +// from its resource ID the same way List does. +func (r *roleBuilder) Get(ctx context.Context, resourceId *v2.ResourceId, parentResourceId *v2.ResourceId) (*v2.Resource, annotations.Annotations, error) { + roleName := resourceId.Resource + if parentResourceId.GetResourceType() == workspaceResourceType.Id { + _, name, found := strings.Cut(resourceId.Resource, ":") + if !found { + return nil, nil, fmt.Errorf("databricks-connector: invalid workspace role resource id: %s", resourceId.Resource) + } + roleName = name + } + + resource, err := roleResource(ctx, roleName, parentResourceId) + if err != nil { + return nil, nil, err + } + + ctxzap.Extract(ctx).Debug("databricks-connector: Get: refreshed role", + zap.String("resource_id", resourceId.Resource), + zap.String("parent_resource_id", parentResourceId.GetResource()), + ) + return resource, nil, nil +} + func (r *roleBuilder) Grant(ctx context.Context, principal *v2.Resource, entitlement *v2.Entitlement) (annotations.Annotations, error) { l := ctxzap.Extract(ctx) diff --git a/pkg/connector/service-principals.go b/pkg/connector/service-principals.go index ccd109ba..a868f054 100644 --- a/pkg/connector/service-principals.go +++ b/pkg/connector/service-principals.go @@ -27,18 +27,19 @@ 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.ResourceType, - "parent_id": parent.Resource, + "parent_type": parent.GetResourceType(), + "parent_id": parent.GetResource(), } options := []rs.ResourceOption{ rs.WithResourceProfile(profile), } + // keep the parent resource id, only if the parent resource is account - if parent.ResourceType == accountResourceType.Id { + if parent.GetResourceType() == accountResourceType.Id { options = append(options, rs.WithParentResourceID(parent)) } @@ -220,6 +221,34 @@ func (s *servicePrincipalBuilder) Grants(ctx context.Context, resource *v2.Resou return rv, nil, nil } +// Get re-fetches a single service principal, used to re-sync it after a RESOURCE_CHANGE event. +func (s *servicePrincipalBuilder) Get(ctx context.Context, resourceId *v2.ResourceId, parentResourceId *v2.ResourceId) (*v2.Resource, annotations.Annotations, error) { + var workspaceId string + if parentResourceId.GetResourceType() == workspaceResourceType.Id { + workspaceId = parentResourceId.Resource + } + + servicePrincipal, rateLimitData, err := s.client.GetServicePrincipal(databricks.WithoutCache(ctx), workspaceId, resourceId.Resource) + annos := annotations.Annotations{} + if rateLimitData != nil { + annos.WithRateLimiting(rateLimitData) + } + if err != nil { + return nil, annos, fmt.Errorf("databricks-connector: failed to get service principal %s: %w", resourceId.Resource, err) + } + + resource, err := s.servicePrincipalResource(ctx, servicePrincipal, parentResourceId) + if err != nil { + return nil, nil, err + } + + ctxzap.Extract(ctx).Debug("databricks-connector: Get: refreshed service principal", + zap.String("resource_id", resourceId.Resource), + zap.String("parent_resource_id", parentResourceId.GetResource()), + ) + return resource, annos, nil +} + func (s *servicePrincipalBuilder) Grant(ctx context.Context, principal *v2.Resource, entitlement *v2.Entitlement) (annotations.Annotations, error) { l := ctxzap.Extract(ctx) diff --git a/pkg/connector/users.go b/pkg/connector/users.go index 5e3f780b..f18bc213 100644 --- a/pkg/connector/users.go +++ b/pkg/connector/users.go @@ -11,6 +11,8 @@ import ( "github.com/conductorone/baton-sdk/pkg/connectorbuilder" 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" ) @@ -46,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, @@ -66,8 +68,9 @@ func (u *userBuilder) userResource(ctx context.Context, user *databricks.User, p rs.WithResourceStatus(status, ""), rs.WithAnnotation(&v2.RawId{Id: user.ID}), } + // keep the parent resource id, only if the parent resource is account - if parent.ResourceType == accountResourceType.Id { + if parent.GetResourceType() == accountResourceType.Id { options = append(options, rs.WithParentResourceID(parent)) } @@ -272,6 +275,34 @@ func (o *userBuilder) CreateAccount(ctx context.Context, accountInfo *v2.Account }, nil, nil, nil } +// Get re-fetches a single user, used to re-sync it after a RESOURCE_CHANGE event. +func (u *userBuilder) Get(ctx context.Context, resourceId *v2.ResourceId, parentResourceId *v2.ResourceId) (*v2.Resource, annotations.Annotations, error) { + var workspaceId string + if parentResourceId.GetResourceType() == workspaceResourceType.Id { + workspaceId = parentResourceId.Resource + } + + user, rateLimitData, err := u.client.GetUser(databricks.WithoutCache(ctx), workspaceId, resourceId.Resource) + annos := annotations.Annotations{} + if rateLimitData != nil { + annos.WithRateLimiting(rateLimitData) + } + if err != nil { + return nil, annos, fmt.Errorf("databricks-connector: failed to get user %s: %w", resourceId.Resource, err) + } + + resource, err := u.userResource(ctx, user, parentResourceId) + if err != nil { + return nil, nil, err + } + + ctxzap.Extract(ctx).Debug("databricks-connector: Get: refreshed user", + zap.String("resource_id", resourceId.Resource), + zap.String("parent_resource_id", parentResourceId.GetResource()), + ) + return resource, annos, nil +} + func (o *userBuilder) Delete(ctx context.Context, resourceId *v2.ResourceId) (annotations.Annotations, error) { _, err := o.client.DeleteUser(ctx, "", resourceId.Resource) if err != nil { diff --git a/pkg/connector/validate_test.go b/pkg/connector/validate_test.go index 225b54e6..28799667 100644 --- a/pkg/connector/validate_test.go +++ b/pkg/connector/validate_test.go @@ -5,9 +5,11 @@ import ( "io" "net/http" "strings" + "sync" "testing" "github.com/conductorone/baton-databricks/pkg/databricks" + "github.com/conductorone/baton-sdk/pkg/pagination" ) // rolesTransport answers the assignable-roles calls Validate makes. failAccount @@ -54,3 +56,114 @@ func TestValidateOAuthAccountCheckFailureReturnsError(t *testing.T) { t.Fatal("Validate: want error on OAuth account check failure, got nil") } } + +// incrementalTransport serves the account calls Validate always makes (roles, one +// workspace) and records every path, so tests can assert whether any incremental-sync +// call (SQL warehouse lookup, statement execution) was made. +type incrementalTransport struct { + mu sync.Mutex + paths []string +} + +func (t *incrementalTransport) RoundTrip(req *http.Request) (*http.Response, error) { + t.mu.Lock() + t.paths = append(t.paths, req.URL.Path) + t.mu.Unlock() + + status, body := http.StatusOK, `{"roles":[]}` + switch { + case strings.HasSuffix(req.URL.Path, "/accounts/acct-1/workspaces"): + body = `[{"workspace_id":1,"workspace_name":"ws1","deployment_name":"ws1"}]` + case strings.Contains(req.URL.Path, "/sql/warehouses/"): + status, body = http.StatusNotFound, `{"error_code":"RESOURCE_DOES_NOT_EXIST","message":"not found"}` + } + return &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + Request: req, + }, nil +} + +func (t *incrementalTransport) hitIncrementalSyncAPI() bool { + t.mu.Lock() + defer t.mu.Unlock() + for _, p := range t.paths { + if strings.Contains(p, "/sql/") { + return true + } + } + return false +} + +// No sql-warehouse-id means incremental sync is off: Validate succeeds without touching +// any warehouse/audit-log API. +func TestValidateWithoutWarehouseIDDisablesIncrementalSync(t *testing.T) { + tr := &incrementalTransport{} + d := newValidateConnector(t, &databricks.NoAuth{}, tr) + + if _, err := d.Validate(context.Background()); err != nil { + t.Fatalf("Validate: want nil with incremental sync disabled, got %v", err) + } + if tr.hitIncrementalSyncAPI() { + t.Errorf("Validate made incremental-sync calls with no sql-warehouse-id: %v", tr.paths) + } +} + +// 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: "one workspace in scope", workspaces: []string{"ws1"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + tr := &incrementalTransport{} + d := newValidateConnector(t, &databricks.NoAuth{}, tr) + d.sqlWarehouseID = "wh-missing" + d.workspaces = tc.workspaces + + if _, err := d.Validate(context.Background()); err != nil { + t.Fatalf("Validate: want nil, got %v", err) + } + if tr.hitIncrementalSyncAPI() { + t.Errorf("Validate made incremental-sync calls: %v", tr.paths) + } + }) + } +} + +// 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, " ", nil) + if err != nil { + t.Fatalf("New: %v", err) + } + if d.sqlWarehouseID != "" { + t.Errorf("sqlWarehouseID = %q, want empty (incremental sync disabled)", d.sqlWarehouseID) + } +} + +// With no sql-warehouse-id the feed is a no-op: no events, no API calls. +func TestListEventsWithoutWarehouseIDIsNoop(t *testing.T) { + tr := &incrementalTransport{} + d := newValidateConnector(t, &databricks.NoAuth{}, tr) + feed := newAuditEventFeed(d.client, nil, "") + + events, state, _, err := feed.ListEvents(context.Background(), nil, &pagination.StreamToken{}) + if err != nil { + t.Fatalf("ListEvents: %v", err) + } + if len(events) != 0 || state.HasMore { + t.Errorf("ListEvents = %d events, HasMore=%v; want none", len(events), state.HasMore) + } + if len(tr.paths) != 0 { + t.Errorf("ListEvents made API calls with incremental sync disabled: %v", tr.paths) + } +} diff --git a/pkg/connector/workspaces.go b/pkg/connector/workspaces.go index 20990ca6..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, } @@ -269,6 +269,31 @@ func (w *workspaceBuilder) Grants(ctx context.Context, resource *v2.Resource, at return rv, &rs.SyncOpResults{NextPageToken: nextPage, Annotations: annos}, nil } +// Get re-fetches a single workspace, used to re-sync it after a RESOURCE_CHANGE event. +// GetWorkspace excludes it the same way ListWorkspaces does, so a deliberately excluded +// workspace naturally comes back not-found rather than needing a separate check here. +func (w *workspaceBuilder) Get(ctx context.Context, resourceId *v2.ResourceId, parentResourceId *v2.ResourceId) (*v2.Resource, annotations.Annotations, error) { + workspace, rateLimitData, err := w.client.GetWorkspace(databricks.WithoutCache(ctx), resourceId.Resource) + annos := annotations.Annotations{} + if rateLimitData != nil { + annos.WithRateLimiting(rateLimitData) + } + if err != nil { + return nil, annos, fmt.Errorf("databricks-connector: failed to get workspace %s: %w", resourceId.Resource, err) + } + + resource, err := workspaceResource(ctx, workspace, parentResourceId) + if err != nil { + return nil, annos, err + } + + ctxzap.Extract(ctx).Debug("databricks-connector: Get: refreshed workspace", + zap.String("resource_id", resourceId.Resource), + zap.String("parent_resource_id", parentResourceId.GetResource()), + ) + return resource, annos, nil +} + // workspacePrincipalID returns the numeric principal ID the permission assignments API expects. // Groups are only synced under the account, so only "account//group/" group IDs of // the given account are accepted; anything else is rejected. @@ -361,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 0678ab24..2ff8722c 100644 --- a/pkg/databricks/client.go +++ b/pkg/databricks/client.go @@ -11,10 +11,10 @@ import ( "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" "go.uber.org/zap" - "google.golang.org/grpc/codes" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" "github.com/conductorone/baton-sdk/pkg/uhttp" + "google.golang.org/grpc/codes" ) const ( @@ -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 { @@ -660,6 +668,30 @@ func (c *Client) ListWorkspaces( return filtered, ratelimitData, nil } +// GetWorkspace finds a single workspace by deployment name from the account's +// workspace list (there is no single-workspace GET endpoint). +func (c *Client) GetWorkspace( + ctx context.Context, + deploymentName string, +) ( + *Workspace, + *v2.RateLimitDescription, + error, +) { + workspaces, ratelimitData, err := c.ListWorkspaces(ctx) + if err != nil { + return nil, ratelimitData, err + } + + for _, w := range workspaces { + if w.DeploymentName == deploymentName { + return &w, ratelimitData, nil + } + } + + return nil, ratelimitData, uhttp.WrapErrors(codes.NotFound, fmt.Sprintf("workspace %s not found", deploymentName)) +} + func (c *Client) ListWorkspaceMembers( ctx context.Context, workspaceId string, @@ -689,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 e22fb6c2..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,10 +90,51 @@ 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 +// callers that need fresh data (statement polling, incremental-sync refreshes). +func WithoutCache(ctx context.Context) context.Context { + return context.WithValue(ctx, noCacheKey{}, true) +} + +func isNoCache(ctx context.Context) bool { + v, _ := ctx.Value(noCacheKey{}).(bool) + return v +} + func (c *Client) Get( ctx context.Context, urlAddress *url.URL, - response interface{}, + response any, + params ...Vars, +) (*v2.RateLimitDescription, error) { + return c.doRequest( + ctx, + urlAddress, + 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( @@ -100,6 +143,7 @@ func (c *Client) Get( http.MethodGet, nil, response, + []uhttp.RequestOption{uhttp.WithNoCache()}, params..., ) } @@ -107,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( @@ -117,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..., ) } @@ -169,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 { @@ -179,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 } @@ -235,9 +260,12 @@ func (c *Client) doRequestWithOptions( if body != nil { options = append(options, uhttp.WithJSONBody(body)) } - options = append(options, extraOptions...) + if method == http.MethodGet && isNoCache(ctx) { + options = append(options, uhttp.WithNoCache()) + } + 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 } @@ -251,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). @@ -273,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 @@ -282,6 +330,7 @@ func (c *Client) doRequestWithOptions( return ratelimitData, &APIError{ StatusCode: resp.StatusCode, + ErrorCode: errorResponse.ErrorCode, Detail: errorResponse.Detail, Message: errorResponse.Message, Err: err, @@ -292,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) - 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...) + req, err := c.prepareRequest(ctx, urlAddress, method, body, nil, params...) 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/sql.go b/pkg/databricks/sql.go new file mode 100644 index 00000000..f7818fc6 --- /dev/null +++ b/pkg/databricks/sql.go @@ -0,0 +1,340 @@ +package databricks + +import ( + "context" + "errors" + "fmt" + "net/http" + "strconv" + "time" + + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "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" +) + +const ( + statementsEndpoint = "/api/2.0/sql/statements" + warehousesEndpoint = "/api/2.0/sql/warehouses" + + statementWaitTimeout = "30s" + + warehouseStateRunning = "RUNNING" + warehouseStateStarting = "STARTING" +) + +// Vars so tests can shorten them. +var ( + statementPollInterval = 2 * time.Second + statementPollMaxWait = 5 * time.Minute +) + +type StatementState string + +const ( + StatementStatePending StatementState = "PENDING" + StatementStateRunning StatementState = "RUNNING" + StatementStateSucceeded StatementState = "SUCCEEDED" + StatementStateFailed StatementState = "FAILED" + StatementStateCanceled StatementState = "CANCELED" + StatementStateClosed StatementState = "CLOSED" +) + +// StatementParameter binds a named parameter referenced as ":name" in a SQL statement. +type StatementParameter struct { + Name string `json:"name"` + Value string `json:"value"` + Type string `json:"type,omitempty"` +} + +type statementRequestBody struct { + WarehouseID string `json:"warehouse_id"` + Statement string `json:"statement"` + WaitTimeout string `json:"wait_timeout,omitempty"` + Format string `json:"format"` + Disposition string `json:"disposition"` + Parameters []StatementParameter `json:"parameters,omitempty"` +} + +type statementError struct { + ErrorCode string `json:"error_code"` + Message string `json:"message"` +} + +type statementStatus struct { + State StatementState `json:"state"` + Error *statementError `json:"error,omitempty"` +} + +type statementManifest struct { + Schema struct { + Columns []struct { + Name string `json:"name"` + } `json:"columns"` + } `json:"schema"` +} + +type statementResultChunk struct { + NextChunkIndex *int `json:"next_chunk_index,omitempty"` + DataArray [][]string `json:"data_array"` +} + +type statementResponse struct { + StatementID string `json:"statement_id"` + Status statementStatus `json:"status"` + Manifest statementManifest `json:"manifest"` + Result statementResultChunk `json:"result"` +} + +// StatementResult is the flattened result of a SQL statement executed via the +// Statement Execution API, with all result chunks already collected. +type StatementResult struct { + Columns []string + Rows [][]string +} + +// ExecuteStatement runs a SQL statement via the Statement Execution API and returns every +// row, along with the rate-limit info from the last call that reported any. +func (c *Client) ExecuteStatement( + ctx context.Context, + workspaceId string, + warehouseId string, + statement string, + params ...StatementParameter, +) (*StatementResult, *v2.RateLimitDescription, error) { + l := ctxzap.Extract(ctx) + u := c.workspaceUrl(workspaceId).JoinPath(statementsEndpoint) + + body := statementRequestBody{ + WarehouseID: warehouseId, + Statement: statement, + WaitTimeout: statementWaitTimeout, + Format: "JSON_ARRAY", + Disposition: "INLINE", + Parameters: params, + } + + var res statementResponse + rateLimit, err := c.Post(ctx, u, body, &res) + if err != nil { + return nil, rateLimit, fmt.Errorf("failed to submit statement: %w", err) + } + l.Debug("ExecuteStatement: submitted sql statement", + zap.String("workspace_id", workspaceId), + zap.String("warehouse_id", warehouseId), + zap.String("statement_id", res.StatementID), + zap.String("state", string(res.Status.State)), + ) + + res, polledRateLimit, err := c.pollStatement(ctx, workspaceId, warehouseId, res) + if polledRateLimit != nil { + rateLimit = polledRateLimit + } + if err != nil { + return nil, rateLimit, err + } + + if res.Status.State != StatementStateSucceeded { + msg := "" + if res.Status.Error != nil { + msg = res.Status.Error.Message + } + // FAILED is usually a query/permission problem; CANCELED/CLOSED means it was stopped. + code := codes.FailedPrecondition + if res.Status.State != StatementStateFailed { + code = codes.Aborted + } + return nil, rateLimit, uhttp.WrapErrors(code, fmt.Sprintf("statement %s did not succeed: state=%s message=%s", res.StatementID, res.Status.State, msg)) + } + + result, resultRateLimit, err := c.collectStatementResult(ctx, workspaceId, res) + if resultRateLimit != nil { + rateLimit = resultRateLimit + } + if err != nil { + return nil, rateLimit, err + } + + l.Debug("ExecuteStatement: sql statement succeeded", + zap.String("statement_id", res.StatementID), + zap.Int("column_count", len(result.Columns)), + zap.Int("row_count", len(result.Rows)), + ) + return result, rateLimit, nil +} + +// pollStatement blocks until the statement reaches a terminal state, for the case of a +// cold warehouse start still running after the initial statementWaitTimeout. Capped at +// statementPollMaxWait so a warehouse stuck PENDING/RUNNING can't hang the caller +// indefinitely; on giving up (or on ctx cancellation) it cancels the statement so it stops +// occupying the warehouse. +func (c *Client) pollStatement(ctx context.Context, workspaceId, warehouseId string, res statementResponse) (statementResponse, *v2.RateLimitDescription, error) { + l := ctxzap.Extract(ctx) + + pollCtx, cancel := context.WithTimeout(ctx, statementPollMaxWait) + defer cancel() + + var rateLimit *v2.RateLimitDescription + for res.Status.State == StatementStatePending || res.Status.State == StatementStateRunning { + select { + case <-pollCtx.Done(): + return res, rateLimit, c.abandonStatement(ctx, workspaceId, warehouseId, res) + case <-time.After(statementPollInterval): + } + + l.Debug("polling databricks sql statement", zap.String("statement_id", res.StatementID), zap.String("state", string(res.Status.State))) + + u := c.workspaceUrl(workspaceId).JoinPath(statementsEndpoint, res.StatementID) + var polled statementResponse + // No cache: the URL is identical on every poll, so a cached PENDING/RUNNING + // response would otherwise be returned until statementPollMaxWait. + polledRateLimit, err := c.Get(WithoutCache(pollCtx), u, &polled) + if polledRateLimit != nil { + rateLimit = polledRateLimit + } + if err != nil { + // The poll deadline can expire mid-request; that's still a timeout. + if pollCtx.Err() != nil { + return res, rateLimit, c.abandonStatement(ctx, workspaceId, warehouseId, res) + } + return res, rateLimit, fmt.Errorf("failed to poll statement %s: %w", res.StatementID, err) + } + res = polled + } + + return res, rateLimit, nil +} + +// abandonStatement cancels a statement that didn't finish in time and returns ctx's error if +// the caller gave up, or else a DeadlineExceeded error saying why the statement didn't finish. +func (c *Client) abandonStatement(ctx context.Context, workspaceId, warehouseId string, res statementResponse) error { + c.cancelStatement(ctx, workspaceId, res.StatementID) + if err := ctx.Err(); err != nil { + return err + } + + l := ctxzap.Extract(ctx) + l.Debug("sql statement did not reach a terminal state before poll timeout, canceled it", + zap.String("statement_id", res.StatementID), + zap.String("state", string(res.Status.State)), + zap.Duration("max_wait", statementPollMaxWait), + ) + + msg := fmt.Sprintf("statement %s did not finish within %s (state=%s)", res.StatementID, statementPollMaxWait, res.Status.State) + switch res.Status.State { + case StatementStateRunning: + msg += ": the query was still running" + case StatementStatePending: + // PENDING means waiting for compute, so the warehouse state says why. + state, err := c.warehouseState(ctx, workspaceId, warehouseId) + switch { + case err != nil: + l.Debug("failed to get sql warehouse state", zap.String("warehouse_id", warehouseId), zap.Error(err)) + case state == "": + case state == warehouseStateRunning: + msg += fmt.Sprintf(": SQL warehouse %s is running but did not pick up the statement; it may be at capacity", warehouseId) + case state == warehouseStateStarting: + msg += fmt.Sprintf(": SQL warehouse %s was still starting; a serverless warehouse starts much faster", warehouseId) + default: + msg += fmt.Sprintf(": SQL warehouse %s is %s and has not started; check its Monitoring tab for start failures", warehouseId, state) + } + default: + // Terminal states never reach here: polling only times out while PENDING or RUNNING. + } + + return uhttp.WrapErrors(codes.DeadlineExceeded, msg) +} + +// warehouseState returns the warehouse's state (e.g. RUNNING, STARTING, STOPPED). Like +// cancelStatement, it ignores ctx's cancellation. +func (c *Client) warehouseState(ctx context.Context, workspaceId, warehouseId string) (string, error) { + reqCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) + defer cancel() + + u := c.workspaceUrl(workspaceId).JoinPath(warehousesEndpoint, warehouseId) + var res struct { + State string `json:"state"` + } + if _, err := c.Get(WithoutCache(reqCtx), u, &res); err != nil { + return "", err + } + return res.State, nil +} + +// cancelStatement best-effort cancels a statement we've given up polling on. The request uses +// ctx's values (for the logger) but not its cancellation, since ctx/pollCtx may already be +// done. Uses /cancel, not DELETE (which only closes an already-terminal statement and +// wouldn't stop one still running). +func (c *Client) cancelStatement(ctx context.Context, workspaceId, statementId string) { + l := ctxzap.Extract(ctx) + + reqCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) + defer cancel() + + u := c.workspaceUrl(workspaceId).JoinPath(statementsEndpoint, statementId, "cancel") + response := struct{}{} + if _, err := c.Post(reqCtx, u, nil, &response); err != nil { + l.Debug("failed to cancel timed-out sql statement", zap.String("statement_id", statementId), zap.Error(err)) + } +} + +func (c *Client) collectStatementResult(ctx context.Context, workspaceId string, res statementResponse) (*StatementResult, *v2.RateLimitDescription, error) { + columns := make([]string, len(res.Manifest.Schema.Columns)) + for i, col := range res.Manifest.Schema.Columns { + columns[i] = col.Name + } + + rows := make([][]string, 0, len(res.Result.DataArray)) + rows = append(rows, res.Result.DataArray...) + + var rateLimit *v2.RateLimitDescription + nextChunk := res.Result.NextChunkIndex + for nextChunk != nil { + u := c.workspaceUrl(workspaceId).JoinPath(statementsEndpoint, res.StatementID, "result", "chunks", strconv.Itoa(*nextChunk)) + var chunk statementResultChunk + chunkRateLimit, err := c.Get(ctx, u, &chunk) + if chunkRateLimit != nil { + rateLimit = chunkRateLimit + } + if err != nil { + return nil, rateLimit, fmt.Errorf("failed to fetch statement result chunk %d: %w", *nextChunk, err) + } + ctxzap.Extract(ctx).Debug("collectStatementResult: fetched sql statement result chunk", + zap.String("statement_id", res.StatementID), + zap.Int("chunk_index", *nextChunk), + zap.Int("row_count", len(chunk.DataArray)), + ) + rows = append(rows, chunk.DataArray...) + nextChunk = chunk.NextChunkIndex + } + + return &StatementResult{Columns: columns, Rows: rows}, rateLimit, nil +} + +// WarehouseExists reports whether warehouseId exists in workspaceId, since SQL warehouses +// are workspace-scoped with no account-level lookup. +func (c *Client) WarehouseExists(ctx context.Context, workspaceId, warehouseId string) (bool, *v2.RateLimitDescription, error) { + u := c.workspaceUrl(workspaceId).JoinPath(warehousesEndpoint, warehouseId) + + var res struct{} + rateLimit, err := c.Get(ctx, u, &res) + if err != nil { + var apiErr *APIError + if errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusNotFound { + ctxzap.Extract(ctx).Debug("WarehouseExists: sql warehouse not found in workspace", + zap.String("workspace_id", workspaceId), + zap.String("warehouse_id", warehouseId), + ) + return false, rateLimit, nil + } + return false, rateLimit, err + } + + ctxzap.Extract(ctx).Debug("WarehouseExists: sql warehouse found in workspace", + zap.String("workspace_id", workspaceId), + zap.String("warehouse_id", warehouseId), + ) + return true, rateLimit, nil +} diff --git a/pkg/databricks/sql_test.go b/pkg/databricks/sql_test.go new file mode 100644 index 00000000..b86d24e3 --- /dev/null +++ b/pkg/databricks/sql_test.go @@ -0,0 +1,130 @@ +package databricks + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "sync/atomic" + "testing" + "time" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +type rewriteHostTransport struct{ target *url.URL } + +func (t rewriteHostTransport) RoundTrip(req *http.Request) (*http.Response, error) { + req = req.Clone(req.Context()) + req.URL.Scheme = t.target.Scheme + req.URL.Host = t.target.Host + return http.DefaultTransport.RoundTrip(req) +} + +func newSQLTestClient(t *testing.T, handler http.HandlerFunc) *Client { + t.Helper() + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + target, err := url.Parse(server.URL) + if err != nil { + t.Fatalf("url.Parse: %v", err) + } + c, err := NewClient(context.Background(), &http.Client{Transport: rewriteHostTransport{target: target}}, + "example.cloud.databricks.com", "accounts.cloud.databricks.com", "acc-1", "", &NoAuth{}, nil) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + return c +} + +func shortenStatementPolling(t *testing.T) { + t.Helper() + interval, maxWait := statementPollInterval, statementPollMaxWait + statementPollInterval, statementPollMaxWait = 10*time.Millisecond, 100*time.Millisecond + t.Cleanup(func() { statementPollInterval, statementPollMaxWait = interval, maxWait }) +} + +func TestExecuteStatementTimeoutExplainsWarehouseState(t *testing.T) { + shortenStatementPolling(t) + + cases := []struct { + name string + statementState StatementState + warehouseState string + want string + }{ + {"stopped warehouse", StatementStatePending, "STOPPED", "SQL warehouse wh-1 is STOPPED and has not started"}, + {"starting warehouse", StatementStatePending, "STARTING", "SQL warehouse wh-1 was still starting"}, + {"running warehouse", StatementStatePending, "RUNNING", "SQL warehouse wh-1 is running but did not pick up the statement"}, + {"running query", StatementStateRunning, "RUNNING", "the query was still running"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var canceled atomic.Bool + c := newSQLTestClient(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case strings.HasSuffix(r.URL.Path, "/cancel"): + canceled.Store(true) + fmt.Fprint(w, `{}`) + case strings.HasSuffix(r.URL.Path, "/sql/warehouses/wh-1"): + fmt.Fprintf(w, `{"id":"wh-1","state":%q}`, tc.warehouseState) + default: + fmt.Fprintf(w, `{"statement_id":"stmt-1","status":{"state":%q}}`, tc.statementState) + } + }) + + _, _, err := c.ExecuteStatement(context.Background(), "dbc-aaa", "wh-1", "SELECT 1") + if err == nil { + t.Fatal("ExecuteStatement() error = nil, want a timeout") + } + if got := status.Code(err); got != codes.DeadlineExceeded { + t.Errorf("error code = %s, want %s", got, codes.DeadlineExceeded) + } + if !strings.Contains(err.Error(), tc.want) { + t.Errorf("error = %q, want it to contain %q", err, tc.want) + } + if !canceled.Load() { + t.Error("statement was not canceled") + } + }) + } +} + +// A poll that is still in flight when the deadline expires must also cancel the statement. +func TestExecuteStatementTimeoutDuringPollCancelsStatement(t *testing.T) { + shortenStatementPolling(t) + + var canceled atomic.Bool + c := newSQLTestClient(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case strings.HasSuffix(r.URL.Path, "/cancel"): + canceled.Store(true) + fmt.Fprint(w, `{}`) + case strings.HasSuffix(r.URL.Path, "/sql/warehouses/wh-1"): + fmt.Fprint(w, `{"id":"wh-1","state":"STOPPED"}`) + case r.Method == http.MethodGet: + select { + case <-r.Context().Done(): + case <-time.After(2 * time.Second): + } + default: + fmt.Fprint(w, `{"statement_id":"stmt-1","status":{"state":"PENDING"}}`) + } + }) + + _, _, err := c.ExecuteStatement(context.Background(), "dbc-aaa", "wh-1", "SELECT 1") + if got := status.Code(err); got != codes.DeadlineExceeded { + t.Fatalf("error = %v, want code %s", err, codes.DeadlineExceeded) + } + if !strings.Contains(err.Error(), "has not started") { + t.Errorf("error = %q, want the warehouse state explained", err) + } + if !canceled.Load() { + t.Error("statement was not canceled") + } +} 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) + } +} diff --git a/test-server/README.md b/test-server/README.md new file mode 100644 index 00000000..4162cd4c --- /dev/null +++ b/test-server/README.md @@ -0,0 +1,287 @@ +# Databricks test server + +A local stand-in for the Databricks APIs that baton-databricks calls from `Validate` and the +incremental sync event feed (`databricks_audit_log`, which reads `system.access.audit` through the +SQL Statement Execution API). It lets you run the event feed end to end without a Databricks +account, locally and in CI. + +Go standard library only, built from the root module. State is in memory, seeded at startup and +reset on every start. Full sync and provisioning endpoints are not implemented yet. + +## Running it + +```bash +# Start the server from the repo root; it writes its CA certificate to ./test-server-ca.pem +go run ./test-server/ + +# Build the connector and run the event feed against it; events are printed to stdout as JSON lines +go build ./cmd/baton-databricks +SSL_CERT_FILE=./test-server-ca.pem ./baton-databricks \ + --hostname localhost:8765 --account-hostname localhost:8765 \ + --account-id test-account \ + --databricks-client-id test-client --databricks-client-secret test-secret \ + --sql-warehouse-id wh-serverless \ + --event-feed true --event-feed-id databricks_audit_log \ + --event-feed-start-at "$(date -u -d '2 days ago' +%Y-%m-%dT%H:%M:%SZ)" +``` + +`--event-feed true` is what selects the event feed mode; `--event-feed-id` alone doesn't. The +connector's logs go to stderr, so stdout holds only events. + +To compare with the expected output (event times depend on when the server started, so +`occurredAt` is dropped): + +```bash +normalize() { jq -c -S 'del(.occurredAt)' "$1" | sort; } +diff <(normalize test-server/testdata/expected-events.jsonl) <(normalize events.jsonl) +``` + +### Flags + +| Flag | Default | Meaning | +|---|---|---| +| `--addr` | `localhost:8765` | Listen address | +| `--ca-out` | `./test-server-ca.pem` | Where to write the generated CA certificate | +| `--warehouse-start-delay` | `0s` | How long a stopped warehouse takes to reach `RUNNING` once a statement starts it | +| `--chunk-rows` | `0` | Rows per result chunk; `0` returns every row in one chunk | + +### TLS and hosts + +The connector always builds `https://` URLs, so the server serves TLS. At startup it creates a +throwaway CA and a certificate for `localhost`, `*.localhost` and each seeded workspace host +(`ws-alpha.localhost`, ...), and writes the CA to `--ca-out`. The connector trusts it through +`SSL_CERT_FILE`. The workspace hosts are listed because OpenSSL, and so curl, doesn't match a +`*.localhost` wildcard; Go does. + +The server routes by `Host` header: + +- `localhost:8765` is the account API and OAuth host (`--account-hostname`). +- `.localhost:8765` is that workspace (the connector builds + `.`). Go resolves any `*.localhost` name to loopback, so no DNS setup + is needed. An unknown deployment name gets a 404. + +**macOS:** Go ignores `SSL_CERT_FILE` on macOS and uses the keychain. Either trust the CA in the +keychain or run the connector in a Linux container. + +### Other warehouses + +- `--sql-warehouse-id wh-pro` reads the same events through a stopped warehouse in `ws-beta`. + The first statement starts it. With `--warehouse-start-delay 35s` the connector's first wait (30s) + times out and it then polls until the warehouse is running. +- `--sql-warehouse-id wh-broken` never starts. The connector gives up after its 5-minute limit, + cancels the statement, and fails with "SQL warehouse wh-broken is STOPPED and has not started". + +## Auth + +| | Value | +|---|---| +| Account ID | `test-account` | +| Client ID / secret | `test-client` / `test-secret` | +| Access token | `test-access-token` | + +The token endpoint takes the client ID and secret as HTTP Basic auth, with form fields +`grant_type=client_credentials` and `scope=all-apis`. Every other endpoint (except `/health`) +needs `Authorization: Bearer test-access-token` and returns 401 +`{"error_code": "UNAUTHENTICATED", ...}` without it. + +## Endpoints + +| Host | Method | Path | Docs | +|---|---|---|---| +| account | POST | `/oidc/accounts/{account_id}/v1/token` | [OAuth M2M](https://docs.databricks.com/aws/en/dev-tools/auth/oauth-m2m) | +| account | GET | `/api/2.0/preview/accounts/{account_id}/access-control/assignable-roles` | [Account access control](https://docs.databricks.com/api/account/accountaccesscontrol/getassignablerolesforresource) | +| account | GET | `/api/2.0/accounts/{account_id}/workspaces` | [List workspaces](https://docs.databricks.com/api/account/workspaces/list) | +| account | GET | `/api/2.0/accounts/{account_id}/scim/v2/Users/{id}` | [Get account user](https://docs.databricks.com/api/account/accountusers/get) | +| account | GET | `/api/2.0/accounts/{account_id}/scim/v2/Groups/{id}` | [Get account group](https://docs.databricks.com/api/account/accountgroups/get) | +| account | GET | `/api/2.0/accounts/{account_id}/scim/v2/ServicePrincipals/{id}` | [Get account service principal](https://docs.databricks.com/api/account/accountserviceprincipals/get) | +| workspace | GET | `/api/2.0/sql/warehouses/{id}` | [Get a warehouse](https://docs.databricks.com/api/workspace/warehouses/get) | +| workspace | POST | `/api/2.0/sql/statements` | [Execute a statement](https://docs.databricks.com/api/workspace/statementexecution/executestatement) | +| workspace | GET | `/api/2.0/sql/statements/{id}` | [Get a statement](https://docs.databricks.com/api/workspace/statementexecution/getstatement) | +| workspace | POST | `/api/2.0/sql/statements/{id}/cancel` | [Cancel a statement](https://docs.databricks.com/api/workspace/statementexecution/cancelexecution) | +| workspace | GET | `/api/2.0/sql/statements/{id}/result/chunks/{n}` | [Get result chunk](https://docs.databricks.com/api/workspace/statementexecution/getstatementresultchunkn) | +| any | GET | `/health` | None: local readiness check | + +Every route in `handlers.go` links its doc page. + +## Behavior and provenance + +Code comments label each behavior: + +- `DOC`: stated in the official Databricks docs (with a link). +- `ASSUMPTION`: not documented; our best guess at the real behavior. +- `DECISION`: a choice made for this server, with no real-world counterpart. + +### From the docs + +- Token endpoint: Basic client credentials, `grant_type=client_credentials`, `scope=all-apis`; + returns `{access_token, token_type: "Bearer", expires_in: 3600}`. +- API errors use `{error_code, message}`; a missing or wrong token is 401 `UNAUTHENTICATED`. +- The workspaces list is a bare JSON array (`workspace_id` int64, `workspace_name`, + `deployment_name`, `aws_region`, `workspace_status`, ...). +- SCIM user, group and service principal objects and their fields. +- Warehouse fields and `state` values (`STARTING`, `RUNNING`, `STOPPED`, ...). A statement on a + warehouse that isn't running stays `PENDING`. +- `wait_timeout` is `0s` or `5s`–`50s` (default `10s`); anything else is 400 + `INVALID_PARAMETER_VALUE`. `on_wait_timeout` is `CONTINUE` (default) or `CANCEL`. +- A statement that finishes within `wait_timeout` returns its result; otherwise the response has + its ID and `PENDING`. GET returns the manifest and first chunk once `SUCCEEDED`. Cancel returns + `{}` and the statement becomes `CANCELED`. +- `data_array` values are strings, nulls are JSON `null`, and `next_chunk_index` is absent on the + last chunk. Column `type_name` values (`STRING`, `TIMESTAMP`, `LONG`, `MAP`, ...). +- Parameters are named `:name` markers with an optional `type` (default `STRING`). +- `system.access.audit` columns; account-level events have `workspace_id = 0`; the + `request_params` keys of each event ([audit log reference](https://docs.databricks.com/aws/en/admin/account-settings/audit-logs)), + including workspace-level `createGroup` without `targetGroupId` and `updatePermissionAssignment` + without `workspace_id`. +- Regionality: "Most audit logs are only available in the region of the workspace" + ([system tables: audit logs](https://docs.databricks.com/aws/en/admin/system-tables/audit-logs)). + Here, a warehouse sees workspace-level rows only from workspaces in its own region, and + account-level rows from every region. + +### What's not from the docs + +| Kind | Behavior | +|---|---| +| `ASSUMPTION` | Token endpoint failures use the RFC 6749 body `{error, error_description}` (401 bad client, 400 bad grant or scope). | +| `ASSUMPTION` | An account ID other than `test-account` gets 404 `RESOURCE_DOES_NOT_EXIST`. | +| `ASSUMPTION` | SCIM 404 body uses the SCIM error schema (`schemas`, `detail`, `status: "404"`). | +| `ASSUMPTION` | Unknown API paths get 404 `ENDPOINT_NOT_FOUND`. | +| `ASSUMPTION` | A warehouse that doesn't exist in the request's workspace is 404 `RESOURCE_DOES_NOT_EXIST`, both for GET warehouse and for `warehouse_id` on a statement. The docs only say it's an error. | +| `ASSUMPTION` | A positional (unnamed) parameter is 400 `INVALID_PARAMETER_VALUE`. | +| `ASSUMPTION` | Submitting a statement starts a stopped warehouse, which is `RUNNING` after `--warehouse-start-delay`. | +| `ASSUMPTION` | `JSON_ARRAY` value formats: `TIMESTAMP` as `2006-01-02T15:04:05.000Z` (UTC, milliseconds), `DATE` as `2006-01-02`, `MAP` as a JSON object string. | +| `ASSUMPTION` | An empty result is one chunk with `row_count: 0` and no `data_array`. | +| `ASSUMPTION` | Unknown statement IDs are 404 `RESOURCE_DOES_NOT_EXIST`; an out-of-range chunk or a chunk of a statement without a result is 400. | +| `ASSUMPTION` | An unbound `:name` marker fails the statement with `[UNBOUND_SQL_PARAMETER]`. Accepted `TIMESTAMP` parameter formats: RFC 3339 or `yyyy-mm-dd hh:mm:ss[.ffffff]`. | +| `DECISION` | Hardcoded credentials, account ID and token. | +| `DECISION` | TLS with a generated CA; routing by `Host` header; `/health`. | +| `DECISION` | Assignable roles returns a fixed list. | +| `DECISION` | Seed warehouses: `wh-serverless` running, `wh-pro` stopped, `wh-broken` never starts (stays `STOPPED` with `health.status: FAILED`). | +| `DECISION` | Statements go straight from `PENDING` to `SUCCEEDED`/`FAILED` once the warehouse runs; there's no visible `RUNNING` state. | +| `DECISION` | One result chunk unless `--chunk-rows` is set; only `format: JSON_ARRAY` and `disposition: INLINE` (no `EXTERNAL_LINKS`); no 12-hour statement expiry. | +| `DECISION` | SQL handling (below). | +| `DECISION` | Seed data is synthetic (below). | + +### SQL handling (`DECISION`) + +The server doesn't run SQL. It recognizes `SELECT FROM system.access.audit` and: + +- binds the typed `:name` parameters; +- applies the predicates the connector's query uses, recognized by shape: + `event_date >= :d`, `(event_time > :t OR (event_time = :t AND event_id > :id))`, + `event_time <= :cutoff`, `service_name IN (...)` and `action_name IN (...)`; +- requires `ORDER BY event_time ASC, event_id ASC` and applies a trailing `LIMIT n`; +- fails any other statement with `FAILED` and "test-server: unsupported statement". + +This is deliberately tied to the shape of the connector's audit query +(`pkg/connector/audit_event_feed.go`). If that query changes shape, update `statements.go`. + +## Seed data + +All IDs and times are made up. Times are relative to server start. + +**Workspaces** + +| `deployment_name` | `workspace_id` | Region | +|---|---|---| +| `ws-alpha` | 7474645012345671 | us-east-1 | +| `ws-beta` | 7474645012345672 | us-east-1 | +| `ws-gamma` | 8828000011112223 | eu-west-1 | + +**Principals** + +| Kind | ID | Name | +|---|---|---| +| User | 4410000000000001 | alice@example.com | +| User | 4410000000000002 | bob@example.com | +| User | 4410000000000003 | carol@example.com (inactive) | +| User | 4410000000000004 | dave@example.com | +| Group | 5520000000000001 | data-engineers | +| Group | 5520000000000002 | analysts | +| Service principal | 6630000000000001 | ci-bot | +| None (deleted principal) | 9990000000000001 | matches nothing | + +**Warehouses** + +| ID | Workspace | State | +|---|---|---| +| `wh-serverless` | `ws-alpha` | `RUNNING` (serverless) | +| `wh-pro` | `ws-beta` | `STOPPED`; starts on the first statement | +| `wh-broken` | `ws-alpha` | `STOPPED`; never starts | + +**Audit rows and the events they should produce** + +Rows the connector should read fall between 26 and 5 hours before server start: older than its +4-hour lag window and inside an `--event-feed-start-at` of 2 days ago. Event IDs are +`/`, and seed event IDs are `e0000001-0000-4000-8000-0000000000NN`. The expected +events, as protojson, are in `testdata/expected-events.jsonl`. Every principal is parented to the +account, as in full sync. + +| # | Time | Workspace | `service_name` / `action_name` | Key `request_params` | Expected events | +|---|---|---|---|---|---| +| 1 | −26h | 0 | accounts / `add` | targetUserId dave | dave changed | +| 2 | −25h | 0 | accounts / `activateUser` | targetUserId bob | bob changed | +| 3 | −24h | 0 | accounts / `deactivateUser` | targetUserId carol | carol changed | +| 4 | −23h | 0 | accounts / `createGroup` | targetGroupId analysts | analysts changed | +| 5 | −22h | 0 | accounts / `addPrincipalToGroup` | data-engineers ← alice | grant: alice gets data-engineers `member` | +| 6 | −21h | 0 | accounts / `addPrincipalToGroup` | data-engineers ← analysts | grant: group analysts gets data-engineers `member` | +| 7 | −20h | 0 | accounts / `addPrincipalsToGroup` | analysts ← ci-bot | grant: service principal ci-bot gets analysts `member` | +| 8 | −19h | 0 | accounts / `removePrincipalFromGroup` | data-engineers ← bob | revoke: bob loses data-engineers `member` | +| 9 | −18h | 0 | accounts / `removePrincipalFromGroup` | data-engineers ← deleted principal | data-engineers changed (member can't be resolved) | +| 10 | −17h | 0 | accounts / `setAccountAdmin` | targetUserId alice | alice changed; `account_admin` role changed | +| 11 | −16h | ws-alpha | accounts / `setAdmin` | targetUserId bob | bob changed; `account_admin` role changed (the connector refreshes it in case `setAdmin` is account admin, which the reference is ambiguous about) | +| 12 | −15h | ws-alpha | accounts / `removeAdmin` | targetUserId bob | bob changed | +| 13 | −14h | ws-alpha | accounts / `createGroup` (workspace-level, no `targetGroupId`) | targetUserId alice | none | +| 14 | −13h | ws-beta | accounts / `delete` | targetUserId carol | carol changed | +| 15 | −12h | 0 | workspace / `addPermissionAssignment` | workspace_id ws-alpha | ws-alpha changed | +| 16 | −11h | 0 | workspace / `updatePermissionAssignment` (no `workspace_id`) | principal_id dave | none | +| 17 | −10h | ws-beta | workspace / `deletePermissionAssignment` | workspace_id ws-beta | ws-beta changed | +| 18 | −9h | ws-gamma | accounts / `add` | targetUserId alice | none: eu-west-1 isn't visible from a us-east-1 warehouse | +| 19 | −8h | ws-alpha | accounts / `updateUser` | targetUserId alice | alice changed; `ws-alpha:allow-cluster-create` and `ws-alpha:allow-instance-pool-create` changed | +| 20 | −7h | 0 | accounts / `updateGroup` | targetGroupId analysts | analysts changed | +| 21 | −6h | 0 | accounts / `removeGroup` | targetGroupId data-engineers | data-engineers changed | +| 22 | −6h (same time as 21) | 0 | accounts / `deactivateUser` | targetUserId dave | dave changed | +| 23 | −5h30m | ws-alpha | accounts / `changeDatabricksWorkspaceAcl` | targetUserId carol | ws-alpha changed; `ws-alpha:workspace-access` changed | +| 24 | −5h | ws-alpha | clusters / `create` | | none: not requested | +| 25 | −2h | 0 | accounts / `add` | targetUserId bob | none: newer than the 4-hour lag cutoff | +| 26 | −72h | 0 | accounts / `add` | targetUserId carol | none: older than the feed start | + +That adds up to 25 events. + +## curl examples + +```bash +CA="--cacert ./test-server-ca.pem" + +curl -s $CA https://localhost:8765/health + +TOKEN=$(curl -s $CA -u test-client:test-secret \ + -d grant_type=client_credentials -d scope=all-apis \ + https://localhost:8765/oidc/accounts/test-account/v1/token | jq -r .access_token) + +curl -s $CA -H "Authorization: Bearer $TOKEN" https://localhost:8765/api/2.0/accounts/test-account/workspaces | jq +curl -s $CA -H "Authorization: Bearer $TOKEN" https://localhost:8765/api/2.0/accounts/test-account/scim/v2/Users/4410000000000001 | jq +curl -s $CA -H "Authorization: Bearer $TOKEN" https://ws-alpha.localhost:8765/api/2.0/sql/warehouses/wh-serverless | jq + +curl -s $CA -H "Authorization: Bearer $TOKEN" https://ws-alpha.localhost:8765/api/2.0/sql/statements \ + -d '{ + "warehouse_id": "wh-serverless", + "statement": "SELECT event_id, event_time, action_name FROM system.access.audit WHERE action_name IN ('\''add'\'') ORDER BY event_time ASC, event_id ASC LIMIT 10", + "wait_timeout": "10s" + }' | jq +``` + +## Tests and CI + +`main_test.go` tests the server in process: auth, routing, `wait_timeout` validation, warehouse +and statement state transitions (including the stopped and broken warehouses), cancel and +`on_wait_timeout: CANCEL`, chunking, regionality and the audit filter. It runs with +`go test ./...`. + +The `test-server` job in `.github/workflows/ci.yaml` starts the server, runs the event feed +against `wh-serverless` and diffs the events with `testdata/expected-events.jsonl`. It then runs +against `wh-broken` and checks for the "has not started" message. That step waits out the +connector's 5-minute statement limit. The job needs no secrets. + +If a diff shows up, either the connector or the expectation is wrong; check the row's documented +meaning in the table above before updating `expected-events.jsonl`. diff --git a/test-server/handlers.go b/test-server/handlers.go new file mode 100644 index 00000000..1b88e348 --- /dev/null +++ b/test-server/handlers.go @@ -0,0 +1,262 @@ +package main + +import ( + "encoding/json" + "log" + "net/http" + "strconv" + "time" +) + +func (s *server) accountRoutes() *http.ServeMux { + mux := http.NewServeMux() + // DOC: https://docs.databricks.com/aws/en/dev-tools/auth/oauth-m2m + mux.HandleFunc("POST /oidc/accounts/{account_id}/v1/token", requireAccount(s.handleToken)) + // DOC: https://docs.databricks.com/api/account/accountaccesscontrol/getassignablerolesforresource + mux.HandleFunc("GET /api/2.0/preview/accounts/{account_id}/access-control/assignable-roles", + requireBearer(requireAccount(s.handleAssignableRoles))) + // DOC: https://docs.databricks.com/api/account/workspaces/list + mux.HandleFunc("GET /api/2.0/accounts/{account_id}/workspaces", requireBearer(requireAccount(s.handleListWorkspaces))) + // DOC: https://docs.databricks.com/api/account/accountusers/get + mux.HandleFunc("GET /api/2.0/accounts/{account_id}/scim/v2/Users/{id}", requireBearer(requireAccount(s.handleGetUser))) + // DOC: https://docs.databricks.com/api/account/accountgroups/get + mux.HandleFunc("GET /api/2.0/accounts/{account_id}/scim/v2/Groups/{id}", requireBearer(requireAccount(s.handleGetGroup))) + // DOC: https://docs.databricks.com/api/account/accountserviceprincipals/get + mux.HandleFunc("GET /api/2.0/accounts/{account_id}/scim/v2/ServicePrincipals/{id}", + requireBearer(requireAccount(s.handleGetServicePrincipal))) + mux.HandleFunc("/", handleNotFound) + return mux +} + +func (s *server) workspaceRoutes() *http.ServeMux { + mux := http.NewServeMux() + // DOC: https://docs.databricks.com/api/workspace/warehouses/get + mux.HandleFunc("GET /api/2.0/sql/warehouses/{id}", requireBearer(s.handleGetWarehouse)) + // DOC: https://docs.databricks.com/api/workspace/statementexecution/executestatement + mux.HandleFunc("POST /api/2.0/sql/statements", requireBearer(s.handleExecuteStatement)) + // DOC: https://docs.databricks.com/api/workspace/statementexecution/getstatement + mux.HandleFunc("GET /api/2.0/sql/statements/{id}", requireBearer(s.handleGetStatement)) + // DOC: https://docs.databricks.com/api/workspace/statementexecution/cancelexecution + mux.HandleFunc("POST /api/2.0/sql/statements/{id}/cancel", requireBearer(s.handleCancelStatement)) + // DOC: https://docs.databricks.com/api/workspace/statementexecution/getstatementresultchunkn + mux.HandleFunc("GET /api/2.0/sql/statements/{id}/result/chunks/{n}", requireBearer(s.handleGetResultChunk)) + mux.HandleFunc("/", handleNotFound) + return mux +} + +// handleNotFound answers any route this server doesn't implement. +// ASSUMPTION: Databricks answers unknown API paths with 404 ENDPOINT_NOT_FOUND. +func handleNotFound(w http.ResponseWriter, r *http.Request) { + writeError(w, http.StatusNotFound, "ENDPOINT_NOT_FOUND", "test-server: no route for "+r.Method+" "+r.URL.Path) +} + +// handleToken implements the OAuth M2M client-credentials token endpoint. +// DOC: HTTP Basic client ID/secret, grant_type=client_credentials, scope=all-apis; returns +// {access_token, token_type: "Bearer", expires_in: 3600}. +// ASSUMPTION: failures use the RFC 6749 error body ({error, error_description}) with 400/401. +func (s *server) handleToken(w http.ResponseWriter, r *http.Request) { + clientID, clientSecret, ok := r.BasicAuth() + if !ok || clientID != testClientID || clientSecret != testClientSecret { + writeJSON(w, http.StatusUnauthorized, map[string]string{ + "error": "invalid_client", + "error_description": "Client authentication failed", + }) + return + } + if err := r.ParseForm(); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid_request", "error_description": err.Error()}) + return + } + if r.PostForm.Get("grant_type") != "client_credentials" { + writeJSON(w, http.StatusBadRequest, map[string]string{ + "error": "unsupported_grant_type", + "error_description": "grant_type must be client_credentials", + }) + return + } + if r.PostForm.Get("scope") != "all-apis" { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid_scope", "error_description": "scope must be all-apis"}) + return + } + writeJSON(w, http.StatusOK, map[string]any{ + "access_token": testAccessToken, + "token_type": "Bearer", + "expires_in": 3600, + "scope": "all-apis", + }) +} + +// handleAssignableRoles returns the roles assignable on a resource. +// DOC: response shape {roles: [{name}]}. DECISION: a fixed list, whatever the resource. +func (s *server) handleAssignableRoles(w http.ResponseWriter, _ *http.Request) { + writeJSON(w, http.StatusOK, map[string]any{ + "roles": []map[string]string{ + {"name": "roles/marketplace.admin"}, + }, + }) +} + +// handleListWorkspaces returns the account's workspaces. +// DOC: a bare JSON array of workspace objects. +func (s *server) handleListWorkspaces(w http.ResponseWriter, _ *http.Request) { + s.mu.Lock() + defer s.mu.Unlock() + + out := make([]map[string]any, 0, len(s.st.workspaces)) + for _, ws := range s.st.workspaces { + out = append(out, map[string]any{ + "account_id": testAccountID, + "workspace_id": ws.ID, + "workspace_name": ws.Name, + "deployment_name": ws.DeploymentName, + "aws_region": ws.Region, + "workspace_status": "RUNNING", + "pricing_tier": "PREMIUM", + "creation_time": s.start.Add(-90 * 24 * time.Hour).UnixMilli(), + }) + } + writeJSON(w, http.StatusOK, out) +} + +// DOC: SCIM user, group and service principal objects, per the account SCIM API pages. +func (s *server) handleGetUser(w http.ResponseWriter, r *http.Request) { + s.mu.Lock() + u, ok := s.st.users[r.PathValue("id")] + s.mu.Unlock() + if !ok { + writeSCIMNotFound(w, "User "+r.PathValue("id")+" not found") + return + } + writeJSON(w, http.StatusOK, map[string]any{ + "schemas": []string{"urn:ietf:params:scim:schemas:core:2.0:User"}, + "id": u.ID, + "userName": u.UserName, + "displayName": u.DisplayName, + "active": u.Active, + "emails": []map[string]any{{"value": u.UserName, "type": "work", "primary": true}}, + }) +} + +func (s *server) handleGetGroup(w http.ResponseWriter, r *http.Request) { + s.mu.Lock() + g, ok := s.st.groups[r.PathValue("id")] + s.mu.Unlock() + if !ok { + writeSCIMNotFound(w, "Group "+r.PathValue("id")+" not found") + return + } + writeJSON(w, http.StatusOK, map[string]any{ + "schemas": []string{"urn:ietf:params:scim:schemas:core:2.0:Group"}, + "id": g.ID, + "displayName": g.DisplayName, + }) +} + +func (s *server) handleGetServicePrincipal(w http.ResponseWriter, r *http.Request) { + s.mu.Lock() + sp, ok := s.st.servicePrincipals[r.PathValue("id")] + s.mu.Unlock() + if !ok { + writeSCIMNotFound(w, "ServicePrincipal "+r.PathValue("id")+" not found") + return + } + writeJSON(w, http.StatusOK, map[string]any{ + "schemas": []string{"urn:ietf:params:scim:schemas:core:2.0:ServicePrincipal"}, + "id": sp.ID, + "applicationId": sp.ApplicationID, + "displayName": sp.DisplayName, + "active": sp.Active, + }) +} + +// writeSCIMNotFound writes a SCIM 404. +// ASSUMPTION: the SCIM error schema (schemas, detail, status); the 404 body is not documented. +func writeSCIMNotFound(w http.ResponseWriter, detail string) { + writeJSON(w, http.StatusNotFound, map[string]any{ + "schemas": []string{"urn:ietf:params:scim:api:messages:2.0:Error"}, + "detail": detail, + "status": "404", + }) +} + +// handleGetWarehouse returns a warehouse in the request's workspace. +// DOC: warehouse fields and state values (STARTING, RUNNING, STOPPING, STOPPED, DELETING, DELETED). +// ASSUMPTION: a warehouse that isn't in this workspace is 404 RESOURCE_DOES_NOT_EXIST. +func (s *server) handleGetWarehouse(w http.ResponseWriter, r *http.Request) { + ws := requestWorkspace(r) + + s.mu.Lock() + defer s.mu.Unlock() + + wh, ok := s.st.warehouses[r.PathValue("id")] + if !ok || wh.DeploymentName != ws.DeploymentName { + writeError(w, http.StatusNotFound, "RESOURCE_DOES_NOT_EXIST", "Warehouse "+r.PathValue("id")+" does not exist.") + return + } + + state := s.warehouseState(wh) + health := map[string]any{"status": "HEALTHY"} + if wh.NeverStarts { + // DECISION: the broken warehouse reports why it can't start. + health = map[string]any{ + "status": "FAILED", + "summary": "test-server: this warehouse never starts", + "failure_reason": map[string]any{ + "code": "INSTANCE_UNREACHABLE", + "type": "CLOUD_FAILURE", + }, + } + } + + writeJSON(w, http.StatusOK, map[string]any{ + "id": wh.ID, + "name": wh.Name, + "warehouse_type": wh.WarehouseType, + "enable_serverless_compute": wh.Serverless, + "cluster_size": "2X-Small", + "min_num_clusters": 1, + "max_num_clusters": 1, + "auto_stop_mins": 10, + "state": state, + "health": health, + "num_clusters": boolToInt(state == warehouseRunning), + "creator_name": "admin@example.com", + "jdbc_url": "jdbc:spark://" + ws.DeploymentName + ".localhost:443/default;transportMode=http;ssl=1;httpPath=/sql/1.0/warehouses/" + wh.ID, + "odbc_params": map[string]any{ + "hostname": ws.DeploymentName + ".localhost", + "path": "/sql/1.0/warehouses/" + wh.ID, + "protocol": "https", + "port": 443, + }, + }) +} + +func boolToInt(b bool) int { + if b { + return 1 + } + return 0 +} + +// apiError is the documented error body: https://docs.databricks.com/api/workspace/introduction. +type apiError struct { + ErrorCode string `json:"error_code"` + Message string `json:"message"` +} + +func writeError(w http.ResponseWriter, status int, code, message string) { + writeJSON(w, status, apiError{ErrorCode: code, Message: message}) +} + +func writeJSON(w http.ResponseWriter, status int, body any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + if err := json.NewEncoder(w).Encode(body); err != nil { + log.Printf("test-server: writing response: %v", err) + } +} + +func parseChunkIndex(s string) (int, bool) { + n, err := strconv.Atoi(s) + return n, err == nil && n >= 0 +} diff --git a/test-server/main.go b/test-server/main.go new file mode 100644 index 00000000..8e23a73d --- /dev/null +++ b/test-server/main.go @@ -0,0 +1,160 @@ +// Command test-server is a local stand-in for the Databricks APIs that baton-databricks calls +// from Validate and the incremental sync event feed (system.access.audit). It serves TLS on +// one address and routes by Host header: "localhost" is the account API and OAuth host, and +// ".localhost" is that workspace. See README.md for usage, endpoints, seed +// data and which behaviors come from the Databricks docs. +package main + +import ( + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "errors" + "flag" + "fmt" + "log" + "math/big" + "net" + "net/http" + "os" + "time" +) + +type config struct { + addr string + caOut string + warehouseStartDelay time.Duration + chunkRows int +} + +func main() { + var cfg config + flag.StringVar(&cfg.addr, "addr", "localhost:8765", "address to listen on") + flag.StringVar(&cfg.caOut, "ca-out", "./test-server-ca.pem", "file to write the generated CA certificate to (PEM)") + flag.DurationVar(&cfg.warehouseStartDelay, "warehouse-start-delay", 0, + "how long a stopped warehouse takes to reach RUNNING after a statement starts it") + flag.IntVar(&cfg.chunkRows, "chunk-rows", 0, "rows per result chunk; 0 returns every row in a single chunk") + flag.Parse() + + if err := run(cfg); err != nil { + fmt.Fprintf(os.Stderr, "test-server: %v\n", err) + os.Exit(1) + } +} + +func run(cfg config) error { + if cfg.chunkRows < 0 { + return errors.New("--chunk-rows must be >= 0") + } + + srv := newServer(cfg, time.Now()) + + hosts := []string{"localhost", "*.localhost"} + for _, ws := range srv.st.workspaces { + hosts = append(hosts, ws.DeploymentName+".localhost") + } + caPEM, cert, err := generateCertificates(hosts) + if err != nil { + return err + } + if err := os.WriteFile(cfg.caOut, caPEM, 0o600); err != nil { + return fmt.Errorf("writing CA certificate: %w", err) + } + httpServer := &http.Server{ + Addr: cfg.addr, + Handler: logRequests(srv), + ReadHeaderTimeout: 10 * time.Second, + TLSConfig: &tls.Config{ + MinVersion: tls.VersionTLS12, + Certificates: []tls.Certificate{cert}, + }, + } + + log.Printf("test-server: listening on https://%s (CA written to %s)", cfg.addr, cfg.caOut) + log.Printf("test-server: account API at https://%s, workspaces at https://.%s", cfg.addr, cfg.addr) + return httpServer.ListenAndServeTLS("", "") +} + +// generateCertificates creates a throwaway CA and a server certificate for hosts (DECISION: the +// connector always builds https:// URLs, so the server needs TLS; clients trust the CA through +// SSL_CERT_FILE). Go accepts the *.localhost wildcard, but OpenSSL (and so curl) doesn't match a +// wildcard that close to the top level, so each seeded workspace host is listed too. +func generateCertificates(hosts []string) ([]byte, tls.Certificate, error) { + now := time.Now() + + caKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + return nil, tls.Certificate{}, fmt.Errorf("generating CA key: %w", err) + } + caTemplate := &x509.Certificate{ + SerialNumber: randomSerial(), + Subject: pkix.Name{CommonName: "baton-databricks test-server CA"}, + NotBefore: now.Add(-time.Hour), + NotAfter: now.Add(7 * 24 * time.Hour), + KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageCRLSign, + BasicConstraintsValid: true, + IsCA: true, + } + caDER, err := x509.CreateCertificate(rand.Reader, caTemplate, caTemplate, &caKey.PublicKey, caKey) + if err != nil { + return nil, tls.Certificate{}, fmt.Errorf("creating CA certificate: %w", err) + } + caCert, err := x509.ParseCertificate(caDER) + if err != nil { + return nil, tls.Certificate{}, fmt.Errorf("parsing CA certificate: %w", err) + } + + leafKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + return nil, tls.Certificate{}, fmt.Errorf("generating server key: %w", err) + } + leafTemplate := &x509.Certificate{ + SerialNumber: randomSerial(), + Subject: pkix.Name{CommonName: "localhost"}, + NotBefore: now.Add(-time.Hour), + NotAfter: now.Add(7 * 24 * time.Hour), + KeyUsage: x509.KeyUsageDigitalSignature, + ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, + DNSNames: hosts, + IPAddresses: []net.IP{net.ParseIP("127.0.0.1"), net.ParseIP("::1")}, + } + leafDER, err := x509.CreateCertificate(rand.Reader, leafTemplate, caCert, &leafKey.PublicKey, caKey) + if err != nil { + return nil, tls.Certificate{}, fmt.Errorf("creating server certificate: %w", err) + } + + caPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: caDER}) + return caPEM, tls.Certificate{Certificate: [][]byte{leafDER, caDER}, PrivateKey: leafKey}, nil +} + +func randomSerial() *big.Int { + serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 62)) + if err != nil { + return big.NewInt(time.Now().UnixNano()) + } + return serial +} + +type statusRecorder struct { + http.ResponseWriter + status int +} + +func (r *statusRecorder) WriteHeader(status int) { + r.status = status + r.ResponseWriter.WriteHeader(status) +} + +func logRequests(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK} + started := time.Now() + next.ServeHTTP(rec, r) + //nolint:gosec // G706: local test-server log; request values are %q-quoted, so they cannot forge log lines. + log.Printf("%q %q -> %d (%s)", r.Method, r.Host+r.URL.Path, rec.status, time.Since(started).Round(time.Millisecond)) + }) +} diff --git a/test-server/main_test.go b/test-server/main_test.go new file mode 100644 index 00000000..525d85ce --- /dev/null +++ b/test-server/main_test.go @@ -0,0 +1,456 @@ +package main + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +var testStart = time.Date(2026, 1, 15, 12, 0, 0, 0, time.UTC) + +// testServer is a server whose clock only moves when the test advances it. +type testServer struct { + *server + clock time.Time +} + +func newTestServer(cfg config) *testServer { + ts := &testServer{server: newServer(cfg, testStart), clock: testStart} + ts.now = func() time.Time { return ts.clock } + return ts +} + +func (ts *testServer) do(t *testing.T, method, url, body string) (int, map[string]any) { + t.Helper() + rec := ts.raw(t, method, url, body, "Bearer "+testAccessToken) + return rec.Code, decodeObject(t, rec) +} + +func (ts *testServer) raw(t *testing.T, method, url, body, authorization string) *httptest.ResponseRecorder { + req := httptest.NewRequestWithContext(t.Context(), method, url, strings.NewReader(body)) + if authorization != "" { + req.Header.Set("Authorization", authorization) + } + if method == http.MethodPost && strings.Contains(url, "/oidc/") { + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + } + rec := httptest.NewRecorder() + ts.ServeHTTP(rec, req) + return rec +} + +func decodeObject(t *testing.T, rec *httptest.ResponseRecorder) map[string]any { + t.Helper() + var out map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatalf("response is not a JSON object: %v: %s", err, rec.Body.String()) + } + return out +} + +const ( + accountURL = "https://localhost:8765" + alphaURL = "https://ws-alpha.localhost:8765" + betaURL = "https://ws-beta.localhost:8765" +) + +// connectorStatement has the shape of the connector's audit query (pkg/connector/audit_event_feed.go). +const connectorStatement = ` + SELECT event_id, event_time, workspace_id, action_name, service_name, request_params + FROM system.access.audit + WHERE event_date >= :start_date + AND (event_time > :start_time OR (event_time = :start_time AND event_id > :start_after_event_id)) + AND event_time <= :lag_cutoff + AND service_name IN ('accounts', 'workspace') + AND action_name IN (%s) + ORDER BY event_time ASC, event_id ASC + LIMIT %d + ` + +const connectorActions = `'activateUser', 'add', 'addPermissionAssignment', 'addPrincipalToGroup', 'addPrincipalsToGroup', + 'changeDatabricksSqlAcl', 'changeDatabricksWorkspaceAcl', 'createGroup', 'deactivateUser', 'delete', + 'deletePermissionAssignment', 'removeAccountAdmin', 'removeAdmin', 'removeGroup', 'removePrincipalFromGroup', + 'removePrincipalsFromGroup', 'setAccountAdmin', 'setAdmin', 'updateGroup', 'updatePermissionAssignment', 'updateUser'` + +func statementRequest(t *testing.T, warehouseID, waitTimeout string, startTime time.Time, startAfter string, limit int) string { + t.Helper() + cutoff := testStart.Add(-4 * time.Hour) + body := map[string]any{ + "warehouse_id": warehouseID, + "statement": fmt.Sprintf(connectorStatement, connectorActions, limit), + "wait_timeout": waitTimeout, + "format": "JSON_ARRAY", + "disposition": "INLINE", + "parameters": []map[string]string{ + {"name": "start_date", "value": startTime.Format(time.DateOnly), "type": "DATE"}, + {"name": "start_time", "value": startTime.Format(time.RFC3339Nano), "type": "TIMESTAMP"}, + {"name": "start_after_event_id", "value": startAfter, "type": "STRING"}, + {"name": "lag_cutoff", "value": cutoff.Format(time.RFC3339Nano), "type": "TIMESTAMP"}, + }, + } + b, err := json.Marshal(body) + if err != nil { + t.Fatal(err) + } + return string(b) +} + +func eventID(n int) string { + return fmt.Sprintf("e0000001-0000-4000-8000-%012d", n) +} + +func eventIDs(t *testing.T, body map[string]any) []string { + t.Helper() + result, _ := body["result"].(map[string]any) + rows, _ := result["data_array"].([]any) + ids := make([]string, 0, len(rows)) + for _, r := range rows { + row, _ := r.([]any) + id, _ := row[0].(string) + ids = append(ids, id) + } + return ids +} + +func stmtState(body map[string]any) string { + status, _ := body["status"].(map[string]any) + s, _ := status["state"].(string) + return s +} + +func TestHealth(t *testing.T) { + ts := newTestServer(config{}) + for _, url := range []string{accountURL + "/health", alphaURL + "/health"} { + if rec := ts.raw(t, http.MethodGet, url, "", ""); rec.Code != http.StatusOK { + t.Errorf("%s: got %d", url, rec.Code) + } + } +} + +func TestToken(t *testing.T) { + ts := newTestServer(config{}) + tokenURL := accountURL + "/oidc/accounts/test-account/v1/token" + form := "grant_type=client_credentials&scope=all-apis" + + req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, tokenURL, strings.NewReader(form)) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.SetBasicAuth(testClientID, testClientSecret) + rec := httptest.NewRecorder() + ts.ServeHTTP(rec, req) + body := decodeObject(t, rec) + if rec.Code != http.StatusOK || body["access_token"] != testAccessToken || body["token_type"] != "Bearer" { + t.Fatalf("good credentials: got %d %v", rec.Code, body) + } + + req = httptest.NewRequestWithContext(t.Context(), http.MethodPost, tokenURL, strings.NewReader(form)) + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.SetBasicAuth(testClientID, "wrong") + rec = httptest.NewRecorder() + ts.ServeHTTP(rec, req) + if rec.Code != http.StatusUnauthorized { + t.Fatalf("bad secret: got %d", rec.Code) + } +} + +func TestBearerRequired(t *testing.T) { + ts := newTestServer(config{}) + for _, auth := range []string{"", "Bearer wrong"} { + rec := ts.raw(t, http.MethodGet, accountURL+"/api/2.0/accounts/test-account/workspaces", "", auth) + body := decodeObject(t, rec) + if rec.Code != http.StatusUnauthorized || body["error_code"] != "UNAUTHENTICATED" { + t.Errorf("auth %q: got %d %v", auth, rec.Code, body) + } + } +} + +func TestAccountEndpoints(t *testing.T) { + ts := newTestServer(config{}) + + rec := ts.raw(t, http.MethodGet, accountURL+"/api/2.0/accounts/test-account/workspaces", "", "Bearer "+testAccessToken) + var workspaces []map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &workspaces); err != nil || len(workspaces) != 3 { + t.Fatalf("workspaces: want a bare array of 3, got %d %s", rec.Code, rec.Body.String()) + } + // Large IDs must come out as exact integers, not floats in scientific notation. + if !strings.Contains(rec.Body.String(), `"workspace_id":7474645012345671`) { + t.Errorf("workspace_id not rendered as an exact integer: %s", rec.Body.String()) + } + + if code, _ := ts.do(t, http.MethodGet, accountURL+"/api/2.0/preview/accounts/test-account/access-control/assignable-roles", ""); code != http.StatusOK { + t.Errorf("assignable roles: got %d", code) + } + if code, _ := ts.do(t, http.MethodGet, accountURL+"/api/2.0/accounts/other-account/workspaces", ""); code != http.StatusNotFound { + t.Errorf("other account: got %d", code) + } + + scim := accountURL + "/api/2.0/accounts/test-account/scim/v2/" + cases := []struct { + path string + want int + }{ + {"Users/" + userAlice, http.StatusOK}, + {"Groups/" + groupAnalysts, http.StatusOK}, + {"ServicePrincipals/" + spCIBot, http.StatusOK}, + {"Users/" + groupAnalysts, http.StatusNotFound}, + {"Users/" + deletedPrincipal, http.StatusNotFound}, + {"Groups/" + deletedPrincipal, http.StatusNotFound}, + {"ServicePrincipals/" + deletedPrincipal, http.StatusNotFound}, + } + for _, c := range cases { + code, body := ts.do(t, http.MethodGet, scim+c.path, "") + if code != c.want { + t.Errorf("%s: got %d, want %d", c.path, code, c.want) + } + if c.want == http.StatusNotFound && body["status"] != "404" { + t.Errorf("%s: want a SCIM error body, got %v", c.path, body) + } + } +} + +func TestRouting(t *testing.T) { + ts := newTestServer(config{}) + if code, _ := ts.do(t, http.MethodGet, "https://ws-nope.localhost:8765/api/2.0/sql/warehouses/wh-serverless", ""); code != http.StatusNotFound { + t.Errorf("unknown workspace host: got %d", code) + } + if code, _ := ts.do(t, http.MethodGet, alphaURL+"/api/2.0/sql/warehouses/wh-serverless", ""); code != http.StatusOK { + t.Errorf("warehouse in its workspace: got %d", code) + } + if code, _ := ts.do(t, http.MethodGet, betaURL+"/api/2.0/sql/warehouses/wh-serverless", ""); code != http.StatusNotFound { + t.Errorf("warehouse from another workspace: got %d", code) + } + if code, _ := ts.do(t, http.MethodGet, accountURL+"/api/2.0/sql/warehouses/wh-serverless", ""); code != http.StatusNotFound { + t.Errorf("workspace API on the account host: got %d", code) + } +} + +func TestWaitTimeoutValidation(t *testing.T) { + ts := newTestServer(config{}) + for _, wait := range []string{"0s", "5s", "50s"} { + code, _ := ts.do(t, http.MethodPost, alphaURL+"/api/2.0/sql/statements", statementRequest(t, "wh-serverless", wait, testStart.Add(-48*time.Hour), "", 10)) + if code != http.StatusOK { + t.Errorf("wait_timeout %s: got %d", wait, code) + } + } + for _, wait := range []string{"1s", "4s", "51s", "10", "5m", "-5s"} { + code, body := ts.do(t, http.MethodPost, alphaURL+"/api/2.0/sql/statements", statementRequest(t, "wh-serverless", wait, testStart.Add(-48*time.Hour), "", 10)) + if code != http.StatusBadRequest || body["error_code"] != "INVALID_PARAMETER_VALUE" { + t.Errorf("wait_timeout %s: got %d %v", wait, code, body) + } + } +} + +func TestStatementErrors(t *testing.T) { + ts := newTestServer(config{}) + url := alphaURL + "/api/2.0/sql/statements" + + code, _ := ts.do(t, http.MethodPost, url, statementRequest(t, "wh-missing", "0s", testStart, "", 10)) + if code != http.StatusNotFound { + t.Errorf("unknown warehouse: got %d", code) + } + code, _ = ts.do(t, http.MethodPost, url, `{"warehouse_id":"wh-serverless","statement":"SELECT ?","parameters":[{"value":"1"}]}`) + if code != http.StatusBadRequest { + t.Errorf("positional parameter: got %d", code) + } + + code, body := ts.do(t, http.MethodPost, url, `{"warehouse_id":"wh-serverless","statement":"SELECT 1","wait_timeout":"0s"}`) + status, _ := body["status"].(map[string]any) + errBody, _ := status["error"].(map[string]any) + msg, _ := errBody["message"].(string) + if code != http.StatusOK || stmtState(body) != statementFailed || !strings.Contains(msg, "test-server: unsupported statement") { + t.Errorf("unsupported statement: got %d %v", code, body) + } + + code, body = ts.do(t, http.MethodPost, url, + `{"warehouse_id":"wh-serverless","statement":"SELECT event_id FROM system.access.audit WHERE event_time > :t ORDER BY event_time, event_id","wait_timeout":"0s"}`) + if code != http.StatusOK || stmtState(body) != statementFailed { + t.Errorf("unbound parameter: got %d %v", code, body) + } + + if code, _ := ts.do(t, http.MethodGet, url+"/no-such-statement", ""); code != http.StatusNotFound { + t.Errorf("unknown statement: got %d", code) + } +} + +func TestAuditQueryRunningWarehouse(t *testing.T) { + ts := newTestServer(config{}) + code, body := ts.do(t, http.MethodPost, alphaURL+"/api/2.0/sql/statements", + statementRequest(t, "wh-serverless", "30s", testStart.Add(-48*time.Hour), "", 1000)) + if code != http.StatusOK || stmtState(body) != statementSucceeded { + t.Fatalf("got %d %v", code, body) + } + + // Visible from us-east-1, inside (start-48h, start-4h], in the requested service/action + // lists: every seed row except 18 (eu-west-1), 24 (clusters), 25 (too new) and 26 (too old). + var want []string + for n := 1; n <= 23; n++ { + if n != 18 { + want = append(want, eventID(n)) + } + } + if got := eventIDs(t, body); strings.Join(got, ",") != strings.Join(want, ",") { + t.Errorf("rows:\n got %v\nwant %v", got, want) + } + + manifest, _ := body["manifest"].(map[string]any) + schema, _ := manifest["schema"].(map[string]any) + columns, _ := schema["columns"].([]any) + if len(columns) != 6 { + t.Errorf("want 6 columns, got %v", columns) + } + result, _ := body["result"].(map[string]any) + if _, ok := result["next_chunk_index"]; ok { + t.Errorf("single chunk must not have next_chunk_index: %v", result) + } + + rows, _ := result["data_array"].([]any) + first, _ := rows[0].([]any) + if first[1] != testStart.Add(-26*time.Hour).Format("2006-01-02T15:04:05.000Z") || first[2] != "0" { + t.Errorf("unexpected value formats: %v", first) + } + var params map[string]string + if s, _ := first[5].(string); json.Unmarshal([]byte(s), ¶ms) != nil || params["targetUserId"] != userDave { + t.Errorf("request_params is not a JSON object string: %v", first[5]) + } +} + +func TestAuditQueryPaging(t *testing.T) { + ts := newTestServer(config{}) + rowTime := func(n int) time.Time { + for _, r := range ts.st.audit { + if r.EventID == eventID(n) { + return r.EventTime + } + } + t.Fatalf("no seed row %d", n) + return time.Time{} + } + + _, body := ts.do(t, http.MethodPost, alphaURL+"/api/2.0/sql/statements", statementRequest(t, "wh-serverless", "10s", rowTime(5), eventID(5), 3)) + if got := eventIDs(t, body); strings.Join(got, ",") != strings.Join([]string{eventID(6), eventID(7), eventID(8)}, ",") { + t.Errorf("page after row 5: got %v", got) + } + + // Rows 21 and 22 share an event_time: the event_id tiebreaker must return only 22. + _, body = ts.do(t, http.MethodPost, alphaURL+"/api/2.0/sql/statements", statementRequest(t, "wh-serverless", "10s", rowTime(21), eventID(21), 1)) + if got := eventIDs(t, body); len(got) != 1 || got[0] != eventID(22) { + t.Errorf("tiebreak after row 21: got %v", got) + } +} + +func TestRegionality(t *testing.T) { + ts := newTestServer(config{}) + stmt := fmt.Sprintf(connectorStatement, connectorActions, 1000) + params := []statementParameter{ + {Name: "start_date", Value: ptr(testStart.Add(-48 * time.Hour).Format(time.DateOnly)), Type: "DATE"}, + {Name: "start_time", Value: ptr(testStart.Add(-48 * time.Hour).Format(time.RFC3339Nano)), Type: "TIMESTAMP"}, + {Name: "start_after_event_id", Value: ptr(""), Type: "STRING"}, + {Name: "lag_cutoff", Value: ptr(testStart.Add(-4 * time.Hour).Format(time.RFC3339Nano)), Type: "TIMESTAMP"}, + } + _, rows, err := ts.st.runAuditQuery(stmt, params, "eu-west-1") + if err != nil { + t.Fatal(err.message) + } + got := map[string]bool{} + for _, r := range rows { + got[*r[0]] = true + } + if !got[eventID(18)] { + t.Error("eu-west-1 should see the ws-gamma row") + } + if !got[eventID(1)] { + t.Error("eu-west-1 should see account-level rows") + } + for _, n := range []int{11, 14, 19} { + if got[eventID(n)] { + t.Errorf("eu-west-1 should not see us-east-1 workspace row %d", n) + } + } +} + +func ptr(s string) *string { return &s } + +func TestStoppedWarehouseColdStart(t *testing.T) { + ts := newTestServer(config{warehouseStartDelay: time.Minute}) + _, body := ts.do(t, http.MethodPost, betaURL+"/api/2.0/sql/statements", statementRequest(t, "wh-pro", "0s", testStart.Add(-48*time.Hour), "", 1000)) + if stmtState(body) != statementPending { + t.Fatalf("want PENDING while the warehouse starts, got %v", body) + } + id, _ := body["statement_id"].(string) + + if _, wh := ts.do(t, http.MethodGet, betaURL+"/api/2.0/sql/warehouses/wh-pro", ""); wh["state"] != warehouseStarting { + t.Errorf("want STARTING, got %v", wh["state"]) + } + if _, polled := ts.do(t, http.MethodGet, betaURL+"/api/2.0/sql/statements/"+id, ""); stmtState(polled) != statementPending { + t.Errorf("want PENDING before the start delay, got %v", polled) + } + + ts.clock = ts.clock.Add(time.Minute) + if _, wh := ts.do(t, http.MethodGet, betaURL+"/api/2.0/sql/warehouses/wh-pro", ""); wh["state"] != warehouseRunning { + t.Errorf("want RUNNING, got %v", wh["state"]) + } + _, polled := ts.do(t, http.MethodGet, betaURL+"/api/2.0/sql/statements/"+id, "") + if stmtState(polled) != statementSucceeded || len(eventIDs(t, polled)) != 22 { + t.Errorf("want SUCCEEDED with 22 rows after the start delay, got %v", polled) + } +} + +func TestBrokenWarehouseAndCancel(t *testing.T) { + ts := newTestServer(config{}) + _, body := ts.do(t, http.MethodPost, alphaURL+"/api/2.0/sql/statements", statementRequest(t, "wh-broken", "0s", testStart.Add(-48*time.Hour), "", 1000)) + id, _ := body["statement_id"].(string) + + ts.clock = ts.clock.Add(time.Hour) + if _, polled := ts.do(t, http.MethodGet, alphaURL+"/api/2.0/sql/statements/"+id, ""); stmtState(polled) != statementPending { + t.Errorf("want PENDING forever, got %v", polled) + } + _, wh := ts.do(t, http.MethodGet, alphaURL+"/api/2.0/sql/warehouses/wh-broken", "") + health, _ := wh["health"].(map[string]any) + if wh["state"] != warehouseStopped || health["status"] != "FAILED" { + t.Errorf("want STOPPED with FAILED health, got %v", wh) + } + + code, cancel := ts.do(t, http.MethodPost, alphaURL+"/api/2.0/sql/statements/"+id+"/cancel", "") + if code != http.StatusOK || len(cancel) != 0 { + t.Errorf("cancel: want 200 {}, got %d %v", code, cancel) + } + if _, polled := ts.do(t, http.MethodGet, alphaURL+"/api/2.0/sql/statements/"+id, ""); stmtState(polled) != statementCanceled { + t.Errorf("want CANCELED after cancel, got %v", polled) + } +} + +func TestOnWaitTimeoutCancel(t *testing.T) { + ts := newTestServer(config{}) + req := strings.Replace(statementRequest(t, "wh-broken", "5s", testStart.Add(-48*time.Hour), "", 1000), + `"wait_timeout":"5s"`, `"wait_timeout":"5s","on_wait_timeout":"CANCEL"`, 1) + if _, body := ts.do(t, http.MethodPost, alphaURL+"/api/2.0/sql/statements", req); stmtState(body) != statementCanceled { + t.Errorf("want CANCELED after wait_timeout with on_wait_timeout=CANCEL, got %v", body) + } +} + +func TestChunks(t *testing.T) { + ts := newTestServer(config{chunkRows: 10}) + _, body := ts.do(t, http.MethodPost, alphaURL+"/api/2.0/sql/statements", statementRequest(t, "wh-serverless", "10s", testStart.Add(-48*time.Hour), "", 1000)) + id, _ := body["statement_id"].(string) + manifest, _ := body["manifest"].(map[string]any) + if manifest["total_chunk_count"] != float64(3) || manifest["total_row_count"] != float64(22) { + t.Fatalf("want 3 chunks / 22 rows, got %v", manifest) + } + + all := eventIDs(t, body) + next, _ := body["result"].(map[string]any)["next_chunk_index"].(float64) + for next > 0 { + _, chunk := ts.do(t, http.MethodGet, fmt.Sprintf("%s/api/2.0/sql/statements/%s/result/chunks/%d", alphaURL, id, int(next)), "") + all = append(all, eventIDs(t, map[string]any{"result": chunk})...) + next, _ = chunk["next_chunk_index"].(float64) + } + if len(all) != 22 { + t.Errorf("want 22 rows across chunks, got %d", len(all)) + } + if code, _ := ts.do(t, http.MethodGet, fmt.Sprintf("%s/api/2.0/sql/statements/%s/result/chunks/3", alphaURL, id), ""); code != http.StatusBadRequest { + t.Errorf("out-of-range chunk: got %d", code) + } +} diff --git a/test-server/seeds.go b/test-server/seeds.go new file mode 100644 index 00000000..5f37ea46 --- /dev/null +++ b/test-server/seeds.go @@ -0,0 +1,150 @@ +package main + +import "time" + +// Seed IDs. DECISION: all made up; numeric IDs are 16 digits like real Databricks IDs. +const ( + wsAlphaID int64 = 7474645012345671 + wsBetaID int64 = 7474645012345672 + wsGammaID int64 = 8828000011112223 + + userAlice = "4410000000000001" + userBob = "4410000000000002" + userCarol = "4410000000000003" + userDave = "4410000000000004" + + groupDataEngineers = "5520000000000001" + groupAnalysts = "5520000000000002" + + spCIBot = "6630000000000001" + + // deletedPrincipal is named in an audit row but matches no user, group or service principal. + deletedPrincipal = "9990000000000001" +) + +// seedState builds the in-memory account. Audit row times are relative to start. +func seedState(start time.Time) *state { + st := &state{ + // DOC: deployment_name is the workspace hostname prefix; aws_region per workspace. + workspaces: []workspace{ + {ID: wsAlphaID, Name: "Alpha", DeploymentName: "ws-alpha", Region: "us-east-1"}, + {ID: wsBetaID, Name: "Beta", DeploymentName: "ws-beta", Region: "us-east-1"}, + {ID: wsGammaID, Name: "Gamma", DeploymentName: "ws-gamma", Region: "eu-west-1"}, + }, + users: map[string]scimUser{ + userAlice: {ID: userAlice, UserName: "alice@example.com", DisplayName: "Alice Anders", Active: true}, + userBob: {ID: userBob, UserName: "bob@example.com", DisplayName: "Bob Brown", Active: true}, + userCarol: {ID: userCarol, UserName: "carol@example.com", DisplayName: "Carol Chen", Active: false}, + userDave: {ID: userDave, UserName: "dave@example.com", DisplayName: "Dave Diaz", Active: true}, + }, + groups: map[string]scimGroup{ + groupDataEngineers: {ID: groupDataEngineers, DisplayName: "data-engineers"}, + groupAnalysts: {ID: groupAnalysts, DisplayName: "analysts"}, + }, + servicePrincipals: map[string]scimServicePrincipal{ + spCIBot: {ID: spCIBot, ApplicationID: "0d1c2b3a-0000-4000-8000-000000000001", DisplayName: "ci-bot", Active: true}, + }, + // DECISION: one running serverless warehouse, one stopped pro warehouse, one that never starts. + warehouses: map[string]*warehouse{ + "wh-serverless": { + ID: "wh-serverless", Name: "Serverless Starter Warehouse", DeploymentName: "ws-alpha", + WarehouseType: "PRO", Serverless: true, InitialState: warehouseRunning, + }, + "wh-pro": { + ID: "wh-pro", Name: "Pro Warehouse", DeploymentName: "ws-beta", + WarehouseType: "PRO", InitialState: warehouseStopped, + }, + "wh-broken": { + ID: "wh-broken", Name: "Broken Warehouse", DeploymentName: "ws-alpha", + WarehouseType: "PRO", InitialState: warehouseStopped, NeverStarts: true, + }, + }, + statements: map[string]*statement{}, + } + st.audit = seedAudit(start) + return st +} + +// seedAudit returns synthetic system.access.audit rows. +// DOC: request_params keys per action follow https://docs.databricks.com/aws/en/admin/account-settings/audit-logs, +// and account-level events have workspace_id 0. +// DECISION: IDs and times are made up. Rows the connector should read fall between 26 and 5 +// hours before start: older than its 4-hour lag window and inside an --event-feed-start-at +// of 2 days ago. Two extra rows sit outside that window on purpose. +func seedAudit(start time.Time) []auditRow { + at := func(h, m int) time.Time { + return start.Add(-time.Duration(h)*time.Hour - time.Duration(m)*time.Minute).UTC().Truncate(time.Millisecond) + } + user := func(endpoint, id, name string) map[string]string { + return map[string]string{"endpoint": endpoint, "targetUserId": id, "targetUserName": name} + } + membership := func(groupID, groupName, memberID, memberName string) map[string]string { + return map[string]string{ + "endpoint": "scim", "targetGroupId": groupID, "targetGroupName": groupName, + "targetUserId": memberID, "targetUserName": memberName, + } + } + tie := at(6, 0) + + return []auditRow{ + {EventID: "e0000001-0000-4000-8000-000000000001", EventTime: at(26, 0), ServiceName: "accounts", ActionName: "add", + RequestParams: user("scim", userDave, "dave@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000002", EventTime: at(25, 0), ServiceName: "accounts", ActionName: "activateUser", + RequestParams: user("scim", userBob, "bob@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000003", EventTime: at(24, 0), ServiceName: "accounts", ActionName: "deactivateUser", + RequestParams: user("scim", userCarol, "carol@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000004", EventTime: at(23, 0), ServiceName: "accounts", ActionName: "createGroup", + RequestParams: map[string]string{"endpoint": "scim", "targetGroupId": groupAnalysts, "targetGroupName": "analysts"}}, + {EventID: "e0000001-0000-4000-8000-000000000005", EventTime: at(22, 0), ServiceName: "accounts", ActionName: "addPrincipalToGroup", + RequestParams: membership(groupDataEngineers, "data-engineers", userAlice, "alice@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000006", EventTime: at(21, 0), ServiceName: "accounts", ActionName: "addPrincipalToGroup", + RequestParams: membership(groupDataEngineers, "data-engineers", groupAnalysts, "analysts")}, + {EventID: "e0000001-0000-4000-8000-000000000007", EventTime: at(20, 0), ServiceName: "accounts", ActionName: "addPrincipalsToGroup", + RequestParams: membership(groupAnalysts, "analysts", spCIBot, "0d1c2b3a-0000-4000-8000-000000000001")}, + {EventID: "e0000001-0000-4000-8000-000000000008", EventTime: at(19, 0), ServiceName: "accounts", ActionName: "removePrincipalFromGroup", + RequestParams: membership(groupDataEngineers, "data-engineers", userBob, "bob@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000009", EventTime: at(18, 0), ServiceName: "accounts", ActionName: "removePrincipalFromGroup", + RequestParams: membership(groupDataEngineers, "data-engineers", deletedPrincipal, "gone@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000010", EventTime: at(17, 0), ServiceName: "accounts", ActionName: "setAccountAdmin", + RequestParams: user("adminConsole", userAlice, "alice@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000011", EventTime: at(16, 0), WorkspaceID: wsAlphaID, ServiceName: "accounts", ActionName: "setAdmin", + RequestParams: user("adminConsole", userBob, "bob@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000012", EventTime: at(15, 0), WorkspaceID: wsAlphaID, ServiceName: "accounts", ActionName: "removeAdmin", + RequestParams: user("adminConsole", userBob, "bob@example.com")}, + // DOC: workspace-level createGroup has no targetGroupId. + {EventID: "e0000001-0000-4000-8000-000000000013", EventTime: at(14, 0), WorkspaceID: wsAlphaID, ServiceName: "accounts", ActionName: "createGroup", + RequestParams: user("scim", userAlice, "alice@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000014", EventTime: at(13, 0), WorkspaceID: wsBetaID, ServiceName: "accounts", ActionName: "delete", + RequestParams: user("scim", userCarol, "carol@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000015", EventTime: at(12, 0), ServiceName: "workspace", ActionName: "addPermissionAssignment", + RequestParams: map[string]string{"principal_id": userDave, "account_id": testAccountID, "workspace_id": "7474645012345671"}}, + // DOC: updatePermissionAssignment has no workspace_id. + {EventID: "e0000001-0000-4000-8000-000000000016", EventTime: at(11, 0), ServiceName: "workspace", ActionName: "updatePermissionAssignment", + RequestParams: map[string]string{"principal_id": userDave, "permissions": `["ADMIN"]`}}, + {EventID: "e0000001-0000-4000-8000-000000000017", EventTime: at(10, 0), WorkspaceID: wsBetaID, ServiceName: "workspace", ActionName: "deletePermissionAssignment", + RequestParams: map[string]string{"principal_id": userBob, "account_id": testAccountID, "workspace_id": "7474645012345672"}}, + // ws-gamma is in eu-west-1, so this row is invisible from a us-east-1 warehouse. + {EventID: "e0000001-0000-4000-8000-000000000018", EventTime: at(9, 0), WorkspaceID: wsGammaID, ServiceName: "accounts", ActionName: "add", + RequestParams: user("scim", userAlice, "alice@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000019", EventTime: at(8, 0), WorkspaceID: wsAlphaID, ServiceName: "accounts", ActionName: "updateUser", + RequestParams: user("scim", userAlice, "alice@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000020", EventTime: at(7, 0), ServiceName: "accounts", ActionName: "updateGroup", + RequestParams: map[string]string{"endpoint": "scim", "targetGroupId": groupAnalysts, "targetGroupName": "analysts"}}, + // Two rows with the same event_time, to exercise the (event_time, event_id) ordering. + {EventID: "e0000001-0000-4000-8000-000000000021", EventTime: tie, ServiceName: "accounts", ActionName: "removeGroup", + RequestParams: map[string]string{"endpoint": "scim", "targetGroupId": groupDataEngineers, "targetGroupName": "data-engineers"}}, + {EventID: "e0000001-0000-4000-8000-000000000022", EventTime: tie, ServiceName: "accounts", ActionName: "deactivateUser", + RequestParams: user("scim", userDave, "dave@example.com")}, + {EventID: "e0000001-0000-4000-8000-000000000023", EventTime: at(5, 30), WorkspaceID: wsAlphaID, ServiceName: "accounts", ActionName: "changeDatabricksWorkspaceAcl", + RequestParams: map[string]string{"shardName": "ws-alpha", "targetUserId": userCarol, "resourceId": "7474645012345671", "aclPermissionSet": "workspace-access"}}, + // Not an action the connector asks for. + {EventID: "e0000001-0000-4000-8000-000000000024", EventTime: at(5, 0), WorkspaceID: wsAlphaID, ServiceName: "clusters", ActionName: "create", + RequestParams: map[string]string{"cluster_name": "etl"}}, + // Outside the window: newer than the connector's 4-hour lag cutoff. + {EventID: "e0000001-0000-4000-8000-000000000025", EventTime: at(2, 0), ServiceName: "accounts", ActionName: "add", + RequestParams: user("scim", userBob, "bob@example.com")}, + // Outside the window: older than an --event-feed-start-at of 2 days ago. + {EventID: "e0000001-0000-4000-8000-000000000026", EventTime: at(72, 0), ServiceName: "accounts", ActionName: "add", + RequestParams: user("scim", userCarol, "carol@example.com")}, + } +} diff --git a/test-server/state.go b/test-server/state.go new file mode 100644 index 00000000..96fef2b2 --- /dev/null +++ b/test-server/state.go @@ -0,0 +1,187 @@ +package main + +import ( + "context" + "net" + "net/http" + "strings" + "sync" + "time" +) + +// DECISION: hardcoded test credentials. +const ( + testAccountID = "test-account" + testClientID = "test-client" + testClientSecret = "test-secret" + testAccessToken = "test-access-token" +) + +type workspace struct { + ID int64 + Name string + DeploymentName string + Region string +} + +type scimUser struct { + ID string + UserName string + DisplayName string + Active bool +} + +type scimGroup struct { + ID string + DisplayName string +} + +type scimServicePrincipal struct { + ID string + ApplicationID string + DisplayName string + Active bool +} + +type warehouse struct { + ID string + Name string + DeploymentName string // workspace hosting the warehouse + WarehouseType string + Serverless bool + InitialState string + // NeverStarts makes start attempts fail: the warehouse stays STOPPED with a FAILED health status. + NeverStarts bool + + startRequestedAt time.Time +} + +type auditRow struct { + EventID string + EventTime time.Time + WorkspaceID int64 + ServiceName string + ActionName string + RequestParams map[string]string +} + +// state is the in-memory account, seeded at startup (see seeds.go) and reset on every start. +type state struct { + workspaces []workspace + users map[string]scimUser + groups map[string]scimGroup + servicePrincipals map[string]scimServicePrincipal + warehouses map[string]*warehouse + audit []auditRow + statements map[string]*statement + nextStatement int +} + +func (st *state) workspaceByDeployment(name string) *workspace { + for i := range st.workspaces { + if st.workspaces[i].DeploymentName == name { + return &st.workspaces[i] + } + } + return nil +} + +func (st *state) workspaceByID(id int64) *workspace { + for i := range st.workspaces { + if st.workspaces[i].ID == id { + return &st.workspaces[i] + } + } + return nil +} + +type server struct { + cfg config + start time.Time + now func() time.Time + + mu sync.Mutex + st *state + + account *http.ServeMux + workspace *http.ServeMux +} + +func newServer(cfg config, start time.Time) *server { + s := &server{ + cfg: cfg, + start: start, + now: time.Now, + st: seedState(start), + } + s.account = s.accountRoutes() + s.workspace = s.workspaceRoutes() + return s +} + +type workspaceKey struct{} + +func requestWorkspace(r *http.Request) *workspace { + ws, _ := r.Context().Value(workspaceKey{}).(*workspace) + return ws +} + +// ServeHTTP routes by Host header (DECISION): localhost is the account host, and +// .localhost is a workspace host, mirroring .. +func (s *server) ServeHTTP(w http.ResponseWriter, r *http.Request) { + // DECISION: local readiness check on every host; not a Databricks endpoint. + if r.URL.Path == "/health" && r.Method == http.MethodGet { + writeJSON(w, http.StatusOK, map[string]string{"status": "ok"}) + return + } + + host := r.Host + if h, _, err := net.SplitHostPort(r.Host); err == nil { + host = h + } + host = strings.ToLower(strings.TrimSuffix(host, ".")) + + switch { + case host == "localhost" || host == "127.0.0.1" || host == "::1": + s.account.ServeHTTP(w, r) + case strings.HasSuffix(host, ".localhost"): + deployment := strings.TrimSuffix(host, ".localhost") + s.mu.Lock() + ws := s.st.workspaceByDeployment(deployment) + s.mu.Unlock() + if ws == nil { + // DECISION: a real unknown deployment name doesn't resolve in DNS at all. + writeError(w, http.StatusNotFound, "ENDPOINT_NOT_FOUND", "test-server: no workspace with deployment name "+deployment) + return + } + s.workspace.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), workspaceKey{}, ws))) + default: + writeError(w, http.StatusNotFound, "ENDPOINT_NOT_FOUND", "test-server: unknown host "+host) + } +} + +// requireBearer enforces the access token on every API call except the token endpoint. +// DOC: https://docs.databricks.com/aws/en/dev-tools/auth/oauth-m2m (Authorization: Bearer ). +// DOC: error body {error_code, message}, https://docs.databricks.com/api/workspace/introduction. +func requireBearer(next http.HandlerFunc) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + token, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ") + if !ok || token != testAccessToken { + writeError(w, http.StatusUnauthorized, "UNAUTHENTICATED", "Credential was not sent or was of an unsupported type for this API.") + return + } + next(w, r) + } +} + +// requireAccount rejects account API paths for any account other than testAccountID. +// ASSUMPTION: the exact status and code for another account's ID are not documented. +func requireAccount(next http.HandlerFunc) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + if r.PathValue("account_id") != testAccountID { + writeError(w, http.StatusNotFound, "RESOURCE_DOES_NOT_EXIST", "Account "+r.PathValue("account_id")+" does not exist.") + return + } + next(w, r) + } +} diff --git a/test-server/statements.go b/test-server/statements.go new file mode 100644 index 00000000..e72ebcf7 --- /dev/null +++ b/test-server/statements.go @@ -0,0 +1,633 @@ +package main + +import ( + "encoding/json" + "fmt" + "net/http" + "regexp" + "sort" + "strconv" + "strings" + "time" +) + +// Warehouse states. DOC: https://docs.databricks.com/api/workspace/warehouses/get +const ( + warehouseStarting = "STARTING" + warehouseRunning = "RUNNING" + warehouseStopped = "STOPPED" +) + +// Statement states. DOC: https://docs.databricks.com/api/workspace/statementexecution/getstatement +const ( + statementPending = "PENDING" + statementSucceeded = "SUCCEEDED" + statementFailed = "FAILED" + statementCanceled = "CANCELED" +) + +// warehouseState returns wh's state now. Callers hold s.mu. +// ASSUMPTION: submitting a statement starts a stopped warehouse, which is RUNNING after +// --warehouse-start-delay. DECISION: a NeverStarts warehouse stays STOPPED. +func (s *server) warehouseState(wh *warehouse) string { + if wh.InitialState != warehouseStopped { + return wh.InitialState + } + if wh.startRequestedAt.IsZero() || wh.NeverStarts { + return warehouseStopped + } + if s.now().Sub(wh.startRequestedAt) >= s.cfg.warehouseStartDelay { + return warehouseRunning + } + return warehouseStarting +} + +type column struct { + Name string `json:"name"` + TypeName string `json:"type_name"` + TypeText string `json:"type_text"` + Position int `json:"position"` +} + +type statement struct { + ID string + DeploymentName string + WarehouseID string + State string + ErrorCode string + ErrorMessage string + + // The outcome, decided at submit time and revealed once the warehouse is running. + outcomeState string + outcomeCode string + outcomeMessage string + columns []column + rows [][]*string +} + +// refresh moves a PENDING statement to its outcome once its warehouse is running. +// DOC: a statement on a warehouse that isn't running stays PENDING, waiting for the warehouse. +// DECISION: statements go straight from PENDING to SUCCEEDED/FAILED (no visible RUNNING). +func (s *server) refresh(stmt *statement) { + if stmt.State != statementPending { + return + } + if s.warehouseState(s.st.warehouses[stmt.WarehouseID]) != warehouseRunning { + return + } + stmt.State = stmt.outcomeState + stmt.ErrorCode = stmt.outcomeCode + stmt.ErrorMessage = stmt.outcomeMessage +} + +func isTerminal(state string) bool { + return state != statementPending +} + +type statementParameter struct { + Name string `json:"name"` + Value *string `json:"value"` + Type string `json:"type"` +} + +type executeRequest struct { + WarehouseID string `json:"warehouse_id"` + Statement string `json:"statement"` + WaitTimeout *string `json:"wait_timeout"` + OnWaitTimeout string `json:"on_wait_timeout"` + Format string `json:"format"` + Disposition string `json:"disposition"` + Parameters []statementParameter `json:"parameters"` +} + +var waitTimeoutRe = regexp.MustCompile(`^(\d+)s$`) + +// parseWaitTimeout validates wait_timeout. +// DOC: 0s, or between 5s and 50s; the default is 10s. +func parseWaitTimeout(v *string) (time.Duration, bool) { + if v == nil { + return 10 * time.Second, true + } + m := waitTimeoutRe.FindStringSubmatch(*v) + if m == nil { + return 0, false + } + n, err := strconv.Atoi(m[1]) + if err != nil || (n != 0 && (n < 5 || n > 50)) { + return 0, false + } + return time.Duration(n) * time.Second, true +} + +// handleExecuteStatement submits a statement. +// DOC: if it finishes within wait_timeout the response has the result; otherwise it has the +// statement ID with PENDING or RUNNING, and on_wait_timeout=CANCEL cancels it instead. +func (s *server) handleExecuteStatement(w http.ResponseWriter, r *http.Request) { + ws := requestWorkspace(r) + + var req executeRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeError(w, http.StatusBadRequest, "MALFORMED_REQUEST", "Could not parse request body: "+err.Error()) + return + } + + wait, ok := parseWaitTimeout(req.WaitTimeout) + if !ok { + writeError(w, http.StatusBadRequest, "INVALID_PARAMETER_VALUE", + "wait_timeout must be 0 seconds (disables wait), or between 5 seconds and 50 seconds.") + return + } + switch req.OnWaitTimeout { + case "", "CONTINUE", "CANCEL": + default: + writeError(w, http.StatusBadRequest, "INVALID_PARAMETER_VALUE", "on_wait_timeout must be CONTINUE or CANCEL.") + return + } + // DECISION: only the JSON_ARRAY format and INLINE disposition are implemented. + if req.Format != "" && req.Format != "JSON_ARRAY" { + writeError(w, http.StatusBadRequest, "INVALID_PARAMETER_VALUE", "test-server: only format JSON_ARRAY is supported") + return + } + if req.Disposition != "" && req.Disposition != "INLINE" { + writeError(w, http.StatusBadRequest, "INVALID_PARAMETER_VALUE", "test-server: only disposition INLINE is supported") + return + } + // DOC: parameters are named (:name) markers. ASSUMPTION: the status and code for a + // positional (unnamed) parameter. + for _, p := range req.Parameters { + if p.Name == "" { + writeError(w, http.StatusBadRequest, "INVALID_PARAMETER_VALUE", "Positional parameters are not supported; use named parameter markers.") + return + } + } + + s.mu.Lock() + wh, ok := s.st.warehouses[req.WarehouseID] + if !ok || wh.DeploymentName != ws.DeploymentName { + s.mu.Unlock() + // DOC: an unknown warehouse_id is an error. ASSUMPTION: its status and code. + writeError(w, http.StatusNotFound, "RESOURCE_DOES_NOT_EXIST", "Warehouse "+req.WarehouseID+" does not exist.") + return + } + if wh.InitialState == warehouseStopped && wh.startRequestedAt.IsZero() { + wh.startRequestedAt = s.now() + } + + s.st.nextStatement++ + stmt := &statement{ + ID: fmt.Sprintf("01f00000-0000-4000-8000-%012d", s.st.nextStatement), + DeploymentName: ws.DeploymentName, + WarehouseID: wh.ID, + State: statementPending, + } + s.st.statements[stmt.ID] = stmt + + region := "" + if hostWS := s.st.workspaceByDeployment(wh.DeploymentName); hostWS != nil { + region = hostWS.Region + } + cols, rows, err := s.st.runAuditQuery(req.Statement, req.Parameters, region) + if err != nil { + stmt.outcomeState = statementFailed + stmt.outcomeCode = err.code + stmt.outcomeMessage = err.message + } else { + stmt.outcomeState = statementSucceeded + stmt.columns = cols + stmt.rows = rows + } + s.refresh(stmt) + s.mu.Unlock() + + deadline := time.NewTimer(wait) + defer deadline.Stop() + tick := time.NewTicker(20 * time.Millisecond) + defer tick.Stop() + for { + s.mu.Lock() + s.refresh(stmt) + if isTerminal(stmt.State) { + body := s.statementBody(stmt) + s.mu.Unlock() + writeJSON(w, http.StatusOK, body) + return + } + s.mu.Unlock() + + select { + case <-r.Context().Done(): + return + case <-deadline.C: + s.mu.Lock() + if req.OnWaitTimeout == "CANCEL" { + stmt.State = statementCanceled + } + body := s.statementBody(stmt) + s.mu.Unlock() + writeJSON(w, http.StatusOK, body) + return + case <-tick.C: + } + } +} + +// lookupStatement finds a statement submitted in the request's workspace. Callers hold s.mu. +// ASSUMPTION: an unknown statement ID is 404 RESOURCE_DOES_NOT_EXIST. +func (s *server) lookupStatement(w http.ResponseWriter, r *http.Request) *statement { + stmt, ok := s.st.statements[r.PathValue("id")] + if !ok || stmt.DeploymentName != requestWorkspace(r).DeploymentName { + writeError(w, http.StatusNotFound, "RESOURCE_DOES_NOT_EXIST", "Statement "+r.PathValue("id")+" does not exist.") + return nil + } + return stmt +} + +// handleGetStatement returns a statement's status. +// DOC: once SUCCEEDED it includes the manifest and the first chunk; terminal states are HTTP 200. +func (s *server) handleGetStatement(w http.ResponseWriter, r *http.Request) { + s.mu.Lock() + defer s.mu.Unlock() + + stmt := s.lookupStatement(w, r) + if stmt == nil { + return + } + s.refresh(stmt) + writeJSON(w, http.StatusOK, s.statementBody(stmt)) +} + +// handleCancelStatement requests cancellation. +// DOC: returns {}; the caller polls the statement to see CANCELED. +func (s *server) handleCancelStatement(w http.ResponseWriter, r *http.Request) { + s.mu.Lock() + defer s.mu.Unlock() + + stmt := s.lookupStatement(w, r) + if stmt == nil { + return + } + s.refresh(stmt) + if !isTerminal(stmt.State) { + stmt.State = statementCanceled + } + writeJSON(w, http.StatusOK, map[string]any{}) +} + +// handleGetResultChunk returns one result chunk of a SUCCEEDED statement. +// DOC: same shape as the result object, with next_chunk_index absent on the last chunk. +// ASSUMPTION: the status and code for an out-of-range chunk or a statement without a result. +func (s *server) handleGetResultChunk(w http.ResponseWriter, r *http.Request) { + s.mu.Lock() + defer s.mu.Unlock() + + stmt := s.lookupStatement(w, r) + if stmt == nil { + return + } + s.refresh(stmt) + if stmt.State != statementSucceeded { + writeError(w, http.StatusBadRequest, "INVALID_STATE", "Statement "+stmt.ID+" has no result: state is "+stmt.State+".") + return + } + n, ok := parseChunkIndex(r.PathValue("n")) + if !ok || n >= s.chunkCount(stmt) { + writeError(w, http.StatusBadRequest, "INVALID_PARAMETER_VALUE", "Chunk index "+r.PathValue("n")+" is out of range.") + return + } + writeJSON(w, http.StatusOK, s.chunkBody(stmt, n)) +} + +// chunkCount is how many chunks stmt's result has. DECISION: one chunk unless --chunk-rows is set. +// ASSUMPTION: an empty result is reported as a single empty chunk. +func (s *server) chunkCount(stmt *statement) int { + if s.cfg.chunkRows <= 0 || len(stmt.rows) == 0 { + return 1 + } + return (len(stmt.rows) + s.cfg.chunkRows - 1) / s.cfg.chunkRows +} + +func (s *server) chunkBounds(stmt *statement, n int) (int, int) { + if s.cfg.chunkRows <= 0 { + return 0, len(stmt.rows) + } + start := n * s.cfg.chunkRows + return start, min(start+s.cfg.chunkRows, len(stmt.rows)) +} + +// chunkBody renders result chunk n. DOC: data_array values are strings and nulls are JSON null. +func (s *server) chunkBody(stmt *statement, n int) map[string]any { + start, end := s.chunkBounds(stmt, n) + body := map[string]any{ + "chunk_index": n, + "row_offset": start, + "row_count": end - start, + } + if end > start { + body["data_array"] = stmt.rows[start:end] + } + if n+1 < s.chunkCount(stmt) { + body["next_chunk_index"] = n + 1 + body["next_chunk_internal_link"] = fmt.Sprintf("/api/2.0/sql/statements/%s/result/chunks/%d", stmt.ID, n+1) + } + return body +} + +// statementBody renders a statement as the Execute/Get responses do. Callers hold s.mu. +func (s *server) statementBody(stmt *statement) map[string]any { + status := map[string]any{"state": stmt.State} + if stmt.ErrorMessage != "" { + status["error"] = map[string]string{"error_code": stmt.ErrorCode, "message": stmt.ErrorMessage} + } + body := map[string]any{"statement_id": stmt.ID, "status": status} + if stmt.State != statementSucceeded { + return body + } + + count := s.chunkCount(stmt) + chunks := make([]map[string]any, 0, count) + for n := range count { + start, end := s.chunkBounds(stmt, n) + chunks = append(chunks, map[string]any{"chunk_index": n, "row_offset": start, "row_count": end - start}) + } + body["manifest"] = map[string]any{ + "format": "JSON_ARRAY", + "schema": map[string]any{ + "column_count": len(stmt.columns), + "columns": stmt.columns, + }, + "total_chunk_count": count, + "chunks": chunks, + "total_row_count": len(stmt.rows), + "truncated": false, + } + body["result"] = s.chunkBody(stmt, 0) + return body +} + +// SQL handling (DECISION): the server doesn't run SQL. It recognizes the shape of the +// connector's query against system.access.audit and evaluates just that. +var ( + auditFromRe = regexp.MustCompile(`(?is)^\s*SELECT\s+(.+?)\s+FROM\s+system\.access\.audit\b`) + paramMarkerRe = regexp.MustCompile(`(?:^|[^:\w]):([A-Za-z_]\w*)`) + eventDateRe = regexp.MustCompile(`(?i)\bevent_date\s*>=\s*:(\w+)`) + afterRe = regexp.MustCompile( + `(?i)\(\s*event_time\s*>\s*:(\w+)\s+OR\s*\(\s*event_time\s*=\s*:(\w+)\s+AND\s+event_id\s*>\s*:(\w+)\s*\)\s*\)`) + cutoffRe = regexp.MustCompile(`(?i)\bevent_time\s*<=\s*:(\w+)`) + serviceListRe = regexp.MustCompile(`(?i)\bservice_name\s+IN\s*\(([^)]*)\)`) + actionListRe = regexp.MustCompile(`(?i)\baction_name\s+IN\s*\(([^)]*)\)`) + orderByRe = regexp.MustCompile(`(?i)\bORDER\s+BY\s+event_time(\s+ASC)?\s*,\s*event_id(\s+ASC)?\b`) + limitRe = regexp.MustCompile(`(?i)\bLIMIT\s+(\d+)\s*;?\s*$`) +) + +// auditColumns are the documented system.access.audit columns this server serves. +// DOC: https://docs.databricks.com/aws/en/admin/system-tables/audit-logs +// DOC: type_name values, https://docs.databricks.com/api/workspace/statementexecution/executestatement +var auditColumns = map[string]column{ + "event_id": {Name: "event_id", TypeName: "STRING", TypeText: "STRING"}, + "event_time": {Name: "event_time", TypeName: "TIMESTAMP", TypeText: "TIMESTAMP"}, + "event_date": {Name: "event_date", TypeName: "DATE", TypeText: "DATE"}, + "workspace_id": {Name: "workspace_id", TypeName: "LONG", TypeText: "BIGINT"}, + "service_name": {Name: "service_name", TypeName: "STRING", TypeText: "STRING"}, + "action_name": {Name: "action_name", TypeName: "STRING", TypeText: "STRING"}, + "request_params": {Name: "request_params", TypeName: "MAP", TypeText: "MAP"}, +} + +type sqlError struct { + code string + message string +} + +func unsupported(detail string) *sqlError { + return &sqlError{code: "BAD_REQUEST", message: "test-server: unsupported statement: " + detail} +} + +type boundParam struct { + typ string + value string + time time.Time +} + +// bindParameters binds typed named parameters. +// DOC: named markers with an optional type; the default type is STRING. +// ASSUMPTION: the accepted TIMESTAMP and DATE literal formats. +func bindParameters(params []statementParameter) (map[string]boundParam, *sqlError) { + bound := make(map[string]boundParam, len(params)) + for _, p := range params { + typ := strings.ToUpper(p.Type) + if typ == "" { + typ = "STRING" + } + if p.Value == nil { + return nil, &sqlError{code: "BAD_REQUEST", message: "test-server: NULL parameter values are not supported"} + } + bp := boundParam{typ: typ, value: *p.Value} + switch typ { + case "STRING": + case "TIMESTAMP": + t, err := time.Parse(time.RFC3339Nano, bp.value) + if err != nil { + t, err = time.Parse("2006-01-02 15:04:05.999999", bp.value) + } + if err != nil { + return nil, &sqlError{code: "INVALID_PARAMETER_VALUE", message: fmt.Sprintf("Cannot parse %q as TIMESTAMP for parameter %s.", bp.value, p.Name)} + } + bp.time = t.UTC() + case "DATE": + t, err := time.Parse(time.DateOnly, bp.value) + if err != nil { + return nil, &sqlError{code: "INVALID_PARAMETER_VALUE", message: fmt.Sprintf("Cannot parse %q as DATE for parameter %s.", bp.value, p.Name)} + } + bp.time = t + default: + return nil, &sqlError{code: "BAD_REQUEST", message: "test-server: unsupported parameter type " + typ} + } + bound[p.Name] = bp + } + return bound, nil +} + +func parseInList(list string) []string { + var values []string + for _, v := range strings.Split(list, ",") { + values = append(values, strings.Trim(strings.TrimSpace(v), "'")) + } + return values +} + +func contains(values []string, v string) bool { + for _, x := range values { + if x == v { + return true + } + } + return false +} + +// timeParam returns a bound TIMESTAMP/DATE parameter. Callers check names against bound first. +func timeParam(bound map[string]boundParam, name string, want string) (time.Time, *sqlError) { + p := bound[name] + if p.typ != want { + return time.Time{}, unsupported(fmt.Sprintf("parameter %s must be bound with type %s", name, want)) + } + return p.time, nil +} + +// runAuditQuery evaluates the connector's audit query against the seed rows visible from a +// warehouse in region. Callers hold s.mu. +// DOC: workspace-level events are only visible from the workspace's region ("Most audit logs +// are only available in the region of the workspace"); account-level events (workspace_id 0) +// are visible from every region. +func (st *state) runAuditQuery(sql string, params []statementParameter, region string) ([]column, [][]*string, *sqlError) { + m := auditFromRe.FindStringSubmatch(sql) + if m == nil { + return nil, nil, unsupported("only SELECT ... FROM system.access.audit is supported") + } + + var cols []column + for i, name := range strings.Split(m[1], ",") { + col, ok := auditColumns[strings.ToLower(strings.TrimSpace(name))] + if !ok { + return nil, nil, unsupported(fmt.Sprintf("column %q", strings.TrimSpace(name))) + } + col.Position = i + cols = append(cols, col) + } + + bound, bindErr := bindParameters(params) + if bindErr != nil { + return nil, nil, bindErr + } + // ASSUMPTION: the error for a marker with no bound parameter (UNBOUND_SQL_PARAMETER). + for _, mm := range paramMarkerRe.FindAllStringSubmatch(sql, -1) { + if _, ok := bound[mm[1]]; !ok { + return nil, nil, &sqlError{code: "BAD_REQUEST", message: fmt.Sprintf("[UNBOUND_SQL_PARAMETER] Found the unbound parameter: %s.", mm[1])} + } + } + + if !orderByRe.MatchString(sql) { + return nil, nil, unsupported("ORDER BY event_time ASC, event_id ASC is required") + } + + var ( + startDate, startTime, cutoff time.Time + startAfterID string + hasStartDate, hasAfter bool + hasCutoff bool + services, actions []string + limit = -1 + err *sqlError + ) + if mm := eventDateRe.FindStringSubmatch(sql); mm != nil { + if startDate, err = timeParam(bound, mm[1], "DATE"); err != nil { + return nil, nil, err + } + hasStartDate = true + } + if mm := afterRe.FindStringSubmatch(sql); mm != nil { + if mm[1] != mm[2] { + return nil, nil, unsupported("the two event_time comparisons must use the same parameter") + } + if startTime, err = timeParam(bound, mm[1], "TIMESTAMP"); err != nil { + return nil, nil, err + } + startAfterID = bound[mm[3]].value + hasAfter = true + } + if mm := cutoffRe.FindStringSubmatch(sql); mm != nil { + if cutoff, err = timeParam(bound, mm[1], "TIMESTAMP"); err != nil { + return nil, nil, err + } + hasCutoff = true + } + if mm := serviceListRe.FindStringSubmatch(sql); mm != nil { + services = parseInList(mm[1]) + } + if mm := actionListRe.FindStringSubmatch(sql); mm != nil { + actions = parseInList(mm[1]) + } + if mm := limitRe.FindStringSubmatch(sql); mm != nil { + limit, _ = strconv.Atoi(mm[1]) + } + + matched := make([]auditRow, 0, len(st.audit)) + for _, row := range st.audit { + if row.WorkspaceID != 0 { + ws := st.workspaceByID(row.WorkspaceID) + if ws == nil || ws.Region != region { + continue + } + } + eventTime := row.EventTime.UTC() + eventDate := time.Date(eventTime.Year(), eventTime.Month(), eventTime.Day(), 0, 0, 0, 0, time.UTC) + if hasStartDate && eventDate.Before(startDate) { + continue + } + afterCursor := eventTime.After(startTime) || (eventTime.Equal(startTime) && row.EventID > startAfterID) + if hasAfter && !afterCursor { + continue + } + if hasCutoff && eventTime.After(cutoff) { + continue + } + if services != nil && !contains(services, row.ServiceName) { + continue + } + if actions != nil && !contains(actions, row.ActionName) { + continue + } + matched = append(matched, row) + } + + sort.Slice(matched, func(i, j int) bool { + if !matched[i].EventTime.Equal(matched[j].EventTime) { + return matched[i].EventTime.Before(matched[j].EventTime) + } + return matched[i].EventID < matched[j].EventID + }) + if limit >= 0 && len(matched) > limit { + matched = matched[:limit] + } + + rows := make([][]*string, 0, len(matched)) + for _, row := range matched { + values := make([]*string, len(cols)) + for i, col := range cols { + values[i] = formatValue(row, col.Name) + } + rows = append(rows, values) + } + return cols, rows, nil +} + +// formatValue renders a column value as JSON_ARRAY does. +// ASSUMPTION: TIMESTAMP as ISO 8601 UTC with milliseconds, DATE as yyyy-mm-dd, and MAP as a +// JSON object string. +func formatValue(row auditRow, name string) *string { + var v string + switch name { + case "event_id": + v = row.EventID + case "event_time": + v = row.EventTime.UTC().Format("2006-01-02T15:04:05.000Z") + case "event_date": + v = row.EventTime.UTC().Format(time.DateOnly) + case "workspace_id": + v = strconv.FormatInt(row.WorkspaceID, 10) + case "service_name": + v = row.ServiceName + case "action_name": + v = row.ActionName + case "request_params": + if row.RequestParams == nil { + return nil + } + b, err := json.Marshal(row.RequestParams) + if err != nil { + return nil + } + v = string(b) + default: + return nil + } + return &v +} diff --git a/test-server/testdata/expected-events.jsonl b/test-server/testdata/expected-events.jsonl new file mode 100644 index 00000000..8a6be962 --- /dev/null +++ b/test-server/testdata/expected-events.jsonl @@ -0,0 +1,25 @@ +{"id":"e0000001-0000-4000-8000-000000000001/0","resourceChangeEvent":{"resourceId":{"resourceType":"user","resource":"4410000000000004"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000002/0","resourceChangeEvent":{"resourceId":{"resourceType":"user","resource":"4410000000000002"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000003/0","resourceChangeEvent":{"resourceId":{"resourceType":"user","resource":"4410000000000003"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000004/0","resourceChangeEvent":{"resourceId":{"resourceType":"group","resource":"account/test-account/group/5520000000000002"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000005/0","createGrantEvent":{"entitlement":{"resource":{"id":{"resourceType":"group","resource":"account/test-account/group/5520000000000001"},"parentResourceId":{"resourceType":"account","resource":"test-account"}},"id":"group:account/test-account/group/5520000000000001:member","displayName":"member","purpose":"PURPOSE_VALUE_ASSIGNMENT","slug":"member"},"principal":{"id":{"resourceType":"user","resource":"4410000000000001"}}}} +{"id":"e0000001-0000-4000-8000-000000000006/0","createGrantEvent":{"entitlement":{"resource":{"id":{"resourceType":"group","resource":"account/test-account/group/5520000000000001"},"parentResourceId":{"resourceType":"account","resource":"test-account"}},"id":"group:account/test-account/group/5520000000000001:member","displayName":"member","purpose":"PURPOSE_VALUE_ASSIGNMENT","slug":"member"},"principal":{"id":{"resourceType":"group","resource":"account/test-account/group/5520000000000002"}}}} +{"id":"e0000001-0000-4000-8000-000000000007/0","createGrantEvent":{"entitlement":{"resource":{"id":{"resourceType":"group","resource":"account/test-account/group/5520000000000002"},"parentResourceId":{"resourceType":"account","resource":"test-account"}},"id":"group:account/test-account/group/5520000000000002:member","displayName":"member","purpose":"PURPOSE_VALUE_ASSIGNMENT","slug":"member"},"principal":{"id":{"resourceType":"service_principal","resource":"6630000000000001"}}}} +{"id":"e0000001-0000-4000-8000-000000000008/0","createRevokeEvent":{"entitlement":{"resource":{"id":{"resourceType":"group","resource":"account/test-account/group/5520000000000001"},"parentResourceId":{"resourceType":"account","resource":"test-account"}},"id":"group:account/test-account/group/5520000000000001:member","displayName":"member","purpose":"PURPOSE_VALUE_ASSIGNMENT","slug":"member"},"principal":{"id":{"resourceType":"user","resource":"4410000000000002"}}}} +{"id":"e0000001-0000-4000-8000-000000000009/0","resourceChangeEvent":{"resourceId":{"resourceType":"group","resource":"account/test-account/group/5520000000000001"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000010/0","resourceChangeEvent":{"resourceId":{"resourceType":"user","resource":"4410000000000001"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000010/1","resourceChangeEvent":{"resourceId":{"resourceType":"role","resource":"account_admin"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000011/0","resourceChangeEvent":{"resourceId":{"resourceType":"user","resource":"4410000000000002"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000011/1","resourceChangeEvent":{"resourceId":{"resourceType":"role","resource":"account_admin"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000012/0","resourceChangeEvent":{"resourceId":{"resourceType":"user","resource":"4410000000000002"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000014/0","resourceChangeEvent":{"resourceId":{"resourceType":"user","resource":"4410000000000003"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000015/0","resourceChangeEvent":{"resourceId":{"resourceType":"workspace","resource":"ws-alpha"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000017/0","resourceChangeEvent":{"resourceId":{"resourceType":"workspace","resource":"ws-beta"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000019/0","resourceChangeEvent":{"resourceId":{"resourceType":"user","resource":"4410000000000001"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000019/1","resourceChangeEvent":{"resourceId":{"resourceType":"role","resource":"ws-alpha:allow-cluster-create"},"parentResourceId":{"resourceType":"workspace","resource":"ws-alpha"}}} +{"id":"e0000001-0000-4000-8000-000000000019/2","resourceChangeEvent":{"resourceId":{"resourceType":"role","resource":"ws-alpha:allow-instance-pool-create"},"parentResourceId":{"resourceType":"workspace","resource":"ws-alpha"}}} +{"id":"e0000001-0000-4000-8000-000000000020/0","resourceChangeEvent":{"resourceId":{"resourceType":"group","resource":"account/test-account/group/5520000000000002"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000021/0","resourceChangeEvent":{"resourceId":{"resourceType":"group","resource":"account/test-account/group/5520000000000001"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000022/0","resourceChangeEvent":{"resourceId":{"resourceType":"user","resource":"4410000000000004"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000023/0","resourceChangeEvent":{"resourceId":{"resourceType":"workspace","resource":"ws-alpha"},"parentResourceId":{"resourceType":"account","resource":"test-account"}}} +{"id":"e0000001-0000-4000-8000-000000000023/1","resourceChangeEvent":{"resourceId":{"resourceType":"role","resource":"ws-alpha:workspace-access"},"parentResourceId":{"resourceType":"workspace","resource":"ws-alpha"}}}