diff --git a/README.md b/README.md index 65c5631c..a39e7aab 100644 --- a/README.md +++ b/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 diff --git a/docs/connector.mdx b/docs/connector.mdx index 6442736f..d9fc54d6 100644 --- a/docs/connector.mdx +++ b/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. +**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. + + Click **Save**. @@ -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. diff --git a/pkg/config/conf.gen.go b/pkg/config/conf.gen.go index a8e30cbc..5eb5af1f 100644 --- a/pkg/config/conf.gen.go +++ b/pkg/config/conf.gen.go @@ -10,6 +10,7 @@ type Snowflake struct { PrivateKeyPath string `mapstructure:"private-key-path"` UserIdentifier string `mapstructure:"user-identifier"` SyncSecrets bool `mapstructure:"sync-secrets"` + SyncTables bool `mapstructure:"sync-tables"` ExcludedDatabases []string `mapstructure:"excluded-databases"` } diff --git a/pkg/config/config.go b/pkg/config/config.go index 64de7d4c..4e992f15 100644 --- a/pkg/config/config.go +++ b/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, } diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 860cdf3c..db88e3c2 100644 --- a/pkg/connector/connector.go +++ b/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 } diff --git a/pkg/connector/databases.go b/pkg/connector/databases.go index d8d4edca..1f7ab74d 100644 --- a/pkg/connector/databases.go +++ b/pkg/connector/databases.go @@ -19,6 +19,7 @@ type databaseBuilder struct { resourceType *v2.ResourceType client *snowflake.Client syncSecrets bool + syncTables bool excludedDatabases map[string]struct{} // uppercase-normalised names to exclude } @@ -26,7 +27,7 @@ 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, } } diff --git a/pkg/connector/tables.go b/pkg/connector/tables.go index 4fbdf71a..e2bb5021 100644 --- a/pkg/connector/tables.go +++ b/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, } } diff --git a/pkg/connector/tables_test.go b/pkg/connector/tables_test.go index 6f95f773..77985a4c 100644 --- a/pkg/connector/tables_test.go +++ b/pkg/connector/tables_test.go @@ -3,15 +3,150 @@ package connector import ( "context" "fmt" + "net/http" + "net/http/httptest" + "sync" "testing" "time" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" rs "github.com/conductorone/baton-sdk/pkg/types/resource" + "github.com/conductorone/baton-sdk/pkg/types/sessions" "github.com/conductorone/baton-snowflake/pkg/snowflake" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +// mapSessionStore duplicates the type in pkg/snowflake/database_test.go — +// Go does not allow importing _test.go symbols across packages. +// IMPORTANT: Keep the two copies in sync if the SessionStore interface changes. +type mapSessionStore struct { + mu sync.Mutex + data map[string][]byte +} + +func newMapSessionStore() *mapSessionStore { + return &mapSessionStore{data: make(map[string][]byte)} +} + +func (m *mapSessionStore) compositeKey(key string, opt []sessions.SessionStoreOption) string { + bag := &sessions.SessionStoreBag{} + for _, o := range opt { + _ = o(context.Background(), bag) + } + if bag.Prefix != "" { + return bag.Prefix + "/" + key + } + return key +} + +func (m *mapSessionStore) Get(_ context.Context, key string, opt ...sessions.SessionStoreOption) ([]byte, bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + v, ok := m.data[m.compositeKey(key, opt)] + return v, ok, nil +} + +func (m *mapSessionStore) GetMany(ctx context.Context, keys []string, opt ...sessions.SessionStoreOption) (map[string][]byte, []string, error) { + result := make(map[string][]byte, len(keys)) + for _, k := range keys { + v, found, err := m.Get(ctx, k, opt...) + if err != nil { + return nil, nil, err + } + if found { + result[k] = v + } + } + return result, nil, nil +} + +func (m *mapSessionStore) Set(_ context.Context, key string, value []byte, opt ...sessions.SessionStoreOption) error { + m.mu.Lock() + defer m.mu.Unlock() + m.data[m.compositeKey(key, opt)] = value + return nil +} + +func (m *mapSessionStore) SetMany(ctx context.Context, values map[string][]byte, opt ...sessions.SessionStoreOption) error { + for k, v := range values { + if err := m.Set(ctx, k, v, opt...); err != nil { + return err + } + } + return nil +} + +func (m *mapSessionStore) Delete(_ context.Context, key string, opt ...sessions.SessionStoreOption) error { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.data, m.compositeKey(key, opt)) + return nil +} + +func (m *mapSessionStore) Clear(_ context.Context, _ ...sessions.SessionStoreOption) error { + m.mu.Lock() + defer m.mu.Unlock() + m.data = make(map[string][]byte) + return nil +} + +func (m *mapSessionStore) GetAll(_ context.Context, _ string, _ ...sessions.SessionStoreOption) (map[string][]byte, string, error) { + m.mu.Lock() + defer m.mu.Unlock() + cp := make(map[string][]byte, len(m.data)) + for k, v := range m.data { + cp[k] = v + } + return cp, "", nil +} + +func newConnectorTestClient(t *testing.T, handler http.Handler) (*snowflake.Client, func()) { + t.Helper() + ts := httptest.NewServer(handler) + client, err := snowflake.New(ts.URL, snowflake.JWTConfig{}, ts.Client()) + require.NoError(t, err) + return client, ts.Close +} + +// JSON response bodies for mock HTTP servers. Column order in rowType matches the +// struct field→column maps in pkg/snowflake/table.go and database.go. +const ( + // 422 with code "002003" (not "003001") — triggers ErrObjectNotFound in ListTableGrants / GetTable. + // Code "003001" means genuine RBAC privilege denial (PermissionDenied); any other 422 soft-skips. + body422ObjectNotFound = `{"code":"002003","message":"Object does not exist or not authorized."}` + + // 200 OK, one grant row with granted_to="APPLICATION" — bypasses the len==0 early-return + // without triggering GetAccountRole or GetUser, and leaves ownerPrincipalID==nil. + // created_on must be a float seconds-since-epoch value (Snowflake TIMESTAMP_LTZ wire format). + body200OneApplicationGrant = `{"resultSetMetadata":{"numRows":1,"rowType":[` + + `{"name":"created_on","type":"timestamp_ltz"},` + + `{"name":"privilege","type":"text"},` + + `{"name":"granted_on","type":"text"},` + + `{"name":"name","type":"text"},` + + `{"name":"granted_to","type":"text"},` + + `{"name":"grantee_name","type":"text"},` + + `{"name":"grant_option","type":"text"},` + + `{"name":"granted_by","type":"text"}` + + `]},"data":[["1704067200.000000000","SELECT","TABLE","MYDB.PUBLIC.MYTABLE","APPLICATION","MYAPP","false","SYSADMIN"]],"statementHandle":"","code":"","message":""}` + + // 200 OK, one STANDARD database row. rowType matches databaseStructFieldToColumnMap. + body200StandardDB = `{"resultSetMetadata":{"numRows":1,"rowType":[` + + `{"name":"name","type":"text"},` + + `{"name":"owner","type":"text"},` + + `{"name":"kind","type":"text"},` + + `{"name":"origin","type":"text"}` + + `]},"data":[["MYDB","SYSADMIN","STANDARD",""]],"statementHandle":"","code":"","message":""}` + + // 200 OK, one SHARED database row. + body200SharedDB = `{"resultSetMetadata":{"numRows":1,"rowType":[` + + `{"name":"name","type":"text"},` + + `{"name":"owner","type":"text"},` + + `{"name":"kind","type":"text"},` + + `{"name":"origin","type":"text"}` + + `]},"data":[["SHAREDDB","SYSADMIN","SHARED",""]],"statementHandle":"","code":"","message":""}` +) + // makeTableResource creates a table resource with profile fields via the real tableResource() function. func makeTableResource(t *testing.T, dbName, schemaName, tableName string) *v2.Resource { t.Helper() @@ -144,3 +279,174 @@ func makePartialProfileResource(t *testing.T, dbName, schemaName, tableName stri require.NoError(t, err) return resource } + +func TestTableBuilderSkipsWhenSyncTablesDisabled(t *testing.T) { + builder := &tableBuilder{syncTables: false} + parentID := &v2.ResourceId{ResourceType: databaseResourceType.Id, Resource: "MYDB"} + resources, results, err := builder.List(context.Background(), parentID, rs.SyncOpAttrs{}) + require.NoError(t, err) + require.Empty(t, resources) + require.NotNil(t, results) +} + +// seedDB writes a STANDARD MYDB entry to ss so isDBSharedOrSystem is answered +// from the cache without an HTTP call. +func seedStandardDB(t *testing.T, ctx context.Context, ss *mapSessionStore) { + t.Helper() + // CacheDatabases only uses the session store, not any Client fields, + // so a zero-value Client is safe to use here. + err := (&snowflake.Client{}).CacheDatabases(ctx, ss, []snowflake.Database{ + {Name: "MYDB", Owner: "SYSADMIN", Kind: "STANDARD"}, + }) + require.NoError(t, err) +} + +// --- Gap 1: ErrObjectNotFound soft-skip --- + +func TestEntitlements_ErrObjectNotFound_SoftSkip(t *testing.T) { + ctx := context.Background() + ss := newMapSessionStore() + seedStandardDB(t, ctx, ss) + + client, cleanup := newConnectorTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnprocessableEntity) + _, _ = w.Write([]byte(body422ObjectNotFound)) + })) + defer cleanup() + + builder := &tableBuilder{client: client, syncTables: true} + resource := makeTableResource(t, "MYDB", "PUBLIC", "MYTABLE") + ents, results, err := builder.Entitlements(ctx, resource, rs.SyncOpAttrs{Session: ss}) + + require.NoError(t, err) + assert.NotNil(t, results) + assert.Nil(t, ents) +} + +func TestGrants_ErrObjectNotFound_ListTableGrants_SoftSkip(t *testing.T) { + ctx := context.Background() + ss := newMapSessionStore() + seedStandardDB(t, ctx, ss) + + client, cleanup := newConnectorTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnprocessableEntity) + _, _ = w.Write([]byte(body422ObjectNotFound)) + })) + defer cleanup() + + builder := &tableBuilder{client: client, syncTables: true} + resource := makeTableResource(t, "MYDB", "PUBLIC", "MYTABLE") + grants, results, err := builder.Grants(ctx, resource, rs.SyncOpAttrs{Session: ss}) + + require.NoError(t, err) + assert.NotNil(t, results) + assert.Nil(t, grants) +} + +func TestGrants_ErrObjectNotFound_GetTable_ReturnsPartialGrants(t *testing.T) { + ctx := context.Background() + ss := newMapSessionStore() + seedStandardDB(t, ctx, ss) + + var callCount int + client, cleanup := newConnectorTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + callCount++ + w.Header().Set("Content-Type", "application/json") + if callCount == 1 { + // ListTableGrants: one APPLICATION grant row — bypasses the len==0 early-return + // without triggering GetAccountRole/GetUser; leaves ownerPrincipalID==nil. + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(body200OneApplicationGrant)) + } else { + // GetTable owner fallback: 422 → ErrObjectNotFound → partial grants returned. + w.WriteHeader(http.StatusUnprocessableEntity) + _, _ = w.Write([]byte(body422ObjectNotFound)) + } + })) + defer cleanup() + + builder := &tableBuilder{client: client, syncTables: true} + resource := makeTableResource(t, "MYDB", "PUBLIC", "MYTABLE") + grants, results, err := builder.Grants(ctx, resource, rs.SyncOpAttrs{Session: ss}) + + require.NoError(t, err) + assert.NotNil(t, results) + // grants is nil because the only grant row has granted_to="APPLICATION", which the + // grant-building switch skips (only ROLE and USER are handled), so nothing is ever appended. + // GetTable then returns ErrObjectNotFound, triggering the early return — which returns + // whatever is in grants at that point (nil here). "Partial" means this return could contain + // a non-empty slice in other fixtures. + assert.Nil(t, grants) + assert.Equal(t, 2, callCount, "expected ListTableGrants + GetTable calls") +} + +// --- Gap 2: isDBSharedOrSystem cache paths --- + +func TestIsDBSharedOrSystem_CacheMiss_StandardDB_ReturnsFalse(t *testing.T) { + ctx := context.Background() + ss := newMapSessionStore() // empty — cache miss forces HTTP call + + client, cleanup := newConnectorTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(body200StandardDB)) + })) + defer cleanup() + + builder := &tableBuilder{client: client} + resource := makeBareResource(t, "MYDB.PUBLIC.MYTABLE") // no profile → fast path skipped + result, err := builder.isDBSharedOrSystem(ctx, ss, resource, "MYDB") + + require.NoError(t, err) + assert.False(t, result) +} + +func TestIsDBSharedOrSystem_CacheMiss_SharedDB_ReturnsTrue(t *testing.T) { + ctx := context.Background() + ss := newMapSessionStore() + + client, cleanup := newConnectorTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(body200SharedDB)) + })) + defer cleanup() + + builder := &tableBuilder{client: client} + resource := makeBareResource(t, "SHAREDDB.PUBLIC.MYTABLE") + result, err := builder.isDBSharedOrSystem(ctx, ss, resource, "SHAREDDB") + + require.NoError(t, err) + assert.True(t, result) +} + +func TestIsDBSharedOrSystem_CacheHit_NoHTTPCall(t *testing.T) { + ctx := context.Background() + ss := newMapSessionStore() + seedStandardDB(t, ctx, ss) + + // &snowflake.Client{} has no transport — any real HTTP call would panic. + builder := &tableBuilder{client: &snowflake.Client{}} + resource := makeBareResource(t, "MYDB.PUBLIC.MYTABLE") + result, err := builder.isDBSharedOrSystem(ctx, ss, resource, "MYDB") + + require.NoError(t, err) + assert.False(t, result) +} + +func TestIsDBSharedOrSystem_FastPath_ProfileFieldTrue(t *testing.T) { + ctx := context.Background() + parentID := &v2.ResourceId{ResourceType: databaseResourceType.Id, Resource: "SHAREDDB"} + table := &snowflake.Table{DatabaseName: "SHAREDDB", SchemaName: "PUBLIC", Name: "T", Kind: "TABLE"} + resource, err := tableResource(ctx, table, parentID, true) // embeds database_is_shared_system: true + require.NoError(t, err) + + // &snowflake.Client{} has no transport — any real HTTP call would panic. + builder := &tableBuilder{client: &snowflake.Client{}} + result, callErr := builder.isDBSharedOrSystem(ctx, nil, resource, "SHAREDDB") + + require.NoError(t, callErr) + assert.True(t, result) +} diff --git a/pkg/snowflake/account_role.go b/pkg/snowflake/account_role.go index 3fbf671a..fcaa1976 100644 --- a/pkg/snowflake/account_role.go +++ b/pkg/snowflake/account_role.go @@ -91,13 +91,7 @@ func (c *Client) ListAccountRoles(ctx context.Context, cursor string, limit int) l := ctxzap.Extract(ctx) l.Debug("ListAccountRoles", zap.String("response.code", response.Code), zap.String("response.message", response.Message)) - req, err = c.GetStatementResponse(ctx, response.StatementHandle) - if err != nil { - return nil, err - } - resp2, err := c.Do(req, uhttp.WithJSONResponse(&response)) - defer closeResponseBody(resp2) - if err != nil { + if err := c.fetchStatementResultIfAsync(ctx, resp1, response.StatementHandle, &response); err != nil { return nil, err } @@ -126,13 +120,7 @@ func (c *Client) ListAccountRoleGrantees(ctx context.Context, roleName string) ( return nil, err } - req, err = c.GetStatementResponse(ctx, response.StatementHandle) - if err != nil { - return nil, err - } - resp2, err := c.Do(req, uhttp.WithJSONResponse(&response)) - defer closeResponseBody(resp2) - if err != nil { + if err := c.fetchStatementResultIfAsync(ctx, resp1, response.StatementHandle, &response); err != nil { return nil, err } @@ -158,7 +146,11 @@ func (c *Client) CacheAccountRoles(ctx context.Context, ss sessions.SessionStore func (c *Client) GetAccountRole(ctx context.Context, ss sessions.SessionStore, roleName string) (*AccountRole, int, error) { if ss != nil { - if cached, found, err := session.GetJSON[*AccountRole](ctx, ss, roleName, accountRoleNamespace); err == nil && found { + cached, found, err := session.GetJSON[*AccountRole](ctx, ss, roleName, accountRoleNamespace) + if err != nil { + ctxzap.Extract(ctx).Debug("account role cache lookup error, falling through to API", + zap.String("role_name", roleName), zap.Error(err)) + } else if found { return cached, http.StatusOK, nil } } diff --git a/pkg/snowflake/client.go b/pkg/snowflake/client.go index 2a7fffb8..411f7bf1 100644 --- a/pkg/snowflake/client.go +++ b/pkg/snowflake/client.go @@ -17,6 +17,7 @@ import ( var ( accountRoleNamespace = sessions.WithPrefix("account_role") + databaseNamespace = sessions.WithPrefix("database") userNamespace = sessions.WithPrefix("user") tableGrantsNamespace = sessions.WithPrefix("table_grants") ) @@ -251,6 +252,22 @@ func Contains[T comparable](ts []T, val T) bool { return false } +// fetchStatementResultIfAsync fires the GET poll only when Snowflake returned 202 +// (query still running). On 200 the full result is already in target from the POST body. +// defer closeResponseBody is used consistently with all other call sites in this package. +func (c *Client) fetchStatementResultIfAsync(ctx context.Context, postResp *http.Response, handle string, target any) error { + if postResp == nil || postResp.StatusCode != http.StatusAccepted { + return nil + } + req, err := c.GetStatementResponse(ctx, handle) + if err != nil { + return err + } + resp, err := c.Do(req, uhttp.WithJSONResponse(target)) + defer closeResponseBody(resp) + return err +} + // closeResponseBody drains and closes the response body if it exists. // This ensures proper resource cleanup and allows connection reuse. func closeResponseBody(resp *http.Response) { diff --git a/pkg/snowflake/client_test.go b/pkg/snowflake/client_test.go new file mode 100644 index 00000000..321bb11f --- /dev/null +++ b/pkg/snowflake/client_test.go @@ -0,0 +1,27 @@ +package snowflake + +import ( + "context" + "net/http" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestFetchStatementResultIfAsync_NilResponse(t *testing.T) { + // &Client{} has a nil StatementsApiUrl. If the nil guard did not fire, + // GetStatementResponse would call c.StatementsApiUrl.String() and panic. + // A clean nil return proves the guard fires before any network call. + c := &Client{} + err := c.fetchStatementResultIfAsync(context.Background(), nil, "", nil) + require.NoError(t, err) +} + +func TestFetchStatementResultIfAsync_SyncResponse(t *testing.T) { + // StatusCode 200 is not 202 — the function must return nil without polling. + // Same nil-transport proof as above applies. + c := &Client{} + postResp := &http.Response{StatusCode: http.StatusOK, Body: http.NoBody} + err := c.fetchStatementResultIfAsync(context.Background(), postResp, "", nil) + require.NoError(t, err) +} diff --git a/pkg/snowflake/database.go b/pkg/snowflake/database.go index a32ced22..47a459b5 100644 --- a/pkg/snowflake/database.go +++ b/pkg/snowflake/database.go @@ -3,8 +3,11 @@ package snowflake import ( "context" "fmt" + "net/http" "strings" + "github.com/conductorone/baton-sdk/pkg/session" + "github.com/conductorone/baton-sdk/pkg/types/sessions" "github.com/conductorone/baton-sdk/pkg/uhttp" "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" "go.uber.org/zap" @@ -43,7 +46,7 @@ func (d *Database) IsSharedOrSystem() bool { return true } kind := strings.ToUpper(strings.TrimSpace(d.Kind)) - return kind == "SHARED" || kind == "APPLICATION" || kind == "IMPORTED DATABASE" + return kind == "SHARED" || kind == "APPLICATION" || kind == "IMPORTED DATABASE" || kind == "CATALOG-LINKED DATABASE" } func (r *ListDatabasesRawResponse) GetDatabases() ([]Database, error) { @@ -83,13 +86,7 @@ func (c *Client) ListDatabases(ctx context.Context, cursor string, limit int) ([ l := ctxzap.Extract(ctx) l.Debug("ListDatabases", zap.String("response.code", response.Code), zap.String("response.message", response.Message)) - req, err = c.GetStatementResponse(ctx, response.StatementHandle) - if err != nil { - return nil, err - } - resp2, err := c.Do(req, uhttp.WithJSONResponse(&response)) - defer closeResponseBody(resp2) - if err != nil { + if err := c.fetchStatementResultIfAsync(ctx, resp1, response.StatementHandle, &response); err != nil { return nil, err } @@ -101,7 +98,28 @@ func (c *Client) ListDatabases(ctx context.Context, cursor string, limit int) ([ return dbs, nil } -func (c *Client) GetDatabase(ctx context.Context, name string) (*Database, int, error) { +func (c *Client) CacheDatabases(ctx context.Context, ss sessions.SessionStore, databases []Database) error { + if ss == nil || len(databases) == 0 { + return nil + } + for i := range databases { + if err := session.SetJSON(ctx, ss, databases[i].Name, &databases[i], databaseNamespace); err != nil { + return err + } + } + return nil +} + +func (c *Client) GetDatabase(ctx context.Context, ss sessions.SessionStore, name string) (*Database, int, error) { + if ss != nil { + cached, found, err := session.GetJSON[*Database](ctx, ss, name, databaseNamespace) + if err != nil { + ctxzap.Extract(ctx).Debug("database cache lookup error, falling through to API", + zap.String("name", name), zap.Error(err)) + } else if found { + return cached, http.StatusOK, nil + } + } queries := []string{ fmt.Sprintf("SHOW DATABASES LIKE '%s' LIMIT 1;", name), } @@ -122,16 +140,36 @@ func (c *Client) GetDatabase(ctx context.Context, name string) (*Database, int, return nil, statusCode, err } + pollStatusCode := resp.StatusCode + if resp.StatusCode == http.StatusAccepted { + req, err = c.GetStatementResponse(ctx, response.StatementHandle) + if err != nil { + return nil, 0, err + } + pollResp, err := c.Do(req, uhttp.WithJSONResponse(&response)) + defer closeResponseBody(pollResp) + if err != nil { + return nil, 0, err + } + if pollResp != nil { + pollStatusCode = pollResp.StatusCode + } + } + databases, err := response.GetDatabases() if err != nil { - return nil, resp.StatusCode, err + return nil, pollStatusCode, err } if len(databases) == 0 { - return nil, resp.StatusCode, fmt.Errorf("database with name %s not found", name) + return nil, pollStatusCode, fmt.Errorf("database with name %s not found", name) } else if len(databases) > 1 { - return nil, resp.StatusCode, fmt.Errorf("expected 1 database with name %s, got %d", name, len(databases)) + return nil, pollStatusCode, fmt.Errorf("expected 1 database with name %s, got %d", name, len(databases)) } - return &databases[0], resp.StatusCode, nil + if ss != nil { + // Write-back is best-effort: a cache miss on a subsequent call just falls through to the API. + _ = session.SetJSON(ctx, ss, name, &databases[0], databaseNamespace) + } + return &databases[0], pollStatusCode, nil } diff --git a/pkg/snowflake/database_test.go b/pkg/snowflake/database_test.go new file mode 100644 index 00000000..243f37db --- /dev/null +++ b/pkg/snowflake/database_test.go @@ -0,0 +1,207 @@ +package snowflake + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "testing" + + "github.com/conductorone/baton-sdk/pkg/types/sessions" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// mapSessionStore is a simple in-memory SessionStore for testing. +// It applies the prefix option by prepending it to the key. +// IMPORTANT: pkg/connector/tables_test.go contains an identical copy (Go does not allow +// importing _test.go symbols across packages). Keep the two in sync. +type mapSessionStore struct { + mu sync.Mutex + data map[string][]byte +} + +func newMapSessionStore() *mapSessionStore { + return &mapSessionStore{data: make(map[string][]byte)} +} + +func (m *mapSessionStore) compositeKey(key string, opt []sessions.SessionStoreOption) string { + bag := &sessions.SessionStoreBag{} + for _, o := range opt { + _ = o(context.Background(), bag) + } + if bag.Prefix != "" { + return bag.Prefix + "/" + key + } + return key +} + +func (m *mapSessionStore) Get(_ context.Context, key string, opt ...sessions.SessionStoreOption) ([]byte, bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + v, ok := m.data[m.compositeKey(key, opt)] + return v, ok, nil +} + +func (m *mapSessionStore) GetMany(ctx context.Context, keys []string, opt ...sessions.SessionStoreOption) (map[string][]byte, []string, error) { + result := make(map[string][]byte, len(keys)) + for _, k := range keys { + v, found, err := m.Get(ctx, k, opt...) + if err != nil { + return nil, nil, err + } + if found { + result[k] = v + } + } + return result, nil, nil +} + +func (m *mapSessionStore) Set(_ context.Context, key string, value []byte, opt ...sessions.SessionStoreOption) error { + m.mu.Lock() + defer m.mu.Unlock() + m.data[m.compositeKey(key, opt)] = value + return nil +} + +func (m *mapSessionStore) SetMany(ctx context.Context, values map[string][]byte, opt ...sessions.SessionStoreOption) error { + for k, v := range values { + if err := m.Set(ctx, k, v, opt...); err != nil { + return err + } + } + return nil +} + +func (m *mapSessionStore) Delete(_ context.Context, key string, opt ...sessions.SessionStoreOption) error { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.data, m.compositeKey(key, opt)) + return nil +} + +func (m *mapSessionStore) Clear(_ context.Context, _ ...sessions.SessionStoreOption) error { + m.mu.Lock() + defer m.mu.Unlock() + m.data = make(map[string][]byte) + return nil +} + +func (m *mapSessionStore) GetAll(_ context.Context, _ string, _ ...sessions.SessionStoreOption) (map[string][]byte, string, error) { + m.mu.Lock() + defer m.mu.Unlock() + cp := make(map[string][]byte, len(m.data)) + for k, v := range m.data { + cp[k] = v + } + return cp, "", nil +} + +// --- IsSharedOrSystem tests --- + +func TestDatabase_IsSharedOrSystem(t *testing.T) { + cases := []struct { + name string + db Database + expect bool + }{ + {"standard", Database{Name: "X", Owner: "SYSADMIN", Kind: "STANDARD"}, false}, + {"shared", Database{Name: "X", Owner: "SYSADMIN", Kind: "SHARED"}, true}, + {"application", Database{Name: "X", Owner: "SYSADMIN", Kind: "APPLICATION"}, true}, + {"imported database", Database{Name: "X", Owner: "SYSADMIN", Kind: "IMPORTED DATABASE"}, true}, + {"catalog-linked uppercase", Database{Name: "X", Owner: "SYSADMIN", Kind: "CATALOG-LINKED DATABASE"}, true}, + {"catalog-linked lowercase", Database{Name: "X", Owner: "SYSADMIN", Kind: "catalog-linked database"}, true}, + {"catalog-linked with spaces", Database{Name: "X", Owner: "SYSADMIN", Kind: " CATALOG-LINKED DATABASE "}, true}, + {"snowflake owner", Database{Name: "X", Owner: "SNOWFLAKE", Kind: "STANDARD"}, true}, + {"snowflake owner lowercase", Database{Name: "X", Owner: "snowflake", Kind: "STANDARD"}, true}, + {"empty owner", Database{Name: "X", Owner: "", Kind: "STANDARD"}, true}, + {"with origin", Database{Name: "X", Owner: "SYSADMIN", Kind: "STANDARD", Origin: "myaccount.myshare"}, true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.expect, tc.db.IsSharedOrSystem()) + }) + } +} + +// --- CacheDatabases / GetDatabase caching tests --- + +func TestCacheDatabases_NilSession(t *testing.T) { + c := &Client{} + err := c.CacheDatabases(context.Background(), nil, []Database{{Name: "DB1"}}) + require.NoError(t, err) +} + +func TestCacheDatabases_Empty(t *testing.T) { + ss := newMapSessionStore() + c := &Client{} + err := c.CacheDatabases(context.Background(), ss, nil) + require.NoError(t, err) + assert.Empty(t, ss.data) +} + +func TestCacheDatabases_PopulatesStore(t *testing.T) { + ctx := context.Background() + ss := newMapSessionStore() + c := &Client{} + + dbs := []Database{ + {Name: "MYDB", Owner: "SYSADMIN", Kind: "STANDARD"}, + {Name: "SHAREDDB", Owner: "SYSADMIN", Kind: "SHARED"}, + } + require.NoError(t, c.CacheDatabases(ctx, ss, dbs)) + + // The store must contain both entries under the database namespace prefix. + assert.Len(t, ss.data, 2) + key := fmt.Sprintf("%s/%s", "database", "MYDB") + assert.Contains(t, ss.data, key) +} + +func TestGetDatabase_CacheMiss(t *testing.T) { + ctx := context.Background() + + var requestCount int + const dbBody = `{"resultSetMetadata":{"numRows":1,"rowType":[{"name":"name","type":"text"},{"name":"owner","type":"text"},{"name":"kind","type":"text"},{"name":"origin","type":"text"}]},"data":[["MYDB","SYSADMIN","STANDARD",""]],"statementHandle":"","code":"","message":""}` + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount++ + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(dbBody)) + })) + defer ts.Close() + + client, err := New(ts.URL, JWTConfig{}, ts.Client()) + require.NoError(t, err) + + ss := newMapSessionStore() + require.NoError(t, client.CacheDatabases(ctx, ss, []Database{{Name: "OTHER", Owner: "SYSADMIN", Kind: "STANDARD"}})) + + db, statusCode, err := client.GetDatabase(ctx, ss, "MYDB") + require.NoError(t, err) + assert.Equal(t, 1, requestCount, "server should have been called exactly once") + require.NotNil(t, db) + assert.Equal(t, "MYDB", db.Name) + assert.Equal(t, http.StatusOK, statusCode) +} + +func TestGetDatabase_CacheHit(t *testing.T) { + ctx := context.Background() + ss := newMapSessionStore() + c := &Client{} + + // Seed the cache via CacheDatabases. + seed := []Database{{Name: "MYDB", Owner: "SYSADMIN", Kind: "STANDARD"}} + require.NoError(t, c.CacheDatabases(ctx, ss, seed)) + + // GetDatabase with a populated cache must return the cached entry without + // hitting the API (Client has no transport, so any real call would panic). + db, statusCode, err := c.GetDatabase(ctx, ss, "MYDB") + require.NoError(t, err) + assert.Equal(t, http.StatusOK, statusCode, "cache hit should return http.StatusOK") + require.NotNil(t, db) + assert.Equal(t, "MYDB", db.Name) + assert.Equal(t, "SYSADMIN", db.Owner) + assert.Equal(t, "STANDARD", db.Kind) // Kind drives IsSharedOrSystem — verify it survives cache round-trip +} + diff --git a/pkg/snowflake/table.go b/pkg/snowflake/table.go index ac96cd7f..30a4f358 100644 --- a/pkg/snowflake/table.go +++ b/pkg/snowflake/table.go @@ -3,6 +3,7 @@ package snowflake import ( "context" "encoding/json" + "errors" "fmt" "net/http" "strings" @@ -18,6 +19,11 @@ import ( "github.com/conductorone/baton-sdk/pkg/uhttp" ) +// ErrObjectNotFound is returned when Snowflake reports a target object does not +// exist or cannot be seen, and the cause is not an RBAC privilege error (code +// 003001). Callers should treat this as a soft-skip condition. +var ErrObjectNotFound = errors.New("baton-snowflake: object does not exist or not authorized") + var schemaStructFieldToColumnMap = map[string]string{ structFieldName: columnName, structFieldDatabaseName: columnDatabaseName, @@ -67,6 +73,9 @@ func (c *Client) ListSchemasInDatabase(ctx context.Context, databaseName string) resp1, err := c.Do(req, uhttp.WithJSONResponse(&response)) defer closeResponseBody(resp1) if err != nil { + // All 422s returned as PermissionDenied. Callers in tables.go propagate this + // error via wrapError — there is no soft-skip at the List level, so returning + // ErrObjectNotFound here would have no observable effect on sync behavior. if resp1 != nil && resp1.StatusCode == http.StatusUnprocessableEntity { l.Debug("Insufficient privileges for SHOW SCHEMAS IN DATABASE", zap.String("database", databaseName)) wrappedErr := fmt.Errorf("baton-snowflake: insufficient privileges for SHOW SCHEMAS IN DATABASE %s: %w", databaseName, err) @@ -75,19 +84,27 @@ func (c *Client) ListSchemasInDatabase(ctx context.Context, databaseName string) return nil, err } - req, err = c.GetStatementResponse(ctx, response.StatementHandle) - if err != nil { - return nil, err - } - resp2, err := c.Do(req, uhttp.WithJSONResponse(&response)) - defer closeResponseBody(resp2) - if err != nil { - if resp2 != nil && resp2.StatusCode == http.StatusUnprocessableEntity { - l.Debug("Insufficient privileges for SHOW SCHEMAS IN DATABASE (statement result)", zap.String("database", databaseName)) - wrappedErr := fmt.Errorf("baton-snowflake: insufficient privileges for SHOW SCHEMAS IN DATABASE %s (statement result): %w", databaseName, err) - return nil, status.Error(codes.PermissionDenied, wrappedErr.Error()) + // Inline async poll: unlike fetchStatementResultIfAsync, this path also wraps 422 as + // PermissionDenied. That extra error handling is why this function does not call the + // shared helper. + if resp1.StatusCode == http.StatusAccepted { + req, err = c.GetStatementResponse(ctx, response.StatementHandle) + if err != nil { + return nil, err + } + resp2, err := c.Do(req, uhttp.WithJSONResponse(&response)) + defer closeResponseBody(resp2) + if err != nil { + // All 422s returned as PermissionDenied. Callers in tables.go propagate this + // error via wrapError — there is no soft-skip at the List level, so returning + // ErrObjectNotFound here would have no observable effect on sync behavior. + if resp2 != nil && resp2.StatusCode == http.StatusUnprocessableEntity { + l.Debug("Insufficient privileges for SHOW SCHEMAS IN DATABASE (statement result)", zap.String("database", databaseName)) + wrappedErr := fmt.Errorf("baton-snowflake: insufficient privileges for SHOW SCHEMAS IN DATABASE %s (statement result): %w", databaseName, err) + return nil, status.Error(codes.PermissionDenied, wrappedErr.Error()) + } + return nil, err } - return nil, err } return response.ListSchemas() @@ -158,6 +175,9 @@ func (c *Client) ListTablesInSchema(ctx context.Context, databaseName, schemaNam resp1, err := c.Do(req, uhttp.WithJSONResponse(&response)) defer closeResponseBody(resp1) if err != nil { + // All 422s returned as PermissionDenied. Callers in tables.go propagate this + // error via wrapError — there is no soft-skip at the List level, so returning + // ErrObjectNotFound here would have no observable effect on sync behavior. if resp1 != nil && resp1.StatusCode == http.StatusUnprocessableEntity { l.Debug("Insufficient privileges for SHOW TABLES IN SCHEMA", zap.String("database", databaseName), zap.String("schema", schemaName)) @@ -167,20 +187,26 @@ func (c *Client) ListTablesInSchema(ctx context.Context, databaseName, schemaNam return nil, "", err } - req, err = c.GetStatementResponse(ctx, response.StatementHandle) - if err != nil { - return nil, "", err - } - resp2, err := c.Do(req, uhttp.WithJSONResponse(&response)) - defer closeResponseBody(resp2) - if err != nil { - if resp2 != nil && resp2.StatusCode == http.StatusUnprocessableEntity { - l.Debug("Insufficient privileges for SHOW TABLES IN SCHEMA (statement result)", - zap.String("database", databaseName), zap.String("schema", schemaName)) - wrappedErr := fmt.Errorf("baton-snowflake: insufficient privileges for SHOW TABLES IN SCHEMA %s.%s (statement result): %w", databaseName, schemaName, err) - return nil, "", status.Error(codes.PermissionDenied, wrappedErr.Error()) + // Inline async poll: see ListSchemasInDatabase for why this does not use fetchStatementResultIfAsync. + if resp1.StatusCode == http.StatusAccepted { + req, err = c.GetStatementResponse(ctx, response.StatementHandle) + if err != nil { + return nil, "", err + } + resp2, err := c.Do(req, uhttp.WithJSONResponse(&response)) + defer closeResponseBody(resp2) + if err != nil { + // All 422s returned as PermissionDenied. Callers in tables.go propagate this + // error via wrapError — there is no soft-skip at the List level, so returning + // ErrObjectNotFound here would have no observable effect on sync behavior. + if resp2 != nil && resp2.StatusCode == http.StatusUnprocessableEntity { + l.Debug("Insufficient privileges for SHOW TABLES IN SCHEMA (statement result)", + zap.String("database", databaseName), zap.String("schema", schemaName)) + wrappedErr := fmt.Errorf("baton-snowflake: insufficient privileges for SHOW TABLES IN SCHEMA %s.%s (statement result): %w", databaseName, schemaName, err) + return nil, "", status.Error(codes.PermissionDenied, wrappedErr.Error()) + } + return nil, "", err } - return nil, "", err } tables, err := response.ListTables() @@ -189,6 +215,8 @@ func (c *Client) ListTablesInSchema(ctx context.Context, databaseName, schemaNam } var nextCursor string + // >= not >: if exactly limit rows returned, there may be more beyond the cursor. + // Fewer than limit means this is definitively the last page. if limit > 0 && len(tables) >= limit { last := tables[len(tables)-1] nextCursor = last.Name @@ -233,18 +261,19 @@ func (c *Client) GetTable(ctx context.Context, database, schema, tableName strin defer closeResponseBody(resp1) if err != nil { if resp1 != nil && resp1.StatusCode == http.StatusUnprocessableEntity { - return nil, nil + // Any 422 is treated as object-not-found rather than decoding for code 003001. + // GetTable is only called as an owner fallback after ListTableGrants already succeeded, + // so a genuine RBAC denial here is extremely unlikely within the same request cycle. + // Even if it occurred, ErrObjectNotFound causes the caller to return partial grants + // with a Warn log — a tolerable degradation, not a silent failure. + // Note: the async path via fetchStatementResultIfAsync also has no 422 handling; + // both paths uniformly soft-skip on any 422 here. + return nil, ErrObjectNotFound } return nil, err } - req, err = c.GetStatementResponse(ctx, response.StatementHandle) - if err != nil { - return nil, err - } - resp2, err := c.Do(req, uhttp.WithJSONResponse(&response)) - defer closeResponseBody(resp2) - if err != nil { + if err := c.fetchStatementResultIfAsync(ctx, resp1, response.StatementHandle, &response); err != nil { return nil, err } @@ -260,7 +289,7 @@ func (c *Client) GetTable(ctx context.Context, database, schema, tableName strin } } - return nil, fmt.Errorf("table %s.%s.%s not found", database, schema, tableName) + return nil, fmt.Errorf("%w: %s.%s.%s", ErrObjectNotFound, database, schema, tableName) } var tableGrantStructFieldToColumnMap = map[string]string{ @@ -320,7 +349,11 @@ func tableGrantsCacheKey(database, schema, tableName, objectKind string) string func (c *Client) ListTableGrants(ctx context.Context, ss sessions.SessionStore, database, schema, tableName, objectKind string) ([]TableGrant, error) { cacheKey := tableGrantsCacheKey(database, schema, tableName, objectKind) if ss != nil { - if cached, found, err := session.GetJSON[[]TableGrant](ctx, ss, cacheKey, tableGrantsNamespace); err == nil && found { + cached, found, err := session.GetJSON[[]TableGrant](ctx, ss, cacheKey, tableGrantsNamespace) + if err != nil { + ctxzap.Extract(ctx).Debug("table grants cache lookup error, falling through to API", + zap.String("cache_key", cacheKey), zap.Error(err)) + } else if found { return cached, nil } } @@ -353,39 +386,56 @@ func (c *Client) ListTableGrants(ctx context.Context, ss sessions.SessionStore, return nil, fmt.Errorf("received 422 but failed to decode response body: %w (request error: %s)", decodeErr, err.Error()) } - // code: 003001 - // message: SQL access control error:\nInsufficient privileges tableRef := fmt.Sprintf("%s.%s.%s", database, schema, tableName) if errMsg.Code == "003001" { + // Genuine RBAC/privilege problem — keep as PermissionDenied so sync fails loudly l.Debug("Insufficient privileges to show grants on table", zap.String("table", tableRef)) - } else { - l.Error(errMsg.Message, zap.String("table", tableRef)) + return nil, status.Errorf(codes.PermissionDenied, "baton-snowflake: insufficient privileges to show grants on table %s: %s", tableRef, errMsg.Message) } - - return nil, status.Errorf(codes.PermissionDenied, "baton-snowflake: insufficient privileges to show grants on table %s: %s", tableRef, errMsg.Message) + // Any other 422: object dropped mid-sync or otherwise no longer accessible + l.Warn("Table no longer exists or not accessible (will soft-skip)", + zap.String("table", tableRef), + zap.String("snowflake_code", errMsg.Code), + zap.String("message", errMsg.Message)) + return nil, ErrObjectNotFound } return nil, err } - if resp != nil { - defer resp.Body.Close() - } + // Close POST response body explicitly — resp is reassigned below for the GET. + closeResponseBody(resp) - req, err = c.GetStatementResponse(ctx, response.StatementHandle) - if err != nil { - return nil, err - } - resp, err = c.Do(req, uhttp.WithJSONResponse(&response)) - if err != nil { - if resp != nil && resp.StatusCode == http.StatusUnprocessableEntity { - l.Debug("Insufficient privileges to show grants on table (statement result)", zap.String("table", fmt.Sprintf("%s.%s.%s", database, schema, tableName))) - wrappedErr := fmt.Errorf("baton-snowflake: insufficient privileges to show grants on table %s.%s.%s (statement result): %w", database, schema, tableName, err) - return nil, status.Error(codes.PermissionDenied, wrappedErr.Error()) + // closeResponseBody drains the body but does not nil resp; StatusCode is still readable. + if resp != nil && resp.StatusCode == http.StatusAccepted { + req, err = c.GetStatementResponse(ctx, response.StatementHandle) + if err != nil { + return nil, err + } + resp, err = c.Do(req, uhttp.WithJSONResponse(&response)) + defer closeResponseBody(resp) + if err != nil { + if resp != nil && resp.StatusCode == http.StatusUnprocessableEntity { + tableRef := fmt.Sprintf("%s.%s.%s", database, schema, tableName) + var errMsg struct { + Code string `json:"code"` + Message string `json:"message"` + } + // resp.Body is safe to read: uhttp.Do replaces the raw stream with + // bytes.NewBuffer(body) before returning, so WithJSONResponse and this + // decode draw from independent copies of the same bytes. + decodeErr := json.NewDecoder(resp.Body).Decode(&errMsg) + if decodeErr != nil { + l.Warn("Failed to decode 422 response body on async poll, treating as object-not-found", + zap.String("table", tableRef), zap.Error(decodeErr)) + } else if errMsg.Code == "003001" { + l.Debug("Insufficient privileges to show grants on table (statement result)", zap.String("table", tableRef)) + return nil, status.Errorf(codes.PermissionDenied, "baton-snowflake: insufficient privileges to show grants on table %s: %s", tableRef, errMsg.Message) + } + l.Warn("Table no longer exists during statement result fetch (will soft-skip)", zap.String("table", tableRef)) + return nil, ErrObjectNotFound + } + return nil, err } - return nil, err - } - if resp != nil { - defer resp.Body.Close() } grants, err := response.GetTableGrants() diff --git a/pkg/snowflake/table_test.go b/pkg/snowflake/table_test.go new file mode 100644 index 00000000..147aea7a --- /dev/null +++ b/pkg/snowflake/table_test.go @@ -0,0 +1,111 @@ +package snowflake + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// newTableGrantsTestClient starts an httptest server that returns a fixed HTTP +// status and JSON body for every request, then builds a Client pointing at it. +// The caller must call the returned cleanup func when done. +func newTableGrantsTestClient(t *testing.T, statusCode int, body string) (*Client, func()) { + t.Helper() + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(statusCode) + _, _ = w.Write([]byte(body)) + })) + client, err := New(ts.URL, JWTConfig{}, ts.Client()) + require.NoError(t, err) + return client, ts.Close +} + +func TestListTableGrants_422_PermissionDenied(t *testing.T) { + body := `{"code":"003001","message":"Insufficient privileges to operate on table"}` + client, cleanup := newTableGrantsTestClient(t, http.StatusUnprocessableEntity, body) + defer cleanup() + + _, err := client.ListTableGrants(context.Background(), nil, "MYDB", "PUBLIC", "MYTABLE", "TABLE") + + require.Error(t, err) + assert.Equal(t, codes.PermissionDenied, status.Code(err)) + assert.False(t, errors.Is(err, ErrObjectNotFound)) +} + +func TestListTableGrants_422_ObjectNotFound(t *testing.T) { + body := `{"code":"002003","message":"Object 'MYDB.PUBLIC.MYTABLE' does not exist or not authorized."}` + client, cleanup := newTableGrantsTestClient(t, http.StatusUnprocessableEntity, body) + defer cleanup() + + _, err := client.ListTableGrants(context.Background(), nil, "MYDB", "PUBLIC", "MYTABLE", "TABLE") + + require.Error(t, err) + assert.True(t, errors.Is(err, ErrObjectNotFound)) + assert.NotEqual(t, codes.PermissionDenied, status.Code(err)) +} + +func TestListTableGrants_200_Success(t *testing.T) { + body := `{"resultSetMetadata":{"numRows":0,"rowType":[]},"data":[],"statementHandle":""}` + client, cleanup := newTableGrantsTestClient(t, http.StatusOK, body) + defer cleanup() + + grants, err := client.ListTableGrants(context.Background(), nil, "MYDB", "PUBLIC", "MYTABLE", "TABLE") + + require.NoError(t, err) + assert.Empty(t, grants) +} + +// newAsyncTableGrantsTestClient starts an httptest server that simulates async execution: +// POST returns 202 with a statement handle; GET on /api/v2/statements/ returns asyncStatusCode and asyncBody. +func newAsyncTableGrantsTestClient(t *testing.T, asyncStatusCode int, asyncBody string) (*Client, func()) { + t.Helper() + mux := http.NewServeMux() + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusAccepted) + _, _ = w.Write([]byte(`{"statementHandle":"test-handle","resultSetMetadata":{"numRows":0,"rowType":[]},"data":[]}`)) + return + } + // GET /api/v2/statements/test-handle + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(asyncStatusCode) + _, _ = w.Write([]byte(asyncBody)) + }) + ts := httptest.NewServer(mux) + client, err := New(ts.URL, JWTConfig{}, ts.Client()) + require.NoError(t, err) + return client, ts.Close +} + +func TestListTableGrants_Async_422_PermissionDenied(t *testing.T) { + body := `{"code":"003001","message":"Insufficient privileges to operate on table"}` + client, cleanup := newAsyncTableGrantsTestClient(t, http.StatusUnprocessableEntity, body) + defer cleanup() + + _, err := client.ListTableGrants(context.Background(), nil, "MYDB", "PUBLIC", "MYTABLE", "TABLE") + + require.Error(t, err) + assert.Equal(t, codes.PermissionDenied, status.Code(err)) + assert.False(t, errors.Is(err, ErrObjectNotFound)) +} + +func TestListTableGrants_Async_422_ObjectNotFound(t *testing.T) { + body := `{"code":"002003","message":"Object 'MYDB.PUBLIC.MYTABLE' does not exist or not authorized."}` + client, cleanup := newAsyncTableGrantsTestClient(t, http.StatusUnprocessableEntity, body) + defer cleanup() + + _, err := client.ListTableGrants(context.Background(), nil, "MYDB", "PUBLIC", "MYTABLE", "TABLE") + + require.Error(t, err) + assert.True(t, errors.Is(err, ErrObjectNotFound)) + assert.NotEqual(t, codes.PermissionDenied, status.Code(err)) +} diff --git a/pkg/snowflake/user.go b/pkg/snowflake/user.go index 4693698d..25849ac8 100644 --- a/pkg/snowflake/user.go +++ b/pkg/snowflake/user.go @@ -11,6 +11,8 @@ import ( "github.com/conductorone/baton-sdk/pkg/session" "github.com/conductorone/baton-sdk/pkg/types/sessions" "github.com/conductorone/baton-sdk/pkg/uhttp" + "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" + "go.uber.org/zap" ) var ( @@ -205,13 +207,7 @@ func (c *Client) ListUsers(ctx context.Context, cursor string, limit int) ([]Use return nil, err } - req, err = c.GetStatementResponse(ctx, response.StatementHandle) - if err != nil { - return nil, err - } - resp2, err := c.Do(req, uhttp.WithJSONResponse(&response)) - defer closeResponseBody(resp2) - if err != nil { + if err := c.fetchStatementResultIfAsync(ctx, resp1, response.StatementHandle, &response); err != nil { return nil, err } @@ -241,7 +237,11 @@ func (c *Client) CacheUsers(ctx context.Context, ss sessions.SessionStore, users func (c *Client) GetUser(ctx context.Context, ss sessions.SessionStore, username string) (*User, int, error) { if ss != nil { - if cached, found, err := session.GetJSON[*User](ctx, ss, username, userNamespace); err == nil && found { + cached, found, err := session.GetJSON[*User](ctx, ss, username, userNamespace) + if err != nil { + ctxzap.Extract(ctx).Debug("user cache lookup error, falling through to API", + zap.String("username", username), zap.Error(err)) + } else if found { return cached, http.StatusOK, nil } }