From 7303c659bc6ad7e4c70f27f6caded3c101bb55f6 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Tue, 18 Aug 2026 17:55:11 -0300 Subject: [PATCH 01/15] feat: incremental sync v1 - scoped implementation --- README.md | 27 ++ baton_capabilities.json | 11 +- config_schema.json | 12 + docs/connector.mdx | 6 + pkg/config/conf.gen.go | 2 + pkg/config/config.go | 13 + pkg/connector/account.go | 10 + pkg/connector/audit_event_feed.go | 393 +++++++++++++++++++++++++ pkg/connector/audit_event_feed_test.go | 247 ++++++++++++++++ pkg/connector/connector.go | 38 ++- pkg/connector/groups.go | 30 ++ pkg/connector/roles.go | 22 ++ pkg/connector/service-principals.go | 25 ++ pkg/connector/users.go | 25 ++ pkg/connector/workspaces.go | 15 + pkg/databricks/client.go | 24 ++ pkg/databricks/sql.go | 181 ++++++++++++ 17 files changed, 1078 insertions(+), 3 deletions(-) create mode 100644 pkg/connector/audit_event_feed.go create mode 100644 pkg/connector/audit_event_feed_test.go create mode 100644 pkg/databricks/sql.go diff --git a/README.md b/README.md index 289d0208..ccd98e8a 100644 --- a/README.md +++ b/README.md @@ -103,6 +103,31 @@ 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: + +- `--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. +- 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 povisioning 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. @@ -139,6 +164,7 @@ Flags: --databricks-client-id string 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 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) + --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) -f, --file string The path to the c1z file to sync with ($BATON_FILE) (default "sync.c1z") -h, --help help for baton-databricks --hostname string The Databricks hostname used to connect to the Databricks API ($BATON_HOSTNAME) (default "cloud.databricks.com") @@ -146,6 +172,7 @@ Flags: --log-level string The log level: debug, info, warn, error ($BATON_LOG_LEVEL) (default "info") -p, --provisioning This must be set in order for provisioning actions to be enabled ($BATON_PROVISIONING) --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. ($BATON_SQL_WAREHOUSE_ID) --ticketing This must be set to enable ticketing support ($BATON_TICKETING) -v, --version version for baton-databricks 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 194a841b..6e615ec3 100644 --- a/config_schema.json +++ b/config_schema.json @@ -145,6 +145,18 @@ "displayName": "Exclude Workspaces", "description": "Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID", "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.", + "stringField": {} } ], "displayName": "Databricks", diff --git a/docs/connector.mdx b/docs/connector.mdx index 88ce529e..f3a0edb3 100644 --- a/docs/connector.mdx +++ b/docs/connector.mdx @@ -18,6 +18,12 @@ sidebarTitle: "Databricks" The Databricks connector supports [automatic account provisioning and deprovisioning](/product/admin/account-provisioning). +### 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. + ## 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 80ebf5b1..241a7ac6 100644 --- a/pkg/config/conf.gen.go +++ b/pkg/config/conf.gen.go @@ -11,6 +11,8 @@ type Databricks struct { Hostname string `mapstructure:"hostname"` 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 fb664cb3..08861070 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -46,6 +46,17 @@ var ( field.WithDescription("Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID"), 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."), + field.WithDisplayName("SQL Warehouse ID"), + ) configFields = []field.SchemaField{ AccountHostnameField, AccountIdField, @@ -54,6 +65,8 @@ var ( HostnameField, BaseURLField, ExcludeWorkspacesField, + EnableIncrementalSyncField, + SQLWarehouseIDField, } ) diff --git a/pkg/connector/account.go b/pkg/connector/account.go index 14dc3d80..e285b2ef 100644 --- a/pkg/connector/account.go +++ b/pkg/connector/account.go @@ -209,6 +209,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..24eba054 --- /dev/null +++ b/pkg/connector/audit_event_feed.go @@ -0,0 +1,393 @@ +package connector + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "strconv" + "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" + + // Databricks audit log delivery can lag up to 24h, so the first poll looks back that far. + auditLogLookback = 24 * 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 +) + +// auditLogActions maps audit log action_name values to the resource type they affect and +// the request_params key holding the native resource ID. Not yet verified against a live workspace. +var auditLogActions = map[string]struct { + resourceType *v2.ResourceType + idParam string +}{ + "createGroup": {groupResourceType, "targetGroupId"}, + "addPrincipalToGroup": {groupResourceType, "targetGroupId"}, + "removePrincipalFromGroup": {groupResourceType, "targetGroupId"}, + "deleteGroup": {groupResourceType, "targetGroupId"}, + "createUser": {userResourceType, "targetUserId"}, + "updateUser": {userResourceType, "targetUserId"}, + "deleteUser": {userResourceType, "targetUserId"}, + "createServicePrincipal": {servicePrincipalResourceType, "targetServicePrincipalId"}, + "updateServicePrincipal": {servicePrincipalResourceType, "targetServicePrincipalId"}, + "deleteServicePrincipal": {servicePrincipalResourceType, "targetServicePrincipalId"}, + "changeDatabricksWorkspaceAcl": {workspaceResourceType, ""}, +} + +func auditLogActionNames() []string { + names := make([]string, 0, len(auditLogActions)) + for name := range auditLogActions { + names = append(names, name) + } + return names +} + +// eventPageCursor is the opaque state persisted between ListEvents calls. StartAt only +// ever advances forward, and LastEventIDs dedupes rows tied exactly on that boundary. +type eventPageCursor struct { + StartAt time.Time `json:"start_at"` + LatestEventSeen time.Time `json:"latest_event_seen"` + LastEventIDs []string `json:"last_event_ids"` +} + +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 when missing or corrupt, so callers +// self-heal by resetting to the lookback default. +func decodeEventCursor(s string) eventPageCursor { + if s == "" { + return eventPageCursor{} + } + + raw, err := base64.StdEncoding.DecodeString(s) + if err != nil { + return eventPageCursor{} + } + + var c eventPageCursor + if err := json.Unmarshal(raw, &c); err != nil { + return eventPageCursor{} + } + + return c +} + +type auditLogRow struct { + EventID string + EventTime time.Time + WorkspaceID int64 + ActionName string + RequestParams map[string]string +} + +type auditEventFeed struct { + client *databricks.Client + enableIncrementalSync bool + sqlWarehouseID string +} + +func newAuditEventFeed(client *databricks.Client, enableIncrementalSync bool, sqlWarehouseID string) *auditEventFeed { + return &auditEventFeed{ + client: client, + enableIncrementalSync: enableIncrementalSync, + sqlWarehouseID: sqlWarehouseID, + } +} + +// 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) + + if !f.enableIncrementalSync { + return nil, &pagination.StreamState{}, nil, nil + } + + cursor := decodeEventCursor(pToken.Cursor) + now := time.Now() + + if cursor.StartAt.IsZero() { + start := now.Add(-auditLogLookback) + if earliestEvent != nil && earliestEvent.AsTime().After(start) { + start = earliestEvent.AsTime() + } + cursor = eventPageCursor{StartAt: start} + } + + workspaces, _, err := f.client.ListWorkspaces(ctx) + 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") + } + + queryWorkspaceId, workspaceLookup := sqlQueryWorkspace(workspaces) + + rows, err := f.queryAuditLog(ctx, queryWorkspaceId, cursor) + if err != nil { + return nil, nil, nil, fmt.Errorf("databricks-connector: failed to query audit log: %w", err) + } + + seen := make(map[string]struct{}, len(cursor.LastEventIDs)) + for _, id := range cursor.LastEventIDs { + seen[id] = struct{}{} + } + + var events []*v2.Event + for _, row := range rows { + if _, ok := seen[row.EventID]; ok { + continue + } + + resourceId, parentResourceId, ok := mapAuditRowToResource(row, f.client.GetAccountId(), workspaceLookup) + if !ok { + l.Debug("databricks-connector: skipping audit row with no resource mapping", + zap.String("action_name", row.ActionName), + zap.String("event_id", row.EventID), + ) + continue + } + + events = append(events, &v2.Event{ + Id: row.EventID, + OccurredAt: timestamppb.New(row.EventTime), + Event: &v2.Event_ResourceChangeEvent{ + ResourceChangeEvent: &v2.ResourceChangeEvent{ + ResourceId: resourceId, + ParentResourceId: 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}, nil, nil +} + +// advanceEventCursor advances only to the last row processed 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 + var latestIDs []string + for _, row := range rows { + switch { + case row.EventTime.After(latest): + latest = row.EventTime + latestIDs = []string{row.EventID} + case row.EventTime.Equal(latest): + latestIDs = append(latestIDs, row.EventID) + } + } + + if hasMore { + return eventPageCursor{StartAt: latest, LatestEventSeen: latest, LastEventIDs: latestIDs} + } + + target := latest.Add(-auditLogTrailingLag) + if len(rows) == 0 { + target = now.Add(-auditLogTrailingLag) + } + if target.Before(cursor.StartAt) { + target = cursor.StartAt + } + + var idsAtTarget []string + if target.Equal(latest) { + idsAtTarget = latestIDs + } + + return eventPageCursor{StartAt: target, LatestEventSeen: latest, LastEventIDs: idsAtTarget} +} + +// mapAuditRowToResource maps an audit row to the Baton resource it affects, returning +// ok=false if the action isn't tracked or the ID/workspace can't be resolved. +func mapAuditRowToResource(row auditLogRow, accountId string, workspaceLookup map[int64]string) (*v2.ResourceId, *v2.ResourceId, bool) { + mapping, ok := auditLogActions[row.ActionName] + if !ok { + return nil, nil, false + } + + accountParent := &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: accountId} + + if mapping.resourceType == workspaceResourceType { + deploymentName, found := workspaceLookup[row.WorkspaceID] + if !found { + return nil, nil, false + } + return &v2.ResourceId{ResourceType: workspaceResourceType.Id, Resource: deploymentName}, accountParent, true + } + + parentResourceId := accountParent + if row.WorkspaceID != 0 { + deploymentName, found := workspaceLookup[row.WorkspaceID] + if !found { + return nil, nil, false + } + parentResourceId = &v2.ResourceId{ResourceType: workspaceResourceType.Id, Resource: deploymentName} + } + + nativeId, ok := row.RequestParams[mapping.idParam] + if !ok || nativeId == "" { + return nil, nil, false + } + + if mapping.resourceType == groupResourceType { + return &v2.ResourceId{ResourceType: groupResourceType.Id, Resource: groupResourceId(context.Background(), nativeId, parentResourceId)}, parentResourceId, true + } + + return &v2.ResourceId{ResourceType: mapping.resourceType.Id, Resource: nativeId}, parentResourceId, true +} + +// sqlQueryWorkspace deterministically picks the workspace used to run the audit log query +// and builds the workspace-ID-to-deployment-name lookup used to resolve audit rows. +func sqlQueryWorkspace(workspaces []databricks.Workspace) (string, map[int64]string) { + best := workspaces[0] + lookup := make(map[int64]string, len(workspaces)) + for _, w := range workspaces { + lookup[int64(w.ID)] = w.DeploymentName + if w.DeploymentName < best.DeploymentName { + best = w + } + } + + return best.DeploymentName, lookup +} + +func (f *auditEventFeed) queryAuditLog(ctx context.Context, workspaceId string, cursor eventPageCursor) ([]auditLogRow, error) { + statement := fmt.Sprintf(` + SELECT event_id, event_time, workspace_id, action_name, request_params + FROM system.access.audit + WHERE event_date >= :start_date + AND event_time >= :start_time + AND action_name IN (%s) + ORDER BY event_time ASC + LIMIT %d + `, quotedInClause(auditLogActionNames()), auditLogPageLimit) + + result, err := f.client.ExecuteStatement( + ctx, + workspaceId, + f.sqlWarehouseID, + statement, + databricks.StatementParameter{Name: "start_date", Value: cursor.StartAt.Format("2006-01-02"), Type: "DATE"}, + databricks.StatementParameter{Name: "start_time", Value: cursor.StartAt.Format(time.RFC3339), Type: "TIMESTAMP"}, + ) + if err != nil { + return nil, err + } + + return parseAuditLogRows(result) +} + +func quotedInClause(values []string) string { + quoted := make([]string, len(values)) + for i, v := range values { + quoted[i] = "'" + v + "'" + } + + out := "" + for i, v := range quoted { + if i > 0 { + out += ", " + } + out += v + } + return out +} + +const ( + colEventID = "event_id" + colEventTime = "event_time" + colWorkspaceID = "workspace_id" + colActionName = "action_name" + colRequestParams = "request_params" +) + +func parseAuditLogRows(result *databricks.StatementResult) ([]auditLogRow, error) { + colIndex := make(map[string]int, len(result.Columns)) + for i, name := range result.Columns { + colIndex[name] = i + } + + for _, name := range []string{colEventID, colEventTime, colWorkspaceID, colActionName, colRequestParams} { + if _, ok := colIndex[name]; !ok { + return nil, fmt.Errorf("audit log query result missing column %q", name) + } + } + + rows := make([]auditLogRow, 0, len(result.Rows)) + for _, r := range result.Rows { + 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 nil, 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 nil, 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 nil, fmt.Errorf("failed to parse request_params %q: %w", v, err) + } + } + + rows = append(rows, auditLogRow{ + EventID: r[colIndex[colEventID]], + EventTime: eventTime, + WorkspaceID: workspaceId, + ActionName: r[colIndex[colActionName]], + RequestParams: requestParams, + }) + } + + return rows, 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..987c81c3 --- /dev/null +++ b/pkg/connector/audit_event_feed_test.go @@ -0,0 +1,247 @@ +package connector + +import ( + "context" + "testing" + "time" + + "github.com/conductorone/baton-databricks/pkg/databricks" + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" +) + +func TestEventCursorRoundTrip(t *testing.T) { + want := eventPageCursor{ + StartAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), + LatestEventSeen: time.Date(2026, 1, 1, 1, 0, 0, 0, time.UTC), + LastEventIDs: []string{"a", "b"}, + } + + encoded, err := encodeEventCursor(want) + if err != nil { + t.Fatalf("encodeEventCursor() error = %v", err) + } + + got := decodeEventCursor(encoded) + if !got.StartAt.Equal(want.StartAt) || !got.LatestEventSeen.Equal(want.LatestEventSeen) || len(got.LastEventIDs) != 2 { + t.Errorf("decodeEventCursor() = %+v, want %+v", got, want) + } +} + +func TestDecodeEventCursorSelfHeals(t *testing.T) { + cases := []string{"", "not-base64!!!", "aW52YWxpZC1qc29u"} // last one is base64("invalid-json") + for _, c := range cases { + got := decodeEventCursor(c) + if !got.StartAt.IsZero() { + t.Errorf("decodeEventCursor(%q) = %+v, want zero-value cursor", c, got) + } + } +} + +func TestAdvanceEventCursorFullPageAdvancesWithoutTrailingLag(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)}, + } + + next := advanceEventCursor(cursor, rows, true, startAt.Add(10*time.Minute)) + + wantStart := startAt.Add(2 * time.Minute) + if !next.StartAt.Equal(wantStart) { + t.Errorf("StartAt = %v, want %v (no trailing lag while more pages remain)", next.StartAt, wantStart) + } + if len(next.LastEventIDs) != 1 || next.LastEventIDs[0] != "2" { + t.Errorf("LastEventIDs = %v, want [2]", next.LastEventIDs) + } +} + +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 len(next.LastEventIDs) != 0 { + t.Errorf("LastEventIDs = %v, want empty", next.LastEventIDs) + } +} + +// 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 len(next.LastEventIDs) != 1 || next.LastEventIDs[0] != "1" { + t.Errorf("LastEventIDs = %v, want [1]", next.LastEventIDs) + } +} + +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" + + cases := []struct { + name string + row auditLogRow + wantOK bool + wantResourceType string + wantResource string + wantParentType string + wantParentID string + }{ + { + name: "account-level group create", + row: auditLogRow{ + ActionName: "createGroup", + WorkspaceID: 0, + RequestParams: map[string]string{"targetGroupId": "g-1"}, + }, + wantOK: true, + wantResourceType: groupResourceType.Id, + wantResource: groupResourceId(context.Background(), "g-1", &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: accountId}), + wantParentType: accountResourceType.Id, + wantParentID: accountId, + }, + { + name: "workspace-scoped acl change", + row: auditLogRow{ + ActionName: "changeDatabricksWorkspaceAcl", + WorkspaceID: 123, + }, + wantOK: true, + wantResourceType: workspaceResourceType.Id, + wantResource: "my-workspace", + wantParentType: accountResourceType.Id, + wantParentID: accountId, + }, + { + name: "unknown action is skipped", + row: auditLogRow{ + ActionName: "someUnityCatalogAction", + }, + wantOK: false, + }, + { + name: "unresolvable workspace is skipped", + row: auditLogRow{ + ActionName: "createUser", + WorkspaceID: 999, + RequestParams: map[string]string{"targetUserId": "u-1"}, + }, + wantOK: false, + }, + { + name: "missing id param is skipped", + row: auditLogRow{ + ActionName: "createUser", + WorkspaceID: 0, + }, + wantOK: false, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + resourceId, parentResourceId, ok := mapAuditRowToResource(tc.row, accountId, workspaceLookup) + if ok != tc.wantOK { + t.Fatalf("ok = %v, want %v", ok, tc.wantOK) + } + if !tc.wantOK { + return + } + if resourceId.ResourceType != tc.wantResourceType || resourceId.Resource != tc.wantResource { + t.Errorf("resourceId = %+v, want type=%s id=%s", resourceId, tc.wantResourceType, tc.wantResource) + } + if parentResourceId.ResourceType != tc.wantParentType || parentResourceId.Resource != tc.wantParentID { + t.Errorf("parentResourceId = %+v, want type=%s id=%s", parentResourceId, tc.wantParentType, tc.wantParentID) + } + }) + } +} + +func TestParseAuditLogRowsDedupesNothingAndParsesFields(t *testing.T) { + result := &databricks.StatementResult{ + Columns: []string{"event_id", "event_time", "workspace_id", "action_name", "request_params"}, + Rows: [][]string{ + {"evt-1", "2026-01-01 00:00:00.000", "123", "createGroup", `{"targetGroupId":"g-1"}`}, + {"evt-2", "2026-01-01T00:01:00Z", "0", "createUser", `{"targetUserId":"u-1"}`}, + }, + } + + rows, err := parseAuditLogRows(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].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(result); err == nil { + t.Error("parseAuditLogRows() error = nil, want error for missing required column") + } +} diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 617e60a0..c529bfbd 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -16,7 +16,9 @@ import ( ) type Databricks struct { - client *databricks.Client + client *databricks.Client + enableIncrementalSync bool + sqlWarehouseID string } // ResourceSyncers returns a ResourceSyncerV2 for each resource type that should be synced from the upstream service. @@ -33,6 +35,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.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) { @@ -136,6 +146,24 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err d.client.UpdateAvailability(isAccAPIAvailable, isWSAPIAvailable) + if d.enableIncrementalSync { + if d.sqlWarehouseID == "" { + return nil, fmt.Errorf("databricks-connector: sql-warehouse-id is required when incremental sync is enabled") + } + + if len(workspaces) == 0 { + return nil, fmt.Errorf("databricks-connector: incremental sync requires at least one workspace to query system.access.audit") + } + + queryWorkspaceId, _ := sqlQueryWorkspace(workspaces) + if err := d.client.ValidateAuditLogAccess(ctx, queryWorkspaceId, d.sqlWarehouseID); 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 } @@ -148,6 +176,8 @@ func New( baseURL string, auth databricks.Auth, excludeWorkspaces []string, + enableIncrementalSync bool, + sqlWarehouseID string, ) (*Databricks, error) { httpClient, err := auth.GetClient(ctx) if err != nil { @@ -160,7 +190,9 @@ func New( } return &Databricks{ - client: client, + client: client, + enableIncrementalSync: enableIncrementalSync, + sqlWarehouseID: sqlWarehouseID, }, nil } @@ -179,6 +211,8 @@ func NewConnector(ctx context.Context, cfg *config.Databricks, opts *cli.Connect cfg.BaseUrl, auth, cfg.DatabricksExcludeWorkspaces, + cfg.EnableIncrementalSync, + cfg.SqlWarehouseId, ) if err != nil { l.Warn("error creating connector", zap.Error(err)) diff --git a/pkg/connector/groups.go b/pkg/connector/groups.go index 4091292d..42f2365f 100644 --- a/pkg/connector/groups.go +++ b/pkg/connector/groups.go @@ -260,6 +260,36 @@ 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. +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.NewGroupMembersAttrVars()) + 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 abed3569..9b8a07f0 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" @@ -294,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 62b92868..5f8c412b 100644 --- a/pkg/connector/service-principals.go +++ b/pkg/connector/service-principals.go @@ -214,6 +214,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..4087bcc1 100644 --- a/pkg/connector/users.go +++ b/pkg/connector/users.go @@ -231,6 +231,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 52ed5d92..d0609e67 100644 --- a/pkg/connector/workspaces.go +++ b/pkg/connector/workspaces.go @@ -178,6 +178,21 @@ 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) { + 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) + } + + 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 fcdd1890..81e954d1 100644 --- a/pkg/databricks/client.go +++ b/pkg/databricks/client.go @@ -578,6 +578,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..ba216ccc --- /dev/null +++ b/pkg/databricks/sql.go @@ -0,0 +1,181 @@ +package databricks + +import ( + "context" + "fmt" + "strconv" + "time" + + "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" + "go.uber.org/zap" +) + +const ( + statementsEndpoint = "/api/2.0/sql/statements" + + statementWaitTimeout = "30s" + statementPollInterval = 2 * time.Second +) + +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. +func (c *Client) ExecuteStatement( + ctx context.Context, + workspaceId string, + warehouseId string, + statement string, + params ...StatementParameter, +) (*StatementResult, 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 + if _, err := c.Post(ctx, u, body, &res); err != nil { + return nil, fmt.Errorf("failed to submit statement: %w", err) + } + + res, err := c.pollStatement(ctx, workspaceId, res) + if err != nil { + return nil, err + } + + switch res.Status.State { + case StatementStateSucceeded: + default: + msg := "" + if res.Status.Error != nil { + msg = res.Status.Error.Message + } + return nil, fmt.Errorf("statement %s did not succeed: state=%s message=%s", res.StatementID, res.Status.State, msg) + } + + return c.collectStatementResult(ctx, workspaceId, res) +} + +// pollStatement blocks until the statement reaches a terminal state, for the case of a +// cold warehouse start still running after the initial statementWaitTimeout. +func (c *Client) pollStatement(ctx context.Context, workspaceId string, res statementResponse) (statementResponse, error) { + l := ctxzap.Extract(ctx) + + for res.Status.State == StatementStatePending || res.Status.State == StatementStateRunning { + select { + case <-ctx.Done(): + return res, ctx.Err() + 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 + if _, err := c.Get(ctx, u, &polled); err != nil { + return res, fmt.Errorf("failed to poll statement %s: %w", res.StatementID, err) + } + res = polled + } + + return res, nil +} + +func (c *Client) collectStatementResult(ctx context.Context, workspaceId string, res statementResponse) (*StatementResult, 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...) + + nextChunk := res.Result.NextChunkIndex + for nextChunk != nil { + u := c.workspaceUrl(workspaceId).JoinPath(statementsEndpoint, res.StatementID, "result", "chunks", strconv.Itoa(*nextChunk)) + var chunk statementResultChunk + if _, err := c.Get(ctx, u, &chunk); err != nil { + return nil, 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}, 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 +} From f4352167c8492af4b1e5e0f5b0caf28bb91d8ccf Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Tue, 18 Aug 2026 21:42:25 -0300 Subject: [PATCH 02/15] feat: extend incremental sync to role/permission-assignment changes Add audit log action mappings for account-admin, workspace-access, and SQL-access role changes, plus a coarser fallback for the cluster-create and instance-pool-create entitlements. mapAuditRowToResource now returns multiple affected resources per audit row so a single action can refresh both a principal and the role(s) it holds. Co-Authored-By: Claude Sonnet 5 --- pkg/connector/audit_event_feed.go | 147 +++++++++++++++++-------- pkg/connector/audit_event_feed_test.go | 94 +++++++++------- pkg/connector/roles.go | 17 +-- 3 files changed, 165 insertions(+), 93 deletions(-) diff --git a/pkg/connector/audit_event_feed.go b/pkg/connector/audit_event_feed.go index 24eba054..fd7ae220 100644 --- a/pkg/connector/audit_event_feed.go +++ b/pkg/connector/audit_event_feed.go @@ -20,8 +20,8 @@ import ( const ( auditEventFeedId = "databricks_audit_log" - // Databricks audit log delivery can lag up to 24h, so the first poll looks back that far. - auditLogLookback = 24 * time.Hour + // 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. @@ -30,23 +30,42 @@ const ( auditLogPageLimit = 1000 ) -// auditLogActions maps audit log action_name values to the resource type they affect and -// the request_params key holding the native resource ID. Not yet verified against a live workspace. -var auditLogActions = map[string]struct { +// 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 -}{ - "createGroup": {groupResourceType, "targetGroupId"}, - "addPrincipalToGroup": {groupResourceType, "targetGroupId"}, - "removePrincipalFromGroup": {groupResourceType, "targetGroupId"}, - "deleteGroup": {groupResourceType, "targetGroupId"}, - "createUser": {userResourceType, "targetUserId"}, - "updateUser": {userResourceType, "targetUserId"}, - "deleteUser": {userResourceType, "targetUserId"}, - "createServicePrincipal": {servicePrincipalResourceType, "targetServicePrincipalId"}, - "updateServicePrincipal": {servicePrincipalResourceType, "targetServicePrincipalId"}, - "deleteServicePrincipal": {servicePrincipalResourceType, "targetServicePrincipalId"}, - "changeDatabricksWorkspaceAcl": {workspaceResourceType, ""}, + accountRole string + roleNames []string +} + +// auditLogActions maps audit log action_name values to the resources they affect. +var auditLogActions = map[string]auditActionMapping{ + "createGroup": {resourceType: groupResourceType, idParam: "targetGroupId"}, + "addPrincipalToGroup": {resourceType: groupResourceType, idParam: "targetGroupId"}, + "removePrincipalFromGroup": {resourceType: groupResourceType, idParam: "targetGroupId"}, + "deleteGroup": {resourceType: groupResourceType, idParam: "targetGroupId"}, + "updateGroup": { + resourceType: groupResourceType, idParam: "targetGroupId", + roleNames: []string{ClusterCreateRole, InstancePoolCreateRole}, + }, + "createUser": {resourceType: userResourceType, idParam: "targetUserId"}, + "updateUser": { + resourceType: userResourceType, idParam: "targetUserId", + roleNames: []string{ClusterCreateRole, InstancePoolCreateRole}, + }, + "deleteUser": {resourceType: userResourceType, idParam: "targetUserId"}, + "createServicePrincipal": {resourceType: servicePrincipalResourceType, idParam: "targetServicePrincipalId"}, + "updateServicePrincipal": { + resourceType: servicePrincipalResourceType, idParam: "targetServicePrincipalId", + roleNames: []string{ClusterCreateRole, InstancePoolCreateRole}, + }, + "deleteServicePrincipal": {resourceType: servicePrincipalResourceType, idParam: "targetServicePrincipalId"}, + "changeDatabricksWorkspaceAcl": {resourceType: workspaceResourceType, roleNames: []string{WorkspaceAccessRole}}, + "changeDatabricksSqlAcl": {roleNames: []string{SQLAccessRole}}, + "setAdmin": {resourceType: userResourceType, idParam: "targetUserId", accountRole: AccountAdminRole}, + "removeAdmin": {resourceType: userResourceType, idParam: "targetUserId", accountRole: AccountAdminRole}, } func auditLogActionNames() []string { @@ -172,8 +191,8 @@ func (f *auditEventFeed) ListEvents( continue } - resourceId, parentResourceId, ok := mapAuditRowToResource(row, f.client.GetAccountId(), workspaceLookup) - if !ok { + affected := mapAuditRowToResource(row, f.client.GetAccountId(), 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), @@ -181,16 +200,18 @@ func (f *auditEventFeed) ListEvents( continue } - events = append(events, &v2.Event{ - Id: row.EventID, - OccurredAt: timestamppb.New(row.EventTime), - Event: &v2.Event_ResourceChangeEvent{ - ResourceChangeEvent: &v2.ResourceChangeEvent{ - ResourceId: resourceId, - ParentResourceId: parentResourceId, + 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 @@ -239,43 +260,75 @@ func advanceEventCursor(cursor eventPageCursor, rows []auditLogRow, hasMore bool return eventPageCursor{StartAt: target, LatestEventSeen: latest, LastEventIDs: idsAtTarget} } -// mapAuditRowToResource maps an audit row to the Baton resource it affects, returning -// ok=false if the action isn't tracked or the ID/workspace can't be resolved. -func mapAuditRowToResource(row auditLogRow, accountId string, workspaceLookup map[int64]string) (*v2.ResourceId, *v2.ResourceId, bool) { +// 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 +// (a principal, an account role, and/or workspace roles), skipping anything unresolvable. +func mapAuditRowToResource(row auditLogRow, accountId string, workspaceLookup map[int64]string) []affectedResource { mapping, ok := auditLogActions[row.ActionName] if !ok { - return nil, nil, false + return nil } accountParent := &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: accountId} - if mapping.resourceType == workspaceResourceType { + var workspaceParent *v2.ResourceId + if row.WorkspaceID != 0 { deploymentName, found := workspaceLookup[row.WorkspaceID] if !found { - return nil, nil, false + return nil } - return &v2.ResourceId{ResourceType: workspaceResourceType.Id, Resource: deploymentName}, accountParent, true + workspaceParent = &v2.ResourceId{ResourceType: workspaceResourceType.Id, Resource: deploymentName} } - parentResourceId := accountParent - if row.WorkspaceID != 0 { - deploymentName, found := workspaceLookup[row.WorkspaceID] - if !found { - return nil, nil, false + 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 workspaceParent != nil { + parent = workspaceParent } - parentResourceId = &v2.ResourceId{ResourceType: workspaceResourceType.Id, Resource: deploymentName} + + 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(context.Background(), nativeId, parent) + } + + affected = append(affected, affectedResource{resourceId: resourceId, parentResourceId: parent}) } - nativeId, ok := row.RequestParams[mapping.idParam] - if !ok || nativeId == "" { - return nil, nil, false + if mapping.accountRole != "" { + affected = append(affected, affectedResource{ + resourceId: &v2.ResourceId{ResourceType: roleResourceType.Id, Resource: roleResourceId(mapping.accountRole, accountParent)}, + parentResourceId: accountParent, + }) } - if mapping.resourceType == groupResourceType { - return &v2.ResourceId{ResourceType: groupResourceType.Id, Resource: groupResourceId(context.Background(), nativeId, parentResourceId)}, parentResourceId, true + 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 &v2.ResourceId{ResourceType: mapping.resourceType.Id, Resource: nativeId}, parentResourceId, true + return affected } // sqlQueryWorkspace deterministically picks the workspace used to run the audit log query diff --git a/pkg/connector/audit_event_feed_test.go b/pkg/connector/audit_event_feed_test.go index 987c81c3..09c92110 100644 --- a/pkg/connector/audit_event_feed_test.go +++ b/pkg/connector/audit_event_feed_test.go @@ -132,14 +132,20 @@ 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 - wantOK bool - wantResourceType string - wantResource string - wantParentType string - wantParentID string + name string + row auditLogRow + want []wantResource }{ { name: "account-level group create", @@ -148,30 +154,48 @@ func TestMapAuditRowToResource(t *testing.T) { WorkspaceID: 0, RequestParams: map[string]string{"targetGroupId": "g-1"}, }, - wantOK: true, - wantResourceType: groupResourceType.Id, - wantResource: groupResourceId(context.Background(), "g-1", &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: accountId}), - wantParentType: accountResourceType.Id, - wantParentID: accountId, + want: []wantResource{ + {groupResourceType.Id, groupResourceId(context.Background(), "g-1", accountParent), accountResourceType.Id, accountId}, + }, }, { - name: "workspace-scoped acl change", + name: "workspace-scoped acl change also refreshes the workspace-access role", row: auditLogRow{ ActionName: "changeDatabricksWorkspaceAcl", WorkspaceID: 123, }, - wantOK: true, - wantResourceType: workspaceResourceType.Id, - wantResource: "my-workspace", - wantParentType: accountResourceType.Id, - wantParentID: accountId, + want: []wantResource{ + {workspaceResourceType.Id, "my-workspace", accountResourceType.Id, accountId}, + {roleResourceType.Id, roleResourceId(WorkspaceAccessRole, workspaceParent), workspaceResourceType.Id, "my-workspace"}, + }, }, { - name: "unknown action is skipped", + name: "setAdmin refreshes the user and the account-admin role", + row: auditLogRow{ + ActionName: "setAdmin", + 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 also refreshes workspace entitlement roles", row: auditLogRow{ - ActionName: "someUnityCatalogAction", + ActionName: "updateUser", + 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"}, }, - wantOK: false, + }, + { + name: "unknown action is skipped", + row: auditLogRow{ActionName: "someUnityCatalogAction"}, }, { name: "unresolvable workspace is skipped", @@ -180,32 +204,26 @@ func TestMapAuditRowToResource(t *testing.T) { WorkspaceID: 999, RequestParams: map[string]string{"targetUserId": "u-1"}, }, - wantOK: false, }, { name: "missing id param is skipped", - row: auditLogRow{ - ActionName: "createUser", - WorkspaceID: 0, - }, - wantOK: false, + row: auditLogRow{ActionName: "createUser", WorkspaceID: 0}, }, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { - resourceId, parentResourceId, ok := mapAuditRowToResource(tc.row, accountId, workspaceLookup) - if ok != tc.wantOK { - t.Fatalf("ok = %v, want %v", ok, tc.wantOK) - } - if !tc.wantOK { - return - } - if resourceId.ResourceType != tc.wantResourceType || resourceId.Resource != tc.wantResource { - t.Errorf("resourceId = %+v, want type=%s id=%s", resourceId, tc.wantResourceType, tc.wantResource) + got := mapAuditRowToResource(tc.row, accountId, workspaceLookup) + if len(got) != len(tc.want) { + t.Fatalf("got %d affected resources, want %d: %+v", len(got), len(tc.want), got) } - if parentResourceId.ResourceType != tc.wantParentType || parentResourceId.Resource != tc.wantParentID { - t.Errorf("parentResourceId = %+v, want type=%s id=%s", parentResourceId, tc.wantParentType, tc.wantParentID) + 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) + } } }) } diff --git a/pkg/connector/roles.go b/pkg/connector/roles.go index 9b8a07f0..d13228aa 100644 --- a/pkg/connector/roles.go +++ b/pkg/connector/roles.go @@ -42,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, } - // 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, From 1d9c118f850fa1d4a806e707108a6e7955b3b152 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Thu, 20 Aug 2026 15:11:54 -0300 Subject: [PATCH 03/15] fix: add event_id tiebreaker to audit log pagination ORDER BY event_time ASC alone gives no deterministic ordering among rows sharing an event_time, so paging via a >= start_time filter plus a remembered ID set can stall forever if a single event_time has >= auditLogPageLimit rows. Order by (event_time, event_id) and page with a composite (event_time, event_id) > predicate instead, so the cursor always advances regardless of how many rows share a timestamp. --- pkg/connector/audit_event_feed.go | 60 ++++++++++++-------------- pkg/connector/audit_event_feed_test.go | 20 ++++----- 2 files changed, 37 insertions(+), 43 deletions(-) diff --git a/pkg/connector/audit_event_feed.go b/pkg/connector/audit_event_feed.go index fd7ae220..b24edad3 100644 --- a/pkg/connector/audit_event_feed.go +++ b/pkg/connector/audit_event_feed.go @@ -76,12 +76,14 @@ func auditLogActionNames() []string { return names } -// eventPageCursor is the opaque state persisted between ListEvents calls. StartAt only -// ever advances forward, and LastEventIDs dedupes rows tied exactly on that boundary. +// 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"` - LatestEventSeen time.Time `json:"latest_event_seen"` - LastEventIDs []string `json:"last_event_ids"` + StartAt time.Time `json:"start_at"` + StartAfterEventID string `json:"start_after_event_id"` + LatestEventSeen time.Time `json:"latest_event_seen"` } func encodeEventCursor(c eventPageCursor) (string, error) { @@ -180,17 +182,8 @@ func (f *auditEventFeed) ListEvents( return nil, nil, nil, fmt.Errorf("databricks-connector: failed to query audit log: %w", err) } - seen := make(map[string]struct{}, len(cursor.LastEventIDs)) - for _, id := range cursor.LastEventIDs { - seen[id] = struct{}{} - } - var events []*v2.Event for _, row := range rows { - if _, ok := seen[row.EventID]; ok { - continue - } - affected := mapAuditRowToResource(row, f.client.GetAccountId(), workspaceLookup) if len(affected) == 0 { l.Debug("databricks-connector: skipping audit row with no resource mapping", @@ -225,23 +218,20 @@ func (f *auditEventFeed) ListEvents( return events, &pagination.StreamState{Cursor: encoded, HasMore: hasMore}, nil, nil } -// advanceEventCursor advances only to the last row processed while a page is full, and -// once drained, trails the newest event seen (or wall-clock time if empty) by auditLogTrailingLag. +// 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 - var latestIDs []string - for _, row := range rows { - switch { - case row.EventTime.After(latest): - latest = row.EventTime - latestIDs = []string{row.EventID} - case row.EventTime.Equal(latest): - latestIDs = append(latestIDs, row.EventID) - } + lastEventID := cursor.StartAfterEventID + if len(rows) > 0 { + last := rows[len(rows)-1] + latest = last.EventTime + lastEventID = last.EventID } if hasMore { - return eventPageCursor{StartAt: latest, LatestEventSeen: latest, LastEventIDs: latestIDs} + return eventPageCursor{StartAt: latest, StartAfterEventID: lastEventID, LatestEventSeen: latest} } target := latest.Add(-auditLogTrailingLag) @@ -252,12 +242,12 @@ func advanceEventCursor(cursor eventPageCursor, rows []auditLogRow, hasMore bool target = cursor.StartAt } - var idsAtTarget []string + startAfterEventID := "" if target.Equal(latest) { - idsAtTarget = latestIDs + startAfterEventID = lastEventID } - return eventPageCursor{StartAt: target, LatestEventSeen: latest, LastEventIDs: idsAtTarget} + return eventPageCursor{StartAt: target, StartAfterEventID: startAfterEventID, LatestEventSeen: latest} } // affectedResource is one resource a mapped audit row's action changed. @@ -347,13 +337,16 @@ func sqlQueryWorkspace(workspaces []databricks.Workspace) (string, map[int64]str } func (f *auditEventFeed) queryAuditLog(ctx context.Context, workspaceId string, cursor eventPageCursor) ([]auditLogRow, 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). statement := fmt.Sprintf(` SELECT event_id, event_time, workspace_id, action_name, request_params FROM system.access.audit WHERE event_date >= :start_date - AND event_time >= :start_time + AND (event_time > :start_time OR (event_time = :start_time AND event_id > :start_after_event_id)) AND action_name IN (%s) - ORDER BY event_time ASC + ORDER BY event_time ASC, event_id ASC LIMIT %d `, quotedInClause(auditLogActionNames()), auditLogPageLimit) @@ -362,8 +355,9 @@ func (f *auditEventFeed) queryAuditLog(ctx context.Context, workspaceId string, workspaceId, f.sqlWarehouseID, statement, - databricks.StatementParameter{Name: "start_date", Value: cursor.StartAt.Format("2006-01-02"), Type: "DATE"}, - databricks.StatementParameter{Name: "start_time", Value: cursor.StartAt.Format(time.RFC3339), Type: "TIMESTAMP"}, + databricks.StatementParameter{Name: "start_date", Value: cursor.StartAt.UTC().Format("2006-01-02"), Type: "DATE"}, + databricks.StatementParameter{Name: "start_time", Value: cursor.StartAt.UTC().Format(time.RFC3339), Type: "TIMESTAMP"}, + databricks.StatementParameter{Name: "start_after_event_id", Value: cursor.StartAfterEventID, Type: "STRING"}, ) if err != nil { return nil, err diff --git a/pkg/connector/audit_event_feed_test.go b/pkg/connector/audit_event_feed_test.go index 09c92110..ccdc745a 100644 --- a/pkg/connector/audit_event_feed_test.go +++ b/pkg/connector/audit_event_feed_test.go @@ -11,9 +11,9 @@ import ( func TestEventCursorRoundTrip(t *testing.T) { want := eventPageCursor{ - StartAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), - LatestEventSeen: time.Date(2026, 1, 1, 1, 0, 0, 0, time.UTC), - LastEventIDs: []string{"a", "b"}, + StartAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), + LatestEventSeen: time.Date(2026, 1, 1, 1, 0, 0, 0, time.UTC), + StartAfterEventID: "b", } encoded, err := encodeEventCursor(want) @@ -22,7 +22,7 @@ func TestEventCursorRoundTrip(t *testing.T) { } got := decodeEventCursor(encoded) - if !got.StartAt.Equal(want.StartAt) || !got.LatestEventSeen.Equal(want.LatestEventSeen) || len(got.LastEventIDs) != 2 { + if !got.StartAt.Equal(want.StartAt) || !got.LatestEventSeen.Equal(want.LatestEventSeen) || got.StartAfterEventID != want.StartAfterEventID { t.Errorf("decodeEventCursor() = %+v, want %+v", got, want) } } @@ -52,8 +52,8 @@ func TestAdvanceEventCursorFullPageAdvancesWithoutTrailingLag(t *testing.T) { if !next.StartAt.Equal(wantStart) { t.Errorf("StartAt = %v, want %v (no trailing lag while more pages remain)", next.StartAt, wantStart) } - if len(next.LastEventIDs) != 1 || next.LastEventIDs[0] != "2" { - t.Errorf("LastEventIDs = %v, want [2]", next.LastEventIDs) + if next.StartAfterEventID != "2" { + t.Errorf("StartAfterEventID = %q, want %q", next.StartAfterEventID, "2") } } @@ -73,8 +73,8 @@ func TestAdvanceEventCursorDrainedPageAppliesTrailingLag(t *testing.T) { 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 len(next.LastEventIDs) != 0 { - t.Errorf("LastEventIDs = %v, want empty", next.LastEventIDs) + if next.StartAfterEventID != "" { + t.Errorf("StartAfterEventID = %q, want empty", next.StartAfterEventID) } } @@ -93,8 +93,8 @@ func TestAdvanceEventCursorTieAtFlooredBoundaryIsRemembered(t *testing.T) { if !next.StartAt.Equal(startAt) { t.Errorf("StartAt = %v, want unchanged %v", next.StartAt, startAt) } - if len(next.LastEventIDs) != 1 || next.LastEventIDs[0] != "1" { - t.Errorf("LastEventIDs = %v, want [1]", next.LastEventIDs) + if next.StartAfterEventID != "1" { + t.Errorf("StartAfterEventID = %q, want %q", next.StartAfterEventID, "1") } } From 22c4e7f1fe4b797347ab15b5df27c6b58e11792d Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Thu, 20 Aug 2026 15:43:09 -0300 Subject: [PATCH 04/15] fix: bound sql statement polling and cancel on give-up pollStatement could block for as long as the caller's context allowed if a warehouse got stuck PENDING/RUNNING (cold start, queued, quota), hanging Validate() indefinitely when incremental sync is enabled. Cap polling at statementPollMaxWait and cancel the statement via DELETE when giving up so it stops occupying the warehouse. --- pkg/databricks/sql.go | 37 +++++++++++++++++++++++++++++++++---- 1 file changed, 33 insertions(+), 4 deletions(-) diff --git a/pkg/databricks/sql.go b/pkg/databricks/sql.go index ba216ccc..3d8034e2 100644 --- a/pkg/databricks/sql.go +++ b/pkg/databricks/sql.go @@ -15,6 +15,7 @@ const ( statementWaitTimeout = "30s" statementPollInterval = 2 * time.Second + statementPollMaxWait = 5 * time.Minute ) type StatementState string @@ -124,14 +125,30 @@ func (c *Client) ExecuteStatement( } // pollStatement blocks until the statement reaches a terminal state, for the case of a -// cold warehouse start still running after the initial statementWaitTimeout. +// 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, error) { l := ctxzap.Extract(ctx) + pollCtx, cancel := context.WithTimeout(ctx, statementPollMaxWait) + defer cancel() + for res.Status.State == StatementStatePending || res.Status.State == StatementStateRunning { select { - case <-ctx.Done(): - return res, ctx.Err() + case <-pollCtx.Done(): + if err := ctx.Err(); err != nil { + c.cancelStatement(workspaceId, res.StatementID) + return res, 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, fmt.Errorf("statement %s did not reach a terminal state within %s", res.StatementID, statementPollMaxWait) case <-time.After(statementPollInterval): } @@ -139,7 +156,7 @@ func (c *Client) pollStatement(ctx context.Context, workspaceId string, res stat u := c.workspaceUrl(workspaceId).JoinPath(statementsEndpoint, res.StatementID) var polled statementResponse - if _, err := c.Get(ctx, u, &polled); err != nil { + if _, err := c.Get(pollCtx, u, &polled); err != nil { return res, fmt.Errorf("failed to poll statement %s: %w", res.StatementID, err) } res = polled @@ -148,6 +165,18 @@ func (c *Client) pollStatement(ctx context.Context, workspaceId string, res stat return res, nil } +// cancelStatement best-effort cancels a statement we've given up polling on, using a fresh +// context since ctx/pollCtx may already be done. +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) + if _, err := c.Delete(ctx, u); 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, error) { columns := make([]string, len(res.Manifest.Schema.Columns)) for i, col := range res.Manifest.Schema.Columns { From 5d31b34506aa7922242384899257fbb9b6bd4c91 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Fri, 21 Aug 2026 11:24:53 -0300 Subject: [PATCH 05/15] fix: nil-safe parent access in user/service-principal/role Get The SDK's GetResource passes request.GetParentResourceId() straight through to Get, which is nil whenever C1 has no parent recorded for the resource (e.g. workspace-token auth, where userResource deliberately omits WithParentResourceID). That nil parent was then dereferenced directly in userResource, servicePrincipalResource, and roleResource, panicking on resync after a RESOURCE_CHANGE event. Use the nil-safe GetResourceType()/GetResource() getters instead, matching the pattern groups.go already used. --- pkg/connector/roles.go | 4 ++-- pkg/connector/service-principals.go | 7 ++++--- pkg/connector/users.go | 3 ++- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/pkg/connector/roles.go b/pkg/connector/roles.go index 49eaaf3c..15f36cb5 100644 --- a/pkg/connector/roles.go +++ b/pkg/connector/roles.go @@ -53,8 +53,8 @@ func roleResourceId(role string, parent *v2.ResourceId) string { func roleResource(ctx context.Context, role string, parent *v2.ResourceId) (*v2.Resource, error) { profile := map[string]interface{}{ "role_name": role, - "parent_type": parent.ResourceType, - "parent_id": parent.Resource, + "parent_type": parent.GetResourceType(), + "parent_id": parent.GetResource(), } roleID := roleResourceId(role, parent) diff --git a/pkg/connector/service-principals.go b/pkg/connector/service-principals.go index 58a99b77..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)) } diff --git a/pkg/connector/users.go b/pkg/connector/users.go index 4087bcc1..f8f0d621 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)) } From f6fbc49e26e86aa59023608950e7a9b258e0b593 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Fri, 21 Aug 2026 11:35:31 -0300 Subject: [PATCH 06/15] fix: mirror token-auth branch in workspaceBuilder.Get Get always called GetWorkspace, which hits the Account API via ListWorkspaces. The Account API is unreachable under workspace-token auth, so any targeted sync of a workspace (advertised via CAPABILITY_TARGETED_SYNC) always failed in that mode, even though List already builds a minimalWorkspaceResource from the configured workspace list to avoid the same call. Get now mirrors that branch. --- pkg/connector/workspaces.go | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/pkg/connector/workspaces.go b/pkg/connector/workspaces.go index 8e0bc63e..9a3e8df0 100644 --- a/pkg/connector/workspaces.go +++ b/pkg/connector/workspaces.go @@ -286,6 +286,20 @@ func (w *workspaceBuilder) Grants(ctx context.Context, resource *v2.Resource, _ // 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) From 590e616b00d4d1ae6574bddd81602d8e52190f60 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Fri, 21 Aug 2026 11:50:22 -0300 Subject: [PATCH 07/15] fix: select audit event parent by Account API availability mapAuditRowToResource picked the user/group/service-principal parent from whether the audit row carried a workspace_id, not from how the resource is actually synced. Users/groups/service principals are only ever synced under the account when the Account API is reachable (accountResource declares them as children only then; groupGrantParent already encodes this rule), so a workspace-scoped row in that mode built an ID that was never synced (e.g. workspace//group/x instead of account//group/x), and the real resource never got refreshed. Select the parent from IsAccountAPIAvailable() instead, matching groupGrantParent. Also thread the real ctx through instead of context.Background(), now that groupResourceId's ctx parameter is actually used for something worth passing correctly. --- pkg/connector/audit_event_feed.go | 15 ++++-- pkg/connector/audit_event_feed_test.go | 64 ++++++++++++++++++++++---- 2 files changed, 66 insertions(+), 13 deletions(-) diff --git a/pkg/connector/audit_event_feed.go b/pkg/connector/audit_event_feed.go index b24edad3..996a3f90 100644 --- a/pkg/connector/audit_event_feed.go +++ b/pkg/connector/audit_event_feed.go @@ -184,7 +184,7 @@ func (f *auditEventFeed) ListEvents( var events []*v2.Event for _, row := range rows { - affected := mapAuditRowToResource(row, f.client.GetAccountId(), workspaceLookup) + 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), @@ -258,7 +258,11 @@ type affectedResource struct { // mapAuditRowToResource maps an audit row to every Baton resource its action affects // (a principal, an account role, and/or workspace roles), skipping anything unresolvable. -func mapAuditRowToResource(row auditLogRow, accountId string, workspaceLookup map[int64]string) []affectedResource { +// The principal's parent mirrors how it's actually synced (see groupGrantParent in +// helpers.go): account when the Account API is reachable, the specific workspace +// otherwise — not whichever scope the audit row happened to occur in. Getting this wrong +// produces a resource ID that was never synced, so the real resource never gets refreshed. +func mapAuditRowToResource(ctx context.Context, row auditLogRow, accountId string, accountAPIAvailable bool, workspaceLookup map[int64]string) []affectedResource { mapping, ok := auditLogActions[row.ActionName] if !ok { return nil @@ -285,7 +289,10 @@ func mapAuditRowToResource(row auditLogRow, accountId string, workspaceLookup ma affected = append(affected, affectedResource{resourceId: workspaceParent, parentResourceId: accountParent}) case mapping.resourceType != nil: parent := accountParent - if workspaceParent != nil { + if !accountAPIAvailable { + if workspaceParent == nil { + return nil + } parent = workspaceParent } @@ -296,7 +303,7 @@ func mapAuditRowToResource(row auditLogRow, accountId string, workspaceLookup ma resourceId := &v2.ResourceId{ResourceType: mapping.resourceType.Id, Resource: nativeId} if mapping.resourceType == groupResourceType { - resourceId.Resource = groupResourceId(context.Background(), nativeId, parent) + resourceId.Resource = groupResourceId(ctx, nativeId, parent) } affected = append(affected, affectedResource{resourceId: resourceId, parentResourceId: parent}) diff --git a/pkg/connector/audit_event_feed_test.go b/pkg/connector/audit_event_feed_test.go index ccdc745a..123d0b62 100644 --- a/pkg/connector/audit_event_feed_test.go +++ b/pkg/connector/audit_event_feed_test.go @@ -143,12 +143,14 @@ func TestMapAuditRowToResource(t *testing.T) { } cases := []struct { - name string - row auditLogRow - want []wantResource + name string + row auditLogRow + accountAPIAvailable bool + want []wantResource }{ { - name: "account-level group create", + name: "account-level group create", + accountAPIAvailable: true, row: auditLogRow{ ActionName: "createGroup", WorkspaceID: 0, @@ -159,7 +161,34 @@ func TestMapAuditRowToResource(t *testing.T) { }, }, { - name: "workspace-scoped acl change also refreshes the workspace-access role", + name: "workspace-scoped group change stays account-parented when the Account API is available", + accountAPIAvailable: true, + row: auditLogRow{ + ActionName: "addPrincipalToGroup", + 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", + 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", WorkspaceID: 123, @@ -170,7 +199,8 @@ func TestMapAuditRowToResource(t *testing.T) { }, }, { - name: "setAdmin refreshes the user and the account-admin role", + name: "setAdmin refreshes the user and the account-admin role", + accountAPIAvailable: true, row: auditLogRow{ ActionName: "setAdmin", RequestParams: map[string]string{"targetUserId": "u-1"}, @@ -181,7 +211,22 @@ func TestMapAuditRowToResource(t *testing.T) { }, }, { - name: "updateUser also refreshes workspace entitlement roles", + name: "updateUser stays account-parented when the Account API is available, but workspace roles still refresh", + accountAPIAvailable: true, + row: auditLogRow{ + ActionName: "updateUser", + 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", WorkspaceID: 123, @@ -198,7 +243,8 @@ func TestMapAuditRowToResource(t *testing.T) { row: auditLogRow{ActionName: "someUnityCatalogAction"}, }, { - name: "unresolvable workspace is skipped", + name: "unresolvable workspace is skipped", + accountAPIAvailable: true, row: auditLogRow{ ActionName: "createUser", WorkspaceID: 999, @@ -213,7 +259,7 @@ func TestMapAuditRowToResource(t *testing.T) { for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { - got := mapAuditRowToResource(tc.row, accountId, workspaceLookup) + 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) } From 16f3b89e2448f9184237dd14782816066efb85a4 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Fri, 21 Aug 2026 12:20:25 -0300 Subject: [PATCH 08/15] fix: don't require the Account API for incremental sync under token auth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validate() and ListEvents() both called ListWorkspaces unconditionally to resolve the audit-log query workspace, but the Account API is unreachable under workspace-token auth — so incremental sync always failed Validate() and every ListEvents poll for token-auth customers. Add resolveSQLWorkspaces, which builds minimal workspaces from the configured deployment names under token auth (mirroring workspaceBuilder.List's token-auth branch) instead of calling ListWorkspaces, and use it in both places. auditEventFeed now carries the configured workspace list to support this. Workspace-scoped audit rows can't be resolved to a deployment name under token auth this way (no numeric workspace ID is ever learned), so they're skipped by mapAuditRowToResource rather than mis-resolved — a known limitation, not a regression, since incremental sync couldn't run under token auth at all before this. --- pkg/connector/audit_event_feed.go | 37 +++++++++++++++++++------- pkg/connector/audit_event_feed_test.go | 27 +++++++++++++++++++ pkg/connector/connector.go | 4 +-- 3 files changed, 56 insertions(+), 12 deletions(-) diff --git a/pkg/connector/audit_event_feed.go b/pkg/connector/audit_event_feed.go index 996a3f90..14d679fc 100644 --- a/pkg/connector/audit_event_feed.go +++ b/pkg/connector/audit_event_feed.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "strconv" + "strings" "time" "github.com/conductorone/baton-databricks/pkg/databricks" @@ -124,13 +125,15 @@ type auditLogRow struct { type auditEventFeed struct { client *databricks.Client + workspaces []string enableIncrementalSync bool sqlWarehouseID string } -func newAuditEventFeed(client *databricks.Client, enableIncrementalSync bool, sqlWarehouseID string) *auditEventFeed { +func newAuditEventFeed(client *databricks.Client, workspaces []string, enableIncrementalSync bool, sqlWarehouseID string) *auditEventFeed { return &auditEventFeed{ client: client, + workspaces: workspaces, enableIncrementalSync: enableIncrementalSync, sqlWarehouseID: sqlWarehouseID, } @@ -167,7 +170,7 @@ func (f *auditEventFeed) ListEvents( cursor = eventPageCursor{StartAt: start} } - workspaces, _, err := f.client.ListWorkspaces(ctx) + 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) } @@ -328,6 +331,27 @@ func mapAuditRowToResource(ctx context.Context, row auditLogRow, accountId strin return affected } +// resolveSQLWorkspaces returns the workspaces available to run the audit-log SQL query +// against. The Account API (ListWorkspaces) is unreachable under workspace-token auth, so +// this builds minimal workspaces from the configured deployment names instead of calling +// it, mirroring workspaceBuilder.List's token-auth branch. Those minimal workspaces have no +// numeric ID (token auth never learns one), so workspace-scoped audit rows can't be +// resolved back to a deployment name via sqlQueryWorkspace's lookup and are skipped by +// mapAuditRowToResource — a known limitation of token auth, not a regression, since +// incremental sync couldn't run under token auth at all before this. +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) + return workspaces, err +} + // sqlQueryWorkspace deterministically picks the workspace used to run the audit log query // and builds the workspace-ID-to-deployment-name lookup used to resolve audit rows. func sqlQueryWorkspace(workspaces []databricks.Workspace) (string, map[int64]string) { @@ -379,14 +403,7 @@ func quotedInClause(values []string) string { quoted[i] = "'" + v + "'" } - out := "" - for i, v := range quoted { - if i > 0 { - out += ", " - } - out += v - } - return out + return strings.Join(quoted, ", ") } const ( diff --git a/pkg/connector/audit_event_feed_test.go b/pkg/connector/audit_event_feed_test.go index 123d0b62..6beaf89c 100644 --- a/pkg/connector/audit_event_feed_test.go +++ b/pkg/connector/audit_event_feed_test.go @@ -2,6 +2,7 @@ package connector import ( "context" + "net/http" "testing" "time" @@ -9,6 +10,32 @@ import ( v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" ) +// 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]) + } + } +} + func TestEventCursorRoundTrip(t *testing.T) { want := eventPageCursor{ StartAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index ade2048c..a88c96ce 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -40,7 +40,7 @@ func (d *Databricks) ResourceSyncers(ctx context.Context) []connectorbuilder.Res // gates its behavior inside ListEvents instead. func (d *Databricks) EventFeeds(ctx context.Context) []connectorbuilder.EventFeed { return []connectorbuilder.EventFeed{ - newAuditEventFeed(d.client, d.enableIncrementalSync, d.sqlWarehouseID), + newAuditEventFeed(d.client, d.workspaces, d.enableIncrementalSync, d.sqlWarehouseID), } } @@ -163,7 +163,7 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err return nil, fmt.Errorf("databricks-connector: sql-warehouse-id is required when incremental sync is enabled") } - auditWorkspaces, _, err := d.client.ListWorkspaces(ctx) + 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) } From 4a14d25df9cd9eeef233429305a4922b00d64f69 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Fri, 21 Aug 2026 12:20:33 -0300 Subject: [PATCH 09/15] refactor: simplify quotedInClause and ExecuteStatement's status check quotedInClause hand-rolled a strings.Join; use it directly. The two-case switch in ExecuteStatement (success vs everything else) reads clearer as a plain if. Cosmetic only, per review feedback. --- pkg/databricks/sql.go | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pkg/databricks/sql.go b/pkg/databricks/sql.go index 3d8034e2..91645a4b 100644 --- a/pkg/databricks/sql.go +++ b/pkg/databricks/sql.go @@ -111,9 +111,7 @@ func (c *Client) ExecuteStatement( return nil, err } - switch res.Status.State { - case StatementStateSucceeded: - default: + if res.Status.State != StatementStateSucceeded { msg := "" if res.Status.Error != nil { msg = res.Status.Error.Message From 62b685835cfe42ae60aca36165b166b0cd4a0dea Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Fri, 21 Aug 2026 13:28:21 -0300 Subject: [PATCH 10/15] feat: add sql-warehouse-workspace to pin the audit-log query workspace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit sqlQueryWorkspace picked whichever workspace sorted alphabetically first by deployment name to run the system.access.audit query, with no way to route it to the workspace that actually hosts sql-warehouse-id. SQL warehouses only exist in one workspace, so querying through the wrong one 404s — this broke deterministically for any account with more than one workspace, unless the warehouse happened to live in the alphabetically-smallest one. Add --sql-warehouse-workspace to pin the deployment name explicitly, validated against the resolved workspace list in both Validate() and ListEvents() via a shared resolveQueryWorkspace helper. When unset and more than one workspace is available, log a Debug line naming the arbitrarily-picked workspace and pointing at the new flag, so it's diagnosable without escalating to Warn/Error for what is, until set, a config gap rather than a connector fault. Regenerated config_schema.json and pkg/config/conf.gen.go; updated README's incremental sync section and flag list. --- README.md | 6 +++ config_schema.json | 6 +++ pkg/config/conf.gen.go | 1 + pkg/config/config.go | 10 +++++ pkg/connector/audit_event_feed.go | 56 ++++++++++++++++++++++++- pkg/connector/audit_event_feed_test.go | 57 ++++++++++++++++++++++++++ pkg/connector/connector.go | 11 ++++- 7 files changed, 143 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index ffa99540..4b6d1a02 100644 --- a/README.md +++ b/README.md @@ -120,6 +120,11 @@ Incremental sync 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. +- `--sql-warehouse-workspace` (or `BATON_SQL_WAREHOUSE_WORKSPACE`), the + deployment name of the workspace that hosts that SQL warehouse. SQL + warehouses only exist in one workspace, so this is required whenever more + than one workspace is available; with only one workspace it's inferred + automatically. - 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 @@ -191,6 +196,7 @@ Flags: --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. ($BATON_SQL_WAREHOUSE_ID) + --sql-warehouse-workspace string Deployment name of the workspace that hosts the SQL warehouse (sql-warehouse-id), since SQL warehouses only exist in one workspace. Required when incremental sync is enabled and more than one workspace is available; if omitted with only one workspace available, that workspace is used automatically. ($BATON_SQL_WAREHOUSE_WORKSPACE) --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) diff --git a/config_schema.json b/config_schema.json index 0e436b4f..677753ef 100644 --- a/config_schema.json +++ b/config_schema.json @@ -175,6 +175,12 @@ "displayName": "SQL Warehouse ID", "description": "ID of the Databricks SQL warehouse used to query system.access.audit. Required when incremental sync is enabled.", "stringField": {} + }, + { + "name": "sql-warehouse-workspace", + "displayName": "SQL Warehouse Workspace", + "description": "Deployment name of the workspace that hosts the SQL warehouse (sql-warehouse-id), since SQL warehouses only exist in one workspace. Required when incremental sync is enabled and more than one workspace is available; if omitted with only one workspace available, that workspace is used automatically.", + "stringField": {} } ], "constraints": [ diff --git a/pkg/config/conf.gen.go b/pkg/config/conf.gen.go index e41c2122..994c3fa9 100644 --- a/pkg/config/conf.gen.go +++ b/pkg/config/conf.gen.go @@ -15,6 +15,7 @@ type Databricks struct { DatabricksExcludeWorkspaces []string `mapstructure:"databricks-exclude-workspaces"` EnableIncrementalSync bool `mapstructure:"enable-incremental-sync"` SqlWarehouseId string `mapstructure:"sql-warehouse-id"` + SqlWarehouseWorkspace string `mapstructure:"sql-warehouse-workspace"` } func (c *Databricks) findFieldByTag(tagValue string) (any, bool) { diff --git a/pkg/config/config.go b/pkg/config/config.go index 8c7dfe75..36ec400f 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -81,6 +81,15 @@ var ( field.WithDescription("ID of the Databricks SQL warehouse used to query system.access.audit. Required when incremental sync is enabled."), field.WithDisplayName("SQL Warehouse ID"), ) + SQLWarehouseWorkspaceField = field.StringField( + "sql-warehouse-workspace", + field.WithDescription( + "Deployment name of the workspace that hosts the SQL warehouse (sql-warehouse-id), since SQL warehouses "+ + "only exist in one workspace. Required when incremental sync is enabled and more than one workspace "+ + "is available; if omitted with only one workspace available, that workspace is used automatically.", + ), + field.WithDisplayName("SQL Warehouse Workspace"), + ) configFields = []field.SchemaField{ AccountHostnameField, AccountIdField, @@ -93,6 +102,7 @@ var ( ExcludeWorkspacesField, EnableIncrementalSyncField, SQLWarehouseIDField, + SQLWarehouseWorkspaceField, } ) diff --git a/pkg/connector/audit_event_feed.go b/pkg/connector/audit_event_feed.go index 14d679fc..2575af53 100644 --- a/pkg/connector/audit_event_feed.go +++ b/pkg/connector/audit_event_feed.go @@ -128,14 +128,22 @@ type auditEventFeed struct { workspaces []string enableIncrementalSync bool sqlWarehouseID string + sqlWarehouseWorkspace string } -func newAuditEventFeed(client *databricks.Client, workspaces []string, enableIncrementalSync bool, sqlWarehouseID string) *auditEventFeed { +func newAuditEventFeed( + client *databricks.Client, + workspaces []string, + enableIncrementalSync bool, + sqlWarehouseID string, + sqlWarehouseWorkspace string, +) *auditEventFeed { return &auditEventFeed{ client: client, workspaces: workspaces, enableIncrementalSync: enableIncrementalSync, sqlWarehouseID: sqlWarehouseID, + sqlWarehouseWorkspace: sqlWarehouseWorkspace, } } @@ -178,7 +186,10 @@ func (f *auditEventFeed) ListEvents( return nil, nil, nil, fmt.Errorf("databricks-connector: no workspace available to query system.access.audit") } - queryWorkspaceId, workspaceLookup := sqlQueryWorkspace(workspaces) + queryWorkspaceId, workspaceLookup, err := resolveQueryWorkspace(ctx, workspaces, f.sqlWarehouseWorkspace) + if err != nil { + return nil, nil, nil, err + } rows, err := f.queryAuditLog(ctx, queryWorkspaceId, cursor) if err != nil { @@ -352,6 +363,47 @@ func resolveSQLWorkspaces(ctx context.Context, client *databricks.Client, config return workspaces, err } +// resolveQueryWorkspace picks the workspace whose SQL warehouse runs the audit-log query, +// and builds the workspace-ID-to-deployment-name lookup used to resolve audit rows. SQL +// warehouses only exist in one workspace, so sqlWarehouseWorkspace should be set to pin the +// workspace that actually hosts sql-warehouse-id; querying the wrong workspace's endpoint +// with that ID 404s. When unset, sqlQueryWorkspace's arbitrary (alphabetically-first) pick +// is used instead, which only happens to be correct when there's a single workspace. +func resolveQueryWorkspace(ctx context.Context, workspaces []databricks.Workspace, sqlWarehouseWorkspace string) (string, map[int64]string, error) { + queryWorkspaceId, lookup := sqlQueryWorkspace(workspaces) + + if sqlWarehouseWorkspace != "" { + found := false + for _, w := range workspaces { + if strings.EqualFold(w.DeploymentName, sqlWarehouseWorkspace) { + queryWorkspaceId = w.DeploymentName + found = true + break + } + } + if !found { + return "", nil, fmt.Errorf( + "databricks-connector: sql-warehouse-workspace %q is not one of the available workspaces", + sqlWarehouseWorkspace, + ) + } + return queryWorkspaceId, lookup, nil + } + + if len(workspaces) > 1 { + ctxzap.Extract(ctx).Debug( + "databricks-connector: sql-warehouse-workspace is not set and more than one workspace is available, so "+ + "the workspace used to query system.access.audit was picked arbitrarily; this will fail if "+ + "sql-warehouse-id does not live in the picked workspace — set sql-warehouse-workspace to the "+ + "deployment name of the workspace that actually hosts the SQL warehouse to fix this deterministically", + zap.String("picked_workspace", queryWorkspaceId), + zap.Int("available_workspace_count", len(workspaces)), + ) + } + + return queryWorkspaceId, lookup, nil +} + // sqlQueryWorkspace deterministically picks the workspace used to run the audit log query // and builds the workspace-ID-to-deployment-name lookup used to resolve audit rows. func sqlQueryWorkspace(workspaces []databricks.Workspace) (string, map[int64]string) { diff --git a/pkg/connector/audit_event_feed_test.go b/pkg/connector/audit_event_feed_test.go index 6beaf89c..0e85ba38 100644 --- a/pkg/connector/audit_event_feed_test.go +++ b/pkg/connector/audit_event_feed_test.go @@ -36,6 +36,63 @@ func TestResolveSQLWorkspacesTokenAuth(t *testing.T) { } } +func TestResolveQueryWorkspace(t *testing.T) { + workspaces := []databricks.Workspace{ + {ID: 1, DeploymentName: "dbc-zzz"}, + {ID: 2, DeploymentName: "dbc-aaa"}, + } + + t.Run("pinned workspace is used regardless of alphabetical order", func(t *testing.T) { + got, lookup, err := resolveQueryWorkspace(context.Background(), workspaces, "dbc-zzz") + if err != nil { + t.Fatalf("resolveQueryWorkspace() error = %v", err) + } + if got != "dbc-zzz" { + t.Errorf("queryWorkspaceId = %q, want %q", got, "dbc-zzz") + } + if lookup[1] != "dbc-zzz" || lookup[2] != "dbc-aaa" { + t.Errorf("lookup = %+v, want ids 1 and 2 mapped to their deployment names", lookup) + } + }) + + t.Run("pinned workspace match is case-insensitive", func(t *testing.T) { + got, _, err := resolveQueryWorkspace(context.Background(), workspaces, "DBC-ZZZ") + if err != nil { + t.Fatalf("resolveQueryWorkspace() error = %v", err) + } + if got != "dbc-zzz" { + t.Errorf("queryWorkspaceId = %q, want %q", got, "dbc-zzz") + } + }) + + t.Run("unknown pinned workspace is a clear config error", func(t *testing.T) { + _, _, err := resolveQueryWorkspace(context.Background(), workspaces, "dbc-does-not-exist") + if err == nil { + t.Fatal("resolveQueryWorkspace() error = nil, want error for unresolvable sql-warehouse-workspace") + } + }) + + t.Run("no pin falls back to the arbitrary alphabetical pick", func(t *testing.T) { + got, _, err := resolveQueryWorkspace(context.Background(), workspaces, "") + if err != nil { + t.Fatalf("resolveQueryWorkspace() error = %v", err) + } + if got != "dbc-aaa" { + t.Errorf("queryWorkspaceId = %q, want %q (sqlQueryWorkspace's default)", got, "dbc-aaa") + } + }) + + t.Run("single workspace needs no pin", func(t *testing.T) { + got, _, err := resolveQueryWorkspace(context.Background(), workspaces[:1], "") + if err != nil { + t.Fatalf("resolveQueryWorkspace() error = %v", err) + } + if got != "dbc-zzz" { + t.Errorf("queryWorkspaceId = %q, want %q", got, "dbc-zzz") + } + }) +} + func TestEventCursorRoundTrip(t *testing.T) { want := eventPageCursor{ StartAt: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index a88c96ce..3188af3d 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -20,6 +20,7 @@ type Databricks struct { workspaces []string enableIncrementalSync bool sqlWarehouseID string + sqlWarehouseWorkspace string } // ResourceSyncers returns a ResourceSyncerV2 for each resource type that should be synced from the upstream service. @@ -40,7 +41,7 @@ func (d *Databricks) ResourceSyncers(ctx context.Context) []connectorbuilder.Res // 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), + newAuditEventFeed(d.client, d.workspaces, d.enableIncrementalSync, d.sqlWarehouseID, d.sqlWarehouseWorkspace), } } @@ -171,7 +172,10 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err return nil, fmt.Errorf("databricks-connector: incremental sync requires at least one workspace to query system.access.audit") } - queryWorkspaceId, _ := sqlQueryWorkspace(auditWorkspaces) + queryWorkspaceId, _, err := resolveQueryWorkspace(ctx, auditWorkspaces, d.sqlWarehouseWorkspace) + if err != nil { + return nil, err + } if err := d.client.ValidateAuditLogAccess(ctx, queryWorkspaceId, d.sqlWarehouseID); err != nil { return nil, fmt.Errorf( "databricks-connector: incremental sync is enabled but the connector cannot query system.access.audit via warehouse %s: %w", @@ -195,6 +199,7 @@ func New( workspaces []string, enableIncrementalSync bool, sqlWarehouseID string, + sqlWarehouseWorkspace string, ) (*Databricks, error) { httpClient, err := auth.GetClient(ctx) if err != nil { @@ -211,6 +216,7 @@ func New( workspaces: workspaces, enableIncrementalSync: enableIncrementalSync, sqlWarehouseID: sqlWarehouseID, + sqlWarehouseWorkspace: sqlWarehouseWorkspace, }, nil } @@ -241,6 +247,7 @@ func NewConnector(ctx context.Context, cfg *config.Databricks, opts *cli.Connect cfg.Workspaces, cfg.EnableIncrementalSync, cfg.SqlWarehouseId, + cfg.SqlWarehouseWorkspace, ) if err != nil { return nil, nil, err From 1484ab67525c31d7609bdd12f7bff10cb7ff0fa1 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Wed, 26 Aug 2026 11:06:46 -0300 Subject: [PATCH 11/15] fix: address remaining audit-log review findings - Use RFC3339Nano for start_time so the (event_time, event_id) tiebreaker keeps sub-second precision, fixing duplicate re-emission at page boundaries. - Reject enable-incremental-sync under workspace token auth at Validate(), since no audit row can ever resolve to a synced resource that way and it would otherwise poll forever for nothing. - Add the three incremental-sync config fields to both field groups so they're selectable in the UI; regenerate config_schema.json. - Propagate rate-limit info from the SQL Statement Execution API through ExecuteStatement into ListEvents' annotations. Co-Authored-By: Claude Sonnet 5 --- config_schema.json | 10 ++++-- pkg/config/config.go | 6 +++- pkg/connector/audit_event_feed.go | 44 ++++++++++++----------- pkg/connector/connector.go | 8 +++++ pkg/databricks/sql.go | 58 +++++++++++++++++++++---------- 5 files changed, 83 insertions(+), 43 deletions(-) diff --git a/config_schema.json b/config_schema.json index 677753ef..173d483b 100644 --- a/config_schema.json +++ b/config_schema.json @@ -217,7 +217,10 @@ "account-hostname", "workspaces", "base-url", - "databricks-exclude-workspaces" + "databricks-exclude-workspaces", + "enable-incremental-sync", + "sql-warehouse-id", + "sql-warehouse-workspace" ], "default": true }, @@ -232,7 +235,10 @@ "hostname", "account-hostname", "base-url", - "databricks-exclude-workspaces" + "databricks-exclude-workspaces", + "enable-incremental-sync", + "sql-warehouse-id", + "sql-warehouse-workspace" ] } ] diff --git a/pkg/config/config.go b/pkg/config/config.go index 36ec400f..12786819 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -124,6 +124,7 @@ var Config = field.NewConfiguration( Fields: []field.SchemaField{ AccountIdField, DatabricksClientIdField, DatabricksClientSecretField, HostnameField, AccountHostnameField, WorkspacesField, BaseURLField, ExcludeWorkspacesField, + EnableIncrementalSyncField, SQLWarehouseIDField, SQLWarehouseWorkspaceField, }, Default: true, }, @@ -131,7 +132,10 @@ var Config = field.NewConfiguration( Name: DatabricksWorkspaceTokenGroup, DisplayName: "Workspace token", HelpText: "Authenticate with a personal access token scoped to each workspace.", - Fields: []field.SchemaField{AccountIdField, WorkspacesField, WorkspaceTokensField, HostnameField, AccountHostnameField, BaseURLField, ExcludeWorkspacesField}, + Fields: []field.SchemaField{ + AccountIdField, WorkspacesField, WorkspaceTokensField, HostnameField, AccountHostnameField, BaseURLField, ExcludeWorkspacesField, + EnableIncrementalSyncField, SQLWarehouseIDField, SQLWarehouseWorkspaceField, + }, Default: false, }, }), diff --git a/pkg/connector/audit_event_feed.go b/pkg/connector/audit_event_feed.go index 2575af53..4d4d0497 100644 --- a/pkg/connector/audit_event_feed.go +++ b/pkg/connector/audit_event_feed.go @@ -162,6 +162,7 @@ func (f *auditEventFeed) ListEvents( 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 @@ -191,9 +192,16 @@ func (f *auditEventFeed) ListEvents( return nil, nil, nil, err } - rows, err := f.queryAuditLog(ctx, queryWorkspaceId, cursor) + rows, rateLimit, err := f.queryAuditLog(ctx, queryWorkspaceId, cursor) if err != nil { - return nil, nil, nil, fmt.Errorf("databricks-connector: failed to query audit log: %w", err) + 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 @@ -229,7 +237,7 @@ func (f *auditEventFeed) ListEvents( return nil, nil, nil, fmt.Errorf("databricks-connector: failed to encode event cursor: %w", err) } - return events, &pagination.StreamState{Cursor: encoded, HasMore: hasMore}, nil, nil + return events, &pagination.StreamState{Cursor: encoded, HasMore: hasMore}, annos, nil } // advanceEventCursor advances only to the last row processed (by the well-ordered @@ -343,13 +351,7 @@ func mapAuditRowToResource(ctx context.Context, row auditLogRow, accountId strin } // resolveSQLWorkspaces returns the workspaces available to run the audit-log SQL query -// against. The Account API (ListWorkspaces) is unreachable under workspace-token auth, so -// this builds minimal workspaces from the configured deployment names instead of calling -// it, mirroring workspaceBuilder.List's token-auth branch. Those minimal workspaces have no -// numeric ID (token auth never learns one), so workspace-scoped audit rows can't be -// resolved back to a deployment name via sqlQueryWorkspace's lookup and are skipped by -// mapAuditRowToResource — a known limitation of token auth, not a regression, since -// incremental sync couldn't run under token auth at all before this. +// 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)) @@ -363,12 +365,9 @@ func resolveSQLWorkspaces(ctx context.Context, client *databricks.Client, config return workspaces, err } -// resolveQueryWorkspace picks the workspace whose SQL warehouse runs the audit-log query, -// and builds the workspace-ID-to-deployment-name lookup used to resolve audit rows. SQL -// warehouses only exist in one workspace, so sqlWarehouseWorkspace should be set to pin the -// workspace that actually hosts sql-warehouse-id; querying the wrong workspace's endpoint -// with that ID 404s. When unset, sqlQueryWorkspace's arbitrary (alphabetically-first) pick -// is used instead, which only happens to be correct when there's a single workspace. +// resolveQueryWorkspace picks the workspace to run the audit-log query against, preferring +// sqlWarehouseWorkspace (SQL warehouses only exist in one workspace) and falling back to +// sqlQueryWorkspace's arbitrary pick otherwise. func resolveQueryWorkspace(ctx context.Context, workspaces []databricks.Workspace, sqlWarehouseWorkspace string) (string, map[int64]string, error) { queryWorkspaceId, lookup := sqlQueryWorkspace(workspaces) @@ -419,7 +418,7 @@ func sqlQueryWorkspace(workspaces []databricks.Workspace) (string, map[int64]str return best.DeploymentName, lookup } -func (f *auditEventFeed) queryAuditLog(ctx context.Context, workspaceId string, cursor eventPageCursor) ([]auditLogRow, error) { +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). @@ -433,20 +432,23 @@ func (f *auditEventFeed) queryAuditLog(ctx context.Context, workspaceId string, LIMIT %d `, quotedInClause(auditLogActionNames()), auditLogPageLimit) - result, err := f.client.ExecuteStatement( + 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"}, - databricks.StatementParameter{Name: "start_time", Value: cursor.StartAt.UTC().Format(time.RFC3339), Type: "TIMESTAMP"}, + // 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, err + return nil, rateLimit, err } - return parseAuditLogRows(result) + rows, err := parseAuditLogRows(result) + return rows, rateLimit, err } func quotedInClause(values []string) string { diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 3188af3d..7cc57021 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -164,6 +164,14 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err 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) diff --git a/pkg/databricks/sql.go b/pkg/databricks/sql.go index 91645a4b..fddd54ab 100644 --- a/pkg/databricks/sql.go +++ b/pkg/databricks/sql.go @@ -6,6 +6,7 @@ import ( "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" ) @@ -82,14 +83,15 @@ type StatementResult struct { Rows [][]string } -// ExecuteStatement runs a SQL statement via the Statement Execution API and returns every row. +// 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, error) { +) (*StatementResult, *v2.RateLimitDescription, error) { u := c.workspaceUrl(workspaceId).JoinPath(statementsEndpoint) body := statementRequestBody{ @@ -102,13 +104,17 @@ func (c *Client) ExecuteStatement( } var res statementResponse - if _, err := c.Post(ctx, u, body, &res); err != nil { - return nil, fmt.Errorf("failed to submit statement: %w", err) + rateLimit, err := c.Post(ctx, u, body, &res) + if err != nil { + return nil, rateLimit, fmt.Errorf("failed to submit statement: %w", err) } - res, err := c.pollStatement(ctx, workspaceId, res) + res, polledRateLimit, err := c.pollStatement(ctx, workspaceId, res) + if polledRateLimit != nil { + rateLimit = polledRateLimit + } if err != nil { - return nil, err + return nil, rateLimit, err } if res.Status.State != StatementStateSucceeded { @@ -116,10 +122,14 @@ func (c *Client) ExecuteStatement( if res.Status.Error != nil { msg = res.Status.Error.Message } - return nil, fmt.Errorf("statement %s did not succeed: state=%s message=%s", res.StatementID, res.Status.State, msg) + return nil, rateLimit, fmt.Errorf("statement %s did not succeed: state=%s message=%s", res.StatementID, res.Status.State, msg) } - return c.collectStatementResult(ctx, workspaceId, res) + 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 @@ -127,18 +137,19 @@ func (c *Client) ExecuteStatement( // 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, error) { +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, err + return res, rateLimit, err } l.Warn("sql statement did not reach a terminal state before poll timeout, canceling", zap.String("statement_id", res.StatementID), @@ -146,7 +157,7 @@ func (c *Client) pollStatement(ctx context.Context, workspaceId string, res stat zap.Duration("max_wait", statementPollMaxWait), ) c.cancelStatement(workspaceId, res.StatementID) - return res, fmt.Errorf("statement %s did not reach a terminal state within %s", res.StatementID, statementPollMaxWait) + return res, rateLimit, fmt.Errorf("statement %s did not reach a terminal state within %s", res.StatementID, statementPollMaxWait) case <-time.After(statementPollInterval): } @@ -154,13 +165,17 @@ func (c *Client) pollStatement(ctx context.Context, workspaceId string, res stat u := c.workspaceUrl(workspaceId).JoinPath(statementsEndpoint, res.StatementID) var polled statementResponse - if _, err := c.Get(pollCtx, u, &polled); err != nil { - return res, fmt.Errorf("failed to poll statement %s: %w", res.StatementID, err) + 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, nil + return res, rateLimit, nil } // cancelStatement best-effort cancels a statement we've given up polling on, using a fresh @@ -175,7 +190,7 @@ func (c *Client) cancelStatement(workspaceId, statementId string) { } } -func (c *Client) collectStatementResult(ctx context.Context, workspaceId string, res statementResponse) (*StatementResult, error) { +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 @@ -184,24 +199,29 @@ func (c *Client) collectStatementResult(ctx context.Context, workspaceId string, 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 - if _, err := c.Get(ctx, u, &chunk); err != nil { - return nil, fmt.Errorf("failed to fetch statement result chunk %d: %w", *nextChunk, err) + 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}, nil + 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 { + 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 From a1835ef707859fc2d292b48b70c5d5f41e42d274 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Wed, 26 Aug 2026 11:46:36 -0300 Subject: [PATCH 12/15] fix: log corrupt event cursors instead of resetting silently decodeEventCursor now logs at Debug when a cursor fails to decode, so the watermark reset is observable instead of invisible. Co-Authored-By: Claude Sonnet 5 --- pkg/connector/audit_event_feed.go | 6 ++++-- pkg/connector/audit_event_feed_test.go | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/pkg/connector/audit_event_feed.go b/pkg/connector/audit_event_feed.go index 4d4d0497..922683df 100644 --- a/pkg/connector/audit_event_feed.go +++ b/pkg/connector/audit_event_feed.go @@ -97,18 +97,20 @@ func encodeEventCursor(c eventPageCursor) (string, error) { // decodeEventCursor returns a zero-value cursor when missing or corrupt, so callers // self-heal by resetting to the lookback default. -func decodeEventCursor(s string) eventPageCursor { +func decodeEventCursor(ctx context.Context, s string) eventPageCursor { if s == "" { return eventPageCursor{} } raw, err := base64.StdEncoding.DecodeString(s) if err != nil { + ctxzap.Extract(ctx).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 { + ctxzap.Extract(ctx).Debug("databricks-connector: corrupt event cursor, resetting to lookback default", zap.Error(err)) return eventPageCursor{} } @@ -168,7 +170,7 @@ func (f *auditEventFeed) ListEvents( return nil, &pagination.StreamState{}, nil, nil } - cursor := decodeEventCursor(pToken.Cursor) + cursor := decodeEventCursor(ctx, pToken.Cursor) now := time.Now() if cursor.StartAt.IsZero() { diff --git a/pkg/connector/audit_event_feed_test.go b/pkg/connector/audit_event_feed_test.go index 0e85ba38..02e51110 100644 --- a/pkg/connector/audit_event_feed_test.go +++ b/pkg/connector/audit_event_feed_test.go @@ -105,7 +105,7 @@ func TestEventCursorRoundTrip(t *testing.T) { t.Fatalf("encodeEventCursor() error = %v", err) } - got := decodeEventCursor(encoded) + got := decodeEventCursor(context.Background(), encoded) if !got.StartAt.Equal(want.StartAt) || !got.LatestEventSeen.Equal(want.LatestEventSeen) || got.StartAfterEventID != want.StartAfterEventID { t.Errorf("decodeEventCursor() = %+v, want %+v", got, want) } @@ -114,7 +114,7 @@ func TestEventCursorRoundTrip(t *testing.T) { func TestDecodeEventCursorSelfHeals(t *testing.T) { cases := []string{"", "not-base64!!!", "aW52YWxpZC1qc29u"} // last one is base64("invalid-json") for _, c := range cases { - got := decodeEventCursor(c) + got := decodeEventCursor(context.Background(), c) if !got.StartAt.IsZero() { t.Errorf("decodeEventCursor(%q) = %+v, want zero-value cursor", c, got) } From ce795ffea4ed15274478e7d07a4faf09045622c9 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Wed, 26 Aug 2026 12:24:50 -0300 Subject: [PATCH 13/15] chore: promote Debug logs to Warn for visibility --- pkg/connector/audit_event_feed.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pkg/connector/audit_event_feed.go b/pkg/connector/audit_event_feed.go index 922683df..39241157 100644 --- a/pkg/connector/audit_event_feed.go +++ b/pkg/connector/audit_event_feed.go @@ -104,13 +104,13 @@ func decodeEventCursor(ctx context.Context, s string) eventPageCursor { raw, err := base64.StdEncoding.DecodeString(s) if err != nil { - ctxzap.Extract(ctx).Debug("databricks-connector: corrupt event cursor, resetting to lookback default", zap.Error(err)) + ctxzap.Extract(ctx).Warn("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 { - ctxzap.Extract(ctx).Debug("databricks-connector: corrupt event cursor, resetting to lookback default", zap.Error(err)) + ctxzap.Extract(ctx).Warn("databricks-connector: corrupt event cursor, resetting to lookback default", zap.Error(err)) return eventPageCursor{} } From a81b45fe945fd60401eefa5aa2e1b6c43c18b3ff Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Thu, 27 Aug 2026 10:49:12 -0300 Subject: [PATCH 14/15] chore: retrigger GitHub jobs From f8ebc73803c90738492377685d5974dac4e84542 Mon Sep 17 00:00:00 2001 From: Javier David Carnelli Date: Thu, 27 Aug 2026 11:34:26 -0300 Subject: [PATCH 15/15] fix: drop incremental-sync fields from the workspace-token config group Validate() rejects enable-incremental-sync under token auth, so offering those fields in the workspace-token group let the UI present an option that can never succeed. Remove them from that group (they stay in oauth2) and regenerate config_schema.json; note the OAuth-only requirement in the README. Co-Authored-By: Claude Sonnet 5 --- README.md | 6 +++++- config_schema.json | 5 +---- pkg/config/config.go | 5 +++-- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 4b6d1a02..d4da30aa 100644 --- a/README.md +++ b/README.md @@ -115,7 +115,11 @@ between full syncs to pick up access changes early, by setting still run as the correctness backstop; incremental sync does not detect deletions, which are only caught by the next full sync. -Incremental sync requires: +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 diff --git a/config_schema.json b/config_schema.json index 173d483b..786aecef 100644 --- a/config_schema.json +++ b/config_schema.json @@ -235,10 +235,7 @@ "hostname", "account-hostname", "base-url", - "databricks-exclude-workspaces", - "enable-incremental-sync", - "sql-warehouse-id", - "sql-warehouse-workspace" + "databricks-exclude-workspaces" ] } ] diff --git a/pkg/config/config.go b/pkg/config/config.go index 12786819..f8e22aab 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -132,11 +132,12 @@ var Config = field.NewConfiguration( Name: DatabricksWorkspaceTokenGroup, DisplayName: "Workspace token", HelpText: "Authenticate with a personal access token scoped to each workspace.", + // 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, BaseURLField, ExcludeWorkspacesField, - EnableIncrementalSyncField, SQLWarehouseIDField, SQLWarehouseWorkspaceField, }, - Default: false, + Default: false, }, }), )