Jeremy/fix/snowflake sync improvements

README.md

@@ -26,6 +26,7 @@ flags or as environment variables via the following variable names:
 | `BATON_PRIVATE_KEY_PATH`    | `--private-key-path`      | Path to private key                              |
 | `BATON_PRIVATE_KEY`         | `--private-key`           | Raw private key value                            |
 | `BATON_EXCLUDED_DATABASES`  | `--excluded-databases`    | Database names to skip during sync (repeatable)  |
+| `BATON_SYNC_TABLES`         | `--sync-tables`           | Set to `false` to skip table sync (default: `true`) |
 
 # Getting Started
 
@@ -172,6 +173,20 @@ baton-snowflake \
 BATON_EXCLUDED_DATABASES="MY_INTERNAL_DB,ANOTHER_DB" baton-snowflake
 ```
 
+### Skipping Table Sync
+
+Use `--sync-tables=false` (or `BATON_SYNC_TABLES=false`) to skip syncing tables entirely. This significantly reduces sync size and duration for accounts with large numbers of tables. Tables are synced by default.
+
+**CLI flag:**
+```bash
+baton-snowflake --sync-tables=false
+```
+
+**Environment variable:**
+```bash
+BATON_SYNC_TABLES=false baton-snowflake
+```
+
 ## brew
 
 ```
@@ -262,6 +277,7 @@ Flags:
 -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)
 --sync-secrets                Enable synchronization of Snowflake secrets. ($BATON_SYNC_SECRETS)
+--sync-tables                 Enable synchronization of Snowflake tables. Set to false to skip. ($BATON_SYNC_TABLES) (default true)
 --ticketing                   This must be set to enable ticketing support ($BATON_TICKETING)
 --user-identifier string      required: User Identifier. ($BATON_USER_IDENTIFIER)
 -v, --version                 version for baton-snowflake

docs/connector.mdx

@@ -135,6 +135,9 @@ In the **RSA Private Key (PEM Format)** field, upload the private key file.
 **Optional.** In the **Excluded Databases** field, enter the names of any Snowflake databases you want to skip during sync. You can add multiple names. Matching is case-insensitive. Excluded databases and all their tables are omitted from every sync.
 </Step>
 <Step>
+**Optional.** Disable **Sync Tables** to skip syncing Snowflake tables entirely. This significantly reduces sync size and duration for accounts with large numbers of tables. Tables are synced by default.
+</Step>
+<Step>
 Click **Save**.
 </Step>
 <Step>
@@ -216,6 +219,9 @@ stringData:
 
   # Optional: comma-separated list of database names to exclude from sync (case-insensitive)
   # BATON_EXCLUDED_DATABASES: "MY_DB,ANOTHER_DB"
