fix: preserve connection contexts in SQL driver wrappers

driver_context_test.go

@@ -0,0 +1,160 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cloudsqlconn_test
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"errors"
+	"fmt"
+	"net/http"
+	"reflect"
+	"testing"
+	"time"
+
+	"cloud.google.com/go/cloudsqlconn"
+	"cloud.google.com/go/cloudsqlconn/internal/mock"
+	"cloud.google.com/go/cloudsqlconn/mysql/mysql"
+	"cloud.google.com/go/cloudsqlconn/postgres/pgxv4"
+	"cloud.google.com/go/cloudsqlconn/postgres/pgxv5"
+	"cloud.google.com/go/cloudsqlconn/sqlserver/mssql"
+	"github.com/google/uuid"
+)
+
+func TestRegisteredDriverConnectionContext(t *testing.T) {
+	for _, drv := range []struct {
+		name     string
+		register func(string, ...cloudsqlconn.Option) (func() error, error)
+	}{
+		{"mysql", mysql.RegisterDriver},
+		{"pgxv5", pgxv5.RegisterDriver},
+		{"pgxv4", pgxv4.RegisterDriver},
+		{"mssql", mssql.RegisterDriver},
+	} {
+		for _, mode := range []string{"cancel", "deadline"} {
+			t.Run(drv.name+"/"+mode, func(t *testing.T) {
+				entered, release := make(chan struct{}, 1), make(chan struct{})
+				transport := driverContextTransport(func(r *http.Request) (*http.Response, error) {
+					select {
+					case entered <- struct{}{}:
+					default:
+					}
+					select {
+					case <-release:
+						return nil, context.Canceled
+					case <-r.Context().Done():
+						return nil, r.Context().Err()
+					}
+				})
+				name := "cloudsql-context-" + uuid.NewString()
+				cleanup, err := drv.register(name,
+					cloudsqlconn.WithTokenSource(mock.EmptyTokenSource{}),
+					cloudsqlconn.WithHTTPClient(&http.Client{Transport: transport}),
+				)
+				if err != nil {
+					t.Fatal(err)
+				}
+				t.Cleanup(func() {
+					close(release)
+					if err := cleanup(); err != nil {
+						t.Error(err)
+					}
+				})
+				dsn := "host=test-project:us-central1:test-instance user=test dbname=test sslmode=disable"
+				switch drv.name {
+				case "mysql":
+					dsn = fmt.Sprintf("test@%s(test-project:us-central1:test-instance)/test", name)
+				case "mssql":
+					dsn = "sqlserver://test:password@localhost?database=test&cloudsql=test-project:us-central1:test-instance"
+				}
+				db, err := sql.Open(name, dsn)
+				if err != nil {
+					t.Fatal(err)
+				}
+				t.Cleanup(func() {
+					if err := db.Close(); err != nil {
+						t.Error(err)
+					}
+				})
+				db.SetMaxOpenConns(1)
+
+				ctx, cancel := context.WithCancel(context.Background())
+				wantErr := context.Canceled
+				if mode == "deadline" {
+					cancel()
+					ctx, cancel = context.WithTimeout(context.Background(), time.Second)
+					wantErr = context.DeadlineExceeded
+				}
+				defer cancel()
+				done := make(chan error, 1)
+				go func() { done <- db.PingContext(ctx) }()
+				select {
+				case <-entered:
+				case err := <-done:
+					t.Fatalf("connection returned before reaching the API: %v", err)
+				case <-time.After(5 * time.Second):
+					t.Fatal("connection did not reach the test API transport")
+				}
+				if mode == "cancel" {
+					cancel()
+				}
+				// The API remains blocked. The caller's context must release the slot.
+				select {
+				case err := <-done:
+					if !errors.Is(err, wantErr) {
+						t.Fatalf("PingContext error = %v, want %v", err, wantErr)
+					}
+				case <-time.After(2 * time.Second):
+					t.Fatal("connection ignored the caller's context")
+				}
+				if stats := db.Stats(); stats.OpenConnections != 0 || stats.InUse != 0 {
+					t.Fatalf("connection attempt retained a pool slot: %+v", stats)
+				}
+
+				driverCtx, ok := db.Driver().(driver.DriverContext)
+				if !ok {
+					t.Fatal("DB.Driver does not implement driver.DriverContext")
+				}
+				wantPackage := "cloud.google.com/go/cloudsqlconn/postgres/pgxv5"
+				switch drv.name {
+				case "mysql":
+					wantPackage = "cloud.google.com/go/cloudsqlconn/mysql/mysql"
+				case "mssql":
+					wantPackage = "cloud.google.com/go/cloudsqlconn/sqlserver/mssql"
+				}
+				if got := reflect.TypeOf(db.Driver()).Elem().PkgPath(); got != wantPackage {
+					t.Fatalf("DB.Driver package = %q, want %q", got, wantPackage)
+				}
+				connector, err := driverCtx.OpenConnector(dsn)
+				if err != nil {
+					t.Fatal(err)
+				}
+				if connector.Driver() != db.Driver() {
+					t.Fatal("connector did not preserve the registered driver")
+				}
+				if _, err := driverCtx.OpenConnector("port=invalid"); err == nil {
+					t.Fatal("OpenConnector accepted an invalid DSN")
+				}
+			})
+		}
+	}
+}
+
+type driverContextTransport func(*http.Request) (*http.Response, error)
+
+func (f driverContextTransport) RoundTrip(r *http.Request) (*http.Response, error) {
+	return f(r)
+}

