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
}
}