Add Unity Catalog resource sync (catalogs, schemas, tables)

baton_capabilities.json

@@ -12,6 +12,17 @@
       ],
       "permissions": {}
     },
+    {
+      "resourceType": {
+        "id": "catalog",
+        "displayName": "Catalog"
+      },
+      "capabilities": [
+        "CAPABILITY_SYNC",
+        "CAPABILITY_PROVISION"
+      ],
+      "permissions": {}
+    },
     {
       "resourceType": {
         "id": "group",
@@ -40,6 +51,17 @@
       ],
       "permissions": {}
     },
+    {
+      "resourceType": {
+        "id": "schema",
+        "displayName": "Schema"
+      },
+      "capabilities": [
+        "CAPABILITY_SYNC",
+        "CAPABILITY_PROVISION"
+      ],
+      "permissions": {}
+    },
     {
       "resourceType": {
         "id": "service_principal",
@@ -54,6 +76,17 @@
       ],
       "permissions": {}
     },
+    {
+      "resourceType": {
+        "id": "table",
+        "displayName": "Table"
+      },
+      "capabilities": [
+        "CAPABILITY_SYNC",
+        "CAPABILITY_PROVISION"
+      ],
+      "permissions": {}
+    },
     {
       "resourceType": {
         "id": "user",

config_schema.json

@@ -163,6 +163,18 @@
       "displayName": "Exclude Workspaces",
       "description": "Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces.",
       "stringSliceField": {}
+    },
+    {
+      "name": "sync-unity-catalog",
+      "displayName": "Sync Unity Catalog",
+      "description": "Sync Unity Catalog catalogs and schemas (and their permission grants) for access reviews.",
+      "boolField": {}
+    },
+    {
+      "name": "sync-unity-catalog-tables",
+      "displayName": "Sync Unity Catalog Tables",
+      "description": "Also sync Unity Catalog tables and their grants. Requires sync-unity-catalog. May be very large in big metastores.",
+      "boolField": {}
     }
   ],
   "constraints": [
@@ -181,6 +193,15 @@
       "secondaryFieldNames": [
         "workspaces"
       ]
+    },
+    {
+      "kind": "CONSTRAINT_KIND_DEPENDENT_ON",
+      "fieldNames": [
+        "sync-unity-catalog-tables"
+      ],
+      "secondaryFieldNames": [
+        "sync-unity-catalog"
+      ]
     }
   ],
   "displayName": "Databricks",
@@ -198,7 +219,9 @@
         "hostname",
         "account-hostname",
         "workspaces",
-        "databricks-exclude-workspaces"
+        "databricks-exclude-workspaces",
+        "sync-unity-catalog",
+        "sync-unity-catalog-tables"
       ],
       "default": true
     }

docs/connector.mdx

@@ -11,9 +11,12 @@ sidebarTitle: "Databricks"
 | Resource | Sync | Provision |
 | :--- | :--- | :--- |
 | Accounts | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> |
+| Catalogs | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> |
 | Groups | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> |
 | Roles | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> |
+| Schemas | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> |
 | Service principals | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> |
+| Tables | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> |
 | Users | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> |
 | Workspaces | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> | <Icon icon="square-check" iconType="solid"  color="#c937ae"/> |
 
@@ -219,6 +222,14 @@ stringData:
   # Mutually exclusive with BATON_WORKSPACES — set one or the other, never both.
   BATON_DATABRICKS_EXCLUDE_WORKSPACES: <workspace-a,workspace-b>
 