mysql/mysql/mysql.go

@@ -82,3 +82,20 @@ type mysqlDriver struct {
 func (d *mysqlDriver) Open(name string) (driver.Conn, error) {
 	return d.d.Open(name)
 }
+
+// OpenConnector implements driver.DriverContext, preserving the caller's context
+// when database/sql opens a connection.
+func (d *mysqlDriver) OpenConnector(name string) (driver.Connector, error) {
+	c, err := d.d.OpenConnector(name)
+	if err != nil {
+		return nil, err
+	}
+	return &mysqlConnector{Connector: c, driver: d}, nil
+}
+
+type mysqlConnector struct {
+	driver.Connector
+	driver *mysqlDriver
+}
+
+func (c *mysqlConnector) Driver() driver.Driver { return c.driver }

postgres/pgxv5/postgres.go

@@ -67,6 +67,27 @@ func (p *pgDriver) Open(name string) (driver.Conn, error) {
 
 }
 
+// OpenConnector implements driver.DriverContext, preserving the caller's context
+// when database/sql opens a connection.
+func (p *pgDriver) OpenConnector(name string) (driver.Connector, error) {
+	dbURI, err := p.dbURI(name)
+	if err != nil {
+		return nil, err
+	}
+	c, err := stdlib.GetDefaultDriver().(driver.DriverContext).OpenConnector(dbURI)
+	if err != nil {
+		return nil, err
+	}
+	return &pgConnector{Connector: c, driver: p}, nil
+}
+
+type pgConnector struct {
+	driver.Connector
+	driver *pgDriver
+}
+
+func (c *pgConnector) Driver() driver.Driver { return c.driver }
+
 // dbURI registers a driver using the provided DSN. If the name has already
 // been registered, dbURI returns the existing registration.
 func (p *pgDriver) dbURI(name string) (string, error) {

sqlserver/mssql/mssql.go

@@ -74,6 +74,16 @@ type sqlserverDriver struct {
 // For details, see
 // https://github.com/microsoft/go-mssqldb#the-connection-string-can-be-specified-in-one-of-three-formats
 func (s *sqlserverDriver) Open(name string) (driver.Conn, error) {
+	c, err := s.OpenConnector(name)
+	if err != nil {
+		return nil, err
+	}
+	return c.Connect(context.Background())
+}
+
+// OpenConnector implements driver.DriverContext, preserving the caller's context
+// when database/sql opens a connection.
+func (s *sqlserverDriver) OpenConnector(name string) (driver.Connector, error) {
 	res, err := msdsn.Parse(name)
 	if err != nil {
 		return nil, err
@@ -87,5 +97,12 @@ func (s *sqlserverDriver) Open(name string) (driver.Conn, error) {
 		d:        s.d,
 		connName: connName,
 	}
-	return c.Connect(context.Background())
+	return &sqlserverConnector{Connector: c, driver: s}, nil
 }
+
+type sqlserverConnector struct {
+	driver.Connector
+	driver *sqlserverDriver
+}
+
+func (c *sqlserverConnector) Driver() driver.Driver { return c.driver }