Make opening driver connections cancellable via a ctx

admin/testadmin/testadmin.go

@@ -134,7 +134,7 @@ func NewWithOptionalRuntime(t *testing.T, startRt bool) *Fixture {
 	}))
 
 	// Initialize mock AI using drivers.Open pattern
-	mockAIHandle, err := drivers.Open("mock_ai", "", "test", map[string]any{}, storage.MustNew(os.TempDir(), nil), activity.NewNoopClient(), logger)
+	mockAIHandle, err := drivers.Open(t.Context(), "mock_ai", "", "test", map[string]any{}, storage.MustNew(os.TempDir(), nil), activity.NewNoopClient(), logger)
 	require.NoError(t, err)
 	t.Cleanup(func() { mockAIHandle.Close() })
 	mockAI, ok := mockAIHandle.AsAI("test")

cli/cmd/admin/start.go

@@ -272,7 +272,7 @@ func StartCmd(ch *cmdutil.Helper) *cobra.Command {
 			default:
 				logger.Fatal("unknown AI driver", zap.String("driver", aiDriver))
 			}
-			aiHandle, err := drivers.Open(aiDriver, "", "", aiConfig, rillstorage.MustNew(os.TempDir(), nil), activity.NewNoopClient(), logger)
+			aiHandle, err := drivers.Open(cmd.Context(), aiDriver, "", "", aiConfig, rillstorage.MustNew(os.TempDir(), nil), activity.NewNoopClient(), logger)
 			if err != nil {
 				logger.Fatal("error creating AI client", zap.Error(err))
 			}

cli/cmd/runtime/install_duckdb_extensions.go

