From 781b98878e2348042e46705575065bdae5233e7c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Benjamin=20Egelund-M=C3=BCller?= Date: Mon, 21 Sep 2026 16:13:47 +0200 Subject: [PATCH 1/5] Make opening driver connections cancellable via a `ctx` --- admin/testadmin/testadmin.go | 2 +- cli/cmd/admin/start.go | 2 +- cli/cmd/runtime/install_duckdb_extensions.go | 2 +- cli/pkg/cmdutil/project.go | 2 +- cli/pkg/local/admin.go | 2 +- runtime/connection_cache.go | 12 +++-- runtime/drivers/admin/admin.go | 2 +- runtime/drivers/athena/athena.go | 2 +- runtime/drivers/athena/olap_test.go | 2 +- runtime/drivers/azure/azure.go | 2 +- runtime/drivers/azure/object_store_test.go | 7 +-- runtime/drivers/bigquery/bigquery.go | 2 +- runtime/drivers/bigquery/olap_test.go | 2 +- runtime/drivers/claude/claude.go | 2 +- runtime/drivers/clickhouse/clickhouse.go | 26 +++++----- runtime/drivers/clickhouse/embed.go | 49 ++++++++++++------- .../clickhouse/information_schema_test.go | 4 +- .../clickhouse/model_executor_self_test.go | 3 +- runtime/drivers/clickhouse/olap_test.go | 30 ++++++------ runtime/drivers/databricks/databricks.go | 2 +- runtime/drivers/databricks/olap_rt_test.go | 3 +- runtime/drivers/databricks/olap_test.go | 3 +- runtime/drivers/drivers.go | 9 ++-- runtime/drivers/drivers_test.go | 2 +- runtime/drivers/druid/druid.go | 10 ++-- runtime/drivers/druid/druid_test.go | 2 +- .../drivers/druid/information_schema_test.go | 2 +- runtime/drivers/druid/sql_driver_test.go | 10 ++-- runtime/drivers/duckdb/config_test.go | 2 +- runtime/drivers/duckdb/crud_test.go | 18 +++---- runtime/drivers/duckdb/duckdb.go | 13 ++--- runtime/drivers/duckdb/duckdb_test.go | 10 ++-- .../duckdb/model_executor_https_self_test.go | 4 +- .../duckdb/model_executor_self_file_test.go | 4 +- .../model_executor_self_objectstore_test.go | 2 +- .../duckdb/model_executor_self_sqlite_test.go | 2 +- .../duckdb/model_executor_self_test.go | 2 +- ...model_executor_sqlstore_self_mysql_test.go | 12 ++--- ...el_executor_sqlstore_self_postgres_test.go | 12 ++--- runtime/drivers/duckdb/motherduck_test.go | 5 +- runtime/drivers/duckdb/olap_test.go | 4 +- runtime/drivers/duckdb/schema_change_test.go | 2 +- runtime/drivers/file/file.go | 4 +- .../file/model_executor_olap_self_test.go | 4 +- runtime/drivers/gcs/gcs.go | 4 +- runtime/drivers/gcs/object_store_test.go | 9 ++-- runtime/drivers/gemini/gemini.go | 4 +- runtime/drivers/https/https.go | 2 +- runtime/drivers/mock/ai/ai.go | 2 +- .../drivers/mock/object_store/object_store.go | 4 +- runtime/drivers/mysql/mysql.go | 2 +- runtime/drivers/mysql/olap_test.go | 3 +- runtime/drivers/openai/openai.go | 2 +- .../drivers/pinot/information_schema_test.go | 2 +- runtime/drivers/pinot/pinot.go | 4 +- runtime/drivers/postgres/olap_test.go | 3 +- runtime/drivers/postgres/postgres.go | 2 +- runtime/drivers/redshift/olap_test.go | 2 +- runtime/drivers/redshift/redshift.go | 2 +- .../s3/mode_executor_olap_to_self_test.go | 2 +- runtime/drivers/s3/object_store_test.go | 5 +- runtime/drivers/s3/s3.go | 2 +- runtime/drivers/salesforce/salesforce.go | 2 +- runtime/drivers/slack/slack.go | 2 +- runtime/drivers/snowflake/olap_test.go | 3 +- runtime/drivers/snowflake/snowflake.go | 2 +- runtime/drivers/sqlite/backups_test.go | 5 +- runtime/drivers/sqlite/migrate_test.go | 3 +- runtime/drivers/sqlite/restore_test.go | 7 +-- runtime/drivers/sqlite/sqlite.go | 11 +++-- runtime/drivers/starrocks/olap_test.go | 2 +- runtime/drivers/starrocks/starrocks.go | 13 +++-- runtime/parser/parser_test.go | 2 +- runtime/server/queries_test.go | 2 +- runtime/storage/storage.go | 7 ++- runtime/testruntime/noopadmin.go | 2 +- 76 files changed, 231 insertions(+), 186 deletions(-) diff --git a/admin/testadmin/testadmin.go b/admin/testadmin/testadmin.go index a0d0f6429e81..e3e063ff46d6 100644 --- a/admin/testadmin/testadmin.go +++ b/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(context.Background(), "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") diff --git a/cli/cmd/admin/start.go b/cli/cmd/admin/start.go index f8521bf3cdd2..eae8a2666339 100644 --- a/cli/cmd/admin/start.go +++ b/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)) } diff --git a/cli/cmd/runtime/install_duckdb_extensions.go b/cli/cmd/runtime/install_duckdb_extensions.go index e447e887ab58..9049a493390b 100644 --- a/cli/cmd/runtime/install_duckdb_extensions.go +++ b/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) } diff --git a/cli/pkg/cmdutil/project.go b/cli/pkg/cmdutil/project.go index b2e0eea72fca..0df9531d9af6 100644 --- a/cli/pkg/cmdutil/project.go +++ b/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 } diff --git a/cli/pkg/local/admin.go b/cli/pkg/local/admin.go index a19cfb9850c5..ae638ef1200a 100644 --- a/cli/pkg/local/admin.go +++ b/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 } diff --git a/runtime/connection_cache.go b/runtime/connection_cache.go index 0c5126d93eee..a3bb8f77be2d 100644 --- a/runtime/connection_cache.go +++ b/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), diff --git a/runtime/drivers/admin/admin.go b/runtime/drivers/admin/admin.go index 04fd766e5721..f6adf4742ad5 100644 --- a/runtime/drivers/admin/admin.go +++ b/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") } diff --git a/runtime/drivers/athena/athena.go b/runtime/drivers/athena/athena.go index d0a4a2e3b492..c469b54d3d50 100644 --- a/runtime/drivers/athena/athena.go +++ b/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") } diff --git a/runtime/drivers/athena/olap_test.go b/runtime/drivers/athena/olap_test.go index 28bca597fc63..5813462db830 100644 --- a/runtime/drivers/athena/olap_test.go +++ b/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(context.Background(), "athena", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/azure/azure.go b/runtime/drivers/azure/azure.go index 228ab0f2cde1..dfa5370dd7a8 100644 --- a/runtime/drivers/azure/azure.go +++ b/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") } diff --git a/runtime/drivers/azure/object_store_test.go b/runtime/drivers/azure/object_store_test.go index 3e5c44fcefc5..d98c9e7c9d97 100644 --- a/runtime/drivers/azure/object_store_test.go +++ b/runtime/drivers/azure/object_store_test.go @@ -1,6 +1,7 @@ package azure_test import ( + "context" "testing" "github.com/rilldata/rill/runtime/drivers" @@ -16,7 +17,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(context.Background(), "azure", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) @@ -43,7 +44,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(context.Background(), "azure", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) @@ -65,7 +66,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(context.Background(), "azure", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/bigquery/bigquery.go b/runtime/drivers/bigquery/bigquery.go index e06ad0ec6e59..cfe2836590c1 100644 --- a/runtime/drivers/bigquery/bigquery.go +++ b/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") } diff --git a/runtime/drivers/bigquery/olap_test.go b/runtime/drivers/bigquery/olap_test.go index a003847d1c21..f620606ba453 100644 --- a/runtime/drivers/bigquery/olap_test.go +++ b/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(context.Background(), "bigquery", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/claude/claude.go b/runtime/drivers/claude/claude.go index 4b56eb5dd654..52ec259ca5f4 100644 --- a/runtime/drivers/claude/claude.go +++ b/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 { diff --git a/runtime/drivers/clickhouse/clickhouse.go b/runtime/drivers/clickhouse/clickhouse.go index 710b5eab47e0..919e337dc828 100644 --- a/runtime/drivers/clickhouse/clickhouse.go +++ b/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,13 @@ 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) } } // group by positional args are supported post 22.7 and we use them heavily in our queries - row := db.QueryRow(` + row := db.QueryRowContext(ctx, ` WITH splitByChar('.', version()) AS parts, toInt32(parts[1]) AS major, @@ -344,7 +344,7 @@ func (d driver) Open(connectorName, instanceID string, config map[string]any, st // 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 _, 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 } @@ -364,7 +364,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 +375,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 +878,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 +918,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 +931,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 +949,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") || diff --git a/runtime/drivers/clickhouse/embed.go b/runtime/drivers/clickhouse/embed.go index f98e8e823eba..e163bf92686c 100644 --- a/runtime/drivers/clickhouse/embed.go +++ b/runtime/drivers/clickhouse/embed.go @@ -60,10 +60,8 @@ func newEmbedClickHouse(tcpPort int, dataDir, tempDir string, logger *zap.Logger // start installs (depending on OS and platform) and starts ClickHouse server. // The destination directory for the ClickHouse binary is .rill/clickhouse. // The function returns the DSN for the ClickHouse server and close function. -// -// TODO: Since this can be a long-running process, we should accept a `ctx`, -// but the `drivers.Open` function currently doesn't propagate that. -func (e *embedClickHouse) start() (*clickhouse.Options, error) { +// The ctx cancels the install and the wait for the server to become ready. +func (e *embedClickHouse) start(ctx context.Context) (*clickhouse.Options, error) { e.mu.Lock() defer e.mu.Unlock() if e.opts != nil { @@ -77,7 +75,7 @@ func (e *embedClickHouse) start() (*clickhouse.Options, error) { return nil, err } - binPath, err := e.install(destDir, e.logger) + binPath, err := e.install(ctx, destDir, e.logger) if err != nil { return nil, err } @@ -95,7 +93,7 @@ func (e *embedClickHouse) start() (*clickhouse.Options, error) { } // Start the ClickHouse server - stderr, err := e.startClickhouse(binPath, configPath) + stderr, err := e.startClickhouse(ctx, binPath, configPath) if err != nil { if !errors.Is(err, errAlreadyRunning) { return nil, err @@ -105,7 +103,7 @@ func (e *embedClickHouse) start() (*clickhouse.Options, error) { if err != nil { return nil, fmt.Errorf("failed to kill existing ClickHouse process: %w", err) } - stderr, err = e.startClickhouse(binPath, configPath) + stderr, err = e.startClickhouse(ctx, binPath, configPath) if err != nil { return nil, err } @@ -164,7 +162,7 @@ func (e *embedClickHouse) start() (*clickhouse.Options, error) { return e.opts, nil } -func (e *embedClickHouse) startClickhouse(binPath, configPath string) (io.ReadCloser, error) { +func (e *embedClickHouse) startClickhouse(ctx context.Context, binPath, configPath string) (io.ReadCloser, error) { e.cmd = exec.Command(binPath, "server", "--config-file", configPath) e.cmd.Stdout = io.Discard @@ -175,7 +173,7 @@ func (e *embedClickHouse) startClickhouse(binPath, configPath string) (io.ReadCl ready := make(chan error, 1) go func() { - err := e.startAndWaitUntilReady(stderr) + err := e.startAndWaitUntilReady(ctx, stderr) ready <- err if err != nil && e.cmd != nil && e.cmd.Process != nil { _ = e.cmd.Process.Kill() @@ -183,8 +181,18 @@ func (e *embedClickHouse) startClickhouse(binPath, configPath string) (io.ReadCl } }() - if err := <-ready; err != nil { - return nil, err + select { + case err := <-ready: + if err != nil { + return nil, err + } + case <-ctx.Done(): + // Kill the server so it doesn't outlive the cancelled open. + // The goroutine above may also kill it, which is harmless. + if e.cmd != nil && e.cmd.Process != nil { + _ = e.cmd.Process.Kill() + } + return nil, ctx.Err() } return stderr, nil } @@ -217,7 +225,7 @@ func (e *embedClickHouse) stop() error { return nil } -func (e *embedClickHouse) install(destDir string, logger *zap.Logger) (string, error) { +func (e *embedClickHouse) install(ctx context.Context, destDir string, logger *zap.Logger) (string, error) { release := "v" + embedVersion + "-stable" destPath := filepath.Join(destDir, embedVersion, "clickhouse") @@ -252,7 +260,7 @@ func (e *embedClickHouse) install(destDir string, logger *zap.Logger) (string, e } url := "https://github.com/ClickHouse/ClickHouse/releases/download/" + release + "/" + fileName logger.Info("Downloading ClickHouse binary", zap.String("url", url), zap.String("dst", destPath)) - if err := downloadFile(destPath, url); err != nil { + if err := downloadFile(ctx, destPath, url); err != nil { return "", fmt.Errorf("error downloading ClickHouse: %w", err) } case "linux": @@ -268,7 +276,7 @@ func (e *embedClickHouse) install(destDir string, logger *zap.Logger) (string, e url := "https://github.com/ClickHouse/ClickHouse/releases/download/" + release + "/" + fileName destTgzPath := filepath.Join(destDir, release, fileName) logger.Info("Downloading ClickHouse binary", zap.String("url", url), zap.String("dst", destTgzPath)) - if err := downloadFile(destTgzPath, url); err != nil { + if err := downloadFile(ctx, destTgzPath, url); err != nil { return "", fmt.Errorf("error downloading ClickHouse: %w", err) } fileToExtract := filepath.Join("clickhouse-common-static-"+embedVersion, "usr", "bin", "clickhouse") @@ -394,7 +402,7 @@ func (e *embedClickHouse) getConfigContent() ([]byte, error) { return config, nil } -func (e *embedClickHouse) startAndWaitUntilReady(stderr io.Reader) error { +func (e *embedClickHouse) startAndWaitUntilReady(ctx context.Context, stderr io.Reader) error { if err := e.cmd.Start(); err != nil { return fmt.Errorf("failed to start clickhouse: %w", err) } @@ -408,6 +416,8 @@ func (e *embedClickHouse) startAndWaitUntilReady(stderr io.Reader) error { select { case <-timer.C: return fmt.Errorf("clickhouse is not ready: timeout") + case <-ctx.Done(): + return ctx.Err() default: if !scanner.Scan() { if scanner.Err() != nil { @@ -434,7 +444,7 @@ func (e *embedClickHouse) startAndWaitUntilReady(stderr io.Reader) error { } } -func downloadFile(path, url string) error { +func downloadFile(ctx context.Context, path, url string) error { if err := os.MkdirAll(filepath.Dir(path), os.ModePerm); err != nil { return err } @@ -445,7 +455,12 @@ func downloadFile(path, url string) error { } defer out.Close() - resp, err := http.Get(url) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, http.NoBody) + if err != nil { + return err + } + + resp, err := http.DefaultClient.Do(req) if err != nil { return err } diff --git a/runtime/drivers/clickhouse/information_schema_test.go b/runtime/drivers/clickhouse/information_schema_test.go index b5b29081ca9e..db08294c1853 100644 --- a/runtime/drivers/clickhouse/information_schema_test.go +++ b/runtime/drivers/clickhouse/information_schema_test.go @@ -17,7 +17,7 @@ import ( func TestInformationSchema(t *testing.T) { testmode.Expensive(t) dsn := testclickhouse.Start(t) - conn, err := drivers.Open("clickhouse", "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "clickhouse", "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) infoSchema, ok := conn.AsInformationSchema() @@ -27,7 +27,7 @@ func TestInformationSchema(t *testing.T) { t.Run("testInformationSchemaAll", func(t *testing.T) { testInformationSchemaAll(t, conn) }) t.Run("testInformationSchemaAllLike", func(t *testing.T) { testInformationSchemaAllLike(t, conn) }) t.Run("testInformationSchemaSystemAllLike", func(t *testing.T) { - conn, err := drivers.Open("clickhouse", "", "default", map[string]any{"dsn": dsn + "/system", "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "clickhouse", "", "default", map[string]any{"dsn": dsn + "/system", "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) testInformationSchemaSystemAllLike(t, conn) }) diff --git a/runtime/drivers/clickhouse/model_executor_self_test.go b/runtime/drivers/clickhouse/model_executor_self_test.go index b86f3ca5a908..bf93494630c9 100644 --- a/runtime/drivers/clickhouse/model_executor_self_test.go +++ b/runtime/drivers/clickhouse/model_executor_self_test.go @@ -1,6 +1,7 @@ package clickhouse_test import ( + "context" "fmt" "regexp" "sort" @@ -37,7 +38,7 @@ func newInstance(t *testing.T, dsn string, opts testruntime.InstanceOptions) (*r t.Helper() database := nonAlphanumeric.ReplaceAllString(t.Name(), "_") - conn, err := drivers.Open("clickhouse", "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "clickhouse", "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer conn.Close() olap, ok := conn.AsOLAP("") diff --git a/runtime/drivers/clickhouse/olap_test.go b/runtime/drivers/clickhouse/olap_test.go index 63ec4b7f648f..bea50d8700b9 100644 --- a/runtime/drivers/clickhouse/olap_test.go +++ b/runtime/drivers/clickhouse/olap_test.go @@ -20,7 +20,7 @@ import ( func TestClickhouseSingle(t *testing.T) { dsn := testclickhouse.Start(t) - conn, err := driver{}.Open("", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer conn.Close() prepareConn(t, conn) @@ -50,7 +50,7 @@ func TestClickhouseCluster(t *testing.T) { dsn, cluster := testclickhouse.StartCluster(t) - conn, err := drivers.Open("clickhouse", "", "default", map[string]any{"dsn": dsn, "cluster": cluster, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "clickhouse", "", "default", map[string]any{"dsn": dsn, "cluster": cluster, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer conn.Close() @@ -422,7 +422,7 @@ func testSyncReplicaNonDefaultDatabase(t *testing.T, olap drivers.OLAPStore, dsn }) // open a connection scoped to the non-default database (the database is taken from the DSN path) - conn, err := drivers.Open("clickhouse", "", "default", map[string]any{"dsn": dsn + "/sync_repl_db", "cluster": cluster, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "clickhouse", "", "default", map[string]any{"dsn": dsn + "/sync_repl_db", "cluster": cluster, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer conn.Close() c := conn.(*Connection) @@ -661,7 +661,7 @@ func TestClickhouseReadWriteMode(t *testing.T) { t.Run("ReadOnlyMode_DisablesModelExecution", func(t *testing.T) { // Test default mode (read-only) with BYODB - conn, err := driver{}.Open("", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer conn.Close() @@ -684,7 +684,7 @@ func TestClickhouseReadWriteMode(t *testing.T) { t.Run("ExplicitReadOnlyMode_DisablesModelExecution", func(t *testing.T) { // Test explicit read-only mode - conn, err := driver{}.Open("", "default", map[string]any{ + conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ "dsn": dsn, "mode": "read", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -710,7 +710,7 @@ func TestClickhouseReadWriteMode(t *testing.T) { t.Run("ReadWriteMode_EnablesModelExecution", func(t *testing.T) { // Test readwrite mode for BYODB - conn, err := driver{}.Open("", "default", map[string]any{ + conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ "dsn": dsn, "mode": "readwrite", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -740,7 +740,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("SeparateReadWriteDSNs", func(t *testing.T) { // Test with both dsn and write_dsn specified - conn, err := driver{}.Open("", "default", map[string]any{ + conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ "dsn": dsn, "write_dsn": dsn, "mode": "readwrite", @@ -777,7 +777,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("OnlyWriteDSN_ShouldFail", func(t *testing.T) { // Test that providing only write_dsn fails - _, err := driver{}.Open("", "default", map[string]any{ + _, err := driver{}.Open(context.Background(), "", "default", map[string]any{ "write_dsn": dsn, "mode": "readwrite", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -787,7 +787,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("DualDSNWithRegularDSN_UsesDualDSN", func(t *testing.T) { // Test that dsn and write_dsn configuration works correctly - conn, err := driver{}.Open("", "default", map[string]any{ + conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ "dsn": dsn, "write_dsn": dsn, "mode": "readwrite", @@ -809,7 +809,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("InvalidDSN_ShouldFail", func(t *testing.T) { // Test that invalid dsn causes failure - _, err := driver{}.Open("", "default", map[string]any{ + _, err := driver{}.Open(context.Background(), "", "default", map[string]any{ "dsn": "invalid-dsn", "write_dsn": dsn, "mode": "readwrite", @@ -820,7 +820,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("InvalidWriteDSN_ShouldFail", func(t *testing.T) { // Test that invalid write_dsn causes failure - _, err := driver{}.Open("", "default", map[string]any{ + _, err := driver{}.Open(context.Background(), "", "default", map[string]any{ "dsn": dsn, "write_dsn": "invalid-dsn", "mode": "readwrite", @@ -831,7 +831,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("SingleDSN_SharedConnection", func(t *testing.T) { // Test that single DSN still uses shared connection (backward compatibility) - conn, err := driver{}.Open("", "default", map[string]any{ + conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ "dsn": dsn, "mode": "readwrite", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -848,7 +848,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("NoConfiguration_ShouldFail", func(t *testing.T) { // Test that providing no valid configuration fails with appropriate error - _, err := driver{}.Open("", "default", map[string]any{ + _, err := driver{}.Open(context.Background(), "", "default", map[string]any{ "mode": "readwrite", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.Error(t, err) @@ -861,7 +861,7 @@ func TestClickhouseDualDSNFunctionality(t *testing.T) { t.Run("ReadWriteOperationsWithDualDSN", func(t *testing.T) { // Test that both read and write operations work with dual DSN setup - conn, err := driver{}.Open("", "default", map[string]any{ + conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ "dsn": dsn, "write_dsn": dsn, "mode": "readwrite", @@ -1103,7 +1103,7 @@ func testEntityTypeRestrictedUser(t *testing.T, olap drivers.OLAPStore, dsn, clu }) restrictedDSN := strings.Replace(dsn, "clickhouse://default@", fmt.Sprintf("clickhouse://%s:%s@", username, password), 1) - handle, err := drivers.Open("clickhouse", "", "restricted", map[string]any{"dsn": restrictedDSN, "cluster": cluster}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := drivers.Open(context.Background(), "clickhouse", "", "restricted", map[string]any{"dsn": restrictedDSN, "cluster": cluster}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer handle.Close() diff --git a/runtime/drivers/databricks/databricks.go b/runtime/drivers/databricks/databricks.go index 0abb62c2854b..37479c28cfc5 100644 --- a/runtime/drivers/databricks/databricks.go +++ b/runtime/drivers/databricks/databricks.go @@ -188,7 +188,7 @@ func rtRequiresSEA(err error) bool { return err != nil && strings.Contains(err.Error(), "not supported for Thrift protocol") } -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("databricks driver can't be shared") } diff --git a/runtime/drivers/databricks/olap_rt_test.go b/runtime/drivers/databricks/olap_rt_test.go index d5ee7520a102..fcbce19b3d79 100644 --- a/runtime/drivers/databricks/olap_rt_test.go +++ b/runtime/drivers/databricks/olap_rt_test.go @@ -3,6 +3,7 @@ package databricks_test import ( + "context" "os" "testing" @@ -91,7 +92,7 @@ func TestOLAP_LakehouseRT(t *testing.T) { func acquireTestDatabricksRT(t *testing.T, dsn string) (drivers.Handle, drivers.OLAPStore) { cfg := map[string]any{"dsn": dsn} - conn, err := drivers.Open("databricks", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "databricks", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/databricks/olap_test.go b/runtime/drivers/databricks/olap_test.go index 11f7b7c0d5c4..06ea92af3348 100644 --- a/runtime/drivers/databricks/olap_test.go +++ b/runtime/drivers/databricks/olap_test.go @@ -1,6 +1,7 @@ package databricks_test import ( + "context" "strings" "testing" "time" @@ -172,7 +173,7 @@ func TestQuerySchema(t *testing.T) { func acquireTestDatabricks(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "databricks") - conn, err := drivers.Open("databricks", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "databricks", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/drivers.go b/runtime/drivers/drivers.go index 88f27ec1f3e4..f4d63157b972 100644 --- a/runtime/drivers/drivers.go +++ b/runtime/drivers/drivers.go @@ -34,13 +34,14 @@ func Register(name string, driver Driver) { // If instanceID is empty, the connection is considered shared and its As...() functions may be invoked with different instance IDs. // If instanceID is not empty, the connection is considered instance-specific and its As...() functions will only be invoked with the same instance ID. // connectorName can be empty if not available. -func Open(driver, connectorName, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (Handle, error) { +// See Driver.Open for details about the ctx. +func Open(ctx context.Context, driver, connectorName, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (Handle, error) { d, ok := Drivers[driver] if !ok { return nil, fmt.Errorf("unknown driver: %s", driver) } - conn, err := d.Open(connectorName, instanceID, config, st, ac, logger) + conn, err := d.Open(ctx, connectorName, instanceID, config, st, ac, logger) if err != nil { return nil, err } @@ -55,7 +56,9 @@ type Driver interface { // Open opens a new handle. // If instanceID is empty, the connection is considered shared and its As...() functions may be invoked with different instance IDs. - Open(connectorName, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (Handle, error) + // The ctx scopes the open itself, such as network dials, handshakes and provisioning of managed resources. + // It is cancelled as soon as Open returns, so implementations must not retain it on the returned handle. + Open(ctx context.Context, connectorName, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (Handle, error) // HasAnonymousSourceAccess returns true if the driver can access the data identified by srcProps without any additional configuration. HasAnonymousSourceAccess(ctx context.Context, srcProps map[string]any, logger *zap.Logger) (bool, error) diff --git a/runtime/drivers/drivers_test.go b/runtime/drivers/drivers_test.go index 10a1a11ade63..ff6a2c0146b3 100644 --- a/runtime/drivers/drivers_test.go +++ b/runtime/drivers/drivers_test.go @@ -32,7 +32,7 @@ func TestAll(t *testing.T) { for _, withDriver := range matrix { err := withDriver(t, func(driver, instanceID string, cfg map[string]any) { // Open - conn, err := drivers.Open(driver, "", instanceID, cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), driver, "", instanceID, cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) require.NotNil(t, conn) diff --git a/runtime/drivers/druid/druid.go b/runtime/drivers/druid/druid.go index bbef301dfa88..c4255c4d196d 100644 --- a/runtime/drivers/druid/druid.go +++ b/runtime/drivers/druid/druid.go @@ -113,7 +113,7 @@ type configProperties struct { // Opens a connection to Apache Druid using HTTP API. // Note that the Druid connection string must have the form "http://user:password@host:port/druid/v2/sql". -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("druid driver can't be shared") } @@ -146,13 +146,13 @@ func (d driver) Open(connectorName, instanceID string, config map[string]any, st } dbx := sqlx.NewDb(db, "druid") - err = dbx.Ping() + err = dbx.PingContext(ctx) if err != nil { return nil, fmt.Errorf("druid: %w", err) } if !conf.SkipVersionCheck { - err = d.checkVersion(dsn) + err = d.checkVersion(ctx, dsn) if err != nil { return nil, fmt.Errorf("druid: %w", err) } @@ -179,7 +179,7 @@ func (d *driver) TertiarySourceConnectors(ctx context.Context, src map[string]an return nil, fmt.Errorf("not implemented") } -func (d driver) checkVersion(dsn string) error { +func (d driver) checkVersion(ctx context.Context, dsn string) error { parsedURL, err := url.Parse(dsn) if err != nil { return err @@ -187,7 +187,7 @@ func (d driver) checkVersion(dsn string) error { parsedURL.Path = "/status" statusURL := parsedURL.String() - req, err := http.NewRequest(http.MethodGet, statusURL, http.NoBody) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, statusURL, http.NoBody) if err != nil { return err } diff --git a/runtime/drivers/druid/druid_test.go b/runtime/drivers/druid/druid_test.go index f266e42b134a..b540542ea51b 100644 --- a/runtime/drivers/druid/druid_test.go +++ b/runtime/drivers/druid/druid_test.go @@ -38,7 +38,7 @@ func TestScan(t *testing.T) { func acquireTestDruid(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "druid") - conn, err := drivers.Open("druid", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "druid", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/druid/information_schema_test.go b/runtime/drivers/druid/information_schema_test.go index 169903295309..337d57bdebfa 100644 --- a/runtime/drivers/druid/information_schema_test.go +++ b/runtime/drivers/druid/information_schema_test.go @@ -17,7 +17,7 @@ func TestInformationSchema(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "druid") - conn, err := drivers.Open("druid", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "druid", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/druid/sql_driver_test.go b/runtime/drivers/druid/sql_driver_test.go index bed1b8098fa7..8f3ebf07ab25 100644 --- a/runtime/drivers/druid/sql_driver_test.go +++ b/runtime/drivers/druid/sql_driver_test.go @@ -19,7 +19,7 @@ import ( */ func Ignore_TestDriver_types(t *testing.T) { driver := &driver{} - handle, err := driver.Open("", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := driver.Open(context.Background(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -54,7 +54,7 @@ func Ignore_TestDriver_types(t *testing.T) { func Ignore_TestDriver_array_type(t *testing.T) { driver := &driver{} - handle, err := driver.Open("", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := driver.Open(context.Background(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -78,7 +78,7 @@ func Ignore_TestDriver_array_type(t *testing.T) { func Ignore_TestDriver_json_type(t *testing.T) { driver := &driver{} - handle, err := driver.Open("", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := driver.Open(context.Background(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -101,7 +101,7 @@ func Ignore_TestDriver_json_type(t *testing.T) { func Ignore_TestDriver_multiple_rows(t *testing.T) { driver := &driver{} - handle, err := driver.Open("", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := driver.Open(context.Background(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -135,7 +135,7 @@ func Ignore_TestDriver_multiple_rows(t *testing.T) { func Ignore_TestDriver_error(t *testing.T) { driver := &driver{} - handle, err := driver.Open("", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := driver.Open(context.Background(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") diff --git a/runtime/drivers/duckdb/config_test.go b/runtime/drivers/duckdb/config_test.go index 4fc9400a92c9..38b021f03466 100644 --- a/runtime/drivers/duckdb/config_test.go +++ b/runtime/drivers/duckdb/config_test.go @@ -40,7 +40,7 @@ func Test_specialCharInPath(t *testing.T) { require.NoError(t, err) dbFile := filepath.Join(path, "st@g3's.db") - conn, err := Driver{}.Open("", "default", map[string]any{"init_sql": fmt.Sprintf("ATTACH %s", safeSQLString(dbFile))}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := Driver{}.Open(context.Background(), "", "default", map[string]any{"init_sql": fmt.Sprintf("ATTACH %s", safeSQLString(dbFile))}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := conn.AsOLAP("") diff --git a/runtime/drivers/duckdb/crud_test.go b/runtime/drivers/duckdb/crud_test.go index 1a3211ea2b6b..966e7e36bc1b 100644 --- a/runtime/drivers/duckdb/crud_test.go +++ b/runtime/drivers/duckdb/crud_test.go @@ -19,7 +19,7 @@ import ( func Test_connection_CreateTableAsSelect(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open("", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) viewConnection := handle.(*connection) require.NoError(t, viewConnection.Migrate(context.Background())) @@ -75,7 +75,7 @@ func Test_connection_CreateTableAsSelect(t *testing.T) { func Test_connection_CreateTableAsSelectMultipleTimes(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open("", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -103,7 +103,7 @@ func Test_connection_CreateTableAsSelectMultipleTimes(t *testing.T) { func Test_connection_DropTable(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open("", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -127,7 +127,7 @@ func Test_connection_DropTable(t *testing.T) { func Test_connection_InsertTableAsSelect_WithAppendStrategy(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open("", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -162,7 +162,7 @@ func Test_connection_InsertTableAsSelect_WithAppendStrategy(t *testing.T) { func Test_connection_InsertTableAsSelect_WithMergeStrategy(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open("", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -213,7 +213,7 @@ func Test_connection_InsertTableAsSelect_WithMergeStrategy(t *testing.T) { func Test_connection_RenameTable(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open("", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -236,7 +236,7 @@ func Test_connection_RenameTable(t *testing.T) { func Test_connection_RenameToExistingTable(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open("", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -261,7 +261,7 @@ func Test_connection_RenameToExistingTable(t *testing.T) { } func Test_connection_RenameToExistingTableOld(t *testing.T) { - handle, err := Driver{}.Open("", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -288,7 +288,7 @@ func Test_connection_RenameToExistingTableOld(t *testing.T) { func Test_connection_CreateTableAsSelectWithComments(t *testing.T) { temp := t.TempDir() require.NoError(t, os.Mkdir(filepath.Join(temp, "default"), fs.ModePerm)) - handle, err := Driver{}.Open("", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) normalConn := handle.(*connection) normalConn.AsOLAP("default") diff --git a/runtime/drivers/duckdb/duckdb.go b/runtime/drivers/duckdb/duckdb.go index de2147dc7267..f3feba6d2d03 100644 --- a/runtime/drivers/duckdb/duckdb.go +++ b/runtime/drivers/duckdb/duckdb.go @@ -154,7 +154,7 @@ type Driver struct { name string } -func (d Driver) Open(connectorName, instanceID string, cfgMap 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, cfgMap map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) { if instanceID == "" { return nil, errors.New("duckdb driver can't be shared") } @@ -178,7 +178,7 @@ func (d Driver) Open(connectorName, instanceID string, cfgMap map[string]any, st // Open remote bucket for backups if configured var remote *blob.Bucket if cfg.EnableBackups { - b, ok, err := st.OpenBucket(context.Background()) + b, ok, err := st.OpenBucket(ctx) if err != nil { return nil, err } @@ -187,8 +187,9 @@ func (d Driver) Open(connectorName, instanceID string, cfgMap map[string]any, st } } - // Create the handle - ctx, cancel := context.WithCancel(context.Background()) + // Create the handle. + // 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{ instanceID: instanceID, connectorName: connectorName, @@ -204,7 +205,7 @@ func (d Driver) Open(connectorName, instanceID string, cfgMap map[string]any, st driverConfig: cfgMap, driverName: d.name, connTimes: make(map[int]time.Time), - ctx: ctx, + ctx: bgctx, cancel: cancel, } @@ -216,7 +217,7 @@ func (d Driver) Open(connectorName, instanceID string, cfgMap map[string]any, st }, connectionsInUse)) // Open the DB - err = c.reopenDB(context.Background()) + err = c.reopenDB(ctx) if err != nil { if remote != nil { _ = remote.Close() diff --git a/runtime/drivers/duckdb/duckdb_test.go b/runtime/drivers/duckdb/duckdb_test.go index a5c4b78d67fb..898fbb17a13c 100644 --- a/runtime/drivers/duckdb/duckdb_test.go +++ b/runtime/drivers/duckdb/duckdb_test.go @@ -16,7 +16,7 @@ import ( func TestNoFatalErr(t *testing.T) { // NOTE: Using this issue to create a fatal error: https://github.com/duckdb/duckdb/issues/7905 - handle, err := Driver{}.Open("", "default", map[string]any{"pool_size": 2}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{"pool_size": 2}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -76,7 +76,7 @@ func TestNoFatalErr(t *testing.T) { func TestNoFatalErrConcurrent(t *testing.T) { // NOTE: Using this issue to create a fatal error: https://github.com/duckdb/duckdb/issues/7905 - handle, err := Driver{}.Open("", "default", map[string]any{"pool_size": 2}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{"pool_size": 2}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -245,7 +245,7 @@ func TestDuckDBModeEnforcement(t *testing.T) { ctx := context.Background() t.Run("read mode blocks model execution", func(t *testing.T) { - handle, err := drivers.Open("duckdb", "", "test", map[string]any{ + handle, err := drivers.Open(context.Background(), "duckdb", "", "test", map[string]any{ "dsn": ":memory:", "mode": "read", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -268,7 +268,7 @@ func TestDuckDBModeEnforcement(t *testing.T) { }) t.Run("readwrite mode allows model execution", func(t *testing.T) { - handle, err := drivers.Open("duckdb", "", "test", map[string]any{ + handle, err := drivers.Open(context.Background(), "duckdb", "", "test", map[string]any{ "dsn": ":memory:", "mode": "readwrite", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -291,7 +291,7 @@ func TestDuckDBModeEnforcement(t *testing.T) { }) t.Run("read mode allows reading", func(t *testing.T) { - handle, err := drivers.Open("duckdb", "", "test", map[string]any{ + handle, err := drivers.Open(context.Background(), "duckdb", "", "test", map[string]any{ "dsn": ":memory:", "mode": "read", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) diff --git a/runtime/drivers/duckdb/model_executor_https_self_test.go b/runtime/drivers/duckdb/model_executor_https_self_test.go index 7afdff95897b..f87e00122cd7 100644 --- a/runtime/drivers/duckdb/model_executor_https_self_test.go +++ b/runtime/drivers/duckdb/model_executor_https_self_test.go @@ -54,10 +54,10 @@ func TestHTTPToDuckDBTransfer(t *testing.T) { })) defer server.Close() // Ensure server shuts down after test - to, err := drivers.Open("duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + to, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open("https", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(context.Background(), "https", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ diff --git a/runtime/drivers/duckdb/model_executor_self_file_test.go b/runtime/drivers/duckdb/model_executor_self_file_test.go index e66eb68692bf..4ed6a0b2cc8c 100644 --- a/runtime/drivers/duckdb/model_executor_self_file_test.go +++ b/runtime/drivers/duckdb/model_executor_self_file_test.go @@ -17,7 +17,7 @@ import ( func TestExecute(t *testing.T) { tempDir := t.TempDir() - duckDB, err := drivers.Open("duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckDB, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer duckDB.Close() @@ -69,7 +69,7 @@ func TestExecute(t *testing.T) { require.NoError(t, err) require.NoError(t, result.Close()) - fileHandle, err := drivers.Open("file", "", "default", map[string]any{}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) + fileHandle, err := drivers.Open(context.Background(), "file", "", "default", map[string]any{}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer fileHandle.Close() diff --git a/runtime/drivers/duckdb/model_executor_self_objectstore_test.go b/runtime/drivers/duckdb/model_executor_self_objectstore_test.go index 765f18e0609f..28ab765231d2 100644 --- a/runtime/drivers/duckdb/model_executor_self_objectstore_test.go +++ b/runtime/drivers/duckdb/model_executor_self_objectstore_test.go @@ -107,7 +107,7 @@ func testExportedObjectExists(t *testing.T, driver string, rt *runtime.Runtime, ctx := context.Background() handle, _, err := rt.AcquireHandle(ctx, id, driver) - conn, err := drivers.Open(driver, "", "default", handle.Config(), storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), driver, "", "default", handle.Config(), storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/duckdb/model_executor_self_sqlite_test.go b/runtime/drivers/duckdb/model_executor_self_sqlite_test.go index 3d23df0eb041..ef18b9a15769 100644 --- a/runtime/drivers/duckdb/model_executor_self_sqlite_test.go +++ b/runtime/drivers/duckdb/model_executor_self_sqlite_test.go @@ -31,7 +31,7 @@ func Test_sqliteToDuckDB_Transfer(t *testing.T) { require.NoError(t, err) db.Close() - to, err := drivers.Open("duckdb", "", "default", map[string]any{"dsn": ":memory:"}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) + to, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{"dsn": ":memory:"}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, _ := to.AsOLAP("") diff --git a/runtime/drivers/duckdb/model_executor_self_test.go b/runtime/drivers/duckdb/model_executor_self_test.go index 2df789122a2c..776aab97cc9a 100644 --- a/runtime/drivers/duckdb/model_executor_self_test.go +++ b/runtime/drivers/duckdb/model_executor_self_test.go @@ -33,7 +33,7 @@ func TestDuckDBToDuckDBTransfer(t *testing.T) { require.NoError(t, err) require.NoError(t, db.Close()) - duckDB, err := drivers.Open("duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckDB, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ diff --git a/runtime/drivers/duckdb/model_executor_sqlstore_self_mysql_test.go b/runtime/drivers/duckdb/model_executor_sqlstore_self_mysql_test.go index c60009efd873..33c32c9efbef 100644 --- a/runtime/drivers/duckdb/model_executor_sqlstore_self_mysql_test.go +++ b/runtime/drivers/duckdb/model_executor_sqlstore_self_mysql_test.go @@ -108,10 +108,10 @@ func TestMySQLToDuckDBTransfer(t *testing.T) { } func mysqlToDuckDB(t *testing.T, dsn string) { - duckDB, err := drivers.Open("duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckDB, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open("mysql", "", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(context.Background(), "mysql", "", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ @@ -205,9 +205,9 @@ func TestMySQLToDuckLakeTransfer(t *testing.T) { } t.Run("retry_after_missing_column", func(t *testing.T) { - duckLake, err := drivers.Open("duckdb", "", "retry", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckLake, err := drivers.Open(context.Background(), "duckdb", "", "retry", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open("mysql", "", "retry", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(context.Background(), "mysql", "", "retry", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) requireDuckLakeRetryAfterFailure(t, duckLake, inputHandle, "mysql", map[string]any{ @@ -224,9 +224,9 @@ func TestMySQLToDuckLakeTransfer(t *testing.T) { } func mysqlToDuckLake(t *testing.T, dsn string, outputConfig map[string]any) { - duckLake, err := drivers.Open("duckdb", "", "default", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckLake, err := drivers.Open(context.Background(), "duckdb", "", "default", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open("mysql", "", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(context.Background(), "mysql", "", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ diff --git a/runtime/drivers/duckdb/model_executor_sqlstore_self_postgres_test.go b/runtime/drivers/duckdb/model_executor_sqlstore_self_postgres_test.go index be1bf72fef89..4c81b5eeca9b 100644 --- a/runtime/drivers/duckdb/model_executor_sqlstore_self_postgres_test.go +++ b/runtime/drivers/duckdb/model_executor_sqlstore_self_postgres_test.go @@ -70,10 +70,10 @@ func TestTransfer(t *testing.T) { } func pgxToDuckDB(t *testing.T, pgdb *sql.DB, dbURL string) { - duckDB, err := drivers.Open("duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckDB, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open("postgres", "", "default", map[string]any{"database_url": dbURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(context.Background(), "postgres", "", "default", map[string]any{"database_url": dbURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ @@ -172,9 +172,9 @@ func TestPostgresToDuckLakeTransfer(t *testing.T) { } t.Run("retry_after_missing_column", func(t *testing.T) { - duckLake, err := drivers.Open("duckdb", "", "retry", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckLake, err := drivers.Open(context.Background(), "duckdb", "", "retry", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open("postgres", "", "retry", map[string]any{"database_url": pg.DatabaseURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(context.Background(), "postgres", "", "retry", map[string]any{"database_url": pg.DatabaseURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) requireDuckLakeRetryAfterFailure(t, duckLake, inputHandle, "postgres", map[string]any{ @@ -190,10 +190,10 @@ func TestPostgresToDuckLakeTransfer(t *testing.T) { } func postgresToDuckLake(t *testing.T, pgdb *sql.DB, dbURL string, outputConfig map[string]any) { - duckLake, err := drivers.Open("duckdb", "", "default", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckLake, err := drivers.Open(context.Background(), "duckdb", "", "default", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open("postgres", "", "default", map[string]any{"database_url": dbURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(context.Background(), "postgres", "", "default", map[string]any{"database_url": dbURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ diff --git a/runtime/drivers/duckdb/motherduck_test.go b/runtime/drivers/duckdb/motherduck_test.go index 5b5d851d2060..b7679d640b31 100644 --- a/runtime/drivers/duckdb/motherduck_test.go +++ b/runtime/drivers/duckdb/motherduck_test.go @@ -1,6 +1,7 @@ package duckdb_test import ( + "context" "testing" "github.com/rilldata/rill/runtime/drivers" @@ -16,7 +17,7 @@ func TestMotherDuckModeEnforcement(t *testing.T) { cfg := testruntime.AcquireConnector(t, "motherduck") cfg["mode"] = "read" - handle, err := drivers.Open("motherduck", "", "test", cfg, + handle, err := drivers.Open(context.Background(), "motherduck", "", "test", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -42,7 +43,7 @@ func TestMotherDuckModeEnforcement(t *testing.T) { cfg := testruntime.AcquireConnector(t, "motherduck") cfg["mode"] = "readwrite" - handle, err := drivers.Open("motherduck", "", "test", cfg, + handle, err := drivers.Open(context.Background(), "motherduck", "", "test", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) diff --git a/runtime/drivers/duckdb/olap_test.go b/runtime/drivers/duckdb/olap_test.go index 678afd5534bd..73b1af5d0d76 100644 --- a/runtime/drivers/duckdb/olap_test.go +++ b/runtime/drivers/duckdb/olap_test.go @@ -209,7 +209,7 @@ func TestClose(t *testing.T) { } func prepareConn(t *testing.T) drivers.Handle { - conn, err := Driver{}.Open("", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := conn.AsOLAP("") @@ -230,7 +230,7 @@ func Test_safeSQLString(t *testing.T) { err := os.Mkdir(path, fs.ModePerm) require.NoError(t, err) - conn, err := Driver{}.Open("", "default", map[string]any{"data_dir": path}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := Driver{}.Open(context.Background(), "", "default", map[string]any{"data_dir": path}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) require.NotNil(t, conn) require.NoError(t, conn.Close()) diff --git a/runtime/drivers/duckdb/schema_change_test.go b/runtime/drivers/duckdb/schema_change_test.go index c5545f163438..f4a2df675c7d 100644 --- a/runtime/drivers/duckdb/schema_change_test.go +++ b/runtime/drivers/duckdb/schema_change_test.go @@ -293,7 +293,7 @@ func TestInsertTableAsSelectEmptyPartitionSkipsSchemaHandling(t *testing.T) { func newSchemaChangeTestConnection(t *testing.T) *connection { t.Helper() - handle, err := Driver{}.Open("", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) diff --git a/runtime/drivers/file/file.go b/runtime/drivers/file/file.go index a7675f822966..8b1424152034 100644 --- a/runtime/drivers/file/file.go +++ b/runtime/drivers/file/file.go @@ -70,7 +70,7 @@ type rillYAML struct { IgnorePaths []string `yaml:"ignore_paths"` } -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(ctx 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("file driver can't be shared") } @@ -103,7 +103,7 @@ func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Cl } // Read rill.yaml and fill in `ignore_paths` - rawYaml, err := c.Get(context.Background(), "/rill.yaml") + rawYaml, err := c.Get(ctx, "/rill.yaml") if err == nil { yml := &rillYAML{} err = yaml.Unmarshal([]byte(rawYaml), yml) diff --git a/runtime/drivers/file/model_executor_olap_self_test.go b/runtime/drivers/file/model_executor_olap_self_test.go index 2fe12422c8dc..246c23b683d7 100644 --- a/runtime/drivers/file/model_executor_olap_self_test.go +++ b/runtime/drivers/file/model_executor_olap_self_test.go @@ -17,7 +17,7 @@ import ( func TestExecute(t *testing.T) { tempDir := t.TempDir() - duckDB, err := drivers.Open("duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckDB, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer duckDB.Close() @@ -74,7 +74,7 @@ func TestExecute(t *testing.T) { require.NoError(t, err) require.NoError(t, result.Close()) - fileHandle, err := drivers.Open("file", "", "default", map[string]any{}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) + fileHandle, err := drivers.Open(context.Background(), "file", "", "default", map[string]any{}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer fileHandle.Close() diff --git a/runtime/drivers/gcs/gcs.go b/runtime/drivers/gcs/gcs.go index 1e8e0a3d72bf..ab393f01ee64 100644 --- a/runtime/drivers/gcs/gcs.go +++ b/runtime/drivers/gcs/gcs.go @@ -102,7 +102,7 @@ func NewConfigProperties(prop map[string]any) (*ConfigProperties, error) { return gcsConfig, nil } -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(ctx 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("gcs driver can't be shared") } @@ -128,7 +128,7 @@ func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Cl if err != nil { return nil, err } - handle, err := drivers.Open("s3", "", instanceID, config, st, ac, logger) + handle, err := drivers.Open(ctx, "s3", "", instanceID, config, st, ac, logger) if err != nil { return nil, fmt.Errorf("failed to open s3 connection for gcs in s3 compatible mode: %w", err) } diff --git a/runtime/drivers/gcs/object_store_test.go b/runtime/drivers/gcs/object_store_test.go index ad5a70dd7e76..6d3495780456 100644 --- a/runtime/drivers/gcs/object_store_test.go +++ b/runtime/drivers/gcs/object_store_test.go @@ -1,6 +1,7 @@ package gcs_test import ( + "context" "testing" "github.com/rilldata/rill/runtime/drivers" @@ -15,7 +16,7 @@ import ( func TestObjectStore(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "gcs") - conn, err := drivers.Open("gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) @@ -54,7 +55,7 @@ func TestObjectStorePathPrefixes(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "gcs") cfg["path_prefixes"] = "gcs://integration-test.rilldata.com/glob_test/" - conn, err := drivers.Open("gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) @@ -71,7 +72,7 @@ func TestObjectStorePathPrefixes(t *testing.T) { func TestObjectStoreHMAC(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "gcs_s3_compat") - conn, err := drivers.Open("gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) @@ -110,7 +111,7 @@ func TestObjectStoreHMACPathPrefixes(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "gcs_s3_compat") cfg["path_prefixes"] = "gcs://integration-test.rilldata.com/glob_test/" - conn, err := drivers.Open("gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/gemini/gemini.go b/runtime/drivers/gemini/gemini.go index 9a492a138767..741a066c96da 100644 --- a/runtime/drivers/gemini/gemini.go +++ b/runtime/drivers/gemini/gemini.go @@ -103,7 +103,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(ctx 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 { @@ -113,7 +113,7 @@ func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Cl if conf.APIKey == "" { return nil, errors.New("api_key is required") } - client, err := genai.NewClient(context.Background(), &genai.ClientConfig{ + client, err := genai.NewClient(ctx, &genai.ClientConfig{ Backend: genai.BackendGeminiAPI, APIKey: conf.APIKey, }) diff --git a/runtime/drivers/https/https.go b/runtime/drivers/https/https.go index 89aede1c3a71..027e9c65301c 100644 --- a/runtime/drivers/https/https.go +++ b/runtime/drivers/https/https.go @@ -97,7 +97,7 @@ func (p *ModelInputProperties) DecodeWithWarnings(props map[string]any) ([]strin return unused, nil } -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("https driver can't be shared") } diff --git a/runtime/drivers/mock/ai/ai.go b/runtime/drivers/mock/ai/ai.go index be0db3e2fe11..617aa1750f32 100644 --- a/runtime/drivers/mock/ai/ai.go +++ b/runtime/drivers/mock/ai/ai.go @@ -48,7 +48,7 @@ func (d driver) TertiarySourceConnectors(ctx context.Context, srcProps map[strin return nil, nil } -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) { toolCallingEnabled, _ := config["enable_tool_calling"].(bool) return &connection{ diff --git a/runtime/drivers/mock/object_store/object_store.go b/runtime/drivers/mock/object_store/object_store.go index 7f44e00fee8c..f1e077dc948f 100644 --- a/runtime/drivers/mock/object_store/object_store.go +++ b/runtime/drivers/mock/object_store/object_store.go @@ -36,14 +36,14 @@ func (driver) Spec() drivers.Spec { } // Open implements drivers.Driver. -func (driver) Open(_, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) { +func (driver) Open(ctx context.Context, _, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) { cfg := &configProperties{} err := mapstructure.WeakDecode(config, cfg) if err != nil { return nil, err } - bucket, err := blob.OpenBucket(context.Background(), "file://"+cfg.Path) + bucket, err := blob.OpenBucket(ctx, "file://"+cfg.Path) if err != nil { return nil, err } diff --git a/runtime/drivers/mysql/mysql.go b/runtime/drivers/mysql/mysql.go index 600cb6701949..52e7100ba670 100644 --- a/runtime/drivers/mysql/mysql.go +++ b/runtime/drivers/mysql/mysql.go @@ -215,7 +215,7 @@ func (c *ConfigProperties) resolveGoFormatDSN() (string, error) { return cfg.FormatDSN(), nil } -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("mysql driver can't be shared") } diff --git a/runtime/drivers/mysql/olap_test.go b/runtime/drivers/mysql/olap_test.go index b0bd76ea1363..e4f4dea3c65c 100644 --- a/runtime/drivers/mysql/olap_test.go +++ b/runtime/drivers/mysql/olap_test.go @@ -1,6 +1,7 @@ package mysql_test import ( + "context" "strings" "testing" @@ -517,7 +518,7 @@ func testLoadDDL(t *testing.T, olap drivers.OLAPStore) { func acquireTestMySQL(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "mysql") - conn, err := drivers.Open("mysql", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "mysql", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/openai/openai.go b/runtime/drivers/openai/openai.go index 2a5d2691e665..e9291eb2c597 100644 --- a/runtime/drivers/openai/openai.go +++ b/runtime/drivers/openai/openai.go @@ -100,7 +100,7 @@ func (d driver) HasAnonymousSourceAccess(ctx context.Context, srcProps map[strin } // 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 { diff --git a/runtime/drivers/pinot/information_schema_test.go b/runtime/drivers/pinot/information_schema_test.go index 20336383dbbd..7862f859c0be 100644 --- a/runtime/drivers/pinot/information_schema_test.go +++ b/runtime/drivers/pinot/information_schema_test.go @@ -33,7 +33,7 @@ var expectedTables = []string{ func TestInformationSchema(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "pinot") - conn, err := drivers.Open("pinot", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "pinot", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/pinot/pinot.go b/runtime/drivers/pinot/pinot.go index 890fc50c281d..6ed487943956 100644 --- a/runtime/drivers/pinot/pinot.go +++ b/runtime/drivers/pinot/pinot.go @@ -134,7 +134,7 @@ func (c *configProperties) validate() error { } // Open a connection to Apache Pinot using HTTP API. -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, fmt.Errorf("pinot driver can't be shared") } @@ -209,7 +209,7 @@ func (d driver) Open(connectorName, instanceID string, config map[string]any, st } dbx := sqlx.NewDb(db, "pinot") - err = dbx.Ping() + err = dbx.PingContext(ctx) if err != nil { return nil, fmt.Errorf("pinot: %w", err) } diff --git a/runtime/drivers/postgres/olap_test.go b/runtime/drivers/postgres/olap_test.go index 26fe837f9adf..f1eba9a0ad35 100644 --- a/runtime/drivers/postgres/olap_test.go +++ b/runtime/drivers/postgres/olap_test.go @@ -1,6 +1,7 @@ package postgres_test import ( + "context" "encoding/json" "fmt" "testing" @@ -425,7 +426,7 @@ func testLoadDDL(t *testing.T, olap drivers.OLAPStore) { func acquireTestPostgres(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "postgres") - conn, err := drivers.Open("postgres", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "postgres", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/postgres/postgres.go b/runtime/drivers/postgres/postgres.go index da2a9e2b797c..726aa54acd8a 100644 --- a/runtime/drivers/postgres/postgres.go +++ b/runtime/drivers/postgres/postgres.go @@ -255,7 +255,7 @@ func (c *ConfigProperties) ResolveDSN() string { return strings.Join(parts, " ") } -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("postgres driver can't be shared") } diff --git a/runtime/drivers/redshift/olap_test.go b/runtime/drivers/redshift/olap_test.go index 5557d667d3ee..0f708cb51495 100644 --- a/runtime/drivers/redshift/olap_test.go +++ b/runtime/drivers/redshift/olap_test.go @@ -366,7 +366,7 @@ func acquireTestRedshift(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg["database"] = "test_db" } cfg["workgroup"] = "integration-test-wg" - conn, err := drivers.Open("redshift", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "redshift", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/redshift/redshift.go b/runtime/drivers/redshift/redshift.go index a1d1af16066c..0c85716c7c78 100644 --- a/runtime/drivers/redshift/redshift.go +++ b/runtime/drivers/redshift/redshift.go @@ -92,7 +92,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("redshift driver can't be shared") } diff --git a/runtime/drivers/s3/mode_executor_olap_to_self_test.go b/runtime/drivers/s3/mode_executor_olap_to_self_test.go index 0e3842dcd4d5..cdf4a763c91e 100644 --- a/runtime/drivers/s3/mode_executor_olap_to_self_test.go +++ b/runtime/drivers/s3/mode_executor_olap_to_self_test.go @@ -140,7 +140,7 @@ func testExportedObjectExists(t *testing.T, driver string, rt *runtime.Runtime, ctx := context.Background() handle, _, err := rt.AcquireHandle(ctx, id, driver) - conn, err := drivers.Open(driver, "", "default", handle.Config(), storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), driver, "", "default", handle.Config(), storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/s3/object_store_test.go b/runtime/drivers/s3/object_store_test.go index bfce339e250c..8925c7dc7510 100644 --- a/runtime/drivers/s3/object_store_test.go +++ b/runtime/drivers/s3/object_store_test.go @@ -1,6 +1,7 @@ package s3_test import ( + "context" "testing" "github.com/rilldata/rill/runtime/drivers" @@ -15,7 +16,7 @@ import ( func TestObjectStore(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "s3") - conn, err := drivers.Open("s3", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "s3", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) @@ -54,7 +55,7 @@ func TestObjectStorePathPrefixes(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "s3") cfg["path_prefixes"] = "s3://integration-test.rilldata.com/glob_test/" - conn, err := drivers.Open("s3", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "s3", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/s3/s3.go b/runtime/drivers/s3/s3.go index 1956e13097ad..de108945b742 100644 --- a/runtime/drivers/s3/s3.go +++ b/runtime/drivers/s3/s3.go @@ -130,7 +130,7 @@ type ConfigProperties struct { } // 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) { if instanceID == "" { return nil, errors.New("s3 driver can't be shared") } diff --git a/runtime/drivers/salesforce/salesforce.go b/runtime/drivers/salesforce/salesforce.go index 7d1518807b9e..5ba63eccca9b 100644 --- a/runtime/drivers/salesforce/salesforce.go +++ b/runtime/drivers/salesforce/salesforce.go @@ -141,7 +141,7 @@ var spec = drivers.Spec{ type driver struct{} -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("salesforce driver can't be shared") } diff --git a/runtime/drivers/slack/slack.go b/runtime/drivers/slack/slack.go index 5e45f3baf55d..992c079c21a4 100644 --- a/runtime/drivers/slack/slack.go +++ b/runtime/drivers/slack/slack.go @@ -36,7 +36,7 @@ func (d driver) Spec() drivers.Spec { return spec } -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, fmt.Errorf("slack driver can't be shared") } diff --git a/runtime/drivers/snowflake/olap_test.go b/runtime/drivers/snowflake/olap_test.go index 0ab31bca2b0f..bdc4362bff96 100644 --- a/runtime/drivers/snowflake/olap_test.go +++ b/runtime/drivers/snowflake/olap_test.go @@ -1,6 +1,7 @@ package snowflake_test import ( + "context" "encoding/json" "strings" "testing" @@ -186,7 +187,7 @@ func TestDryRun(t *testing.T) { func acquireTestSnowflake(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "snowflake") - conn, err := drivers.Open("snowflake", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "snowflake", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/snowflake/snowflake.go b/runtime/drivers/snowflake/snowflake.go index cd62dca9c87f..2004584b0b0d 100644 --- a/runtime/drivers/snowflake/snowflake.go +++ b/runtime/drivers/snowflake/snowflake.go @@ -208,7 +208,7 @@ func (c *configProperties) resolveDSN() (string, error) { return gosnowflake.DSN(cfg) } -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("snowflake driver can't be shared") } diff --git a/runtime/drivers/sqlite/backups_test.go b/runtime/drivers/sqlite/backups_test.go index 2c2604198f55..6c690a0d6cfc 100644 --- a/runtime/drivers/sqlite/backups_test.go +++ b/runtime/drivers/sqlite/backups_test.go @@ -1,6 +1,7 @@ package sqlite import ( + "context" "fmt" "path/filepath" "strings" @@ -34,7 +35,7 @@ func TestBackup(t *testing.T) { "id": "test-backup", "backups_enable": true, } - h, err := driver{}.Open("", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) + h, err := driver{}.Open(context.Background(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer h.Close() err = h.Migrate(t.Context()) @@ -141,7 +142,7 @@ func TestDBFilePath(t *testing.T) { for idx, tc := range cases { t.Run(fmt.Sprintf("case-%d", idx), func(t *testing.T) { cfg := map[string]any{"dsn": tc.dsn} - h, err := driver{}.Open("", "", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + h, err := driver{}.Open(context.Background(), "", "", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer h.Close() diff --git a/runtime/drivers/sqlite/migrate_test.go b/runtime/drivers/sqlite/migrate_test.go index eaa2f22b0417..60eef2b8a441 100644 --- a/runtime/drivers/sqlite/migrate_test.go +++ b/runtime/drivers/sqlite/migrate_test.go @@ -1,6 +1,7 @@ package sqlite import ( + "context" "path/filepath" "testing" "time" @@ -22,7 +23,7 @@ func TestDeleteExpiredAISessions(t *testing.T) { cfg := map[string]any{"dsn": dbPath} // Open the database, run migrations, and seed test data. - h, err := driver{}.Open("", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) + h, err := driver{}.Open(context.Background(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer h.Close() require.NoError(t, h.Migrate(t.Context())) diff --git a/runtime/drivers/sqlite/restore_test.go b/runtime/drivers/sqlite/restore_test.go index 5f4131699df2..0517de8d03a3 100644 --- a/runtime/drivers/sqlite/restore_test.go +++ b/runtime/drivers/sqlite/restore_test.go @@ -1,6 +1,7 @@ package sqlite import ( + "context" "fmt" "os" "path/filepath" @@ -33,7 +34,7 @@ func TestBackupAndRestore(t *testing.T) { } // Create a database with an instance in it, then back it up. - h, err := driver{}.Open("", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) + h, err := driver{}.Open(context.Background(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) require.NoError(t, h.Migrate(t.Context())) registry, ok := h.AsRegistry() @@ -47,7 +48,7 @@ func TestBackupAndRestore(t *testing.T) { require.NoError(t, restoreBackup(t.Context(), bucket, dbPath, zap.NewNop())) // Reopen the database and check the instance is back. - h, err = driver{}.Open("", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) + h, err = driver{}.Open(context.Background(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer h.Close() require.NoError(t, h.Migrate(t.Context())) @@ -124,7 +125,7 @@ func TestShouldRestoreBackup(t *testing.T) { require.Equal(t, dsn, dbPath) // Once migrated, the database has data and must not be overwritten by a restore. - h, err := driver{}.Open("", "", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + h, err := driver{}.Open(context.Background(), "", "", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) require.NoError(t, h.Migrate(t.Context())) require.NoError(t, h.Close()) diff --git a/runtime/drivers/sqlite/sqlite.go b/runtime/drivers/sqlite/sqlite.go index d33e31b3c025..b33d7a523a17 100644 --- a/runtime/drivers/sqlite/sqlite.go +++ b/runtime/drivers/sqlite/sqlite.go @@ -38,7 +38,7 @@ type configProperties struct { BackupsEnable bool `mapstructure:"backups_enable"` } -func (d driver) Open(_, _ string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) { +func (d driver) Open(ctx context.Context, _, _ string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) { // Parse config conf := &configProperties{} err := mapstructure.WeakDecode(config, conf) @@ -62,7 +62,7 @@ func (d driver) Open(_, _ string, config map[string]any, st *storage.Client, ac // This must run before the handle is opened below because it replaces the database file, // which an already-open SQLite connection would not observe. if conf.ID != "" && conf.BackupsEnable { - err := restoreBackupIfEmpty(context.Background(), st, conf.ID, conf.DSN, logger) + err := restoreBackupIfEmpty(ctx, st, conf.ID, conf.DSN, logger) if err != nil { return nil, fmt.Errorf("sqlite: failed to restore backup %q: %w", conf.ID, err) } @@ -76,13 +76,14 @@ func (d driver) Open(_, _ string, config map[string]any, st *storage.Client, ac dbx := sqlx.NewDb(db, "sqlite") db.SetMaxOpenConns(1) - // Create the handle - ctx, cancel := context.WithCancel(context.Background()) + // Create the handle. + // 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()) h := &connection{ db: dbx, logger: logger, config: config, - ctx: ctx, + ctx: bgctx, cancel: cancel, storage: st, backupID: conf.ID, diff --git a/runtime/drivers/starrocks/olap_test.go b/runtime/drivers/starrocks/olap_test.go index ff8c02080356..f2844557f185 100644 --- a/runtime/drivers/starrocks/olap_test.go +++ b/runtime/drivers/starrocks/olap_test.go @@ -21,7 +21,7 @@ func TestStarRocksOLAP(t *testing.T) { dsn := teststarrocks.StartWithData(t) - conn, err := driver{}.Open("", "default", map[string]any{ + conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ "dsn": dsn, }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) diff --git a/runtime/drivers/starrocks/starrocks.go b/runtime/drivers/starrocks/starrocks.go index 5524130338c5..cf171b18ad89 100644 --- a/runtime/drivers/starrocks/starrocks.go +++ b/runtime/drivers/starrocks/starrocks.go @@ -165,7 +165,7 @@ const ( defaultPort = 9030 ) -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(ctx 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("starrocks driver: instance ID is required") } @@ -194,7 +194,7 @@ func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Cl // Open database connection immediately in drivers.Open // This ensures the connection is established and validated upfront - if err := conn.initDB(); err != nil { + if err := conn.initDB(ctx); err != nil { return nil, fmt.Errorf("failed to initialize database connection: %w", err) } @@ -328,7 +328,7 @@ func (c *connection) AsModelManager(instanceID string) (drivers.ModelManager, er // initDB initializes the database connection. // Called during drivers.Open to establish connection upfront. -func (c *connection) initDB() error { +func (c *connection) initDB(ctx context.Context) error { dsn := c.buildDSN() db, err := sqlx.Open("mysql", dsn) @@ -344,10 +344,9 @@ func (c *connection) initDB() error { db.SetConnMaxLifetime(30 * time.Minute) db.SetConnMaxIdleTime(5 * time.Minute) - // Test connection with an independent context to prevent premature cancellation - // Use a context with sufficient timeout (30 seconds) instead of the request context - // This prevents 499 errors when the frontend request is cancelled quickly - pingCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + // Cap the connection test and validation queries below at 30 seconds. + // The ctx scopes the open as a whole; it is not the request context, so it won't be cancelled by a frontend request going away. + pingCtx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() if err := db.PingContext(pingCtx); err != nil { db.Close() diff --git a/runtime/parser/parser_test.go b/runtime/parser/parser_test.go index 2ec375d81b4a..185dd8512278 100644 --- a/runtime/parser/parser_test.go +++ b/runtime/parser/parser_test.go @@ -3167,7 +3167,7 @@ func requireResourcesAndErrors(t testing.TB, p *Parser, wantResources []*Resourc func makeRepo(t testing.TB, files map[string]string) drivers.RepoStore { root := t.TempDir() - handle, err := drivers.Open("file", "", "default", map[string]any{"dsn": root}, storage.MustNew(root, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := drivers.Open(context.Background(), "file", "", "default", map[string]any{"dsn": root}, storage.MustNew(root, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) repo, ok := handle.AsRepoStore("") diff --git a/runtime/server/queries_test.go b/runtime/server/queries_test.go index b2672b4645fa..f52a17e3d0c7 100644 --- a/runtime/server/queries_test.go +++ b/runtime/server/queries_test.go @@ -133,7 +133,7 @@ func TestServer_UpdateLimit_UNION(t *testing.T) { } func prepareOLAPStore(t *testing.T) drivers.OLAPStore { - conn, err := drivers.Open("duckdb", "", "default", map[string]any{"dsn": ":memory:?access_mode=read_write", "pool_size": 4}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{"dsn": ":memory:?access_mode=read_write", "pool_size": 4}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := conn.AsOLAP("") require.True(t, ok) diff --git a/runtime/storage/storage.go b/runtime/storage/storage.go index 04e4dca11db5..49381c4e815f 100644 --- a/runtime/storage/storage.go +++ b/runtime/storage/storage.go @@ -130,6 +130,9 @@ func (c *Client) RandomTempDir(pattern string, elem ...string) (string, error) { return path, nil } +// OpenBucket opens the client's external bucket, scoped to the client's prefixes plus any additional elem. +// It returns false if no bucket is configured. +// The ctx scopes the open itself; the returned bucket does not retain it and stays usable after it is cancelled. func (c *Client) OpenBucket(ctx context.Context, elem ...string) (*blob.Bucket, bool, error) { if c.bucketConfig == nil { return nil, false, nil @@ -167,7 +170,9 @@ func (c *Client) path(base string, elem ...string) string { } func (c *Client) newGCPClient(ctx context.Context) (*gcp.HTTPClient, error) { - creds, err := gcputil.Credentials(ctx, c.bucketConfig.GoogleApplicationCredentialsJSON, false) + // Detach cancellation: the credentials capture the ctx and reuse it to refresh tokens for as long as the bucket is used. + // Some credential types bind that ctx to the refresh request, so a cancelled ctx would permanently break a long-lived bucket. + creds, err := gcputil.Credentials(context.WithoutCancel(ctx), c.bucketConfig.GoogleApplicationCredentialsJSON, false) if err != nil { return nil, err } diff --git a/runtime/testruntime/noopadmin.go b/runtime/testruntime/noopadmin.go index 94a13900b81a..2129b66e0e5e 100644 --- a/runtime/testruntime/noopadmin.go +++ b/runtime/testruntime/noopadmin.go @@ -52,7 +52,7 @@ func (n *noopAdminService) HasAnonymousSourceAccess(ctx context.Context, srcProp } // Open implements [drivers.Driver]. -func (n *noopAdminService) Open(connectorName, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) { +func (n *noopAdminService) Open(_ context.Context, connectorName, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) { return n, nil } From 2e986810f91888891e366a1935d38f825a68da22 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Benjamin=20Egelund-M=C3=BCller?= Date: Wed, 23 Sep 2026 09:56:30 +0200 Subject: [PATCH 2/5] Don't let DuckDB per-connection init queries observe the open ctx --- runtime/drivers/duckdb/open_ctx_test.go | 45 +++++++++++++++++++++++++ runtime/pkg/rduckdb/db.go | 7 ++-- runtime/pkg/rduckdb/generic.go | 7 ++-- 3 files changed, 55 insertions(+), 4 deletions(-) create mode 100644 runtime/drivers/duckdb/open_ctx_test.go diff --git a/runtime/drivers/duckdb/open_ctx_test.go b/runtime/drivers/duckdb/open_ctx_test.go new file mode 100644 index 000000000000..c1599fcd8980 --- /dev/null +++ b/runtime/drivers/duckdb/open_ctx_test.go @@ -0,0 +1,45 @@ +package duckdb + +import ( + "context" + "testing" + + "github.com/rilldata/rill/runtime/drivers" + "github.com/rilldata/rill/runtime/pkg/activity" + "github.com/rilldata/rill/runtime/storage" + "github.com/stretchr/testify/require" + "go.uber.org/zap" + "golang.org/x/sync/errgroup" +) + +// TestOpenCtxNotRetained verifies that the handle does not retain the ctx passed to Open. +// The connection cache cancels that ctx as soon as Open returns, +// so anything that outlives Open (notably the per-connection init queries, which run for every connection the pool opens later) +// must not observe its cancellation. +func TestOpenCtxNotRetained(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + + handle, err := Driver{}.Open(ctx, "", "default", map[string]any{"data_dir": t.TempDir(), "pool_size": 4}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + require.NoError(t, err) + defer handle.Close() + require.NoError(t, handle.Migrate(ctx)) + + // Mimic the connection cache, which cancels the open ctx immediately after Open returns. + cancel() + + olap, ok := handle.AsOLAP("default") + require.True(t, ok) + + // Query concurrently so the pool has to open connections it didn't already have at the time of the cancellation. + var g errgroup.Group + for range 8 { + g.Go(func() error { + res, err := olap.Query(context.Background(), &drivers.Statement{Query: "SELECT 1"}) + if err != nil { + return err + } + return res.Close() + }) + } + require.NoError(t, g.Wait()) +} diff --git a/runtime/pkg/rduckdb/db.go b/runtime/pkg/rduckdb/db.go index ce7a1d575223..8f5631f0beda 100644 --- a/runtime/pkg/rduckdb/db.go +++ b/runtime/pkg/rduckdb/db.go @@ -893,12 +893,15 @@ func (d *db) openDBAndAttach(ctx context.Context, uri, ignoreTable string, initQ } // Rebuild DuckDB DSN (which should be "path?key=val&...") // this is required since spaces and other special characters are valid in db file path but invalid and hence encoded in URL + // The init queries run for every new connection the pool opens, for as long as the returned DB is in use. + // So they must not use the caller's ctx, which may be cancelled while the DB is still open. + connCtx := d.ctx connector, err := duckdb.NewConnector(generateDSN(dsn.Path, query.Encode()), func(execer driver.ExecerContext) error { for _, qry := range d.opts.ConnInitQueries { - _, err := execer.ExecContext(ctx, qry, nil) + _, err := execer.ExecContext(connCtx, qry, nil) if err != nil && strings.Contains(err.Error(), "Failed to download extension") { // Retry using another mirror. Based on: https://github.com/duckdb/duckdb/issues/9378 - _, err = execer.ExecContext(ctx, qry+" FROM 'http://nightly-extensions.duckdb.org'", nil) + _, err = execer.ExecContext(connCtx, qry+" FROM 'http://nightly-extensions.duckdb.org'", nil) } if err != nil { return err diff --git a/runtime/pkg/rduckdb/generic.go b/runtime/pkg/rduckdb/generic.go index 89f833b7ad44..72873f1cc29a 100644 --- a/runtime/pkg/rduckdb/generic.go +++ b/runtime/pkg/rduckdb/generic.go @@ -99,12 +99,15 @@ func NewGeneric(ctx context.Context, opts *GenericOptions) (res DB, dbErr error) } localFileName := filepath.Join(opts.LocalDataDir, "main"+uuid.NewString()[:8]+".db") + // The init queries run for every new connection the pool opens, for as long as the returned DB is in use. + // So they must not observe cancellation of the caller's ctx, which may be cancelled while the DB is still open. + connCtx := context.WithoutCancel(ctx) connector, err := duckdb.NewConnector(dsnForLocalDuckDB(localFileName, opts.Settings), func(execer driver.ExecerContext) error { for _, qry := range opts.ConnInitQueries { - _, err := execer.ExecContext(ctx, qry, nil) + _, err := execer.ExecContext(connCtx, qry, nil) if err != nil && strings.Contains(err.Error(), "Failed to download extension") { // Retry using another mirror. Based on: https://github.com/duckdb/duckdb/issues/9378 - _, err = execer.ExecContext(ctx, qry+" FROM 'http://nightly-extensions.duckdb.org'", nil) + _, err = execer.ExecContext(connCtx, qry+" FROM 'http://nightly-extensions.duckdb.org'", nil) } if err != nil { return err From 461ab506ec40cdc3d6b1906780ff262f625fdd11 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Benjamin=20Egelund-M=C3=BCller?= Date: Wed, 23 Sep 2026 10:21:38 +0200 Subject: [PATCH 3/5] Don't let the open ctx deadline break readonly ClickHouse connections --- runtime/drivers/clickhouse/clickhouse.go | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/runtime/drivers/clickhouse/clickhouse.go b/runtime/drivers/clickhouse/clickhouse.go index 919e337dc828..d94c8f1bfe00 100644 --- a/runtime/drivers/clickhouse/clickhouse.go +++ b/runtime/drivers/clickhouse/clickhouse.go @@ -322,8 +322,10 @@ func (d driver) Open(ctx context.Context, connectorName, instanceID string, conf return nil, fmt.Errorf("failed to open write connection: %w", err) } } - // group by positional args are supported post 22.7 and we use them heavily in our queries - row := db.QueryRowContext(ctx, ` + // group by positional args are supported post 22.7 and we use them heavily in our queries. + // Note the statements below strip the ctx deadline: see the note in openHandle. + stmtCtx := contextWithoutDeadline(ctx) + row := db.QueryRowContext(stmtCtx, ` WITH splitByChar('.', version()) AS parts, toInt32(parts[1]) AS major, @@ -344,7 +346,7 @@ func (d driver) Open(ctx context.Context, connectorName, instanceID string, conf // whether the cluster mode supports modifying query settings. This setting // has no practical use for our purposes. supportSettings := true - if _, err := db.ExecContext(ctx, "SET show_table_uuid_in_table_create_query_if_not_nil = 1"); err != nil { + if _, err := db.ExecContext(stmtCtx, "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 } @@ -929,9 +931,13 @@ func openHandle(ctx context.Context, instanceID string, conf *configProperties, opts.DialTimeout = time.Second * 60 } - // Open the connection + // Open the connection. + // The clickhouse-go driver derives a 'max_execution_time' setting from the ctx deadline, + // which a readonly server rejects; we can't know yet whether this one is readonly. + // So strip the deadline while keeping cancellation: the dial and read timeouts in opts already bound the ping. + pingCtx := contextWithoutDeadline(ctx) db := sqlx.NewDb(otelsql.OpenDB(clickhouse.Connector(opts)), "clickhouse") - err := db.PingContext(ctx) + err := db.PingContext(pingCtx) if err != nil { // Detect SSL/TLS mismatch (common causes: "read: EOF" or TLS Alert [21]) if strings.Contains(err.Error(), "EOF") || @@ -949,7 +955,7 @@ func openHandle(ctx context.Context, instanceID string, conf *configProperties, // 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.PingContext(ctx) + err := db.PingContext(pingCtx) if err != nil { // Detect SSL/TLS mismatch (common causes: "read: EOF" or \x15 means TLS Alert [21]"]) if strings.Contains(err.Error(), "EOF") || From 6490f52f4552c121e60d73b90e99277341911478 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Benjamin=20Egelund-M=C3=BCller?= Date: Wed, 30 Sep 2026 17:31:39 +0200 Subject: [PATCH 4/5] Use t.Context() in tests and share the readonly ctx rule in ClickHouse --- admin/testadmin/testadmin.go | 2 +- runtime/drivers/athena/olap_test.go | 2 +- runtime/drivers/azure/object_store_test.go | 7 ++- runtime/drivers/bigquery/olap_test.go | 2 +- runtime/drivers/clickhouse/clickhouse.go | 42 +++++++++-------- .../clickhouse/information_schema_test.go | 4 +- .../clickhouse/model_executor_self_test.go | 3 +- runtime/drivers/clickhouse/olap.go | 26 +++++------ runtime/drivers/clickhouse/olap_test.go | 30 ++++++------- runtime/drivers/databricks/olap_rt_test.go | 3 +- runtime/drivers/databricks/olap_test.go | 3 +- runtime/drivers/drivers_test.go | 2 +- runtime/drivers/druid/druid_test.go | 2 +- .../drivers/druid/information_schema_test.go | 2 +- runtime/drivers/druid/sql_driver_test.go | 10 ++--- runtime/drivers/duckdb/config_test.go | 2 +- runtime/drivers/duckdb/crud_test.go | 18 ++++---- runtime/drivers/duckdb/duckdb_test.go | 10 ++--- .../duckdb/model_executor_https_self_test.go | 4 +- .../duckdb/model_executor_self_file_test.go | 4 +- .../model_executor_self_objectstore_test.go | 2 +- .../duckdb/model_executor_self_sqlite_test.go | 2 +- .../duckdb/model_executor_self_test.go | 2 +- ...model_executor_sqlstore_self_mysql_test.go | 12 ++--- ...el_executor_sqlstore_self_postgres_test.go | 12 ++--- runtime/drivers/duckdb/motherduck_test.go | 5 +-- runtime/drivers/duckdb/olap_test.go | 4 +- runtime/drivers/duckdb/open_ctx_test.go | 45 ------------------- runtime/drivers/duckdb/schema_change_test.go | 2 +- .../file/model_executor_olap_self_test.go | 4 +- runtime/drivers/gcs/object_store_test.go | 9 ++-- runtime/drivers/mysql/olap_test.go | 3 +- .../drivers/pinot/information_schema_test.go | 2 +- runtime/drivers/postgres/olap_test.go | 3 +- runtime/drivers/redshift/olap_test.go | 2 +- .../s3/mode_executor_olap_to_self_test.go | 2 +- runtime/drivers/s3/object_store_test.go | 5 +-- runtime/drivers/snowflake/olap_test.go | 3 +- runtime/drivers/sqlite/backups_test.go | 5 +-- runtime/drivers/sqlite/migrate_test.go | 3 +- runtime/drivers/sqlite/restore_test.go | 7 ++- runtime/drivers/starrocks/olap_test.go | 2 +- runtime/parser/parser_test.go | 2 +- runtime/server/queries_test.go | 2 +- runtime/storage/storage.go | 6 ++- 45 files changed, 132 insertions(+), 192 deletions(-) delete mode 100644 runtime/drivers/duckdb/open_ctx_test.go diff --git a/admin/testadmin/testadmin.go b/admin/testadmin/testadmin.go index e3e063ff46d6..a33faa60c4f3 100644 --- a/admin/testadmin/testadmin.go +++ b/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(context.Background(), "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") diff --git a/runtime/drivers/athena/olap_test.go b/runtime/drivers/athena/olap_test.go index 5813462db830..25ec8e0d63cd 100644 --- a/runtime/drivers/athena/olap_test.go +++ b/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(context.Background(), "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() }) diff --git a/runtime/drivers/azure/object_store_test.go b/runtime/drivers/azure/object_store_test.go index d98c9e7c9d97..506368ba1ec5 100644 --- a/runtime/drivers/azure/object_store_test.go +++ b/runtime/drivers/azure/object_store_test.go @@ -1,7 +1,6 @@ package azure_test import ( - "context" "testing" "github.com/rilldata/rill/runtime/drivers" @@ -17,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(context.Background(), "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() }) @@ -44,7 +43,7 @@ func TestObjectStoreCloud(t *testing.T) { func TestObjectStore(t *testing.T) { cfg := testruntime.AcquireConnector(t, "azure") - conn, err := drivers.Open(context.Background(), "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() }) @@ -66,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(context.Background(), "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() }) diff --git a/runtime/drivers/bigquery/olap_test.go b/runtime/drivers/bigquery/olap_test.go index f620606ba453..01d2347b7c4a 100644 --- a/runtime/drivers/bigquery/olap_test.go +++ b/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(context.Background(), "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() }) diff --git a/runtime/drivers/clickhouse/clickhouse.go b/runtime/drivers/clickhouse/clickhouse.go index d94c8f1bfe00..e0483a9e8bf5 100644 --- a/runtime/drivers/clickhouse/clickhouse.go +++ b/runtime/drivers/clickhouse/clickhouse.go @@ -322,10 +322,23 @@ func (d driver) Open(ctx context.Context, connectorName, instanceID string, conf return nil, fmt.Errorf("failed to open write connection: %w", err) } } - // group by positional args are supported post 22.7 and we use them heavily in our queries. - // Note the statements below strip the ctx deadline: see the note in openHandle. - stmtCtx := contextWithoutDeadline(ctx) - row := db.QueryRowContext(stmtCtx, ` + // 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.QueryRowContext(settingsCtx(ctx, supportSettings), ` WITH splitByChar('.', version()) AS parts, toInt32(parts[1]) AS major, @@ -341,17 +354,6 @@ func (d driver) Open(ctx context.Context, connectorName, instanceID string, conf 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.ExecContext(stmtCtx, "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 { @@ -931,13 +933,9 @@ func openHandle(ctx context.Context, instanceID string, conf *configProperties, opts.DialTimeout = time.Second * 60 } - // Open the connection. - // The clickhouse-go driver derives a 'max_execution_time' setting from the ctx deadline, - // which a readonly server rejects; we can't know yet whether this one is readonly. - // So strip the deadline while keeping cancellation: the dial and read timeouts in opts already bound the ping. - pingCtx := contextWithoutDeadline(ctx) + // Open the connection db := sqlx.NewDb(otelsql.OpenDB(clickhouse.Connector(opts)), "clickhouse") - err := db.PingContext(pingCtx) + 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") || @@ -955,7 +953,7 @@ func openHandle(ctx context.Context, instanceID string, conf *configProperties, // 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.PingContext(pingCtx) + 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") || diff --git a/runtime/drivers/clickhouse/information_schema_test.go b/runtime/drivers/clickhouse/information_schema_test.go index db08294c1853..c6079cd6e0ea 100644 --- a/runtime/drivers/clickhouse/information_schema_test.go +++ b/runtime/drivers/clickhouse/information_schema_test.go @@ -17,7 +17,7 @@ import ( func TestInformationSchema(t *testing.T) { testmode.Expensive(t) dsn := testclickhouse.Start(t) - conn, err := drivers.Open(context.Background(), "clickhouse", "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "clickhouse", "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) infoSchema, ok := conn.AsInformationSchema() @@ -27,7 +27,7 @@ func TestInformationSchema(t *testing.T) { t.Run("testInformationSchemaAll", func(t *testing.T) { testInformationSchemaAll(t, conn) }) t.Run("testInformationSchemaAllLike", func(t *testing.T) { testInformationSchemaAllLike(t, conn) }) t.Run("testInformationSchemaSystemAllLike", func(t *testing.T) { - conn, err := drivers.Open(context.Background(), "clickhouse", "", "default", map[string]any{"dsn": dsn + "/system", "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "clickhouse", "", "default", map[string]any{"dsn": dsn + "/system", "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) testInformationSchemaSystemAllLike(t, conn) }) diff --git a/runtime/drivers/clickhouse/model_executor_self_test.go b/runtime/drivers/clickhouse/model_executor_self_test.go index bf93494630c9..673f5f440398 100644 --- a/runtime/drivers/clickhouse/model_executor_self_test.go +++ b/runtime/drivers/clickhouse/model_executor_self_test.go @@ -1,7 +1,6 @@ package clickhouse_test import ( - "context" "fmt" "regexp" "sort" @@ -38,7 +37,7 @@ func newInstance(t *testing.T, dsn string, opts testruntime.InstanceOptions) (*r t.Helper() database := nonAlphanumeric.ReplaceAllString(t.Name(), "_") - conn, err := drivers.Open(context.Background(), "clickhouse", "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "clickhouse", "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer conn.Close() olap, ok := conn.AsOLAP("") diff --git a/runtime/drivers/clickhouse/olap.go b/runtime/drivers/clickhouse/olap.go index 042b293bbcb8..044d2a515664 100644 --- a/runtime/drivers/clickhouse/olap.go +++ b/runtime/drivers/clickhouse/olap.go @@ -507,27 +507,25 @@ type SQLConn struct { } func (sc *SQLConn) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) { - if sc.supportSettings { - return sc.Conn.ExecContext(ctx, query, args...) - } - ctx2 := contextWithoutDeadline(ctx) - return sc.Conn.ExecContext(ctx2, query, args...) + return sc.Conn.ExecContext(settingsCtx(ctx, sc.supportSettings), query, args...) } func (sc *SQLConn) QueryxContext(ctx context.Context, query string, args ...any) (*sqlx.Rows, error) { - if sc.supportSettings { - return sc.Conn.QueryxContext(ctx, query, args...) - } - ctx2 := contextWithoutDeadline(ctx) - return sc.Conn.QueryxContext(ctx2, query, args...) + return sc.Conn.QueryxContext(settingsCtx(ctx, sc.supportSettings), query, args...) } func (sc *SQLConn) QueryRowContext(ctx context.Context, query string, args ...any) *sqlx.Row { - if sc.supportSettings { - return sc.Conn.QueryRowxContext(ctx, query, args...) + return sc.Conn.QueryRowxContext(settingsCtx(ctx, sc.supportSettings), query, args...) +} + +// settingsCtx returns the ctx to use for a statement sent to a server that may not support modifying settings. +// When the server is readonly, it strips the deadline (keeping cancellation), +// since clickhouse-go would otherwise derive a `max_execution_time` setting from the deadline, which the server rejects. +func settingsCtx(ctx context.Context, supportSettings bool) context.Context { + if supportSettings { + return ctx } - ctx2 := contextWithoutDeadline(ctx) - return sc.Conn.QueryRowxContext(ctx2, query, args...) + return contextWithoutDeadline(ctx) } func contextWithoutDeadline(parent context.Context) context.Context { diff --git a/runtime/drivers/clickhouse/olap_test.go b/runtime/drivers/clickhouse/olap_test.go index bea50d8700b9..2b7bc2dab94b 100644 --- a/runtime/drivers/clickhouse/olap_test.go +++ b/runtime/drivers/clickhouse/olap_test.go @@ -20,7 +20,7 @@ import ( func TestClickhouseSingle(t *testing.T) { dsn := testclickhouse.Start(t) - conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := driver{}.Open(t.Context(), "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer conn.Close() prepareConn(t, conn) @@ -50,7 +50,7 @@ func TestClickhouseCluster(t *testing.T) { dsn, cluster := testclickhouse.StartCluster(t) - conn, err := drivers.Open(context.Background(), "clickhouse", "", "default", map[string]any{"dsn": dsn, "cluster": cluster, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "clickhouse", "", "default", map[string]any{"dsn": dsn, "cluster": cluster, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer conn.Close() @@ -422,7 +422,7 @@ func testSyncReplicaNonDefaultDatabase(t *testing.T, olap drivers.OLAPStore, dsn }) // open a connection scoped to the non-default database (the database is taken from the DSN path) - conn, err := drivers.Open(context.Background(), "clickhouse", "", "default", map[string]any{"dsn": dsn + "/sync_repl_db", "cluster": cluster, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "clickhouse", "", "default", map[string]any{"dsn": dsn + "/sync_repl_db", "cluster": cluster, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer conn.Close() c := conn.(*Connection) @@ -661,7 +661,7 @@ func TestClickhouseReadWriteMode(t *testing.T) { t.Run("ReadOnlyMode_DisablesModelExecution", func(t *testing.T) { // Test default mode (read-only) with BYODB - conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := driver{}.Open(t.Context(), "", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer conn.Close() @@ -684,7 +684,7 @@ func TestClickhouseReadWriteMode(t *testing.T) { t.Run("ExplicitReadOnlyMode_DisablesModelExecution", func(t *testing.T) { // Test explicit read-only mode - conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ + conn, err := driver{}.Open(t.Context(), "", "default", map[string]any{ "dsn": dsn, "mode": "read", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -710,7 +710,7 @@ func TestClickhouseReadWriteMode(t *testing.T) { t.Run("ReadWriteMode_EnablesModelExecution", func(t *testing.T) { // Test readwrite mode for BYODB - conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ + conn, err := driver{}.Open(t.Context(), "", "default", map[string]any{ "dsn": dsn, "mode": "readwrite", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -740,7 +740,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("SeparateReadWriteDSNs", func(t *testing.T) { // Test with both dsn and write_dsn specified - conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ + conn, err := driver{}.Open(t.Context(), "", "default", map[string]any{ "dsn": dsn, "write_dsn": dsn, "mode": "readwrite", @@ -777,7 +777,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("OnlyWriteDSN_ShouldFail", func(t *testing.T) { // Test that providing only write_dsn fails - _, err := driver{}.Open(context.Background(), "", "default", map[string]any{ + _, err := driver{}.Open(t.Context(), "", "default", map[string]any{ "write_dsn": dsn, "mode": "readwrite", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -787,7 +787,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("DualDSNWithRegularDSN_UsesDualDSN", func(t *testing.T) { // Test that dsn and write_dsn configuration works correctly - conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ + conn, err := driver{}.Open(t.Context(), "", "default", map[string]any{ "dsn": dsn, "write_dsn": dsn, "mode": "readwrite", @@ -809,7 +809,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("InvalidDSN_ShouldFail", func(t *testing.T) { // Test that invalid dsn causes failure - _, err := driver{}.Open(context.Background(), "", "default", map[string]any{ + _, err := driver{}.Open(t.Context(), "", "default", map[string]any{ "dsn": "invalid-dsn", "write_dsn": dsn, "mode": "readwrite", @@ -820,7 +820,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("InvalidWriteDSN_ShouldFail", func(t *testing.T) { // Test that invalid write_dsn causes failure - _, err := driver{}.Open(context.Background(), "", "default", map[string]any{ + _, err := driver{}.Open(t.Context(), "", "default", map[string]any{ "dsn": dsn, "write_dsn": "invalid-dsn", "mode": "readwrite", @@ -831,7 +831,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("SingleDSN_SharedConnection", func(t *testing.T) { // Test that single DSN still uses shared connection (backward compatibility) - conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ + conn, err := driver{}.Open(t.Context(), "", "default", map[string]any{ "dsn": dsn, "mode": "readwrite", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -848,7 +848,7 @@ func TestClickhouseDualDSN(t *testing.T) { t.Run("NoConfiguration_ShouldFail", func(t *testing.T) { // Test that providing no valid configuration fails with appropriate error - _, err := driver{}.Open(context.Background(), "", "default", map[string]any{ + _, err := driver{}.Open(t.Context(), "", "default", map[string]any{ "mode": "readwrite", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.Error(t, err) @@ -861,7 +861,7 @@ func TestClickhouseDualDSNFunctionality(t *testing.T) { t.Run("ReadWriteOperationsWithDualDSN", func(t *testing.T) { // Test that both read and write operations work with dual DSN setup - conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ + conn, err := driver{}.Open(t.Context(), "", "default", map[string]any{ "dsn": dsn, "write_dsn": dsn, "mode": "readwrite", @@ -1103,7 +1103,7 @@ func testEntityTypeRestrictedUser(t *testing.T, olap drivers.OLAPStore, dsn, clu }) restrictedDSN := strings.Replace(dsn, "clickhouse://default@", fmt.Sprintf("clickhouse://%s:%s@", username, password), 1) - handle, err := drivers.Open(context.Background(), "clickhouse", "", "restricted", map[string]any{"dsn": restrictedDSN, "cluster": cluster}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := drivers.Open(t.Context(), "clickhouse", "", "restricted", map[string]any{"dsn": restrictedDSN, "cluster": cluster}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer handle.Close() diff --git a/runtime/drivers/databricks/olap_rt_test.go b/runtime/drivers/databricks/olap_rt_test.go index fcbce19b3d79..c68badd70b4c 100644 --- a/runtime/drivers/databricks/olap_rt_test.go +++ b/runtime/drivers/databricks/olap_rt_test.go @@ -3,7 +3,6 @@ package databricks_test import ( - "context" "os" "testing" @@ -92,7 +91,7 @@ func TestOLAP_LakehouseRT(t *testing.T) { func acquireTestDatabricksRT(t *testing.T, dsn string) (drivers.Handle, drivers.OLAPStore) { cfg := map[string]any{"dsn": dsn} - conn, err := drivers.Open(context.Background(), "databricks", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "databricks", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/databricks/olap_test.go b/runtime/drivers/databricks/olap_test.go index 06ea92af3348..1add3fe9480d 100644 --- a/runtime/drivers/databricks/olap_test.go +++ b/runtime/drivers/databricks/olap_test.go @@ -1,7 +1,6 @@ package databricks_test import ( - "context" "strings" "testing" "time" @@ -173,7 +172,7 @@ func TestQuerySchema(t *testing.T) { func acquireTestDatabricks(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "databricks") - conn, err := drivers.Open(context.Background(), "databricks", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "databricks", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/drivers_test.go b/runtime/drivers/drivers_test.go index ff6a2c0146b3..a48c897580d7 100644 --- a/runtime/drivers/drivers_test.go +++ b/runtime/drivers/drivers_test.go @@ -32,7 +32,7 @@ func TestAll(t *testing.T) { for _, withDriver := range matrix { err := withDriver(t, func(driver, instanceID string, cfg map[string]any) { // Open - conn, err := drivers.Open(context.Background(), driver, "", instanceID, cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), driver, "", instanceID, cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) require.NotNil(t, conn) diff --git a/runtime/drivers/druid/druid_test.go b/runtime/drivers/druid/druid_test.go index b540542ea51b..3f4a84ec7adc 100644 --- a/runtime/drivers/druid/druid_test.go +++ b/runtime/drivers/druid/druid_test.go @@ -38,7 +38,7 @@ func TestScan(t *testing.T) { func acquireTestDruid(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "druid") - conn, err := drivers.Open(context.Background(), "druid", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "druid", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/druid/information_schema_test.go b/runtime/drivers/druid/information_schema_test.go index 337d57bdebfa..8e1474995f78 100644 --- a/runtime/drivers/druid/information_schema_test.go +++ b/runtime/drivers/druid/information_schema_test.go @@ -17,7 +17,7 @@ func TestInformationSchema(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "druid") - conn, err := drivers.Open(context.Background(), "druid", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "druid", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/druid/sql_driver_test.go b/runtime/drivers/druid/sql_driver_test.go index 8f3ebf07ab25..49fd064d8506 100644 --- a/runtime/drivers/druid/sql_driver_test.go +++ b/runtime/drivers/druid/sql_driver_test.go @@ -19,7 +19,7 @@ import ( */ func Ignore_TestDriver_types(t *testing.T) { driver := &driver{} - handle, err := driver.Open(context.Background(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := driver.Open(t.Context(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -54,7 +54,7 @@ func Ignore_TestDriver_types(t *testing.T) { func Ignore_TestDriver_array_type(t *testing.T) { driver := &driver{} - handle, err := driver.Open(context.Background(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := driver.Open(t.Context(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -78,7 +78,7 @@ func Ignore_TestDriver_array_type(t *testing.T) { func Ignore_TestDriver_json_type(t *testing.T) { driver := &driver{} - handle, err := driver.Open(context.Background(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := driver.Open(t.Context(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -101,7 +101,7 @@ func Ignore_TestDriver_json_type(t *testing.T) { func Ignore_TestDriver_multiple_rows(t *testing.T) { driver := &driver{} - handle, err := driver.Open(context.Background(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := driver.Open(t.Context(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -135,7 +135,7 @@ func Ignore_TestDriver_multiple_rows(t *testing.T) { func Ignore_TestDriver_error(t *testing.T) { driver := &driver{} - handle, err := driver.Open(context.Background(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := driver.Open(t.Context(), "", "default", map[string]any{"pool_size": 2, "dsn": "http://localhost:8888/druid/v2/sql"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") diff --git a/runtime/drivers/duckdb/config_test.go b/runtime/drivers/duckdb/config_test.go index 38b021f03466..63f8a9af63bf 100644 --- a/runtime/drivers/duckdb/config_test.go +++ b/runtime/drivers/duckdb/config_test.go @@ -40,7 +40,7 @@ func Test_specialCharInPath(t *testing.T) { require.NoError(t, err) dbFile := filepath.Join(path, "st@g3's.db") - conn, err := Driver{}.Open(context.Background(), "", "default", map[string]any{"init_sql": fmt.Sprintf("ATTACH %s", safeSQLString(dbFile))}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := Driver{}.Open(t.Context(), "", "default", map[string]any{"init_sql": fmt.Sprintf("ATTACH %s", safeSQLString(dbFile))}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := conn.AsOLAP("") diff --git a/runtime/drivers/duckdb/crud_test.go b/runtime/drivers/duckdb/crud_test.go index 966e7e36bc1b..5fd52ec754cb 100644 --- a/runtime/drivers/duckdb/crud_test.go +++ b/runtime/drivers/duckdb/crud_test.go @@ -19,7 +19,7 @@ import ( func Test_connection_CreateTableAsSelect(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) viewConnection := handle.(*connection) require.NoError(t, viewConnection.Migrate(context.Background())) @@ -75,7 +75,7 @@ func Test_connection_CreateTableAsSelect(t *testing.T) { func Test_connection_CreateTableAsSelectMultipleTimes(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -103,7 +103,7 @@ func Test_connection_CreateTableAsSelectMultipleTimes(t *testing.T) { func Test_connection_DropTable(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -127,7 +127,7 @@ func Test_connection_DropTable(t *testing.T) { func Test_connection_InsertTableAsSelect_WithAppendStrategy(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -162,7 +162,7 @@ func Test_connection_InsertTableAsSelect_WithAppendStrategy(t *testing.T) { func Test_connection_InsertTableAsSelect_WithMergeStrategy(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -213,7 +213,7 @@ func Test_connection_InsertTableAsSelect_WithMergeStrategy(t *testing.T) { func Test_connection_RenameTable(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -236,7 +236,7 @@ func Test_connection_RenameTable(t *testing.T) { func Test_connection_RenameToExistingTable(t *testing.T) { temp := t.TempDir() - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -261,7 +261,7 @@ func Test_connection_RenameToExistingTable(t *testing.T) { } func Test_connection_RenameToExistingTableOld(t *testing.T) { - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) @@ -288,7 +288,7 @@ func Test_connection_RenameToExistingTableOld(t *testing.T) { func Test_connection_CreateTableAsSelectWithComments(t *testing.T) { temp := t.TempDir() require.NoError(t, os.Mkdir(filepath.Join(temp, "default"), fs.ModePerm)) - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{}, storage.MustNew(temp, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) normalConn := handle.(*connection) normalConn.AsOLAP("default") diff --git a/runtime/drivers/duckdb/duckdb_test.go b/runtime/drivers/duckdb/duckdb_test.go index 898fbb17a13c..8ca8d2633b66 100644 --- a/runtime/drivers/duckdb/duckdb_test.go +++ b/runtime/drivers/duckdb/duckdb_test.go @@ -16,7 +16,7 @@ import ( func TestNoFatalErr(t *testing.T) { // NOTE: Using this issue to create a fatal error: https://github.com/duckdb/duckdb/issues/7905 - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{"pool_size": 2}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{"pool_size": 2}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -76,7 +76,7 @@ func TestNoFatalErr(t *testing.T) { func TestNoFatalErrConcurrent(t *testing.T) { // NOTE: Using this issue to create a fatal error: https://github.com/duckdb/duckdb/issues/7905 - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{"pool_size": 2}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{"pool_size": 2}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := handle.AsOLAP("") @@ -245,7 +245,7 @@ func TestDuckDBModeEnforcement(t *testing.T) { ctx := context.Background() t.Run("read mode blocks model execution", func(t *testing.T) { - handle, err := drivers.Open(context.Background(), "duckdb", "", "test", map[string]any{ + handle, err := drivers.Open(t.Context(), "duckdb", "", "test", map[string]any{ "dsn": ":memory:", "mode": "read", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -268,7 +268,7 @@ func TestDuckDBModeEnforcement(t *testing.T) { }) t.Run("readwrite mode allows model execution", func(t *testing.T) { - handle, err := drivers.Open(context.Background(), "duckdb", "", "test", map[string]any{ + handle, err := drivers.Open(t.Context(), "duckdb", "", "test", map[string]any{ "dsn": ":memory:", "mode": "readwrite", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -291,7 +291,7 @@ func TestDuckDBModeEnforcement(t *testing.T) { }) t.Run("read mode allows reading", func(t *testing.T) { - handle, err := drivers.Open(context.Background(), "duckdb", "", "test", map[string]any{ + handle, err := drivers.Open(t.Context(), "duckdb", "", "test", map[string]any{ "dsn": ":memory:", "mode": "read", }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) diff --git a/runtime/drivers/duckdb/model_executor_https_self_test.go b/runtime/drivers/duckdb/model_executor_https_self_test.go index f87e00122cd7..653731d69b9c 100644 --- a/runtime/drivers/duckdb/model_executor_https_self_test.go +++ b/runtime/drivers/duckdb/model_executor_https_self_test.go @@ -54,10 +54,10 @@ func TestHTTPToDuckDBTransfer(t *testing.T) { })) defer server.Close() // Ensure server shuts down after test - to, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + to, err := drivers.Open(t.Context(), "duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open(context.Background(), "https", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(t.Context(), "https", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ diff --git a/runtime/drivers/duckdb/model_executor_self_file_test.go b/runtime/drivers/duckdb/model_executor_self_file_test.go index 4ed6a0b2cc8c..d44b389027a3 100644 --- a/runtime/drivers/duckdb/model_executor_self_file_test.go +++ b/runtime/drivers/duckdb/model_executor_self_file_test.go @@ -17,7 +17,7 @@ import ( func TestExecute(t *testing.T) { tempDir := t.TempDir() - duckDB, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckDB, err := drivers.Open(t.Context(), "duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer duckDB.Close() @@ -69,7 +69,7 @@ func TestExecute(t *testing.T) { require.NoError(t, err) require.NoError(t, result.Close()) - fileHandle, err := drivers.Open(context.Background(), "file", "", "default", map[string]any{}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) + fileHandle, err := drivers.Open(t.Context(), "file", "", "default", map[string]any{}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer fileHandle.Close() diff --git a/runtime/drivers/duckdb/model_executor_self_objectstore_test.go b/runtime/drivers/duckdb/model_executor_self_objectstore_test.go index 28ab765231d2..e0b3e1adaa68 100644 --- a/runtime/drivers/duckdb/model_executor_self_objectstore_test.go +++ b/runtime/drivers/duckdb/model_executor_self_objectstore_test.go @@ -107,7 +107,7 @@ func testExportedObjectExists(t *testing.T, driver string, rt *runtime.Runtime, ctx := context.Background() handle, _, err := rt.AcquireHandle(ctx, id, driver) - conn, err := drivers.Open(context.Background(), driver, "", "default", handle.Config(), storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), driver, "", "default", handle.Config(), storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/duckdb/model_executor_self_sqlite_test.go b/runtime/drivers/duckdb/model_executor_self_sqlite_test.go index ef18b9a15769..2e796d04feac 100644 --- a/runtime/drivers/duckdb/model_executor_self_sqlite_test.go +++ b/runtime/drivers/duckdb/model_executor_self_sqlite_test.go @@ -31,7 +31,7 @@ func Test_sqliteToDuckDB_Transfer(t *testing.T) { require.NoError(t, err) db.Close() - to, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{"dsn": ":memory:"}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) + to, err := drivers.Open(t.Context(), "duckdb", "", "default", map[string]any{"dsn": ":memory:"}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, _ := to.AsOLAP("") diff --git a/runtime/drivers/duckdb/model_executor_self_test.go b/runtime/drivers/duckdb/model_executor_self_test.go index 776aab97cc9a..75a64d367e50 100644 --- a/runtime/drivers/duckdb/model_executor_self_test.go +++ b/runtime/drivers/duckdb/model_executor_self_test.go @@ -33,7 +33,7 @@ func TestDuckDBToDuckDBTransfer(t *testing.T) { require.NoError(t, err) require.NoError(t, db.Close()) - duckDB, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckDB, err := drivers.Open(t.Context(), "duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ diff --git a/runtime/drivers/duckdb/model_executor_sqlstore_self_mysql_test.go b/runtime/drivers/duckdb/model_executor_sqlstore_self_mysql_test.go index 33c32c9efbef..89eda1b07e22 100644 --- a/runtime/drivers/duckdb/model_executor_sqlstore_self_mysql_test.go +++ b/runtime/drivers/duckdb/model_executor_sqlstore_self_mysql_test.go @@ -108,10 +108,10 @@ func TestMySQLToDuckDBTransfer(t *testing.T) { } func mysqlToDuckDB(t *testing.T, dsn string) { - duckDB, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckDB, err := drivers.Open(t.Context(), "duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open(context.Background(), "mysql", "", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(t.Context(), "mysql", "", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ @@ -205,9 +205,9 @@ func TestMySQLToDuckLakeTransfer(t *testing.T) { } t.Run("retry_after_missing_column", func(t *testing.T) { - duckLake, err := drivers.Open(context.Background(), "duckdb", "", "retry", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckLake, err := drivers.Open(t.Context(), "duckdb", "", "retry", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open(context.Background(), "mysql", "", "retry", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(t.Context(), "mysql", "", "retry", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) requireDuckLakeRetryAfterFailure(t, duckLake, inputHandle, "mysql", map[string]any{ @@ -224,9 +224,9 @@ func TestMySQLToDuckLakeTransfer(t *testing.T) { } func mysqlToDuckLake(t *testing.T, dsn string, outputConfig map[string]any) { - duckLake, err := drivers.Open(context.Background(), "duckdb", "", "default", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckLake, err := drivers.Open(t.Context(), "duckdb", "", "default", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open(context.Background(), "mysql", "", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(t.Context(), "mysql", "", "default", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ diff --git a/runtime/drivers/duckdb/model_executor_sqlstore_self_postgres_test.go b/runtime/drivers/duckdb/model_executor_sqlstore_self_postgres_test.go index 4c81b5eeca9b..5869900e3b44 100644 --- a/runtime/drivers/duckdb/model_executor_sqlstore_self_postgres_test.go +++ b/runtime/drivers/duckdb/model_executor_sqlstore_self_postgres_test.go @@ -70,10 +70,10 @@ func TestTransfer(t *testing.T) { } func pgxToDuckDB(t *testing.T, pgdb *sql.DB, dbURL string) { - duckDB, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckDB, err := drivers.Open(t.Context(), "duckdb", "", "default", map[string]any{"data_dir": t.TempDir()}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open(context.Background(), "postgres", "", "default", map[string]any{"database_url": dbURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(t.Context(), "postgres", "", "default", map[string]any{"database_url": dbURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ @@ -172,9 +172,9 @@ func TestPostgresToDuckLakeTransfer(t *testing.T) { } t.Run("retry_after_missing_column", func(t *testing.T) { - duckLake, err := drivers.Open(context.Background(), "duckdb", "", "retry", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckLake, err := drivers.Open(t.Context(), "duckdb", "", "retry", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open(context.Background(), "postgres", "", "retry", map[string]any{"database_url": pg.DatabaseURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(t.Context(), "postgres", "", "retry", map[string]any{"database_url": pg.DatabaseURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) requireDuckLakeRetryAfterFailure(t, duckLake, inputHandle, "postgres", map[string]any{ @@ -190,10 +190,10 @@ func TestPostgresToDuckLakeTransfer(t *testing.T) { } func postgresToDuckLake(t *testing.T, pgdb *sql.DB, dbURL string, outputConfig map[string]any) { - duckLake, err := drivers.Open(context.Background(), "duckdb", "", "default", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckLake, err := drivers.Open(t.Context(), "duckdb", "", "default", outputConfig, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) - inputHandle, err := drivers.Open(context.Background(), "postgres", "", "default", map[string]any{"database_url": dbURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + inputHandle, err := drivers.Open(t.Context(), "postgres", "", "default", map[string]any{"database_url": dbURL}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) opts := &drivers.ModelExecutorOptions{ diff --git a/runtime/drivers/duckdb/motherduck_test.go b/runtime/drivers/duckdb/motherduck_test.go index b7679d640b31..97f299630662 100644 --- a/runtime/drivers/duckdb/motherduck_test.go +++ b/runtime/drivers/duckdb/motherduck_test.go @@ -1,7 +1,6 @@ package duckdb_test import ( - "context" "testing" "github.com/rilldata/rill/runtime/drivers" @@ -17,7 +16,7 @@ func TestMotherDuckModeEnforcement(t *testing.T) { cfg := testruntime.AcquireConnector(t, "motherduck") cfg["mode"] = "read" - handle, err := drivers.Open(context.Background(), "motherduck", "", "test", cfg, + handle, err := drivers.Open(t.Context(), "motherduck", "", "test", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) @@ -43,7 +42,7 @@ func TestMotherDuckModeEnforcement(t *testing.T) { cfg := testruntime.AcquireConnector(t, "motherduck") cfg["mode"] = "readwrite" - handle, err := drivers.Open(context.Background(), "motherduck", "", "test", cfg, + handle, err := drivers.Open(t.Context(), "motherduck", "", "test", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) diff --git a/runtime/drivers/duckdb/olap_test.go b/runtime/drivers/duckdb/olap_test.go index 73b1af5d0d76..817ad8138ac7 100644 --- a/runtime/drivers/duckdb/olap_test.go +++ b/runtime/drivers/duckdb/olap_test.go @@ -209,7 +209,7 @@ func TestClose(t *testing.T) { } func prepareConn(t *testing.T) drivers.Handle { - conn, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := Driver{}.Open(t.Context(), "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := conn.AsOLAP("") @@ -230,7 +230,7 @@ func Test_safeSQLString(t *testing.T) { err := os.Mkdir(path, fs.ModePerm) require.NoError(t, err) - conn, err := Driver{}.Open(context.Background(), "", "default", map[string]any{"data_dir": path}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := Driver{}.Open(t.Context(), "", "default", map[string]any{"data_dir": path}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) require.NotNil(t, conn) require.NoError(t, conn.Close()) diff --git a/runtime/drivers/duckdb/open_ctx_test.go b/runtime/drivers/duckdb/open_ctx_test.go deleted file mode 100644 index c1599fcd8980..000000000000 --- a/runtime/drivers/duckdb/open_ctx_test.go +++ /dev/null @@ -1,45 +0,0 @@ -package duckdb - -import ( - "context" - "testing" - - "github.com/rilldata/rill/runtime/drivers" - "github.com/rilldata/rill/runtime/pkg/activity" - "github.com/rilldata/rill/runtime/storage" - "github.com/stretchr/testify/require" - "go.uber.org/zap" - "golang.org/x/sync/errgroup" -) - -// TestOpenCtxNotRetained verifies that the handle does not retain the ctx passed to Open. -// The connection cache cancels that ctx as soon as Open returns, -// so anything that outlives Open (notably the per-connection init queries, which run for every connection the pool opens later) -// must not observe its cancellation. -func TestOpenCtxNotRetained(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - - handle, err := Driver{}.Open(ctx, "", "default", map[string]any{"data_dir": t.TempDir(), "pool_size": 4}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) - require.NoError(t, err) - defer handle.Close() - require.NoError(t, handle.Migrate(ctx)) - - // Mimic the connection cache, which cancels the open ctx immediately after Open returns. - cancel() - - olap, ok := handle.AsOLAP("default") - require.True(t, ok) - - // Query concurrently so the pool has to open connections it didn't already have at the time of the cancellation. - var g errgroup.Group - for range 8 { - g.Go(func() error { - res, err := olap.Query(context.Background(), &drivers.Statement{Query: "SELECT 1"}) - if err != nil { - return err - } - return res.Close() - }) - } - require.NoError(t, g.Wait()) -} diff --git a/runtime/drivers/duckdb/schema_change_test.go b/runtime/drivers/duckdb/schema_change_test.go index f4a2df675c7d..05d0c4027df8 100644 --- a/runtime/drivers/duckdb/schema_change_test.go +++ b/runtime/drivers/duckdb/schema_change_test.go @@ -293,7 +293,7 @@ func TestInsertTableAsSelectEmptyPartitionSkipsSchemaHandling(t *testing.T) { func newSchemaChangeTestConnection(t *testing.T) *connection { t.Helper() - handle, err := Driver{}.Open(context.Background(), "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := Driver{}.Open(t.Context(), "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) c := handle.(*connection) require.NoError(t, c.Migrate(context.Background())) diff --git a/runtime/drivers/file/model_executor_olap_self_test.go b/runtime/drivers/file/model_executor_olap_self_test.go index 246c23b683d7..1749958774f2 100644 --- a/runtime/drivers/file/model_executor_olap_self_test.go +++ b/runtime/drivers/file/model_executor_olap_self_test.go @@ -17,7 +17,7 @@ import ( func TestExecute(t *testing.T) { tempDir := t.TempDir() - duckDB, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + duckDB, err := drivers.Open(t.Context(), "duckdb", "", "default", map[string]any{}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer duckDB.Close() @@ -74,7 +74,7 @@ func TestExecute(t *testing.T) { require.NoError(t, err) require.NoError(t, result.Close()) - fileHandle, err := drivers.Open(context.Background(), "file", "", "default", map[string]any{}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) + fileHandle, err := drivers.Open(t.Context(), "file", "", "default", map[string]any{}, storage.MustNew(tempDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer fileHandle.Close() diff --git a/runtime/drivers/gcs/object_store_test.go b/runtime/drivers/gcs/object_store_test.go index 6d3495780456..fc8293b6f1c9 100644 --- a/runtime/drivers/gcs/object_store_test.go +++ b/runtime/drivers/gcs/object_store_test.go @@ -1,7 +1,6 @@ package gcs_test import ( - "context" "testing" "github.com/rilldata/rill/runtime/drivers" @@ -16,7 +15,7 @@ import ( func TestObjectStore(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "gcs") - conn, err := drivers.Open(context.Background(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) @@ -55,7 +54,7 @@ func TestObjectStorePathPrefixes(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "gcs") cfg["path_prefixes"] = "gcs://integration-test.rilldata.com/glob_test/" - conn, err := drivers.Open(context.Background(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) @@ -72,7 +71,7 @@ func TestObjectStorePathPrefixes(t *testing.T) { func TestObjectStoreHMAC(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "gcs_s3_compat") - conn, err := drivers.Open(context.Background(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) @@ -111,7 +110,7 @@ func TestObjectStoreHMACPathPrefixes(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "gcs_s3_compat") cfg["path_prefixes"] = "gcs://integration-test.rilldata.com/glob_test/" - conn, err := drivers.Open(context.Background(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "gcs", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/mysql/olap_test.go b/runtime/drivers/mysql/olap_test.go index e4f4dea3c65c..b1400032281a 100644 --- a/runtime/drivers/mysql/olap_test.go +++ b/runtime/drivers/mysql/olap_test.go @@ -1,7 +1,6 @@ package mysql_test import ( - "context" "strings" "testing" @@ -518,7 +517,7 @@ func testLoadDDL(t *testing.T, olap drivers.OLAPStore) { func acquireTestMySQL(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "mysql") - conn, err := drivers.Open(context.Background(), "mysql", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "mysql", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/pinot/information_schema_test.go b/runtime/drivers/pinot/information_schema_test.go index 7862f859c0be..0017bb3ffcea 100644 --- a/runtime/drivers/pinot/information_schema_test.go +++ b/runtime/drivers/pinot/information_schema_test.go @@ -33,7 +33,7 @@ var expectedTables = []string{ func TestInformationSchema(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "pinot") - conn, err := drivers.Open(context.Background(), "pinot", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "pinot", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/postgres/olap_test.go b/runtime/drivers/postgres/olap_test.go index f1eba9a0ad35..d09d79e9beaf 100644 --- a/runtime/drivers/postgres/olap_test.go +++ b/runtime/drivers/postgres/olap_test.go @@ -1,7 +1,6 @@ package postgres_test import ( - "context" "encoding/json" "fmt" "testing" @@ -426,7 +425,7 @@ func testLoadDDL(t *testing.T, olap drivers.OLAPStore) { func acquireTestPostgres(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "postgres") - conn, err := drivers.Open(context.Background(), "postgres", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "postgres", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/redshift/olap_test.go b/runtime/drivers/redshift/olap_test.go index 0f708cb51495..d97e3d2b98a4 100644 --- a/runtime/drivers/redshift/olap_test.go +++ b/runtime/drivers/redshift/olap_test.go @@ -366,7 +366,7 @@ func acquireTestRedshift(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg["database"] = "test_db" } cfg["workgroup"] = "integration-test-wg" - conn, err := drivers.Open(context.Background(), "redshift", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "redshift", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/s3/mode_executor_olap_to_self_test.go b/runtime/drivers/s3/mode_executor_olap_to_self_test.go index cdf4a763c91e..187b44e578b7 100644 --- a/runtime/drivers/s3/mode_executor_olap_to_self_test.go +++ b/runtime/drivers/s3/mode_executor_olap_to_self_test.go @@ -140,7 +140,7 @@ func testExportedObjectExists(t *testing.T, driver string, rt *runtime.Runtime, ctx := context.Background() handle, _, err := rt.AcquireHandle(ctx, id, driver) - conn, err := drivers.Open(context.Background(), driver, "", "default", handle.Config(), storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), driver, "", "default", handle.Config(), storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/s3/object_store_test.go b/runtime/drivers/s3/object_store_test.go index 8925c7dc7510..582dd5667322 100644 --- a/runtime/drivers/s3/object_store_test.go +++ b/runtime/drivers/s3/object_store_test.go @@ -1,7 +1,6 @@ package s3_test import ( - "context" "testing" "github.com/rilldata/rill/runtime/drivers" @@ -16,7 +15,7 @@ import ( func TestObjectStore(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "s3") - conn, err := drivers.Open(context.Background(), "s3", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "s3", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) @@ -55,7 +54,7 @@ func TestObjectStorePathPrefixes(t *testing.T) { testmode.Expensive(t) cfg := testruntime.AcquireConnector(t, "s3") cfg["path_prefixes"] = "s3://integration-test.rilldata.com/glob_test/" - conn, err := drivers.Open(context.Background(), "s3", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "s3", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/snowflake/olap_test.go b/runtime/drivers/snowflake/olap_test.go index bdc4362bff96..2d46efce525f 100644 --- a/runtime/drivers/snowflake/olap_test.go +++ b/runtime/drivers/snowflake/olap_test.go @@ -1,7 +1,6 @@ package snowflake_test import ( - "context" "encoding/json" "strings" "testing" @@ -187,7 +186,7 @@ func TestDryRun(t *testing.T) { func acquireTestSnowflake(t *testing.T) (drivers.Handle, drivers.OLAPStore) { cfg := testruntime.AcquireConnector(t, "snowflake") - conn, err := drivers.Open(context.Background(), "snowflake", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "snowflake", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) t.Cleanup(func() { conn.Close() }) diff --git a/runtime/drivers/sqlite/backups_test.go b/runtime/drivers/sqlite/backups_test.go index 6c690a0d6cfc..4bdb436cf457 100644 --- a/runtime/drivers/sqlite/backups_test.go +++ b/runtime/drivers/sqlite/backups_test.go @@ -1,7 +1,6 @@ package sqlite import ( - "context" "fmt" "path/filepath" "strings" @@ -35,7 +34,7 @@ func TestBackup(t *testing.T) { "id": "test-backup", "backups_enable": true, } - h, err := driver{}.Open(context.Background(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) + h, err := driver{}.Open(t.Context(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer h.Close() err = h.Migrate(t.Context()) @@ -142,7 +141,7 @@ func TestDBFilePath(t *testing.T) { for idx, tc := range cases { t.Run(fmt.Sprintf("case-%d", idx), func(t *testing.T) { cfg := map[string]any{"dsn": tc.dsn} - h, err := driver{}.Open(context.Background(), "", "", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + h, err := driver{}.Open(t.Context(), "", "", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer h.Close() diff --git a/runtime/drivers/sqlite/migrate_test.go b/runtime/drivers/sqlite/migrate_test.go index 60eef2b8a441..d0f1369ff624 100644 --- a/runtime/drivers/sqlite/migrate_test.go +++ b/runtime/drivers/sqlite/migrate_test.go @@ -1,7 +1,6 @@ package sqlite import ( - "context" "path/filepath" "testing" "time" @@ -23,7 +22,7 @@ func TestDeleteExpiredAISessions(t *testing.T) { cfg := map[string]any{"dsn": dbPath} // Open the database, run migrations, and seed test data. - h, err := driver{}.Open(context.Background(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) + h, err := driver{}.Open(t.Context(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer h.Close() require.NoError(t, h.Migrate(t.Context())) diff --git a/runtime/drivers/sqlite/restore_test.go b/runtime/drivers/sqlite/restore_test.go index 0517de8d03a3..20a412011740 100644 --- a/runtime/drivers/sqlite/restore_test.go +++ b/runtime/drivers/sqlite/restore_test.go @@ -1,7 +1,6 @@ package sqlite import ( - "context" "fmt" "os" "path/filepath" @@ -34,7 +33,7 @@ func TestBackupAndRestore(t *testing.T) { } // Create a database with an instance in it, then back it up. - h, err := driver{}.Open(context.Background(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) + h, err := driver{}.Open(t.Context(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) require.NoError(t, h.Migrate(t.Context())) registry, ok := h.AsRegistry() @@ -48,7 +47,7 @@ func TestBackupAndRestore(t *testing.T) { require.NoError(t, restoreBackup(t.Context(), bucket, dbPath, zap.NewNop())) // Reopen the database and check the instance is back. - h, err = driver{}.Open(context.Background(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) + h, err = driver{}.Open(t.Context(), "", "", cfg, storage.MustNew(storageDir, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) defer h.Close() require.NoError(t, h.Migrate(t.Context())) @@ -125,7 +124,7 @@ func TestShouldRestoreBackup(t *testing.T) { require.Equal(t, dsn, dbPath) // Once migrated, the database has data and must not be overwritten by a restore. - h, err := driver{}.Open(context.Background(), "", "", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + h, err := driver{}.Open(t.Context(), "", "", map[string]any{"dsn": dsn}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) require.NoError(t, h.Migrate(t.Context())) require.NoError(t, h.Close()) diff --git a/runtime/drivers/starrocks/olap_test.go b/runtime/drivers/starrocks/olap_test.go index f2844557f185..3aaba36dadee 100644 --- a/runtime/drivers/starrocks/olap_test.go +++ b/runtime/drivers/starrocks/olap_test.go @@ -21,7 +21,7 @@ func TestStarRocksOLAP(t *testing.T) { dsn := teststarrocks.StartWithData(t) - conn, err := driver{}.Open(context.Background(), "", "default", map[string]any{ + conn, err := driver{}.Open(t.Context(), "", "default", map[string]any{ "dsn": dsn, }, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) diff --git a/runtime/parser/parser_test.go b/runtime/parser/parser_test.go index 185dd8512278..d009010e10a3 100644 --- a/runtime/parser/parser_test.go +++ b/runtime/parser/parser_test.go @@ -3167,7 +3167,7 @@ func requireResourcesAndErrors(t testing.TB, p *Parser, wantResources []*Resourc func makeRepo(t testing.TB, files map[string]string) drivers.RepoStore { root := t.TempDir() - handle, err := drivers.Open(context.Background(), "file", "", "default", map[string]any{"dsn": root}, storage.MustNew(root, nil), activity.NewNoopClient(), zap.NewNop()) + handle, err := drivers.Open(t.Context(), "file", "", "default", map[string]any{"dsn": root}, storage.MustNew(root, nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) repo, ok := handle.AsRepoStore("") diff --git a/runtime/server/queries_test.go b/runtime/server/queries_test.go index f52a17e3d0c7..f910ca12c791 100644 --- a/runtime/server/queries_test.go +++ b/runtime/server/queries_test.go @@ -133,7 +133,7 @@ func TestServer_UpdateLimit_UNION(t *testing.T) { } func prepareOLAPStore(t *testing.T) drivers.OLAPStore { - conn, err := drivers.Open(context.Background(), "duckdb", "", "default", map[string]any{"dsn": ":memory:?access_mode=read_write", "pool_size": 4}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + conn, err := drivers.Open(t.Context(), "duckdb", "", "default", map[string]any{"dsn": ":memory:?access_mode=read_write", "pool_size": 4}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) require.NoError(t, err) olap, ok := conn.AsOLAP("") require.True(t, ok) diff --git a/runtime/storage/storage.go b/runtime/storage/storage.go index 49381c4e815f..244400ed45c3 100644 --- a/runtime/storage/storage.go +++ b/runtime/storage/storage.go @@ -170,8 +170,10 @@ func (c *Client) path(base string, elem ...string) string { } func (c *Client) newGCPClient(ctx context.Context) (*gcp.HTTPClient, error) { - // Detach cancellation: the credentials capture the ctx and reuse it to refresh tokens for as long as the bucket is used. - // Some credential types bind that ctx to the refresh request, so a cancelled ctx would permanently break a long-lived bucket. + // The credentials capture the ctx and reuse it to refresh tokens for as long as the bucket is used, + // and some credential types bind it to the refresh request, so a cancelled ctx would permanently break a long-lived bucket. + // Detaching cancellation costs nothing here: the call does no I/O on this path. + // It only parses the JSON into a lazily-refreshing token source, and never reaches the metadata server since allowHostAccess is false. creds, err := gcputil.Credentials(context.WithoutCancel(ctx), c.bucketConfig.GoogleApplicationCredentialsJSON, false) if err != nil { return nil, err From fa337c2bfc7bdb9650f64cc6b1867eab13398cbf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Benjamin=20Egelund-M=C3=BCller?= Date: Thu, 1 Oct 2026 17:32:28 +0200 Subject: [PATCH 5/5] Scope the cached BigQuery client's ctx to the handle, not the first request --- runtime/drivers/bigquery/bigquery.go | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/runtime/drivers/bigquery/bigquery.go b/runtime/drivers/bigquery/bigquery.go index cfe2836590c1..77e9dc73eece 100644 --- a/runtime/drivers/bigquery/bigquery.go +++ b/runtime/drivers/bigquery/bigquery.go @@ -78,10 +78,14 @@ func (d driver) Open(_ context.Context, _, instanceID string, config map[string] 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