From 16f22803e74ceac3ea5a5a905794b8fc77b769ee Mon Sep 17 00:00:00 2001 From: Sahil Muthoo Date: Fri, 28 Jun 2024 12:00:38 +1000 Subject: [PATCH] wip --- pkg/dbconn/conn.go | 1 + pkg/dbconn/conn_test.go | 32 ++++++++++++++++++++++++++++---- 2 files changed, 29 insertions(+), 4 deletions(-) diff --git a/pkg/dbconn/conn.go b/pkg/dbconn/conn.go index 6c604c27..58f147d1 100644 --- a/pkg/dbconn/conn.go +++ b/pkg/dbconn/conn.go @@ -3032,6 +3032,7 @@ func newDSN(dsn string, config *DBConfig) (string, error) { // character_set_client, character_set_connection, character_set_results ops = append(ops, fmt.Sprintf("%s=%s", "charset", "binary")) ops = append(ops, fmt.Sprintf("%s=%s", "collation", "binary")) + ops = append(ops, fmt.Sprintf("%s=%s", "rejectReadOnly", "true")) dsn = fmt.Sprintf("%s?%s", dsn, strings.Join(ops, "&")) return dsn, nil } diff --git a/pkg/dbconn/conn_test.go b/pkg/dbconn/conn_test.go index 6f6010d6..c9f53c2f 100644 --- a/pkg/dbconn/conn_test.go +++ b/pkg/dbconn/conn_test.go @@ -13,25 +13,25 @@ func TestNewDSN(t *testing.T) { dsn := "root:password@tcp(127.0.0.1:3306)/test" resp, err := newDSN(dsn, NewDBConfig()) assert.NoError(t, err) - assert.Equal(t, "root:password@tcp(127.0.0.1:3306)/test?sql_mode=%22%22&time_zone=%22%2B00%3A00%22&innodb_lock_wait_timeout=3&lock_wait_timeout=30&range_optimizer_max_mem_size=0&transaction_isolation=%22read-committed%22&charset=binary&collation=binary", resp) + assert.Equal(t, "root:password@tcp(127.0.0.1:3306)/test?sql_mode=%22%22&time_zone=%22%2B00%3A00%22&innodb_lock_wait_timeout=3&lock_wait_timeout=30&range_optimizer_max_mem_size=0&transaction_isolation=%22read-committed%22&charset=binary&collation=binary&rejectReadOnly=true", resp) // Also without TLS options dsn = "root:password@tcp(mydbhost.internal:3306)/test" resp, err = newDSN(dsn, NewDBConfig()) assert.NoError(t, err) - assert.Equal(t, "root:password@tcp(mydbhost.internal:3306)/test?sql_mode=%22%22&time_zone=%22%2B00%3A00%22&innodb_lock_wait_timeout=3&lock_wait_timeout=30&range_optimizer_max_mem_size=0&transaction_isolation=%22read-committed%22&charset=binary&collation=binary", resp) // unchanged + assert.Equal(t, "root:password@tcp(mydbhost.internal:3306)/test?sql_mode=%22%22&time_zone=%22%2B00%3A00%22&innodb_lock_wait_timeout=3&lock_wait_timeout=30&range_optimizer_max_mem_size=0&transaction_isolation=%22read-committed%22&charset=binary&collation=binary&rejectReadOnly=true", resp) // unchanged // However, if it is RDS - it will be changed. dsn = "root:password@tcp(tern-001.cluster-ro-ckxxxxxxvm.us-west-2.rds.amazonaws.com)/test" resp, err = newDSN(dsn, NewDBConfig()) assert.NoError(t, err) - assert.Equal(t, "root:password@tcp(tern-001.cluster-ro-ckxxxxxxvm.us-west-2.rds.amazonaws.com)/test?tls=rds&sql_mode=%22%22&time_zone=%22%2B00%3A00%22&innodb_lock_wait_timeout=3&lock_wait_timeout=30&range_optimizer_max_mem_size=0&transaction_isolation=%22read-committed%22&charset=binary&collation=binary", resp) + assert.Equal(t, "root:password@tcp(tern-001.cluster-ro-ckxxxxxxvm.us-west-2.rds.amazonaws.com)/test?tls=rds&sql_mode=%22%22&time_zone=%22%2B00%3A00%22&innodb_lock_wait_timeout=3&lock_wait_timeout=30&range_optimizer_max_mem_size=0&transaction_isolation=%22read-committed%22&charset=binary&collation=binary&rejectReadOnly=true", resp) // This is with optional port too dsn = "root:password@tcp(tern-001.cluster-ro-ckxxxxxxvm.us-west-2.rds.amazonaws.com:12345)/test" resp, err = newDSN(dsn, NewDBConfig()) assert.NoError(t, err) - assert.Equal(t, "root:password@tcp(tern-001.cluster-ro-ckxxxxxxvm.us-west-2.rds.amazonaws.com:12345)/test?tls=rds&sql_mode=%22%22&time_zone=%22%2B00%3A00%22&innodb_lock_wait_timeout=3&lock_wait_timeout=30&range_optimizer_max_mem_size=0&transaction_isolation=%22read-committed%22&charset=binary&collation=binary", resp) + assert.Equal(t, "root:password@tcp(tern-001.cluster-ro-ckxxxxxxvm.us-west-2.rds.amazonaws.com:12345)/test?tls=rds&sql_mode=%22%22&time_zone=%22%2B00%3A00%22&innodb_lock_wait_timeout=3&lock_wait_timeout=30&range_optimizer_max_mem_size=0&transaction_isolation=%22read-committed%22&charset=binary&collation=binary&rejectReadOnly=true", resp) // Invalid DSN, can't parse. dsn = "invalid" @@ -59,3 +59,27 @@ func TestNewConn(t *testing.T) { assert.Error(t, err) assert.Nil(t, db) } + +func TestNewConnRejectsReadOnlyConnections(t *testing.T) { + testutils.RunSQL(t, "DROP TABLE IF EXISTS t1") + testutils.RunSQL(t, "CREATE TABLE t1 (a INT NOT NULL, b INT, c INT, PRIMARY KEY (a))") + + db, err := New(testutils.DSN(), NewDBConfig()) + assert.NoError(t, err) + defer db.Close() + _, err = db.Exec("set session transaction read only") + assert.NoError(t, err) + + // This would error, but `database/sql` automatically retries on a + // new connection (because we enforce `rejectReadOnly=true`) which is not read-only, and eventually succeed. + // See also: + // docs: https://github.com/go-sql-driver/mysql?tab=readme-ov-file#rejectreadonly + // rejectReadOnly test in `go-sql-driver/mysql`: https://github.com/go-sql-driver/mysql/blob/52c1917d99904701db2b0e4f14baffa948009cd7/driver_test.go#L2270-L2301 + _, err = db.Exec("insert into t1 values (1, 2, 3)") + assert.NoError(t, err) + + var count int + err = db.QueryRow("select count(*) from t1 where a=1").Scan(&count) + assert.NoError(t, err) + assert.Equal(t, 1, count) +}