integration-tests/go-sql-server-driver: Move max-connections testing to go-sql-server-driver.

go/go.mod

@@ -18,7 +18,7 @@ require (
 	github.com/dustin/go-humanize v1.0.1
 	github.com/fatih/color v1.13.0
 	github.com/flynn-archive/go-shlex v0.0.0-20150515145356-3f9db97f8568
-	github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d
+	github.com/go-sql-driver/mysql v1.9.1
 	github.com/gocraft/dbr/v2 v2.7.2
 	github.com/golang/snappy v0.0.4
 	github.com/google/uuid v1.6.0

go/go.sum

@@ -297,8 +297,8 @@ github.com/go-sql-driver/mysql v1.4.0/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG
 github.com/go-sql-driver/mysql v1.4.1/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w=
 github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
 github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
-github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d h1:QQP1nE4qh5aHTGvI1LgOFxZYVxYoGeMfbNHikogPyoA=
-github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
+github.com/go-sql-driver/mysql v1.9.1 h1:FrjNGn/BsJQjVRuSa8CBrM5BWA9BWoXXat3KrtSb/iI=
+github.com/go-sql-driver/mysql v1.9.1/go.mod h1:qn46aNg1333BRMNU69Lq93t8du/dwxI64Gl8i5p1WMU=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/gobwas/glob v0.2.3/go.mod h1:d3Ez4x06l9bZtSvzIay5+Yzi0fmZzPgnTbPcKjJAkT8=
 github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=

go/libraries/doltcore/dtestutils/sql_server_driver/cmd.go

@@ -17,7 +17,9 @@ package sql_server_driver
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"database/sql"
+	"database/sql/driver"
 	"fmt"
 	"io"
 	"log"
@@ -28,6 +30,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/go-sql-driver/mysql"
 	_ "github.com/go-sql-driver/mysql"
 )
 
@@ -379,23 +382,66 @@ func (s *SqlServer) Restart(newargs *[]string, newenvs *[]string) error {
 }
 
 func (s *SqlServer) DB(c Connection) (*sql.DB, error) {
-	var pass string
+	connector, err := s.Connector(c)
+	if err != nil {
+		return nil, err
+	}
+	return OpenDB(connector)
+}
+
+// If a test needs to circumvent the database/sql connection pool
+// it can use the raw MySQL connector.
+func (s *SqlServer) Connector(c Connection) (driver.Connector, error) {
 	pass, err := c.Password()
 	if err != nil {
 		return nil, err
 	}
-	return ConnectDB(c.User, pass, s.DBName, "127.0.0.1", s.Port, c.DriverParams)
+	dsn := GetDSN(c.User, pass, s.DBName, "127.0.0.1", s.Port, c.DriverParams)
+	cfg, err := mysql.ParseDSN(dsn)
+	if err != nil {
+		return nil, err
+	}
+	// See the comment on WithConnectRetriesDisabled for why we do this.
+	cfg.Apply(mysql.BeforeConnect(func(ctx context.Context, cfg *mysql.Config) error {
+		// TODO: This could be more robust if we sniffed it on first connect.
+		const numAttemptsGoLibraryMakes = 3
+		if attempt, ok := incrementConnectRetryAttempts(ctx); ok && attempt < numAttemptsGoLibraryMakes {
+			return driver.ErrBadConn
+		}
+		return nil
+	}))
+	return mysql.NewConnector(cfg)
 }
 
