chore: move connection variables to connection state

Move individual connection variables to the generic connection state.

Author: olavloite

client_side_statement_test.go

@@ -31,7 +31,7 @@ import (
 )
 
 func TestStatementExecutor_StartBatchDdl(t *testing.T) {
-	c := &conn{retryAborts: true, logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
+	c := &conn{logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
 	s := &statementExecutor{}
 	ctx := context.Background()
 
@@ -62,7 +62,7 @@ func TestStatementExecutor_StartBatchDdl(t *testing.T) {
 }
 
 func TestStatementExecutor_StartBatchDml(t *testing.T) {
-	c := &conn{retryAborts: true, logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
+	c := &conn{logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
 	s := &statementExecutor{}
 	ctx := context.Background()
 
@@ -99,7 +99,7 @@ func TestStatementExecutor_StartBatchDml(t *testing.T) {
 }
 
 func TestStatementExecutor_RetryAbortsInternally(t *testing.T) {
-	c := &conn{retryAborts: true, logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
+	c := &conn{logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
 	s := &statementExecutor{}
 	ctx := context.Background()
 	for i, test := range []struct {
@@ -283,7 +283,7 @@ func TestStatementExecutor_ReadOnlyStaleness(t *testing.T) {
 func TestShowCommitTimestamp(t *testing.T) {
 	t.Parallel()
 
-	c := &conn{retryAborts: true, logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
+	c := &conn{logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
 	s := &statementExecutor{}
 	ctx := context.Background()
 
@@ -329,7 +329,7 @@ func TestShowCommitTimestamp(t *testing.T) {
 }
 
 func TestStatementExecutor_ExcludeTxnFromChangeStreams(t *testing.T) {
-	c := &conn{retryAborts: true, logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
+	c := &conn{logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
 	s := &statementExecutor{}
 	ctx := context.Background()
 	for i, test := range []struct {
@@ -458,7 +458,7 @@ func TestStatementExecutor_SetTransactionTag(t *testing.T) {
 		{"", "tag-with-missing-opening-quote'", true},
 		{"", "'tag-with-missing-closing-quote", true},
 	} {
-		c := &conn{retryAborts: true, logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
+		c := &conn{logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
 		s := &statementExecutor{}
 
 		it, err := s.ShowTransactionTag(ctx, c, "", ExecOptions{}, nil)
@@ -518,7 +518,7 @@ func TestStatementExecutor_SetTransactionTag(t *testing.T) {
 
 func TestStatementExecutor_UsesExecOptions(t *testing.T) {
 	ctx := context.Background()
-	c := &conn{retryAborts: true, logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
+	c := &conn{logger: noopLogger, state: createInitialConnectionState(connectionstate.TypeNonTransactional, map[string]connectionstate.ConnectionPropertyValue{})}
 	s := &statementExecutor{}
 
 	it, err := s.ShowTransactionTag(ctx, c, "", ExecOptions{DecodeOption: DecodeOptionProto, ReturnResultSetMetadata: true, ReturnResultSetStats: true}, nil)

conn.go

@@ -225,7 +225,6 @@ type conn struct {
 	resetForRetry  bool
 	commitResponse *spanner.CommitResponse
 	database       string
-	retryAborts    bool
 
 	execSingleQuery              func(ctx context.Context, c *spanner.Client, statement spanner.Statement, bound spanner.TimestampBound, options ExecOptions) *spanner.RowIterator
 	execSingleQueryTransactional func(ctx context.Context, c *spanner.Client, statement spanner.Statement, options ExecOptions) (rowIterator, *spanner.CommitResponse, error)
@@ -236,25 +235,6 @@ type conn struct {
 	state *connectionstate.ConnectionState
 	// batch is the currently active DDL or DML batch on this connection.
 	batch *batch
-	// autoBatchDml determines whether DML statements should automatically
-	// be batched and sent to Spanner when a non-DML statement is encountered.
-	autoBatchDml bool
-	// autoBatchDmlUpdateCount determines the update count that is returned for
-	// DML statements that are executed when autoBatchDml is true.
-	autoBatchDmlUpdateCount int64
-	// autoBatchDmlUpdateCountVerification enables/disables the verification
-	// that the update count that was returned for automatically batched DML
-	// statements was correct.
-	autoBatchDmlUpdateCountVerification bool
-
-	// readOnlyStaleness is used for queries in autocommit mode and for read-only transactions.
-	readOnlyStaleness spanner.TimestampBound
-	// isolationLevel determines the default isolation level that is used for read/write
-	// transactions on this connection. This default is ignored if the BeginTx function is
-	// called with an isolation level other than sql.LevelDefault.
-	isolationLevel sql.IsolationLevel
-	// beginTransactionOption determines the default transactions start mode.
-	beginTransactionOption spanner.BeginTransactionOption
 
 	// execOptions are applied to the next statement or transaction that is executed
 	// on this connection. It can also be set by passing it in as an argument to
@@ -289,7 +269,7 @@ func (c *conn) CommitResponse() (commitResponse *spanner.CommitResponse, err err
 }
 
 func (c *conn) RetryAbortsInternally() bool {
-	return c.retryAborts
+	return propertyRetryAbortsInternally.GetValueOrDefault(c.state)
 }
 
 func (c *conn) SetRetryAbortsInternally(retry bool) error {
@@ -299,9 +279,13 @@ func (c *conn) SetRetryAbortsInternally(retry bool) error {
 
 func (c *conn) setRetryAbortsInternally(retry bool) (driver.Result, error) {
 	if c.inTransaction() {
-		return c.tx.setRetryAbortsInternally(retry)
+		if _, err := c.tx.setRetryAbortsInternally(retry); err != nil {
+			return nil, err
+		}
+	}
+	if err := propertyRetryAbortsInternally.SetValue(c.state, retry, connectionstate.ContextUser); err != nil {
+		return nil, err
 	}
-	c.retryAborts = retry
 	return driver.ResultNoRows, nil
 }
 
@@ -325,7 +309,7 @@ func (c *conn) setAutocommitDMLMode(mode AutocommitDMLMode) (driver.Result, erro
 }
 
 func (c *conn) ReadOnlyStaleness() spanner.TimestampBound {
-	return c.readOnlyStaleness
+	return propertyReadOnlyStaleness.GetValueOrDefault(c.state)
 }
 
 func (c *conn) SetReadOnlyStaleness(staleness spanner.TimestampBound) error {
@@ -334,17 +318,18 @@ func (c *conn) SetReadOnlyStaleness(staleness spanner.TimestampBound) error {
 }
 
 func (c *conn) setReadOnlyStaleness(staleness spanner.TimestampBound) (driver.Result, error) {
-	c.readOnlyStaleness = staleness
+	if err := propertyReadOnlyStaleness.SetValue(c.state, staleness, connectionstate.ContextUser); err != nil {
+		return nil, err
+	}
 	return driver.ResultNoRows, nil
 }
 
 func (c *conn) IsolationLevel() sql.IsolationLevel {
-	return c.isolationLevel
+	return propertyIsolationLevel.GetValueOrDefault(c.state)
 }
 
 func (c *conn) SetIsolationLevel(level sql.IsolationLevel) error {
-	c.isolationLevel = level
-	return nil
+	return propertyIsolationLevel.SetValue(c.state, level, connectionstate.ContextUser)
 }
 
 func (c *conn) MaxCommitDelay() time.Duration {
@@ -419,30 +404,27 @@ func (c *conn) setStatementTag(statementTag string) (driver.Result, error) {
 }
 
 func (c *conn) AutoBatchDml() bool {
-	return c.autoBatchDml
+	return propertyAutoBatchDml.GetValueOrDefault(c.state)
 }
 
 func (c *conn) SetAutoBatchDml(autoBatch bool) error {
-	c.autoBatchDml = autoBatch
-	return nil
+	return propertyAutoBatchDml.SetValue(c.state, autoBatch, connectionstate.ContextUser)
 }
 
 func (c *conn) AutoBatchDmlUpdateCount() int64 {
-	return c.autoBatchDmlUpdateCount
+	return propertyAutoBatchDmlUpdateCount.GetValueOrDefault(c.state)
 }
 
 func (c *conn) SetAutoBatchDmlUpdateCount(updateCount int64) error {
-	c.autoBatchDmlUpdateCount = updateCount
-	return nil
+	return propertyAutoBatchDmlUpdateCount.SetValue(c.state, updateCount, connectionstate.ContextUser)
 }
 
 func (c *conn) AutoBatchDmlUpdateCountVerification() bool {
-	return c.autoBatchDmlUpdateCountVerification
+	return propertyAutoBatchDmlUpdateCountVerification.GetValueOrDefault(c.state)
 }
 
 func (c *conn) SetAutoBatchDmlUpdateCountVerification(verify bool) error {
-	c.autoBatchDmlUpdateCountVerification = verify
-	return nil
+	return propertyAutoBatchDmlUpdateCountVerification.SetValue(c.state, verify, connectionstate.ContextUser)
 }
 
 func (c *conn) StartBatchDDL() error {
@@ -683,16 +665,8 @@ func (c *conn) ResetSession(_ context.Context) error {
 	}
 	c.commitResponse = nil
 	c.batch = nil
-	c.autoBatchDml = c.connector.connectorConfig.AutoBatchDml
-	c.autoBatchDmlUpdateCount = c.connector.connectorConfig.AutoBatchDmlUpdateCount
-	c.autoBatchDmlUpdateCountVerification = !c.connector.connectorConfig.DisableAutoBatchDmlUpdateCountVerification
-	c.retryAborts = c.connector.retryAbortsInternally
-	c.isolationLevel = c.connector.connectorConfig.IsolationLevel
-	c.beginTransactionOption = c.connector.connectorConfig.BeginTransactionOption
 
 	_ = c.state.Reset(connectionstate.ContextUser)
-	// TODO: Reset the following fields to the connector default
-	c.readOnlyStaleness = spanner.TimestampBound{}
 	c.execOptions = ExecOptions{
 		DecodeToNativeArrays: c.connector.connectorConfig.DecodeToNativeArrays,
 	}
@@ -817,7 +791,7 @@ func (c *conn) queryContext(ctx context.Context, query string, execOptions ExecO
 			// The statement was either detected as being a query, or potentially not recognized at all.
 			// In that case, just default to using a single-use read-only transaction and let Spanner
 			// return an error if the statement is not suited for that type of transaction.
-			iter = &readOnlyRowIterator{c.execSingleQuery(ctx, c.client, stmt, c.readOnlyStaleness, execOptions)}
+			iter = &readOnlyRowIterator{c.execSingleQuery(ctx, c.client, stmt, c.ReadOnlyStaleness(), execOptions)}
 		}
 	} else {
 		if execOptions.PartitionedQueryOptions.PartitionQuery {
@@ -875,7 +849,7 @@ func (c *conn) execContext(ctx context.Context, query string, execOptions ExecOp
 	}
 
 	// Start an automatic DML batch.
-	if c.autoBatchDml && !c.inBatch() && c.inReadWriteTransaction() {
+	if c.AutoBatchDml() && !c.inBatch() && c.inReadWriteTransaction() {
 		if _, err := c.startBatchDML( /* automatic = */ true); err != nil {
 			return nil, err
 		}
@@ -964,19 +938,19 @@ func (c *conn) getTransactionOptions() ReadWriteTransactionOptions {
 	}()
 	txOpts := ReadWriteTransactionOptions{
 		TransactionOptions:     c.execOptions.TransactionOptions,
-		DisableInternalRetries: !c.retryAborts,
+		DisableInternalRetries: !c.RetryAbortsInternally(),
 	}
 	// Only use the default isolation level from the connection if the ExecOptions
 	// did not contain a more specific isolation level.
 	if txOpts.TransactionOptions.IsolationLevel == spannerpb.TransactionOptions_ISOLATION_LEVEL_UNSPECIFIED {
 		// This should never really return an error, but we check just to be absolutely sure.
-		level, err := toProtoIsolationLevel(c.isolationLevel)
+		level, err := toProtoIsolationLevel(c.IsolationLevel())
 		if err == nil {
 			txOpts.TransactionOptions.IsolationLevel = level
 		}
 	}
 	if txOpts.TransactionOptions.BeginTransactionOption == spanner.DefaultBeginTransaction {
-		txOpts.TransactionOptions.BeginTransactionOption = c.convertDefaultBeginTransactionOption(c.beginTransactionOption)
+		txOpts.TransactionOptions.BeginTransactionOption = c.convertDefaultBeginTransactionOption(propertyBeginTransactionOption.GetValueOrDefault(c.state))
 	}
 	return txOpts
 }
@@ -992,7 +966,7 @@ func (c *conn) getReadOnlyTransactionOptions() ReadOnlyTransactionOptions {
 		opts.BeginTransactionOption = c.convertDefaultBeginTransactionOption(opts.BeginTransactionOption)
 		return opts
 	}
-	return ReadOnlyTransactionOptions{TimestampBound: c.readOnlyStaleness, BeginTransactionOption: c.convertDefaultBeginTransactionOption(c.beginTransactionOption)}
+	return ReadOnlyTransactionOptions{TimestampBound: c.ReadOnlyStaleness(), BeginTransactionOption: c.convertDefaultBeginTransactionOption(propertyBeginTransactionOption.GetValueOrDefault(c.state))}
 }
 
 func (c *conn) withTempBatchReadOnlyTransactionOptions(options *BatchReadOnlyTransactionOptions) {
@@ -1004,7 +978,7 @@ func (c *conn) getBatchReadOnlyTransactionOptions() BatchReadOnlyTransactionOpti
 		defer func() { c.tempBatchReadOnlyTransactionOptions = nil }()
 		return *c.tempBatchReadOnlyTransactionOptions
 	}
-	return BatchReadOnlyTransactionOptions{TimestampBound: c.readOnlyStaleness}
+	return BatchReadOnlyTransactionOptions{TimestampBound: c.ReadOnlyStaleness()}
 }
 
 func (c *conn) Begin() (driver.Tx, error) {
@@ -1134,10 +1108,10 @@ func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, e
 
 func (c *conn) convertDefaultBeginTransactionOption(opt spanner.BeginTransactionOption) spanner.BeginTransactionOption {
 	if opt == spanner.DefaultBeginTransaction {
-		if c.beginTransactionOption == spanner.DefaultBeginTransaction {
+		if propertyBeginTransactionOption.GetValueOrDefault(c.state) == spanner.DefaultBeginTransaction {
 			return spanner.InlinedBeginTransaction
 		}
-		return c.beginTransactionOption
+		return propertyBeginTransactionOption.GetValueOrDefault(c.state)
 	}
 	return opt
 }

conn_with_mockserver_test.go

@@ -551,3 +551,86 @@ func TestSetAutocommitDMLMode(t *testing.T) {
 		})
 	}
 }
+
+func TestConfigureConnectionProperty(t *testing.T) {
+	t.Parallel()
+
+	type config int
+	const (
+		configConnString config = iota
+		configConnectorConfig
+		configConnectorConfigAndParams
+	)
+
+	for _, cfg := range []config{configConnString, configConnectorConfig, configConnectorConfigAndParams} {
+		var db *sql.DB
+		var teardown func()
+		switch cfg {
+		case configConnString:
+			db, _, teardown = setupTestDBConnectionWithParams(t, "isolationLevel=repeatable_read")
+		case configConnectorConfig:
+			db, _, teardown = setupTestDBConnectionWithConnectorConfig(t, ConnectorConfig{
+				Project:        "p",
+				Instance:       "i",
+				Database:       "d",
+				IsolationLevel: sql.LevelRepeatableRead,
+			})
+		case configConnectorConfigAndParams:
+			db, _, teardown = setupTestDBConnectionWithConnectorConfig(t, ConnectorConfig{
+				Project:        "p",
+				Instance:       "i",
+				Database:       "d",
+				IsolationLevel: sql.LevelSerializable,
+				// The configuration in the params (which come from the connection string)
+				// take precedence over the configuration in the ConnectorConfig.
+				Params: map[string]string{
+					"isolation_level": "repeatable_read",
+				},
+			})
+		default:
+			t.Fatalf("invalid config: %v", cfg)
+		}
+		defer teardown()
+		// Limit the number of open connections to 1 to ensure that the test re-uses the same connection.
+		db.SetMaxOpenConns(1)
+
+		// Repeat twice to ensure that we get a connection that has been reset
+		// when getting a fresh connection from the pool
+		for range 2 {
+			sqlConn, err := db.Conn(context.Background())
+			if err != nil {
+				t.Fatal(err)
+			}
+			_ = sqlConn.Raw(func(driverConn interface{}) error {
+				c, _ := driverConn.(SpannerConn)
+				if g, w := c.IsolationLevel(), sql.LevelRepeatableRead; g != w {
+					t.Fatalf("isolation level mismatch\n Got: %v\nWant: %v", g, w)
+				}
+				if err := c.SetIsolationLevel(sql.LevelSerializable); err != nil {
+					t.Fatal(err)
+				}
+				if g, w := c.IsolationLevel(), sql.LevelSerializable; g != w {
+					t.Fatalf("isolation level mismatch\n Got: %v\nWant: %v", g, w)
+				}
+				// Reset the connection manually (this should also happen automatically when the connection is closed).
+				sc := c.(*conn)
+				if err := sc.ResetSession(context.Background()); err != nil {
+					t.Fatal(err)
+				}
+				// Verify that the isolation level is reset to the value it had when the connection was created.
+				if g, w := c.IsolationLevel(), sql.LevelRepeatableRead; g != w {
+					t.Fatalf("isolation level mismatch\n Got: %v\nWant: %v", g, w)
+				}
+				// Set the isolation level back to serializable in order to ensure that it is also correctly reset
+				// when the connection is closed.
+				if err := c.SetIsolationLevel(sql.LevelSerializable); err != nil {
+					t.Fatal(err)
+				}
+				return nil
+			})
+			if err := sqlConn.Close(); err != nil {
+				t.Fatal(err)
+			}
+		}
+	}
+}