From 49ca1515183dfb277b064265a6a999b32f61b0ba Mon Sep 17 00:00:00 2001
From: Jeremy Johnstone <j.johnstone@draftkings.com>
Date: Wed, 27 May 2026 16:04:27 -0400
Subject: [PATCH 1/3] fix: improve table sync resilience and refactor async
 polling

Introduces ErrObjectNotFound to distinguish RBAC denials (error code
003001 -> PermissionDenied) from tables dropped mid-sync (other 422s ->
soft-skip with warning). Entitlements() and Grants() now return partial
results instead of failing the entire sync when a table disappears
between List and Grants phases.

Extracts fetchStatementResultIfAsync() helper to deduplicate async
polling logic (Snowflake 202 Accepted pattern) across ListUsers,
ListAccountRoles, ListAccountRoleGrantees, ListSchemasInDatabase,
ListTablesInSchema, and GetDatabase.

Fixes: EIT-4535
---
 pkg/connector/tables.go       |  15 +++
 pkg/snowflake/account_role.go |  22 ++---
 pkg/snowflake/client.go       |  17 ++++
 pkg/snowflake/client_test.go  |  27 ++++++
 pkg/snowflake/table.go        | 166 ++++++++++++++++++++++------------
 pkg/snowflake/table_test.go   | 111 +++++++++++++++++++++++
 pkg/snowflake/user.go         |  16 ++--
 7 files changed, 293 insertions(+), 81 deletions(-)
 create mode 100644 pkg/snowflake/client_test.go
 create mode 100644 pkg/snowflake/table_test.go

diff --git a/pkg/connector/tables.go b/pkg/connector/tables.go
index 4fbdf71a..cd0016d3 100644
--- a/pkg/connector/tables.go
+++ b/pkg/connector/tables.go
@@ -277,6 +277,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))
 	}
 
@@ -317,6 +322,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 +401,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" {
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/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/<handle> 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
 		}
 	}

From 9f9a1a7fc3ee9f831688d6ed8ba61110929e0d5a Mon Sep 17 00:00:00 2001
From: Jeremy Johnstone <j.johnstone@draftkings.com>
Date: Wed, 27 May 2026 16:04:52 -0400
Subject: [PATCH 2/3] fix: add database caching to reduce redundant Snowflake
 queries

Adds CacheDatabases() to bulk-seed the session cache after List(),
and updates GetDatabase() to check cache before making an API call.
Cache errors fall through to the API rather than failing the call.

Also adds CATALOG-LINKED DATABASE to IsSharedOrSystem() kind check
and plumbs SessionStore through the table builder's isDBSharedOrSystem()
helper to enable cache hits during table grant evaluation.

Fixes: EIT-4535
---
 pkg/connector/databases.go     |  23 +++-
 pkg/snowflake/database.go      |  64 +++++++---
 pkg/snowflake/database_test.go | 207 +++++++++++++++++++++++++++++++++
 3 files changed, 276 insertions(+), 18 deletions(-)
 create mode 100644 pkg/snowflake/database_test.go

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

From e0d6fddbfd4d19f4bb860879cf1f6f6a0a93318c Mon Sep 17 00:00:00 2001
From: Jeremy Johnstone <j.johnstone@draftkings.com>
Date: Wed, 27 May 2026 16:05:29 -0400
Subject: [PATCH 3/3] feat: add --sync-tables flag to skip table
 synchronization

Adds an optional --sync-tables boolean flag (default: true). When set
to false, tableBuilder.List() returns early and databaseResource() omits
the child table resource type annotation, skipping all table sync work.

Useful for accounts with very large numbers of tables where table
grants are not needed, significantly reducing sync duration and output
size.

Fixes: EIT-4535
---
 README.md                    |  16 ++
 docs/connector.mdx           |   6 +
 pkg/config/conf.gen.go       |   1 +
 pkg/config/config.go         |   7 +
 pkg/connector/connector.go   |  12 +-
 pkg/connector/tables.go      |  31 +++-
 pkg/connector/tables_test.go | 306 +++++++++++++++++++++++++++++++++++
 7 files changed, 364 insertions(+), 15 deletions(-)

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.
 </Step>
 <Step>
+**Optional.** Disable **Sync Tables** to skip syncing Snowflake tables entirely. This significantly reduces sync size and duration for accounts with large numbers of tables. Tables are synced by default.
+</Step>
+<Step>
 Click **Save**.
 </Step>
 <Step>
@@ -216,6 +219,9 @@ stringData:
 
   # Optional: comma-separated list of database names to exclude from sync (case-insensitive)
   # BATON_EXCLUDED_DATABASES: "MY_DB,ANOTHER_DB"
+
+  # Optional: set to false to skip syncing tables (reduces sync size significantly)
+  # BATON_SYNC_TABLES: false
 ```
 
 See the connector's README or run `--help` to see all available configuration flags and environment variables.
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/tables.go b/pkg/connector/tables.go
index cd0016d3..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
 	}
@@ -306,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
 	}
@@ -430,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)
+}