-func ConnectDB(user, password, name, host string, port int, driverParams map[string]string) (*sql.DB, error) {
+func GetDSN(user, password, name, host string, port int, driverParams map[string]string) string {
 	params := make(url.Values)
 	params.Set("allowAllFiles", "true")
 	params.Set("tls", "preferred")
 	for k, v := range driverParams {
 		params.Set(k, v)
 	}
 	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?%s", user, password, host, port, name, params.Encode())
+	return dsn
+}
 
+func OpenDB(connector driver.Connector) (*sql.DB, error) {
+	db := sql.OpenDB(connector)
+	var err error
+	for i := 0; i < ConnectAttempts; i++ {
+		err = db.Ping()
+		if err == nil {
+			return db, nil
+		}
+		time.Sleep(RetrySleepDuration)
+	}
+	if err != nil {
+		return nil, err
+	}
+	return db, nil
+}
+
+func ConnectDB(user, password, name, host string, port int, params map[string]string) (*sql.DB, error) {
+	dsn := GetDSN(user, password, name, host, port, params)
 	db, err := sql.Open("mysql", dsn)
 	if err != nil {
 		return nil, err
@@ -412,3 +458,41 @@ func ConnectDB(user, password, name, host string, port int, driverParams map[str
 	}
 	return db, nil
 }
+
+type connectRetryAttemptKeyType int
+
+var connectRetryAttemptKey connectRetryAttemptKeyType
+
+// The database/sql package in Go takes connections out of a
+// connection pool or opens a new connection to the database. It has
+// logic in it that looks for driver.ErrBadConn responses when it is
+// opening a connection and automatically retries those connections a
+// fixed number of times before actually surfacing the error to the
+// caller. This behavior interferes with some testing we want to do
+// against the behavior of the server.
+//
+// WithConnectionRetriesDisabled is a hack which circumvents these
+// retries. It works by embedding a counter into the returned context,
+// which should then be passed to *sql.DB.Conn(). An interceptor has
+// been installed on the |mysql.Connector| which looks for this
+// counter and fast-fails the first few calls into the driver. That
+// way the first call which goes through to the driver is the last and
+// final retry from *sql.DB.
+func WithConnectRetriesDisabled(ctx context.Context) context.Context {
+	return context.WithValue(ctx, connectRetryAttemptKey, &retryAttempt{})
+}
+
+type retryAttempt struct {
+	attempt int
+}
+
+func incrementConnectRetryAttempts(ctx context.Context) (int, bool) {
+	v := ctx.Value(connectRetryAttemptKey)
+	if v != nil {
+		if v, ok := v.(*retryAttempt); ok {
+			v.attempt += 1
+			return v.attempt, true
+		}
+	}
+	return 0, false
+}

integration-tests/bats/sql-server.bats

@@ -2115,167 +2115,3 @@ EOF
     [[ "$output" =~ "br3  | true" ]] || false
     [[ "$output" =~ "main | false" ]] || false
 }
-
-@test "sql-server: test --max-connections and --back-log flags" {
-    cd repo1
-    # 3 connections allowed, everything else immediate failure.
-    start_sql_server_with_args --max-connections 3 --back-log 0
-
-    # Start 3 long running connections.
-    set +e # errors expected.
-    pids=()
-    for i in {1..3}; do
-      dolt sql &
-      pids+=($!)
-    done
-
-    # Give the first 3 connections time to start/connect
-    sleep 5
-
-    run dolt sql -q "select(1);"
-    [ $status -ne 0 ]
-    [[ "$output" =~ "bad connection" ]] || false
-
-    # Kill all the running shells
-    for pid in "${pids[@]}"; do
-      kill -9 "$pid"
-    done
-}
-
-@test "sql-server: test --max-connections 3 flag" {
-    cd repo1
-
-    dolt sql -q "CREATE TABLE test_table (
-        id INT AUTO_INCREMENT PRIMARY KEY,
-        str VARCHAR(20)
-    );"
-    dolt commit -A -m "create test_table"
-
-    # 3 connections allowed, default back-log (50).
-    start_sql_server_with_args --max-connections 3
-
-    set +e # errors expected.
-
-    # Start 3 long running connections.
-    pids=()
-    for i in {1..3}; do
-      dolt sql &
-      pids+=($!)
-    done
-
-    sleep 1 # give all connections a chance to start
-
-    # These jobs will wait until there is a connection available. verify
-    # by ensuring the inserts complete.
-    dolt sql -q "insert into test_table (str) values ('test4223');" &
-    second_to_last_pid=$!
-    dolt sql -q "insert into test_table (str) values ('test9119');" &
-    last_pid=$!
-
-    sleep 1
-
-    # Now have two jobs waiting for a connection.
-
-    # Kill all the running shells
-    for pid in "${pids[@]}"; do
-      kill -9 "$pid"
-    done
-
-    wait "$second_to_last_pid"
-    wait "$last_pid"
-
-    run dolt sql -q "select * from test_table;"
-    [ $status -eq 0 ]
-    [[ "$output" =~ "test4223" ]] || false
-    [[ "$output" =~ "test9119" ]] || false
-}
-
-@test "sql-server: test --max-connections 3 and --back-log 1 flags" {
-    cd repo1
-
-    dolt sql -q "CREATE TABLE test_table (
-        id INT AUTO_INCREMENT PRIMARY KEY,
-        str VARCHAR(20)
-    );"
-    dolt commit -A -m "create test_table"
-
-    # 3 connections allowed, 1 connection in back-log.
-    start_sql_server_with_args --max-connections 3 --back-log 1
-
-    set +e # errors expected.
-    # Start 3 long running connections.
-    pids=()
-    for i in {1..3}; do
-      dolt sql &
-      pids+=($!)
-    done
-
-    # This job will wait until there is a connection available. verify
-    # by ensuring the insert completes.
-    dolt sql -q "insert into test_table (str) values ('test4223');" &
-    last_pid=$!
-
-    sleep 1 # give all connections a chance to start
-
-    # This operation should fail immediately.
-    run dolt sql -q "insert into test_table (str) values ('testxxxx');"
-    [ "$status" -ne 0 ]
-    [[ "$stderr" =~ "bad connection" ]] || false
-
-    # Kill all the running shells
-    for pid in "${pids[@]}"; do
-      kill -9 "$pid"
-    done
-
-    wait "$last_pid"
-
-    run dolt sql -q "select * from test_table;"
-    [ $status -eq 0 ]
-    [[ "$output" =~ "test4223" ]] || false
-    [[ ! "$output" =~ "testxxxx" ]] || false
-}
-
-# bats test_tags=no_lambda
-@test "sql-server: test --max-connections-timeout 10s and --max-connections 3 flags" {
-    skiponwindows "mysql client required"
-
-    cd repo1
-
-    export PORT=$( definePORT )
-
-    # Default is 60s, but I don't want to extend the test time for this.
-    start_sql_server_with_args_no_port --max-connections-timeout=10s --max-connections=3 --port=$PORT
-
-    # For this test we use the mysql client, which doesn't retry connections, but also seems to exit early
-    # if we don't give it a query. So we use a sleep to hang the three connections. These will be killed, so
-    # they won't run the full 30s.
-    pids=()
-    for i in {1..3}; do
-      mysql -h 127.0.0.1 -P $PORT -u root -D repo1 -e "select sleep(30)" &
-      pids+=($!)
-    done
-
-    sleep 3
-
-    # Attempt to connect with a fourth connection - should fail after 10s. `run` caused a lot of variance in the test
-    # time, so we avoid using it here.
-    start_time=$(date +%s)
-    set +e
-    mysql -h 127.0.0.1 -P $PORT -u root > $BATS_TMPDIR/mysql.out
-    status=$?
-    end_time=$(date +%s)
-    [ $status -ne 0 ]
-
-    # There is a high amount of variance in the time it takes to fail, so we just check that it is within a range
-    # with a pretty high upper bound.
-    elapsed_time=$((end_time - start_time))
-    [[ $elapsed_time -lt 15 ]] || false
-    [[ $elapsed_time -gt 9 ]] || false
-
-    run cat $BATS_TMPDIR/mysql.out
-    [[ "$output" =~ "Lost connection to MySQL server at 'reading initial communication packet'" ]] || false
-
-    for pid in "${pids[@]}"; do
-      kill "$pid" 2>/dev/null
-    done
-}