+
+  # Optional: set to false to skip syncing tables (reduces sync size significantly)
+  # BATON_SYNC_TABLES: false
 ```
 
 See the connector's README or run `--help` to see all available configuration flags and environment variables.

pkg/config/config.go

@@ -50,6 +50,12 @@ var (
 		field.WithDisplayName("Excluded Databases"),
 		field.WithDescription("Database names to exclude from sync (case-insensitive). Can be specified multiple times. When set, matching databases and all their tables are skipped entirely."),
 	)
+	SyncTables = field.BoolField(
+		"sync-tables",
+		field.WithDisplayName("Sync Tables"),
+		field.WithDescription("Enable synchronization of Snowflake tables. Set to false to skip table syncing and significantly reduce sync size. Defaults to true."),
+		field.WithDefaultValue(true),
+	)
 
 	fieldRelationships = []field.SchemaFieldRelationship{
 		field.FieldsMutuallyExclusive(
@@ -69,6 +75,7 @@ var (
 		PrivateKeyPathField,
 		UserIdentifierField,
 		SyncSecrets,
+		SyncTables,
 		ExcludedDatabases,
 	}
 

pkg/connector/connector.go

@@ -18,6 +18,7 @@ import (
 type Connector struct {
 	Client            *snowflake.Client
 	syncSecrets       bool
+	syncTables        bool
 	excludedDatabases []string
 }
 
@@ -26,16 +27,12 @@ func (d *Connector) ResourceSyncers(ctx context.Context) []connectorbuilder.Reso
 	builders := []connectorbuilder.ResourceSyncerV2{
 		newUserBuilder(d.Client, d.syncSecrets),
 		newAccountRoleBuilder(d.Client),
-		newDatabaseBuilder(d.Client, d.syncSecrets, d.excludedDatabases),
-		newTableBuilder(d.Client),
+		newDatabaseBuilder(d.Client, d.syncSecrets, d.syncTables, d.excludedDatabases),
+		newTableBuilder(d.Client, d.syncTables),
 	}
 
 	if d.syncSecrets {
-		builders = append(
-			builders,
-			newSecretBuilder(d.Client),
-			newRsaBuilder(d.Client),
-		)
+		builders = append(builders, newSecretBuilder(d.Client), newRsaBuilder(d.Client))
 	}
 
 	return builders
@@ -239,6 +236,7 @@ func New(ctx context.Context, cfg *config.Snowflake, _ *cli.ConnectorOpts) (conn
 	return &Connector{
 		Client:            client,
 		syncSecrets:       cfg.SyncSecrets,
+		syncTables:        cfg.SyncTables,
 		excludedDatabases: cfg.ExcludedDatabases,
 	}, nil, nil
 }

pkg/connector/databases.go

@@ -19,14 +19,15 @@ type databaseBuilder struct {
 	resourceType      *v2.ResourceType
 	client            *snowflake.Client
 	syncSecrets       bool
+	syncTables        bool
 	excludedDatabases map[string]struct{} // uppercase-normalised names to exclude
 }
 
 func (o *databaseBuilder) ResourceType(ctx context.Context) *v2.ResourceType {
 	return databaseResourceType
 }
 
-func databaseResource(database *snowflake.Database, syncSecrets bool) (*v2.Resource, error) {
+func databaseResource(database *snowflake.Database, syncSecrets bool, syncTables bool) (*v2.Resource, error) {
 	profile := map[string]interface{}{
 		profileKeyName:        database.Name,
 		"kind":                database.Kind,
@@ -38,7 +39,10 @@ func databaseResource(database *snowflake.Database, syncSecrets bool) (*v2.Resou
 		rs.WithAppProfile(profile),
 	}
 
-	opts := []rs.ResourceOption{rs.WithAnnotation(&v2.ChildResourceType{ResourceTypeId: tableResourceType.Id})}
+	var opts []rs.ResourceOption
+	if syncTables {
+		opts = append(opts, rs.WithAnnotation(&v2.ChildResourceType{ResourceTypeId: tableResourceType.Id}))
+	}
 	if syncSecrets {
 		opts = append(opts, rs.WithAnnotation(&v2.ChildResourceType{ResourceTypeId: secretResourceType.Id}))
 	}
@@ -68,12 +72,19 @@ func (o *databaseBuilder) List(ctx context.Context, parentResourceID *v2.Resourc
 		return nil, nil, wrapError(err, "failed to list databases")
 	}
 
+	// Seed database cache for both databaseBuilder.Grants and tableBuilder.isDBSharedOrSystem
+	// (via GetDatabase), both of which run regardless of syncTables.
+	// Do not guard this with syncTables.
+	if err := o.client.CacheDatabases(ctx, opts.Session, databases); err != nil {
+		return nil, nil, wrapError(err, "failed to seed database cache")
+	}
+
 	var resources []*v2.Resource
 	for _, database := range databases {
 		if _, excluded := o.excludedDatabases[strings.ToUpper(database.Name)]; excluded {
 			continue
 		}
-		resource, err := databaseResource(&database, o.syncSecrets) // #nosec G601
+		resource, err := databaseResource(&database, o.syncSecrets, o.syncTables) // #nosec G601
 		if err != nil {
 			return nil, nil, wrapError(err, "failed to create database resource")
 		}
@@ -108,7 +119,8 @@ func (o *databaseBuilder) Entitlements(_ context.Context, resource *v2.Resource,
 }
 
 func (o *databaseBuilder) Grants(ctx context.Context, resource *v2.Resource, opts rs.SyncOpAttrs) ([]*v2.Grant, *rs.SyncOpResults, error) {
-	database, _, err := o.client.GetDatabase(ctx, resource.Id.Resource)
+	// Uses the database cache seeded during List to avoid a redundant Snowflake query.
+	database, _, err := o.client.GetDatabase(ctx, opts.Session, resource.Id.Resource)
 	if err != nil {
 		return nil, nil, wrapError(err, "failed to get database")
 	}
@@ -138,7 +150,7 @@ func (o *databaseBuilder) Grants(ctx context.Context, resource *v2.Resource, opt
 	return grants, nil, nil
 }
 
-func newDatabaseBuilder(client *snowflake.Client, syncSecrets bool, excludedDatabases []string) *databaseBuilder {
+func newDatabaseBuilder(client *snowflake.Client, syncSecrets bool, syncTables bool, excludedDatabases []string) *databaseBuilder {
 	excluded := make(map[string]struct{}, len(excludedDatabases))
 	for _, name := range excludedDatabases {
 		excluded[strings.ToUpper(name)] = struct{}{}
@@ -147,6 +159,7 @@ func newDatabaseBuilder(client *snowflake.Client, syncSecrets bool, excludedData
 		resourceType:      databaseResourceType,
 		client:            client,
 		syncSecrets:       syncSecrets,
+		syncTables:        syncTables,
 		excludedDatabases: excluded,
 	}
 }

pkg/connector/tables.go

@@ -2,14 +2,19 @@ package connector
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strings"
 
+	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
+	"go.uber.org/zap"
+
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
 	"github.com/conductorone/baton-sdk/pkg/pagination"
 	ent "github.com/conductorone/baton-sdk/pkg/types/entitlement"
 	"github.com/conductorone/baton-sdk/pkg/types/grant"
 	rs "github.com/conductorone/baton-sdk/pkg/types/resource"
+	"github.com/conductorone/baton-sdk/pkg/types/sessions"
 	"github.com/conductorone/baton-snowflake/pkg/snowflake"
 )
 
@@ -59,7 +64,7 @@ func getObjectKind(resource *v2.Resource) string {
 	return defaultObjectKind
 }
 
-func (o *tableBuilder) isDBSharedOrSystem(ctx context.Context, resource *v2.Resource, databaseName string) (bool, error) {
+func (o *tableBuilder) isDBSharedOrSystem(ctx context.Context, ss sessions.SessionStore, resource *v2.Resource, databaseName string) (bool, error) {
 	if v := getTableProfileField(resource, "database_is_shared_system"); v != nil {
 		switch val := v.(type) {
 		case bool:
@@ -70,7 +75,7 @@ func (o *tableBuilder) isDBSharedOrSystem(ctx context.Context, resource *v2.Reso
 			return val == "true" || val == "1", nil
 		}
 	}
-	db, statusCode, err := o.client.GetDatabase(ctx, databaseName)
+	db, statusCode, err := o.client.GetDatabase(ctx, ss, databaseName)
 	if snowflake.IsUnprocessableEntity(statusCode, err) {
 		return true, nil
 	}
@@ -84,7 +89,8 @@ func (o *tableBuilder) isDBSharedOrSystem(ctx context.Context, resource *v2.Reso
 }
 
 type tableBuilder struct {
-	client *snowflake.Client
+	client     *snowflake.Client
+	syncTables bool
 }
 
 func (o *tableBuilder) ResourceType(ctx context.Context) *v2.ResourceType {
@@ -129,6 +135,9 @@ func (o *tableBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId
 	if parentResourceID == nil {
 		return nil, &rs.SyncOpResults{}, nil
 	}
+	if !o.syncTables {
+		return nil, &rs.SyncOpResults{}, nil
+	}
 
 	if parentResourceID.ResourceType != databaseResourceType.Id {
 		return nil, nil, wrapError(fmt.Errorf("invalid parent resource type: %s", parentResourceID.ResourceType), "invalid parent resource type")
@@ -149,7 +158,7 @@ func (o *tableBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId
 	// Encoding isSharedOrSystemDB in ResourceTypeID avoids re-querying the
 	// database on every subsequent page.
 	if bag.Current() == nil {
-		parentDB, statusCode, err := o.client.GetDatabase(ctx, databaseName)
+		parentDB, statusCode, err := o.client.GetDatabase(ctx, opts.Session, databaseName)
 		if err != nil && !snowflake.IsUnprocessableEntity(statusCode, err) {
 			return nil, nil, wrapError(err, "failed to get parent database")
 		}
@@ -260,13 +269,16 @@ func grantsContainPrincipal(grants []*v2.Grant, principalID *v2.ResourceId, enti
 }
 
 func (o *tableBuilder) Entitlements(ctx context.Context, resource *v2.Resource, opts rs.SyncOpAttrs) ([]*v2.Entitlement, *rs.SyncOpResults, error) {
+	// No syncTables guard here: the SDK calls Entitlements from the C1Z store (persistent
+	// across syncs), not just from the current List return. Skipping would silently leave
+	// stale grants if syncTables was enabled on a previous run.
 	databaseName, schemaName, tableName, err := parseTableResourceID(resource)
 	if err != nil {
 		return nil, nil, err
 	}
 	var rv []*v2.Entitlement
 
-	isSharedOrSystem, err := o.isDBSharedOrSystem(ctx, resource, databaseName)
+	isSharedOrSystem, err := o.isDBSharedOrSystem(ctx, opts.Session, resource, databaseName)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -277,6 +289,11 @@ func (o *tableBuilder) Entitlements(ctx context.Context, resource *v2.Resource,
 	objectKind := getObjectKind(resource)
 	tableGrants, err := o.client.ListTableGrants(ctx, opts.Session, databaseName, schemaName, tableName, objectKind)
 	if err != nil {
+		if errors.Is(err, snowflake.ErrObjectNotFound) {
+			ctxzap.Extract(ctx).Warn("table no longer exists during entitlements phase, skipping",
+				zap.String("table", resource.Id.Resource))
+			return nil, &rs.SyncOpResults{}, nil
+		}
 		return nil, nil, wrapError(err, fmt.Sprintf("failed to list table grants for %s", resource.Id.Resource))
 	}
 
@@ -301,12 +318,14 @@ func (o *tableBuilder) Entitlements(ctx context.Context, resource *v2.Resource,
 }
 
 func (o *tableBuilder) Grants(ctx context.Context, resource *v2.Resource, opts rs.SyncOpAttrs) ([]*v2.Grant, *rs.SyncOpResults, error) {
+	// No syncTables guard here: same reasoning as Entitlements — the SDK iterates from
+	// the persistent C1Z store, so a guard would silently leave stale grants from prior syncs.
 	databaseName, schemaName, tableName, err := parseTableResourceID(resource)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	isSharedOrSystem, err := o.isDBSharedOrSystem(ctx, resource, databaseName)
+	isSharedOrSystem, err := o.isDBSharedOrSystem(ctx, opts.Session, resource, databaseName)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -317,6 +336,11 @@ func (o *tableBuilder) Grants(ctx context.Context, resource *v2.Resource, opts r
 	objectKind := getObjectKind(resource)
 	tableGrants, err := o.client.ListTableGrants(ctx, opts.Session, databaseName, schemaName, tableName, objectKind)
 	if err != nil {
+		if errors.Is(err, snowflake.ErrObjectNotFound) {
+			ctxzap.Extract(ctx).Warn("table no longer exists during grants phase, skipping",
+				zap.String("table", resource.Id.Resource))
+			return nil, &rs.SyncOpResults{}, nil
+		}
 		return nil, nil, wrapError(err, "failed to list table grants")
 	}
 	if len(tableGrants) == 0 {
@@ -391,6 +415,11 @@ func (o *tableBuilder) Grants(ctx context.Context, resource *v2.Resource, opts r
 	if ownerPrincipalID == nil {
 		table, err := o.client.GetTable(ctx, databaseName, schemaName, tableName)
 		if err != nil {
+			if errors.Is(err, snowflake.ErrObjectNotFound) {
+				ctxzap.Extract(ctx).Warn("table disappeared before owner fallback, returning partial grants",
+					zap.String("table", resource.Id.Resource))
+				return grants, &rs.SyncOpResults{}, nil
+			}
 			return nil, nil, wrapError(err, "failed to get table for owner fallback")
 		}
 		if table != nil && table.Owner != "" && table.Owner != "SNOWFLAKE" {
@@ -415,8 +444,9 @@ func (o *tableBuilder) Grants(ctx context.Context, resource *v2.Resource, opts r
 	return grants, &rs.SyncOpResults{}, nil
 }
 
-func newTableBuilder(client *snowflake.Client) *tableBuilder {
+func newTableBuilder(client *snowflake.Client, syncTables bool) *tableBuilder {
 	return &tableBuilder{
-		client: client,
+		client:     client,
+		syncTables: syncTables,
 	}
 }