+  # Optional: sync Unity Catalog catalogs and schemas (and their grants) for
+  # data access reviews. Off by default.
+  # BATON_SYNC_UNITY_CATALOG: true
+
+  # Optional: also sync Unity Catalog tables and their grants. Requires
+  # BATON_SYNC_UNITY_CATALOG. Can be very large in big metastores.
+  # BATON_SYNC_UNITY_CATALOG_TABLES: true
+
   # Optional: include if you want C1 to provision access using this connector
   BATON_PROVISIONING: true
 ```

pkg/config/config.go

@@ -70,6 +70,18 @@ var (
 		field.WithDescription("Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces."),
 		field.WithDisplayName("Exclude Workspaces"),
 	)
+	SyncUnityCatalogField = field.BoolField(
+		"sync-unity-catalog",
+		field.WithDescription("Sync Unity Catalog catalogs and schemas (and their permission grants) for access reviews."),
+		field.WithDefaultValue(false),
+		field.WithDisplayName("Sync Unity Catalog"),
+	)
+	SyncUnityCatalogTablesField = field.BoolField(
+		"sync-unity-catalog-tables",
+		field.WithDescription("Also sync Unity Catalog tables and their grants. Requires sync-unity-catalog. May be very large in big metastores."),
+		field.WithDefaultValue(false),
+		field.WithDisplayName("Sync Unity Catalog Tables"),
+	)
 	configFields = []field.SchemaField{
 		AccountHostnameField,
 		AccountIdField,
@@ -80,6 +92,8 @@ var (
 		WorkspaceTokensField,
 		BaseURLField,
 		ExcludeWorkspacesField,
+		SyncUnityCatalogField,
+		SyncUnityCatalogTablesField,
 	}
 )
 
@@ -92,6 +106,7 @@ var Config = field.NewConfiguration(
 	field.WithConstraints(
 		field.FieldsMutuallyExclusive(WorkspacesField, ExcludeWorkspacesField),
 		field.FieldsDependentOn([]field.SchemaField{WorkspaceTokensField}, []field.SchemaField{WorkspacesField}),
+		field.FieldsDependentOn([]field.SchemaField{SyncUnityCatalogTablesField}, []field.SchemaField{SyncUnityCatalogField}),
 	),
 	field.WithFieldGroups([]field.SchemaFieldGroup{
 		{
@@ -101,6 +116,7 @@ var Config = field.NewConfiguration(
 			Fields: []field.SchemaField{
 				AccountIdField, DatabricksClientIdField, DatabricksClientSecretField,
 				HostnameField, AccountHostnameField, WorkspacesField, ExcludeWorkspacesField,
+				SyncUnityCatalogField, SyncUnityCatalogTablesField,
 			},
 			Default: true,
 		},

pkg/connector/connector.go

@@ -16,19 +16,25 @@ import (
 )
 
 type Databricks struct {
-	client     *databricks.Client
-	workspaces []string
+	client                 *databricks.Client
+	workspaces             []string
+	syncUnityCatalog       bool
+	syncUnityCatalogTables bool
 }
 
 // ResourceSyncers returns a ResourceSyncerV2 for each resource type that should be synced from the upstream service.
 func (d *Databricks) ResourceSyncers(ctx context.Context) []connectorbuilder.ResourceSyncerV2 {
+	resolver := newUCResolver(d.client)
 	syncers := []connectorbuilder.ResourceSyncerV2{
 		newAccountBuilder(d.client),
 		newGroupBuilder(d.client),
 		newServicePrincipalBuilder(d.client),
 		newUserBuilder(d.client),
-		newWorkspaceBuilder(d.client, d.workspaces),
+		newWorkspaceBuilder(d.client, d.workspaces, d.syncUnityCatalog),
 		newRoleBuilder(d.client),
+		newCatalogBuilder(resolver),
+		newSchemaBuilder(resolver, d.syncUnityCatalogTables),
+		newTableBuilder(resolver),
 	}
 
 	return syncers
@@ -174,6 +180,8 @@ func New(
 	auth databricks.Auth,
 	excludeWorkspaces []string,
 	workspaces []string,
+	syncUnityCatalog bool,
+	syncUnityCatalogTables bool,
 ) (*Databricks, error) {
 	httpClient, err := auth.GetClient(ctx)
 	if err != nil {
@@ -186,8 +194,10 @@ func New(
 	}
 
 	return &Databricks{
-		client:     client,
-		workspaces: workspaces,
+		client:                 client,
+		workspaces:             workspaces,
+		syncUnityCatalog:       syncUnityCatalog,
+		syncUnityCatalogTables: syncUnityCatalogTables,
 	}, nil
 }
 
@@ -216,6 +226,8 @@ func NewConnector(ctx context.Context, cfg *config.Databricks, opts *cli.Connect
 		auth,
 		cfg.DatabricksExcludeWorkspaces,
 		cfg.Workspaces,
+		cfg.SyncUnityCatalog,
+		cfg.SyncUnityCatalogTables,
 	)
 	if err != nil {
 		return nil, nil, err

pkg/connector/pagination.go

@@ -58,3 +58,23 @@ func prepareNextToken(page uint, pageTotal int, total uint) string {
 
 	return token
 }
+
+// parseCursorToken initializes a pagination Bag for cursor-based (opaque
+// string token) APIs such as Unity Catalog. Unlike parsePageToken it returns
+// the current cursor string directly instead of a numeric offset. Advance with
+// bag.NextToken(next) where next is the API's next_page_token ("" ends).
+func parseCursorToken(i string, resourceID *v2.ResourceId) (*pagination.Bag, string, error) {
+	b := &pagination.Bag{}
+	if err := b.Unmarshal(i); err != nil {
+		return nil, "", err
+	}
+
+	if b.Current() == nil {
+		b.Push(pagination.PageState{
+			ResourceTypeID: resourceID.ResourceType,
+			ResourceID:     resourceID.Resource,
+		})
+	}
+
+	return b, b.PageToken(), nil
+}

pkg/connector/resource_types.go

@@ -46,4 +46,22 @@ var (
 		Id:          "account",
 		DisplayName: "Account",
 	}
+
+	// The catalog resource type is a Unity Catalog top-level data container.
+	catalogResourceType = &v2.ResourceType{
+		Id:          "catalog",
+		DisplayName: "Catalog",
+	}
+
+	// The schema resource type is a Unity Catalog schema within a catalog.
+	schemaResourceType = &v2.ResourceType{
+		Id:          "schema",
+		DisplayName: "Schema",
+	}
+
+	// The table resource type is a Unity Catalog table or view within a schema.
+	tableResourceType = &v2.ResourceType{
+		Id:          "table",
+		DisplayName: "Table",
+	}
 )