diff --git a/conn.go b/conn.go new file mode 100644 index 0000000..86e4dae --- /dev/null +++ b/conn.go @@ -0,0 +1,50 @@ +package qube + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "fmt" +) + +type Conn struct { + db *sql.DB + raw *sql.Conn +} + +func (conn *Conn) Exec(query string, args ...any) (sql.Result, error) { + // Avoid "bad connection". + return conn.withRetry(context.Background(), query, args...) +} + +func (conn *Conn) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) { + // Avoid "bad connection". + return conn.withRetry(ctx, query, args...) +} + +func (conn *Conn) withRetry(ctx context.Context, query string, args ...any) (sql.Result, error) { + res, err := conn.raw.ExecContext(ctx, query, args...) + + if errors.Is(err, driver.ErrBadConn) || errors.Is(err, sql.ErrConnDone) { + if !errors.Is(err, sql.ErrConnDone) { + conn.raw.Close() + } + + raw, err := conn.db.Conn(ctx) + + if err != nil { + return nil, fmt.Errorf("failed to reopen DB connection (%w)", err) + } + + conn.raw = raw + return conn.raw.ExecContext(ctx, query, args...) + } + + return res, err +} + +func (conn *Conn) Close() { + conn.raw.Close() + conn.db.Close() +} diff --git a/db_config.go b/db_config.go index 2555669..3511f1d 100644 --- a/db_config.go +++ b/db_config.go @@ -92,20 +92,32 @@ func (config *DBConfig) OpenDBWithPing(autoCommit bool) (DBIface, error) { return nil, fmt.Errorf("failed to ping DB (%w)", err) } + rawConn, err := db.Conn(context.Background()) + + if err != nil { + db.Close() + return nil, fmt.Errorf("failed to get a single DB connection (%w)", err) + } + + conn := &Conn{ + db: db, + raw: rawConn, + } + if config.Driver == DBDriverMySQL { if autoCommit { - _, err = db.Exec("set autocommit = 1") + _, err = conn.Exec("set autocommit = 1") } else { - _, err = db.Exec("set autocommit = 0") + _, err = conn.Exec("set autocommit = 0") } if err != nil { - db.Close() + conn.Close() return nil, fmt.Errorf("failed to disable autocommit (%w)", err) } } - return db, nil + return conn, nil } func (cfg *DBConfig) getMySQLConnector() (driver.Connector, error) { diff --git a/db_iface.go b/db_iface.go index fef1172..e01933e 100644 --- a/db_iface.go +++ b/db_iface.go @@ -8,5 +8,5 @@ import ( type DBIface interface { Exec(query string, args ...any) (sql.Result, error) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) - Close() error + Close() } diff --git a/go.mod b/go.mod index ed2b32b..6da03db 100644 --- a/go.mod +++ b/go.mod @@ -10,9 +10,9 @@ require ( github.com/creack/pty v1.1.24 github.com/go-sql-driver/mysql v1.9.3 github.com/google/uuid v1.6.0 - github.com/jackc/pgx/v5 v5.8.0 + github.com/jackc/pgx/v5 v5.9.1 github.com/jamiealquiza/tachymeter v2.0.0+incompatible - github.com/klauspost/compress v1.18.4 + github.com/klauspost/compress v1.18.5 github.com/mattn/go-isatty v0.0.20 github.com/neilotoole/jsoncolor v0.7.2 github.com/shirou/gopsutil/v4 v4.26.2 diff --git a/go.sum b/go.sum index b070e4b..3b3f9da 100644 --- a/go.sum +++ b/go.sum @@ -66,14 +66,14 @@ github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsI github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= -github.com/jackc/pgx/v5 v5.8.0 h1:TYPDoleBBme0xGSAX3/+NujXXtpZn9HBONkQC7IEZSo= -github.com/jackc/pgx/v5 v5.8.0/go.mod h1:QVeDInX2m9VyzvNeiCJVjCkNFqzsNb43204HshNSZKw= +github.com/jackc/pgx/v5 v5.9.1 h1:uwrxJXBnx76nyISkhr33kQLlUqjv7et7b9FjCen/tdc= +github.com/jackc/pgx/v5 v5.9.1/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/jamiealquiza/tachymeter v2.0.0+incompatible h1:mGiF1DGo8l6vnGT8FXNNcIXht/YmjzfraiUprXYwJ6g= github.com/jamiealquiza/tachymeter v2.0.0+incompatible/go.mod h1:Ayf6zPZKEnLsc3winWEXJRkTBhdHo58HODAu1oFJkYU= -github.com/klauspost/compress v1.18.4 h1:RPhnKRAQ4Fh8zU2FY/6ZFDwTVTxgJ/EMydqSTzE9a2c= -github.com/klauspost/compress v1.18.4/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4= +github.com/klauspost/compress v1.18.5 h1:/h1gH5Ce+VWNLSWqPzOVn6XBO+vJbCNGvjoaGBFW2IE= +github.com/klauspost/compress v1.18.5/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0= github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= diff --git a/null_db.go b/null_db.go index 6cbdac2..8dfc400 100644 --- a/null_db.go +++ b/null_db.go @@ -21,6 +21,5 @@ func (db *NullDB) ExecContext(ctx context.Context, query string, args ...any) (s return nil, nil } -func (db *NullDB) Close() error { - return nil +func (db *NullDB) Close() { } diff --git a/null_db_test.go b/null_db_test.go index 3f452cd..fe01a53 100644 --- a/null_db_test.go +++ b/null_db_test.go @@ -23,8 +23,7 @@ func Test_NullDB(t *testing.T) { _, err = db.ExecContext(context.Background(), "select 2") require.NoError(err) - err = db.Close() - require.NoError(err) + db.Close() assert.Equal("select 1\nselect 2\n", buf.String()) } diff --git a/task_test.go b/task_test.go index b83a6b3..da48f44 100644 --- a/task_test.go +++ b/task_test.go @@ -338,14 +338,12 @@ func TestAcc_Task_MySQLSession(t *testing.T) { db, err := sql.Open("mysql", testDSN_MySQL) require.NoError(err) + _, err = db.Exec("drop table if exists qube_test") + require.NoError(err) _, err = db.Exec("create table qube_test (data text not null)") require.NoError(err) - t.Cleanup(func() { - _, err := db.Exec("drop table qube_test") - require.NoError(err) - db.Close() - }) + t.Cleanup(func() { db.Close() }) f, _ := os.CreateTemp("", "") defer os.Remove(f.Name()) @@ -404,14 +402,12 @@ func TestAcc_Task_PostgreSQLSession(t *testing.T) { db, err := sql.Open("pgx", testDSN_PostgreSQL) require.NoError(err) + _, err = db.Exec("drop table if exists qube_test") + require.NoError(err) _, err = db.Exec("create table qube_test (data text not null)") require.NoError(err) - t.Cleanup(func() { - _, err := db.Exec("drop table qube_test") - require.NoError(err) - db.Close() - }) + t.Cleanup(func() { db.Close() }) f, _ := os.CreateTemp("", "") defer os.Remove(f.Name()) @@ -470,14 +466,12 @@ func TestAcc_Task_MySQLTx(t *testing.T) { db, err := sql.Open("mysql", testDSN_MySQL) require.NoError(err) + _, err = db.Exec("drop table if exists qube_test") + require.NoError(err) _, err = db.Exec("create table qube_test (data text not null)") require.NoError(err) - t.Cleanup(func() { - _, err := db.Exec("drop table qube_test") - require.NoError(err) - db.Close() - }) + t.Cleanup(func() { db.Close() }) f, _ := os.CreateTemp("", "") defer os.Remove(f.Name()) @@ -537,14 +531,12 @@ func TestAcc_Task_PostgreSQLTx(t *testing.T) { db, err := sql.Open("pgx", testDSN_PostgreSQL) require.NoError(err) + _, err = db.Exec("drop table if exists qube_test") + require.NoError(err) _, err = db.Exec("create table qube_test (data text not null)") require.NoError(err) - t.Cleanup(func() { - _, err := db.Exec("drop table qube_test") - require.NoError(err) - db.Close() - }) + t.Cleanup(func() { db.Close() }) f, _ := os.CreateTemp("", "") defer os.Remove(f.Name())