diff --git a/README.md b/README.md index 5a018e06..f05a7dd0 100644 --- a/README.md +++ b/README.md @@ -115,6 +115,36 @@ To instead exclude specific workspaces from the sync, pass them to the list. Each entry can be a workspace name, deployment name, or numeric workspace ID. Excluded workspaces and their roles are skipped entirely. +## 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, by setting +`--enable-incremental-sync` (or `BATON_ENABLE_INCREMENTAL_SYNC`). Full syncs +still run as the correctness backstop; incremental sync does not detect +deletions, which are only caught by the next full sync. + +Incremental sync requires OAuth2 (service principal) authentication — it's not +available with workspace tokens, since the Account API needed to resolve audit +events back to synced resources is unreachable that way. + +Incremental sync also requires: + +- `--sql-warehouse-id` (or `BATON_SQL_WAREHOUSE_ID`), the ID of a Databricks SQL + warehouse the connector can use to query the `system.access.audit` table. A + small serverless warehouse is recommended to minimize cold-start latency. + The connector automatically discovers which workspace hosts it. +- A one-time setup performed by a Databricks admin, which the connector cannot + do on its own: + - An account admin must [enable the `access` system + schema](https://docs.databricks.com/en/admin/system-tables/index.html) for + the account's Unity Catalog metastore. + - A metastore admin must grant `SELECT` on `system.access` to the service + principal or user the connector authenticates as. + +Once enabled, ongoing polling only needs that `SELECT` grant plus warehouse +access; no further elevated privilege is required. + ## Group provisioning limitations provisioning of account groups from a workspace token is not supported, if you need to provision groups you can only do it using the client-id and client-secret flow, this is due to the fact that the Databricks API does not allow provisioning of groups from a workspace token. @@ -153,7 +183,8 @@ Flags: --client-secret string The client secret used to authenticate with ConductorOne ($BATON_CLIENT_SECRET) --databricks-client-id string required: The Databricks service principal's client ID used to connect to the Databricks Account and Workspace API ($BATON_DATABRICKS_CLIENT_ID) --databricks-client-secret string required: The Databricks service principal's client secret used to connect to the Databricks Account and Workspace API ($BATON_DATABRICKS_CLIENT_SECRET) - --databricks-exclude-workspaces strings Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID ($BATON_DATABRICKS_EXCLUDE_WORKSPACES) + --databricks-exclude-workspaces strings Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces. ($BATON_DATABRICKS_EXCLUDE_WORKSPACES) + --enable-incremental-sync Poll a Databricks audit-log event feed between full syncs to pick up access changes early. Deletions are still only caught by the next full sync. ($BATON_ENABLE_INCREMENTAL_SYNC) --external-resource-c1z string The path to the c1z file to sync external baton resources with ($BATON_EXTERNAL_RESOURCE_C1Z) --external-resource-entitlement-id-filter string The entitlement that external users, groups must have access to sync external baton resources ($BATON_EXTERNAL_RESOURCE_ENTITLEMENT_ID_FILTER) --external-resource-traits strings Resource type traits (e.g. "user", "group", "app") to sync and match from the external resource c1z. When unset the matcher falls back to user and group; passing this flag replaces the full set rather than adding to it. ($BATON_EXTERNAL_RESOURCE_TRAITS) @@ -173,6 +204,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 ID of the Databricks SQL warehouse used to query system.access.audit. Required when incremental sync is enabled; the workspace hosting it is discovered automatically. ($BATON_SQL_WAREHOUSE_ID) --storage-engine string The storage engine to use when opening the sync c1z file: sqlite or pebble. Leave unset to use the baton-sdk default. ($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) @@ -181,7 +213,7 @@ Flags: -v, --version version for baton-databricks --workers int The number of sync workers to use. -1 for auto-detect, 0 for sequential, >0 for parallel ($BATON_WORKERS) --workspace-tokens strings required: The Databricks personal access tokens scoped to specific workspaces used to connect to the Databricks Workspace API ($BATON_WORKSPACE_TOKENS) - --workspaces strings Limit syncing to the specified workspaces, by deployment name, not workspace ID. Required when using workspace tokens, in the same order as workspace-tokens. ($BATON_WORKSPACES) + --workspaces strings Limit syncing to the specified workspaces, by deployment name, not workspace ID. Required when using workspace tokens, in the same order as workspace-tokens. Mutually exclusive with databricks-exclude-workspaces. ($BATON_WORKSPACES) Use "baton-databricks [command] --help" for more information about a command. ``` diff --git a/baton_capabilities.json b/baton_capabilities.json index 91985cec..57cfd894 100644 --- a/baton_capabilities.json +++ b/baton_capabilities.json @@ -8,6 +8,7 @@ }, "capabilities": [ "CAPABILITY_SYNC", + "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], "permissions": {} @@ -22,6 +23,7 @@ }, "capabilities": [ "CAPABILITY_SYNC", + "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], "permissions": {} @@ -36,6 +38,7 @@ }, "capabilities": [ "CAPABILITY_SYNC", + "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], "permissions": {} @@ -50,6 +53,7 @@ }, "capabilities": [ "CAPABILITY_SYNC", + "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], "permissions": {} @@ -69,6 +73,7 @@ }, "capabilities": [ "CAPABILITY_SYNC", + "CAPABILITY_TARGETED_SYNC", "CAPABILITY_ACCOUNT_PROVISIONING", "CAPABILITY_RESOURCE_DELETE" ], @@ -84,6 +89,7 @@ }, "capabilities": [ "CAPABILITY_SYNC", + "CAPABILITY_TARGETED_SYNC", "CAPABILITY_PROVISION" ], "permissions": {} @@ -93,7 +99,10 @@ "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 3e6dbc4f..9d4d6a4a 100644 --- a/config_schema.json +++ b/config_schema.json @@ -163,6 +163,18 @@ "displayName": "Exclude Workspaces", "description": "Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces.", "stringSliceField": {} + }, + { + "name": "enable-incremental-sync", + "displayName": "Enable Incremental Sync", + "description": "Poll a Databricks audit-log event feed between full syncs to pick up access changes early. Deletions are still only caught by the next full sync.", + "boolField": {} + }, + { + "name": "sql-warehouse-id", + "displayName": "SQL Warehouse ID", + "description": "ID of the Databricks SQL warehouse used to query system.access.audit. Required when incremental sync is enabled; the workspace hosting it is discovered automatically.", + "stringField": {} } ], "constraints": [ @@ -198,7 +210,9 @@ "hostname", "account-hostname", "workspaces", - "databricks-exclude-workspaces" + "databricks-exclude-workspaces", + "enable-incremental-sync", + "sql-warehouse-id" ], "default": true }, diff --git a/docs/connector.mdx b/docs/connector.mdx index 5307fd67..311b6b48 100644 --- a/docs/connector.mdx +++ b/docs/connector.mdx @@ -23,6 +23,14 @@ The Databricks connector supports [automatic account provisioning and deprovisio Provisioning **account groups** requires OAuth authentication. It is not available when authenticating with a workspace token, because the Databricks API does not allow provisioning account groups from a workspace token. Workspace-scoped groups can still be provisioned with a workspace token. +### 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 access is only reflected after the next full sync. + +Turning this on requires a small Databricks SQL warehouse and a one-time setup step performed by a Databricks admin (granting the connector read access to Databricks' `system.access` activity log). Ask your connector operator or C1 support contact to enable it for you. + +Incremental sync requires OAuth2 (service principal) authentication — it's not available with workspace tokens, since the Account API needed to resolve audit events back to synced resources is unreachable that way. + ## Gather Databricks credentials Configuring the connector requires you to pass in credentials generated in Databricks. Gather these credentials before you move on. diff --git a/pkg/config/conf.gen.go b/pkg/config/conf.gen.go index 6ca5facf..e41c2122 100644 --- a/pkg/config/conf.gen.go +++ b/pkg/config/conf.gen.go @@ -13,6 +13,8 @@ type Databricks struct { WorkspaceTokens []string `mapstructure:"workspace-tokens"` BaseUrl string `mapstructure:"base-url"` DatabricksExcludeWorkspaces []string `mapstructure:"databricks-exclude-workspaces"` + EnableIncrementalSync bool `mapstructure:"enable-incremental-sync"` + 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 f6790134..e5c420fc 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -70,6 +70,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"), ) + EnableIncrementalSyncField = field.BoolField( + "enable-incremental-sync", + field.WithDescription("Poll a Databricks audit-log event feed between full syncs to pick up access changes early. Deletions are still only caught by the next full sync."), + field.WithDisplayName("Enable Incremental Sync"), + field.WithDefaultValue(false), + ) + SQLWarehouseIDField = field.StringField( + "sql-warehouse-id", + field.WithDescription( + "ID of the Databricks SQL warehouse used to query system.access.audit. Required when incremental "+ + "sync is enabled; the workspace hosting it is discovered automatically.", + ), + field.WithDisplayName("SQL Warehouse ID"), + ) configFields = []field.SchemaField{ AccountHostnameField, AccountIdField, @@ -80,6 +94,8 @@ var ( WorkspaceTokensField, BaseURLField, ExcludeWorkspacesField, + EnableIncrementalSyncField, + SQLWarehouseIDField, } ) @@ -101,6 +117,7 @@ var Config = field.NewConfiguration( Fields: []field.SchemaField{ AccountIdField, DatabricksClientIdField, DatabricksClientSecretField, HostnameField, AccountHostnameField, WorkspacesField, ExcludeWorkspacesField, + EnableIncrementalSyncField, SQLWarehouseIDField, }, Default: true, }, @@ -110,6 +127,8 @@ var Config = field.NewConfiguration( HelpText: "Authenticate with a personal access token scoped to each workspace. " + "Does not sync account-level data (account entitlements and grants, and " + "workspace-membership entitlements); use OAuth for full account coverage.", + // Incremental sync requires the Account API, which workspace tokens can't reach + // (see Validate) — omitted here so the UI doesn't offer an option that can never work. Fields: []field.SchemaField{AccountIdField, WorkspacesField, WorkspaceTokensField, HostnameField, AccountHostnameField}, Default: false, }, diff --git a/pkg/connector/account.go b/pkg/connector/account.go index 96aa8f2c..4b0a4e3f 100644 --- a/pkg/connector/account.go +++ b/pkg/connector/account.go @@ -216,6 +216,16 @@ 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 + } + + 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..f0d4c451 --- /dev/null +++ b/pkg/connector/audit_event_feed.go @@ -0,0 +1,624 @@ +package connector + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "sort" + "strconv" + "strings" + "sync" + "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" + "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" + "go.uber.org/zap" + "google.golang.org/protobuf/types/known/timestamppb" +) + +const ( + auditEventFeedId = "databricks_audit_log" + + // The first poll looks back this far since there's no prior watermark yet. + auditLogLookback = 1 * time.Hour + + // 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 + + auditLogPageLimit = 1000 + + // auditLogRetention mirrors system.access.audit's documented 365-day retention; a cursor + // older than this can no longer be satisfied by the table and is treated as stale. + auditLogRetention = 365 * 24 * time.Hour + + auditServiceAccounts = "accounts" +) + +// 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. +type auditActionMapping struct { + resourceType *v2.ResourceType + idParam string + accountRole string + roleNames []string +} + +// 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"}, + {auditServiceAccounts, "addPrincipalToGroup"}: {resourceType: groupResourceType, idParam: "targetGroupId"}, + {auditServiceAccounts, "removePrincipalFromGroup"}: {resourceType: groupResourceType, idParam: "targetGroupId"}, + {auditServiceAccounts, "addPrincipalsToGroup"}: {resourceType: groupResourceType, idParam: "targetGroupId"}, + {auditServiceAccounts, "removePrincipalsFromGroup"}: {resourceType: groupResourceType, idParam: "targetGroupId"}, + {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, "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}}, + {auditServiceAccounts, "setAdmin"}: {resourceType: userResourceType, idParam: "targetUserId", accountRole: AccountAdminRole}, + // removeAdmin revokes *workspace* admin, not account admin, so only the user is refreshed. + {auditServiceAccounts, "removeAdmin"}: {resourceType: userResourceType, idParam: "targetUserId"}, +} + +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. Corrupt/missing is routine and logged at Debug; a stale-but-valid +// cursor indicates a real data gap and is logged at Warn. +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.Warn("databricks-connector: event cursor is older than system.access.audit's 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 + enableIncrementalSync bool + sqlWarehouseID string + + // queryWorkspaceID caches which workspace hosts sqlWarehouseID across polls. + queryWorkspaceMu sync.Mutex + queryWorkspaceID string +} + +func newAuditEventFeed( + client *databricks.Client, + workspaces []string, + enableIncrementalSync bool, + sqlWarehouseID string, +) *auditEventFeed { + return &auditEventFeed{ + client: client, + workspaces: workspaces, + enableIncrementalSync: enableIncrementalSync, + sqlWarehouseID: sqlWarehouseID, + } +} + +// resolveQueryWorkspaceID returns which workspace hosts f.sqlWarehouseID, resolving it +// once via resolveWarehouseWorkspace rather than probing on every poll. +func (f *auditEventFeed) resolveQueryWorkspaceID(ctx context.Context, workspaces []databricks.Workspace) (string, *v2.RateLimitDescription, error) { + f.queryWorkspaceMu.Lock() + defer f.queryWorkspaceMu.Unlock() + + if f.queryWorkspaceID != "" { + return f.queryWorkspaceID, nil, nil + } + + id, rateLimit, err := resolveWarehouseWorkspace(ctx, f.client, workspaces, f.sqlWarehouseID) + if err != nil { + return "", rateLimit, err + } + + f.queryWorkspaceID = id + return id, rateLimit, nil +} + +// EventFeedMetadata is registered unconditionally; enable-incremental-sync gates behavior +// inside ListEvents instead, to avoid confusing "feed not found" errors when it's off. +func (f *auditEventFeed) EventFeedMetadata(ctx context.Context) *v2.EventFeedMetadata { + return &v2.EventFeedMetadata{ + Id: auditEventFeedId, + SupportedEventTypes: []v2.EventType{v2.EventType_EVENT_TYPE_RESOURCE_CHANGE}, + } +} + +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.enableIncrementalSync { + return nil, &pagination.StreamState{}, nil, nil + } + + now := time.Now() + cursor := decodeEventCursor(ctx, pToken.Cursor, now) + + if cursor.StartAt.IsZero() { + start := now.Add(-auditLogLookback) + if earliestEvent != nil { + start = earliestEvent.AsTime() + } + cursor = eventPageCursor{StartAt: start} + } + + workspaces, err := resolveSQLWorkspaces(ctx, f.client, f.workspaces) + if err != nil { + return nil, nil, nil, fmt.Errorf("databricks-connector: failed to list workspaces: %w", err) + } + if len(workspaces) == 0 { + return nil, nil, nil, fmt.Errorf("databricks-connector: no workspace available to query system.access.audit") + } + + workspaceLookup := make(map[int64]string, len(workspaces)) + for _, w := range workspaces { + workspaceLookup[int64(w.ID)] = w.DeploymentName + } + + queryWorkspaceId, rateLimit, err := f.resolveQueryWorkspaceID(ctx, workspaces) + if rateLimit != nil { + annos.WithRateLimiting(rateLimit) + } + if err != nil { + return nil, nil, annos, err + } + + rows, rateLimit, err := f.queryAuditLog(ctx, queryWorkspaceId, cursor) + 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 + for _, row := range rows { + affected := mapAuditRowToResource(ctx, row, f.client.GetAccountId(), f.client.IsAccountAPIAvailable(), workspaceLookup) + if len(affected) == 0 { + l.Debug("databricks-connector: skipping audit row with no resource mapping", + zap.String("action_name", row.ActionName), + zap.String("event_id", row.EventID), + ) + continue + } + + for i, a := range affected { + events = append(events, &v2.Event{ + Id: fmt.Sprintf("%s/%d", row.EventID, i), + OccurredAt: timestamppb.New(row.EventTime), + Event: &v2.Event_ResourceChangeEvent{ + ResourceChangeEvent: &v2.ResourceChangeEvent{ + ResourceId: a.resourceId, + ParentResourceId: a.parentResourceId, + }, + }, + }) + } + } + + hasMore := len(rows) >= auditLogPageLimit + nextCursor := advanceEventCursor(cursor, rows, hasMore, now) + + encoded, err := encodeEventCursor(nextCursor) + if err != nil { + return nil, nil, nil, fmt.Errorf("databricks-connector: failed to encode event cursor: %w", err) + } + + return events, &pagination.StreamState{Cursor: encoded, HasMore: hasMore}, annos, nil +} + +// advanceEventCursor advances only to the last row processed (by the well-ordered +// (event_time, event_id) boundary) while a page is full, and once drained, trails the +// newest event seen (or wall-clock time if empty) by auditLogTrailingLag. +func advanceEventCursor(cursor eventPageCursor, rows []auditLogRow, hasMore bool, now time.Time) eventPageCursor { + latest := cursor.StartAt + lastEventID := cursor.StartAfterEventID + if len(rows) > 0 { + last := rows[len(rows)-1] + latest = last.EventTime + lastEventID = last.EventID + } + + if hasMore { + startAt := latest + startAfterEventID := lastEventID + // Clamp intra-page advances to the trailing-lag boundary; otherwise the never-regress + // floor below would permanently defeat auditLogTrailingLag for this burst. + if laggedFloor := now.Add(-auditLogTrailingLag); startAt.After(laggedFloor) { + startAt = laggedFloor + startAfterEventID = "" + } + return eventPageCursor{StartAt: startAt, StartAfterEventID: startAfterEventID} + } + + target := latest.Add(-auditLogTrailingLag) + if len(rows) == 0 { + target = now.Add(-auditLogTrailingLag) + } + if target.Before(cursor.StartAt) { + target = cursor.StartAt + } + + startAfterEventID := "" + if target.Equal(latest) { + startAfterEventID = lastEventID + } + + return eventPageCursor{StartAt: target, StartAfterEventID: startAfterEventID} +} + +// affectedResource is one resource a mapped audit row's action changed. +type affectedResource struct { + resourceId *v2.ResourceId + parentResourceId *v2.ResourceId +} + +// mapAuditRowToResource maps an audit row to every Baton resource its action affects, skipping +// anything unresolvable. The principal's parent mirrors how it's actually synced (see +// groupGrantParent in helpers.go), not the scope the audit row occurred in. +func mapAuditRowToResource(ctx context.Context, row auditLogRow, accountId string, accountAPIAvailable bool, workspaceLookup map[int64]string) []affectedResource { + mapping, ok := auditLogActions[auditActionKey{Service: row.ServiceName, Action: row.ActionName}] + if !ok { + return nil + } + + accountParent := &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: accountId} + + var workspaceParent *v2.ResourceId + if row.WorkspaceID != 0 { + deploymentName, found := workspaceLookup[row.WorkspaceID] + if !found { + return nil + } + workspaceParent = &v2.ResourceId{ResourceType: workspaceResourceType.Id, Resource: deploymentName} + } + + var affected []affectedResource + + switch { + case mapping.resourceType == workspaceResourceType: + if workspaceParent == nil { + return nil + } + affected = append(affected, affectedResource{resourceId: workspaceParent, parentResourceId: accountParent}) + case mapping.resourceType != nil: + parent := accountParent + if !accountAPIAvailable { + if workspaceParent == nil { + return nil + } + parent = workspaceParent + } + + nativeId, ok := row.RequestParams[mapping.idParam] + if !ok || nativeId == "" { + return nil + } + + resourceId := &v2.ResourceId{ResourceType: mapping.resourceType.Id, Resource: nativeId} + if mapping.resourceType == groupResourceType { + resourceId.Resource = groupResourceId(ctx, nativeId, parent) + } + + affected = append(affected, affectedResource{resourceId: resourceId, parentResourceId: parent}) + } + + 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 +} + +// resolveSQLWorkspaces returns the workspaces available to run the audit-log SQL query +// against, without calling the Account API under token auth (unreachable there). +func resolveSQLWorkspaces(ctx context.Context, client *databricks.Client, configuredWorkspaces []string) ([]databricks.Workspace, error) { + if client.IsTokenAuth() { + workspaces := make([]databricks.Workspace, 0, len(configuredWorkspaces)) + for _, name := range configuredWorkspaces { + workspaces = append(workspaces, databricks.Workspace{DeploymentName: name}) + } + return workspaces, nil + } + + workspaces, _, err := client.ListWorkspaces(ctx) + if err != nil { + return nil, err + } + + if len(configuredWorkspaces) == 0 { + return workspaces, nil + } + + 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, nil +} + +// 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) { + if len(workspaces) == 1 { + return workspaces[0].DeploymentName, nil, nil + } + + var rateLimit *v2.RateLimitDescription + for _, w := range workspaces { + found, rl, err := client.WarehouseExists(ctx, w.DeploymentName, warehouseId) + if rl != nil { + rateLimit = rl + } + if err != nil { + 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 + } + } + + return "", rateLimit, fmt.Errorf( + "databricks-connector: sql-warehouse-id %q was not found in any of the %d available workspaces", + warehouseId, len(workspaces), + ) +} + +func (f *auditEventFeed) queryAuditLog(ctx context.Context, workspaceId string, cursor eventPageCursor) ([]auditLogRow, *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 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"}, + ) + if err != nil { + return nil, rateLimit, err + } + + rows, err := parseAuditLogRows(ctx, result) + return rows, rateLimit, err +} + +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. +func parseAuditLogRows(ctx context.Context, result *databricks.StatementResult) ([]auditLogRow, 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, 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.Warn("databricks-connector: skipping malformed audit log row", zap.Error(err)) + continue + } + rows = append(rows, row) + } + + return rows, 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) + } + + eventTime, err := time.Parse("2006-01-02 15:04:05.999", r[colIndex[colEventTime]]) + if err != nil { + eventTime, err = time.Parse(time.RFC3339, r[colIndex[colEventTime]]) + if err != nil { + return auditLogRow{}, fmt.Errorf("failed to parse event_time %q: %w", r[colIndex[colEventTime]], 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: r[colIndex[colEventID]], + EventTime: eventTime, + 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..7bd7db59 --- /dev/null +++ b/pkg/connector/audit_event_feed_test.go @@ -0,0 +1,708 @@ +package connector + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strconv" + "strings" + "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" +) + +// TestResolveSQLWorkspacesTokenAuth ensures the audit-log workspace lookup never calls the +// Account API under workspace-token auth (unreachable in that mode), building minimal +// workspaces from the configured deployment names instead. +func TestResolveSQLWorkspacesTokenAuth(t *testing.T) { + auth := databricks.NewTokenAuth([]string{"dbc-1", "dbc-2"}, []string{"token-1", "token-2"}) + client, err := databricks.NewClient(context.Background(), &http.Client{}, "example.cloud.databricks.com", "accounts.cloud.databricks.com", "", "", auth, nil) + if err != nil { + t.Fatalf("NewClient() error = %v", err) + } + + got, err := resolveSQLWorkspaces(context.Background(), client, []string{"dbc-1", "dbc-2"}) + if err != nil { + t.Fatalf("resolveSQLWorkspaces() error = %v", err) + } + + want := []databricks.Workspace{{DeploymentName: "dbc-1"}, {DeploymentName: "dbc-2"}} + if len(got) != len(want) { + t.Fatalf("got %d workspaces, want %d: %+v", len(got), len(want), got) + } + for i := range want { + if got[i].DeploymentName != want[i].DeploymentName || got[i].ID != 0 { + t.Errorf("[%d] = %+v, want %+v", i, got[i], want[i]) + } + } +} + +// 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.NewTokenAuth(nil, nil) + 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 needs no probe", func(t *testing.T) { + client := newProbeTestClient(t, func(w http.ResponseWriter, r *http.Request) { + t.Fatalf("unexpected request %s %s: a single workspace should skip probing", r.Method, r.URL.Path) + }) + 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") + } + }) + + 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("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 +// system.access.audit's 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) + } +} + +// TestAdvanceEventCursorFullPageAdvancesPastLagWindow verifies that a full page of rows +// older than the trailing-lag boundary still advances StartAt to the last row processed. +func TestAdvanceEventCursorFullPageAdvancesPastLagWindow(t *testing.T) { + startAt := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + cursor := eventPageCursor{StartAt: startAt} + + rows := []auditLogRow{ + {EventID: "1", EventTime: startAt.Add(1 * time.Minute)}, + {EventID: "2", EventTime: startAt.Add(2 * time.Minute)}, + } + + // now is far enough past the rows that startAt+2min is still older than + // now-auditLogTrailingLag, so the lag clamp shouldn't kick in. + now := startAt.Add(2*time.Minute + auditLogTrailingLag + time.Hour) + + next := advanceEventCursor(cursor, rows, true, now) + + wantStart := startAt.Add(2 * time.Minute) + if !next.StartAt.Equal(wantStart) { + t.Errorf("StartAt = %v, want %v (rows are older than the lag window, so no clamp)", next.StartAt, wantStart) + } + if next.StartAfterEventID != "2" { + t.Errorf("StartAfterEventID = %q, want %q", next.StartAfterEventID, "2") + } +} + +// TestAdvanceEventCursorFullPageClampsToLagWindow verifies intra-page paging never +// advances StartAt past now-auditLogTrailingLag when the rows are within the lag window. +func TestAdvanceEventCursorFullPageClampsToLagWindow(t *testing.T) { + startAt := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + cursor := eventPageCursor{StartAt: startAt} + + rows := []auditLogRow{ + {EventID: "1", EventTime: startAt.Add(1 * time.Minute)}, + {EventID: "2", EventTime: startAt.Add(2 * time.Minute)}, + } + + // now is close to the rows' timestamps, so startAt+2min falls inside the lag window. + now := startAt.Add(10 * time.Minute) + + next := advanceEventCursor(cursor, rows, true, now) + + wantStart := now.Add(-auditLogTrailingLag) + if !next.StartAt.Equal(wantStart) { + t.Errorf("StartAt = %v, want %v (clamped to the trailing-lag boundary)", next.StartAt, wantStart) + } + if next.StartAfterEventID != "" { + t.Errorf("StartAfterEventID = %q, want empty (clamped boundary doesn't tie to a real row)", next.StartAfterEventID) + } + + // Once the burst drains, the trailing lag must still apply going forward. + drainedLatest := startAt.Add(3 * time.Hour) + drainedRows := []auditLogRow{{EventID: "3", EventTime: drainedLatest}} + drainedNow := drainedLatest.Add(5 * time.Minute) + drained := advanceEventCursor(next, drainedRows, false, drainedNow) + + wantDrainedStart := drainedLatest.Add(-auditLogTrailingLag) + if !drained.StartAt.Equal(wantDrainedStart) { + t.Errorf("drained StartAt = %v, want %v (trailing lag re-applied after drain)", drained.StartAt, wantDrainedStart) + } + if !drained.StartAt.After(next.StartAt) { + t.Errorf("drained StartAt = %v did not advance past the clamped intra-page StartAt %v", drained.StartAt, next.StartAt) + } +} + +func TestAdvanceEventCursorDrainedPageAppliesTrailingLag(t *testing.T) { + startAt := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + cursor := eventPageCursor{StartAt: startAt} + latest := startAt.Add(5 * time.Hour) + + rows := []auditLogRow{ + {EventID: "1", EventTime: latest}, + } + + next := advanceEventCursor(cursor, rows, false, latest) + + wantStart := latest.Add(-auditLogTrailingLag) + if !next.StartAt.Equal(wantStart) { + t.Errorf("StartAt = %v, want %v", next.StartAt, wantStart) + } + // The trailing lag pushes the boundary well before the only row seen, so nothing ties. + if next.StartAfterEventID != "" { + t.Errorf("StartAfterEventID = %q, want empty", next.StartAfterEventID) + } +} + +// TestAdvanceEventCursorTieAtFlooredBoundaryIsRemembered covers a row landing exactly on +// the floored StartAt boundary, which would otherwise be re-fetched forever. +func TestAdvanceEventCursorTieAtFlooredBoundaryIsRemembered(t *testing.T) { + startAt := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + cursor := eventPageCursor{StartAt: startAt} + + rows := []auditLogRow{ + {EventID: "1", EventTime: startAt}, + } + + next := advanceEventCursor(cursor, rows, false, startAt) + + if !next.StartAt.Equal(startAt) { + t.Errorf("StartAt = %v, want unchanged %v", next.StartAt, startAt) + } + if next.StartAfterEventID != "1" { + t.Errorf("StartAfterEventID = %q, want %q", next.StartAfterEventID, "1") + } +} + +func TestAdvanceEventCursorNeverRegresses(t *testing.T) { + startAt := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC) + cursor := eventPageCursor{StartAt: startAt} + + // "now" is barely past startAt, so subtracting the trailing lag would regress. + now := startAt.Add(1 * time.Minute) + + next := advanceEventCursor(cursor, nil, false, now) + + if next.StartAt.Before(cursor.StartAt) { + t.Errorf("StartAt regressed: got %v, was %v", next.StartAt, cursor.StartAt) + } + if !next.StartAt.Equal(cursor.StartAt) { + t.Errorf("StartAt = %v, want unchanged %v", next.StartAt, cursor.StartAt) + } +} + +func TestAdvanceEventCursorEmptyWindowTrailsWallClock(t *testing.T) { + startAt := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + cursor := eventPageCursor{StartAt: startAt} + now := startAt.Add(10 * time.Hour) + + next := advanceEventCursor(cursor, nil, false, now) + + wantStart := now.Add(-auditLogTrailingLag) + if !next.StartAt.Equal(wantStart) { + t.Errorf("StartAt = %v, want %v", next.StartAt, wantStart) + } +} + +func TestMapAuditRowToResource(t *testing.T) { + 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 + accountAPIAvailable bool + want []wantResource + }{ + { + name: "account-level group create", + accountAPIAvailable: true, + 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 when the Account API is available", + accountAPIAvailable: true, + row: auditLogRow{ + ActionName: "addPrincipalToGroup", + 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 group change is workspace-parented under token auth", + accountAPIAvailable: false, + row: auditLogRow{ + ActionName: "addPrincipalToGroup", + ServiceName: "accounts", + WorkspaceID: 123, + RequestParams: map[string]string{"targetGroupId": "g-1"}, + }, + want: []wantResource{ + {groupResourceType.Id, groupResourceId(context.Background(), "g-1", workspaceParent), workspaceResourceType.Id, "my-workspace"}, + }, + }, + { + name: "workspace-scoped acl change also refreshes the workspace-access role", + accountAPIAvailable: true, + 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", + accountAPIAvailable: true, + 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: "updateUser stays account-parented when the Account API is available, but workspace roles still refresh", + accountAPIAvailable: true, + 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: "updateUser is workspace-parented under token auth", + accountAPIAvailable: false, + row: auditLogRow{ + ActionName: "updateUser", + ServiceName: "accounts", + WorkspaceID: 123, + RequestParams: map[string]string{"targetUserId": "u-1"}, + }, + want: []wantResource{ + {userResourceType.Id, "u-1", workspaceResourceType.Id, "my-workspace"}, + {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", + accountAPIAvailable: true, + 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 := mapAuditRowToResource(context.Background(), tc.row, accountId, tc.accountAPIAvailable, workspaceLookup) + 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 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") + } +} + +// TestAdvanceEventCursorLargeTiedBurstMakesProgress verifies more than auditLogPageLimit rows +// sharing one event_time still make progress via the (event_time, event_id) cursor. +func TestAdvanceEventCursorLargeTiedBurstMakesProgress(t *testing.T) { + tied := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + cursor := eventPageCursor{StartAt: tied.Add(-time.Minute)} + + rows := make([]auditLogRow, auditLogPageLimit) + for i := range rows { + rows[i] = auditLogRow{EventID: fmt.Sprintf("evt-%04d", i), EventTime: tied} + } + + // now is far past the lag window so the intra-page clamp doesn't interfere. + now := tied.Add(auditLogTrailingLag + time.Hour) + + next := advanceEventCursor(cursor, rows, true, now) + if !next.StartAt.Equal(tied) { + t.Fatalf("StartAt = %v, want %v (advances to the tied timestamp, not stuck before it)", next.StartAt, tied) + } + lastID := rows[len(rows)-1].EventID + if next.StartAfterEventID != lastID { + t.Fatalf("StartAfterEventID = %q, want %q (last row in the tied burst)", next.StartAfterEventID, lastID) + } + + // A later, non-tied page must advance the cursor past the tied burst. + followUpLatest := tied.Add(5 * time.Hour) + followUp := []auditLogRow{{EventID: "evt-1000", EventTime: followUpLatest}} + drained := advanceEventCursor(next, followUp, false, followUpLatest.Add(5*time.Minute)) + if !drained.StartAt.After(tied) { + t.Errorf("drained StartAt = %v did not advance past the tied burst's timestamp %v", drained.StartAt, tied) + } +} + +// 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) +} + +// 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) { + 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.NewTokenAuth([]string{"ws1"}, []string{"token-1"}) + 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. + client.UpdateAvailability(true, true) + + feed := newAuditEventFeed(client, []string{"ws1"}, true, "wh-1") + + events, streamState, annos, err := feed.ListEvents(context.Background(), nil, &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) + } +} diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 10996986..ea27d846 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "io" + "time" "github.com/conductorone/baton-databricks/pkg/config" "github.com/conductorone/baton-databricks/pkg/databricks" @@ -15,9 +16,15 @@ import ( "go.uber.org/zap" ) +// validateAuditLogAccessTimeout bounds the one-off audit-log probe in Validate() so a cold +// warehouse doesn't hang credential validation for minutes. +const validateAuditLogAccessTimeout = 90 * time.Second + type Databricks struct { - client *databricks.Client - workspaces []string + client *databricks.Client + workspaces []string + enableIncrementalSync bool + sqlWarehouseID string } // ResourceSyncers returns a ResourceSyncerV2 for each resource type that should be synced from the upstream service. @@ -34,6 +41,14 @@ func (d *Databricks) ResourceSyncers(ctx context.Context) []connectorbuilder.Res return syncers } +// EventFeeds registers the audit-log event feed unconditionally; enable-incremental-sync +// gates its behavior inside ListEvents instead. +func (d *Databricks) EventFeeds(ctx context.Context) []connectorbuilder.EventFeed { + return []connectorbuilder.EventFeed{ + newAuditEventFeed(d.client, d.workspaces, d.enableIncrementalSync, 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) { @@ -161,6 +176,42 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err ) } + if d.enableIncrementalSync { + if d.sqlWarehouseID == "" { + return nil, fmt.Errorf("databricks-connector: sql-warehouse-id is required when incremental sync is enabled") + } + + // Under token auth no numeric workspace ID is ever learned, so audit rows can never + // be resolved back to a synced resource (see mapAuditRowToResource) — every poll + // would silently produce zero events. Fail loudly instead of polling forever for + // nothing. + if d.client.IsTokenAuth() { + return nil, fmt.Errorf("databricks-connector: incremental sync is not supported with workspace token auth") + } + + auditWorkspaces, err := resolveSQLWorkspaces(ctx, d.client, d.workspaces) + if err != nil { + return nil, fmt.Errorf("databricks-connector: incremental sync requires the account API to list workspaces: %w", err) + } + if len(auditWorkspaces) == 0 { + return nil, fmt.Errorf("databricks-connector: incremental sync requires at least one workspace to query system.access.audit") + } + + queryWorkspaceId, _, err := resolveWarehouseWorkspace(ctx, d.client, auditWorkspaces, d.sqlWarehouseID) + if err != nil { + return nil, err + } + validateCtx, cancel := context.WithTimeout(ctx, validateAuditLogAccessTimeout) + err = d.client.ValidateAuditLogAccess(validateCtx, queryWorkspaceId, d.sqlWarehouseID) + cancel() + if err != nil { + return nil, fmt.Errorf( + "databricks-connector: incremental sync is enabled but the connector cannot query system.access.audit via warehouse %s: %w", + d.sqlWarehouseID, err, + ) + } + } + return nil, nil } @@ -174,6 +225,8 @@ func New( auth databricks.Auth, excludeWorkspaces []string, workspaces []string, + enableIncrementalSync bool, + sqlWarehouseID string, ) (*Databricks, error) { httpClient, err := auth.GetClient(ctx) if err != nil { @@ -186,8 +239,10 @@ func New( } return &Databricks{ - client: client, - workspaces: workspaces, + client: client, + workspaces: workspaces, + enableIncrementalSync: enableIncrementalSync, + sqlWarehouseID: sqlWarehouseID, }, nil } @@ -216,6 +271,8 @@ func NewConnector(ctx context.Context, cfg *config.Databricks, opts *cli.Connect auth, cfg.DatabricksExcludeWorkspaces, cfg.Workspaces, + cfg.EnableIncrementalSync, + cfg.SqlWarehouseId, ) if err != nil { return nil, nil, err diff --git a/pkg/connector/groups.go b/pkg/connector/groups.go index 7d223f8b..2c193cb5 100644 --- a/pkg/connector/groups.go +++ b/pkg/connector/groups.go @@ -272,6 +272,37 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.S return rv, &rs.SyncOpResults{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(ctx, workspaceId, groupId.Resource, databricks.NewGroupAttrVars()) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err) + } + + annos := annotations.Annotations{} + if rateLimitData != nil { + annos.WithRateLimiting(rateLimitData) + } + + resource, err := groupResource(ctx, group, parentResourceId) + if err != nil { + return nil, nil, err + } + + 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/roles.go b/pkg/connector/roles.go index 8fb0ec1f..15f36cb5 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" @@ -41,21 +42,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{}{ "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, @@ -293,6 +295,27 @@ 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 + } + + 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 a51d383c..12005188 100644 --- a/pkg/connector/service-principals.go +++ b/pkg/connector/service-principals.go @@ -29,15 +29,16 @@ func (s *servicePrincipalBuilder) servicePrincipalResource(ctx context.Context, profile := map[string]interface{}{ "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)) } @@ -215,6 +216,31 @@ 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(ctx, workspaceId, resourceId.Resource) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to get service principal %s: %w", resourceId.Resource, err) + } + + annos := annotations.Annotations{} + if rateLimitData != nil { + annos.WithRateLimiting(rateLimitData) + } + + resource, err := s.servicePrincipalResource(ctx, servicePrincipal, parentResourceId) + if err != nil { + return nil, nil, err + } + + 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 775ada76..171eec7c 100644 --- a/pkg/connector/users.go +++ b/pkg/connector/users.go @@ -61,8 +61,9 @@ func (u *userBuilder) userResource(ctx context.Context, user *databricks.User, p rs.WithResourceProfile(profile), rs.WithResourceStatus(status, ""), } + // 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)) } @@ -231,6 +232,31 @@ 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(ctx, workspaceId, resourceId.Resource) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to get user %s: %w", resourceId.Resource, err) + } + + annos := annotations.Annotations{} + if rateLimitData != nil { + annos.WithRateLimiting(rateLimitData) + } + + resource, err := u.userResource(ctx, user, parentResourceId) + if err != nil { + return nil, nil, err + } + + 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/workspaces.go b/pkg/connector/workspaces.go index 8a594192..5e9db86a 100644 --- a/pkg/connector/workspaces.go +++ b/pkg/connector/workspaces.go @@ -284,6 +284,43 @@ func (w *workspaceBuilder) Grants(ctx context.Context, resource *v2.Resource, _ return rv, nil, nil } +// Get re-fetches a single workspace, used to re-sync it after a RESOURCE_CHANGE event. +func (w *workspaceBuilder) Get(ctx context.Context, resourceId *v2.ResourceId, parentResourceId *v2.ResourceId) (*v2.Resource, annotations.Annotations, error) { + if w.client.IsTokenAuth() { + if _, ok := w.workspaces[resourceId.Resource]; !ok { + return nil, nil, fmt.Errorf("databricks-connector: workspace %s is not configured", resourceId.Resource) + } + + ws := &databricks.Workspace{DeploymentName: resourceId.Resource} + resource, err := minimalWorkspaceResource(ctx, ws, parentResourceId) + if err != nil { + return nil, nil, err + } + + return resource, nil, nil + } + + workspace, _, err := w.client.GetWorkspace(ctx, resourceId.Resource) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to get workspace %s: %w", resourceId.Resource, err) + } + + // Mirror List's --workspaces allowlist so a targeted re-sync can't resurrect a + // workspace that was deliberately excluded from the configured set. + if len(w.workspaces) > 0 { + if _, ok := matchConfiguredWorkspace(w.workspaces, workspace.DeploymentName, workspace.Name, strconv.Itoa(workspace.ID)); !ok { + return nil, nil, fmt.Errorf("databricks-connector: workspace %s is not configured", resourceId.Resource) + } + } + + resource, err := workspaceResource(ctx, workspace, parentResourceId) + if err != nil { + return nil, nil, err + } + + return resource, nil, nil +} + func (w *workspaceBuilder) Grant(ctx context.Context, principal *v2.Resource, entitlement *v2.Entitlement) (annotations.Annotations, error) { l := ctxzap.Extract(ctx) diff --git a/pkg/databricks/client.go b/pkg/databricks/client.go index 5882e8fb..5c099e88 100644 --- a/pkg/databricks/client.go +++ b/pkg/databricks/client.go @@ -608,6 +608,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, fmt.Errorf("workspace %s not found", deploymentName) +} + func (c *Client) ListWorkspaceMembers( ctx context.Context, workspaceId string, diff --git a/pkg/databricks/sql.go b/pkg/databricks/sql.go new file mode 100644 index 00000000..f6520ee8 --- /dev/null +++ b/pkg/databricks/sql.go @@ -0,0 +1,251 @@ +package databricks + +import ( + "context" + "errors" + "fmt" + "net/http" + "strconv" + "time" + + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" + "go.uber.org/zap" +) + +const ( + statementsEndpoint = "/api/2.0/sql/statements" + warehousesEndpoint = "/api/2.0/sql/warehouses" + + statementWaitTimeout = "30s" + 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) { + 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) + } + + res, polledRateLimit, err := c.pollStatement(ctx, workspaceId, 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 + } + return nil, rateLimit, fmt.Errorf("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 + } + return result, rateLimit, err +} + +// 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 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(): + if err := ctx.Err(); err != nil { + c.cancelStatement(workspaceId, res.StatementID) + return res, rateLimit, err + } + l.Warn("sql statement did not reach a terminal state before poll timeout, canceling", + zap.String("statement_id", res.StatementID), + zap.String("state", string(res.Status.State)), + zap.Duration("max_wait", statementPollMaxWait), + ) + c.cancelStatement(workspaceId, res.StatementID) + return res, rateLimit, fmt.Errorf("statement %s did not reach a terminal state within %s", res.StatementID, statementPollMaxWait) + 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 + polledRateLimit, err := c.Get(pollCtx, u, &polled) + if polledRateLimit != nil { + rateLimit = polledRateLimit + } + if err != nil { + return res, rateLimit, fmt.Errorf("failed to poll statement %s: %w", res.StatementID, err) + } + res = polled + } + + return res, rateLimit, nil +} + +// cancelStatement best-effort cancels a statement we've given up polling on, using a fresh +// context 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(workspaceId, statementId string) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + u := c.workspaceUrl(workspaceId).JoinPath(statementsEndpoint, statementId, "cancel") + response := struct{}{} + if _, err := c.Post(ctx, u, nil, &response); err != nil { + ctxzap.Extract(ctx).Warn("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) + } + rows = append(rows, chunk.DataArray...) + nextChunk = chunk.NextChunkIndex + } + + return &StatementResult{Columns: columns, Rows: rows}, rateLimit, nil +} + +// ValidateAuditLogAccess confirms the configured warehouse can query system.access.audit, +// which requires a one-time SELECT grant from a metastore admin (see README). +func (c *Client) ValidateAuditLogAccess(ctx context.Context, workspaceId, warehouseId string) error { + if _, _, err := c.ExecuteStatement(ctx, workspaceId, warehouseId, "SELECT 1 FROM system.access.audit LIMIT 1"); err != nil { + return fmt.Errorf("failed to query system.access.audit: %w", err) + } + return 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 { + return false, rateLimit, nil + } + return false, rateLimit, err + } + + return true, rateLimit, nil +}