@@ -19,7 +19,7 @@ func InstallDuckDBExtensionsCmd(ch *cmdutil.Helper) *cobra.Command {
 		Use: "install-duckdb-extensions",
 		RunE: func(cmd *cobra.Command, args []string) error {
 			cfg := map[string]any{"dsn": ":memory:"} // In-memory
-			h, err := drivers.Open("duckdb", "", "default", cfg, storage.MustNew(os.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
+			h, err := drivers.Open(cmd.Context(), "duckdb", "", "default", cfg, storage.MustNew(os.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
 			if err != nil {
 				return fmt.Errorf("failed to open ephemeral duckdb: %w", err)
 			}

cli/pkg/cmdutil/project.go

@@ -16,7 +16,7 @@ import (
 // RepoForProjectPath creates an ad-hoc drivers.RepoStore for a local project file path
 func RepoForProjectPath(path string) (drivers.RepoStore, string, error) {
 	instanceID := "default"
-	repoHandle, err := drivers.Open("file", "", instanceID, map[string]any{"dsn": path}, storage.MustNew(os.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
+	repoHandle, err := drivers.Open(context.Background(), "file", "", instanceID, map[string]any{"dsn": path}, storage.MustNew(os.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
 	if err != nil {
 		return nil, "", err
 	}

cli/pkg/local/admin.go

@@ -165,7 +165,7 @@ func (l *localAdminService) HasAnonymousSourceAccess(ctx context.Context, srcPro
 }
 
 // Open implements [drivers.Driver].
-func (l *localAdminService) Open(connectorName, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
+func (l *localAdminService) Open(_ context.Context, connectorName, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
 	return l, nil
 }
 

runtime/connection_cache.go

@@ -166,9 +166,15 @@ func (r *Runtime) openAndMigrate(ctx context.Context, cfg cachedConnectionConfig
 	}
 
 	r.Logger.Debug("opening connection", zap.String("instance_id", cfg.instanceID), zap.String("driver", cfg.driver), zap.String("name", cfg.name), zap.Bool("provision", cfg.provision))
-	handle, err := drivers.Open(cfg.driver, cfg.name, cfg.instanceID, cfg.config, storage, activityClient, logger)
-	if err == nil && ctx.Err() != nil {
-		err = fmt.Errorf("timed out while opening driver %q", cfg.driver)
+	handle, err := drivers.Open(ctx, cfg.driver, cfg.name, cfg.instanceID, cfg.config, storage, activityClient, logger)
+	if ctx.Err() != nil {
+		if err == nil {
+			// The driver didn't observe the cancellation and opened anyway. Discard the handle since we're returning an error.
+			handle.Close()
+			err = fmt.Errorf("timed out while opening driver %q", cfg.driver)
+		} else {
+			err = fmt.Errorf("timed out while opening driver %q: %w", cfg.driver, err)
+		}
 	}
 	r.activity.Record(ctx, activity.EventTypeLog, "connection_open",
 		attribute.String("instance_id", cfg.instanceID),

runtime/drivers/admin/admin.go

@@ -48,7 +48,7 @@ type configProperties struct {
 	ProjectID   string `mapstructure:"project_id"`
 }
 
-func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
+func (d driver) Open(_ context.Context, _, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
 	if instanceID == "" {
 		return nil, errors.New("admin driver can't be shared")
 	}

runtime/drivers/athena/athena.go

@@ -80,7 +80,7 @@ type configProperties struct {
 	AllowHostAccess bool   `mapstructure:"allow_host_access"`
 }
 
-func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
+func (d driver) Open(_ context.Context, _, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
 	if instanceID == "" {
 		return nil, errors.New("athena driver can't be shared")
 	}

runtime/drivers/athena/olap_test.go

@@ -457,7 +457,7 @@ func TestExec(t *testing.T) {
 func acquireTestAthena(t *testing.T) (drivers.Handle, drivers.OLAPStore) {
 	cfg := testruntime.AcquireConnector(t, "athena")
 	cfg["output_location"] = "s3://integration-test.rilldata.com/athena/"
-	conn, err := drivers.Open("athena", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
+	conn, err := drivers.Open(t.Context(), "athena", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
 	require.NoError(t, err)
 	t.Cleanup(func() { conn.Close() })
 

runtime/drivers/azure/azure.go

@@ -135,7 +135,7 @@ func (c *ConfigProperties) GetConnectionString() string {
 	return dsn
 }
 
-func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
+func (d driver) Open(_ context.Context, _, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
 	if instanceID == "" {
 		return nil, errors.New("azure driver can't be shared")
 	}

runtime/drivers/azure/object_store_test.go

@@ -16,7 +16,7 @@ func TestObjectStoreCloud(t *testing.T) {
 	testmode.Expensive(t)
 	// using azure cloud for these test because azurite does not support startFrom
 	cfg := testruntime.AcquireConnector(t, "azure_cloud")
-	conn, err := drivers.Open("azure", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
+	conn, err := drivers.Open(t.Context(), "azure", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
 	require.NoError(t, err)
 	t.Cleanup(func() { conn.Close() })
 
@@ -43,7 +43,7 @@ func TestObjectStoreCloud(t *testing.T) {
 
 func TestObjectStore(t *testing.T) {
 	cfg := testruntime.AcquireConnector(t, "azure")
-	conn, err := drivers.Open("azure", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
+	conn, err := drivers.Open(t.Context(), "azure", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
 	require.NoError(t, err)
 	t.Cleanup(func() { conn.Close() })
 
@@ -65,7 +65,7 @@ func TestObjectStore(t *testing.T) {
 func TestObjectStorePathPrefixes(t *testing.T) {
 	cfg := testruntime.AcquireConnector(t, "azure")
 	cfg["path_prefixes"] = "azure://integration-test/glob_test/"
-	conn, err := drivers.Open("azure", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
+	conn, err := drivers.Open(t.Context(), "azure", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
 	require.NoError(t, err)
 	t.Cleanup(func() { conn.Close() })
 

runtime/drivers/bigquery/bigquery.go

@@ -67,7 +67,7 @@ type configProperties struct {
 	LogQueries bool `mapstructure:"log_queries"`
 }
 
-func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
+func (d driver) Open(_ context.Context, _, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
 	if instanceID == "" {
 		return nil, errors.New("bigquery driver can't be shared")
 	}
@@ -78,10 +78,14 @@ func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Cl
 		return nil, err
 	}
 
+	// Note the handle's ctx tracks the handle's lifetime, so it must not derive from the ctx passed to Open.
+	bgctx, cancel := context.WithCancel(context.Background())
 	conn := &Connection{
 		config:   conf,
 		storage:  st,
 		logger:   logger,
+		ctx:      bgctx,
+		cancel:   cancel,
 		clientMu: semaphore.NewWeighted(1),
 	}
 	return conn, nil
@@ -105,6 +109,12 @@ type Connection struct {
 	storage *storage.Client
 	logger  *zap.Logger
 
+	// ctx tracks the handle's lifetime and is cancelled by Close.
+	// The cached client's credentials capture it and reuse it to refresh tokens,
+	// so it must outlive any single request that happens to open the client.
+	ctx    context.Context
+	cancel context.CancelFunc
+
 	client    *bigquery.Client // lazily populated using getClient
 	clientErr error
 	clientMu  *semaphore.Weighted
@@ -114,11 +124,11 @@ var _ drivers.Handle = &Connection{}
 
 // Ping implements drivers.Handle.
 func (c *Connection) Ping(ctx context.Context) error {
+	// NOTE: Not closing the client here since getClient caches it for the handle's lifetime. Close does that.
 	client, err := c.getClient(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to create client: %w", err)
 	}
-	defer client.Close()
 
 	// Run a simple query to verify connection
 	q := client.Query("SELECT 1")
@@ -144,6 +154,7 @@ func (c *Connection) Config() map[string]any {
 
 // Close implements drivers.Connection.
 func (c *Connection) Close() error {
+	c.cancel()
 	if c.client != nil {
 		return c.client.Close()
 	}
@@ -244,7 +255,9 @@ func (c *Connection) getClient(ctx context.Context) (*bigquery.Client, error) {
 	if c.client != nil || c.clientErr != nil {
 		return c.client, c.clientErr
 	}
-	client, err := c.createClient(ctx, "")
+	// The client is cached for the handle's lifetime, so it's created with the handle's ctx instead of the caller's.
+	// Its credentials reuse that ctx to refresh tokens, which a cancelled request ctx would permanently break.
+	client, err := c.createClient(c.ctx, "")
 	if err != nil {
 		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
 			// don't cache context errors

runtime/drivers/bigquery/olap_test.go

@@ -393,7 +393,7 @@ func TestAllDatatypesRowCount(t *testing.T) {
 
 func acquireTestBigQuery(t *testing.T) (drivers.Handle, drivers.OLAPStore) {
 	cfg := testruntime.AcquireConnector(t, "bigquery")
-	conn, err := drivers.Open("bigquery", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
+	conn, err := drivers.Open(t.Context(), "bigquery", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
 	require.NoError(t, err)
 	t.Cleanup(func() { conn.Close() })
 

runtime/drivers/claude/claude.go

@@ -83,7 +83,7 @@ func (d driver) Spec() drivers.Spec {
 }
 
 // Open implements drivers.Driver.
-func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
+func (d driver) Open(_ context.Context, _, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
 	conf := &configProperties{}
 	err := mapstructure.WeakDecode(config, conf)
 	if err != nil {

runtime/drivers/clickhouse/clickhouse.go

@@ -221,7 +221,7 @@ func (c *configProperties) validate() error {
 
 // Open connects to Clickhouse using std API.
 // Connection string format : https://github.com/ClickHouse/clickhouse-go?tab=readme-ov-file#dsn
-func (d driver) Open(connectorName, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
+func (d driver) Open(ctx context.Context, connectorName, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {
 	if instanceID == "" {
 		return nil, errors.New("clickhouse driver can't be shared")
 	}
@@ -296,7 +296,7 @@ func (d driver) Open(connectorName, instanceID string, config map[string]any, st
 		if err != nil {
 			return nil, err
 		}
-		opts, err = embed.start()
+		opts, err = embed.start(ctx)
 		if err != nil {
 			return nil, err
 		}
@@ -305,7 +305,7 @@ func (d driver) Open(connectorName, instanceID string, config map[string]any, st
 	}
 
 	// Open the main database connection
-	db, err := openHandle(instanceID, conf, opts, logger)
+	db, err := openHandle(ctx, instanceID, conf, opts, logger)
 	if err != nil {
 		return nil, err
 	}
@@ -317,13 +317,28 @@ func (d driver) Open(connectorName, instanceID string, config map[string]any, st
 		if err != nil {
 			return nil, fmt.Errorf("failed to parse write DSN: %w", err)
 		}
-		writeDB, err = openHandle(instanceID, conf, writeOpts, logger)
+		writeDB, err = openHandle(ctx, instanceID, conf, writeOpts, logger)
 		if err != nil {
 			return nil, fmt.Errorf("failed to open write connection: %w", err)
 		}
 	}
+	// Using the harmless, non–side-effecting setting
+	// `show_table_uuid_in_table_create_query_if_not_nil` as a probe to check
+	// whether the cluster mode supports modifying query settings. This setting
+	// has no practical use for our purposes.
+	// It runs before the other statements below because they need to know the result.
+	// The probe itself is safe to run on the raw ctx:
+	// if the ctx has a deadline, clickhouse-go derives a `max_execution_time` setting from it,
+	// which a readonly server rejects with the same error we're matching on here.
+	supportSettings := true
+	if _, err := db.ExecContext(ctx, "SET show_table_uuid_in_table_create_query_if_not_nil = 1"); err != nil {
+		if strings.Contains(err.Error(), "Cannot modify") && strings.Contains(err.Error(), "setting in readonly mode") {
+			supportSettings = false
+		}
+	}
+
 	// group by positional args are supported post 22.7 and we use them heavily in our queries
-	row := db.QueryRow(`
+	row := db.QueryRowContext(settingsCtx(ctx, supportSettings), `
         WITH
             splitByChar('.', version()) AS parts,
             toInt32(parts[1]) AS major,
@@ -339,17 +354,6 @@ func (d driver) Open(connectorName, instanceID string, config map[string]any, st
 		return nil, fmt.Errorf("clickhouse version must be 22.7 or higher")
 	}
 
-	// Using the harmless, non–side-effecting setting
-	// `show_table_uuid_in_table_create_query_if_not_nil` as a probe to check
-	// whether the cluster mode supports modifying query settings. This setting
-	// has no practical use for our purposes.
-	supportSettings := true
-	if _, err := db.Exec("SET show_table_uuid_in_table_create_query_if_not_nil = 1"); err != nil {
-		if strings.Contains(err.Error(), "Cannot modify") && strings.Contains(err.Error(), "setting in readonly mode") {
-			supportSettings = false
-		}
-	}
-
 	// Compute OLAP queue size
 	var olapSemSize int
 	if conf.MaxOpenConns < 1 {
@@ -364,7 +368,8 @@ func (d driver) Open(connectorName, instanceID string, config map[string]any, st
 		olapSemSize = 1
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	// Note the handle's ctx tracks the handle's lifetime, so it must not derive from the ctx passed to Open.
+	bgctx, cancel := context.WithCancel(context.Background())
 	c := &Connection{
 		readDB:          db,
 		writeDB:         writeDB,
@@ -374,7 +379,7 @@ func (d driver) Open(connectorName, instanceID string, config map[string]any, st
 		instanceID:      instanceID,
 		connectorName:   connectorName,
 		supportSettings: supportSettings,
-		ctx:             ctx,
+		ctx:             bgctx,
 		cancel:          cancel,
 		metaSem:         semaphore.NewWeighted(1),
 		olapSem:         priorityqueue.NewSemaphore(olapSemSize),
@@ -877,7 +882,7 @@ func (c *Connection) checkBillingTableExists(ctx context.Context, cluster string
 	return existsEverywhere, nil
 }
 
-func openHandle(instanceID string, conf *configProperties, opts *clickhouse.Options, logger *zap.Logger) (*sqlx.DB, error) {
+func openHandle(ctx context.Context, instanceID string, conf *configProperties, opts *clickhouse.Options, logger *zap.Logger) (*sqlx.DB, error) {
 	// Apply certain options from conf that are not set in the DSN.
 	if conf.MaxIdleConns != 0 {
 		opts.MaxIdleConns = conf.MaxIdleConns
@@ -917,7 +922,8 @@ func openHandle(instanceID string, conf *configProperties, opts *clickhouse.Opti
 	// It prevents invalid host/port combinations from proceeding to db.Ping, which uses a longer timeout to handle scale-to-zero scenarios.
 	if conf.Host != "" && conf.Port != 0 {
 		target := net.JoinHostPort(conf.Host, fmt.Sprintf("%d", conf.Port))
-		conn, err := net.DialTimeout("tcp", target, 10*time.Second)
+		dialer := net.Dialer{Timeout: 10 * time.Second}
+		conn, err := dialer.DialContext(ctx, "tcp", target)
 		if err != nil {
 			return nil, fmt.Errorf("please check that the host and port are correct %s: %w", target, err)
 		}
@@ -929,7 +935,7 @@ func openHandle(instanceID string, conf *configProperties, opts *clickhouse.Opti
 
 	// Open the connection
 	db := sqlx.NewDb(otelsql.OpenDB(clickhouse.Connector(opts)), "clickhouse")
-	err := db.Ping()
+	err := db.PingContext(ctx)
 	if err != nil {
 		// Detect SSL/TLS mismatch (common causes: "read: EOF" or TLS Alert [21])
 		if strings.Contains(err.Error(), "EOF") ||
@@ -947,7 +953,7 @@ func openHandle(instanceID string, conf *configProperties, opts *clickhouse.Opti
 		// may be the port is http, also try with http protocol if DSN is not provided
 		opts.Protocol = clickhouse.HTTP
 		db = sqlx.NewDb(otelsql.OpenDB(clickhouse.Connector(opts)), "clickhouse")
-		err := db.Ping()
+		err := db.PingContext(ctx)
 		if err != nil {
 			// Detect SSL/TLS mismatch (common causes: "read: EOF" or  \x15 means TLS Alert [21]"])
 			if strings.Contains(err.Error(), "EOF") ||