integration-tests/go-sql-server-driver/go.mod

@@ -15,7 +15,7 @@ require (
 	filippo.io/edwards25519 v1.1.0 // indirect
 	github.com/creasty/defaults v1.6.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d // indirect
+	github.com/go-sql-driver/mysql v1.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	golang.org/x/crypto v0.35.0 // indirect
 	golang.org/x/sys v0.30.0 // indirect

integration-tests/go-sql-server-driver/go.sum

@@ -4,8 +4,8 @@ github.com/creasty/defaults v1.6.0 h1:ltuE9cfphUtlrBeomuu8PEyISTXnxqkBIoQfXgv7BS
 github.com/creasty/defaults v1.6.0/go.mod h1:iGzKe6pbEHnpMPtfDXZEr0NVxWnPTjb1bbDy08fPzYM=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d h1:QQP1nE4qh5aHTGvI1LgOFxZYVxYoGeMfbNHikogPyoA=
-github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
+github.com/go-sql-driver/mysql v1.9.1 h1:FrjNGn/BsJQjVRuSa8CBrM5BWA9BWoXXat3KrtSb/iI=
+github.com/go-sql-driver/mysql v1.9.1/go.mod h1:qn46aNg1333BRMNU69Lq93t8du/dwxI64Gl8i5p1WMU=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=