feat: parse SET TRANSACTION statements (#549)

* feat: parse SET TRANSACTION statements

Parse SET TRANSACTION statements and translate these to SET LOCAL statements.
SET TRANSACTION may only be executed in a transaction block, and can only be
used for a specific, limited set of connection properties. The syntax is
specified by the SQL standard and PostgreSQL.

See also https://www.postgresql.org/docs/current/sql-set-transaction.html

This change only adds partial support. The following features will be added
in future changes:
1. SET TRANSACTION READ {WRITE | ONLY} is not picked up by the driver, as
   the type of transaction is set directly when BeginTx is called. A
   refactor of this transaction handling is needed to be able to pick up
   SET TRANSACTION READ ONLY / SET TRANSACTION READ WRITE statements that
   are executed after BeginTx has been called.
2. PostgreSQL allows multiple transaction modes to be set in a single
   SET TRANSACTION statement. E.g. the following is allowed:
   SET TRANSACTION READ WRITE, ISOLATION LEVEL SERIALIZABLE
   The current implementation only supports one transaction mode per
   SET statement.

* feat: support multiple transaction options in one statement

Author: olavloite

conn.go

@@ -327,13 +327,18 @@ func (c *conn) showConnectionVariable(identifier parser.Identifier) (any, bool,
 	return c.state.GetValue(extension, name)
 }
 
-func (c *conn) setConnectionVariable(identifier parser.Identifier, value string, local bool) error {
+func (c *conn) setConnectionVariable(identifier parser.Identifier, value string, local bool, transaction bool) error {
+	if transaction && !local {
+		// When transaction == true, then local must also be true.
+		// We should never hit this condition, as this is an indication of a bug in the driver code.
+		return status.Errorf(codes.FailedPrecondition, "transaction properties must be set as a local value")
+	}
 	extension, name, err := toExtensionAndName(identifier)
 	if err != nil {
 		return err
 	}
 	if local {
-		return c.state.SetLocalValue(extension, name, value)
+		return c.state.SetLocalValue(extension, name, value, transaction)
 	}
 	return c.state.SetValue(extension, name, value, connectionstate.ContextUser)
 }
@@ -1188,6 +1193,8 @@ func (c *conn) beginTx(ctx context.Context, driverOpts driver.TxOptions, closeFu
 		}
 	}()
 
+	// TODO: Delay the actual determination of the transaction type until the first query.
+	//       This is required in order to support SET TRANSACTION READ {ONLY | WRITE}
 	readOnlyTxOpts := c.getReadOnlyTransactionOptions()
 	batchReadOnlyTxOpts := c.getBatchReadOnlyTransactionOptions()
 	if c.inTransaction() {

connection_properties.go

@@ -228,6 +228,24 @@ var propertyDecodeToNativeArrays = createConnectionProperty(
 // Transaction connection properties.
 // ------------------------------------------------------------------------------------------------
 
+var propertyTransactionReadOnly = createConnectionProperty(
+	"transaction_read_only",
+	"transaction_read_only is the default read-only mode for transactions on this connection.",
+	false,
+	false,
+	nil,
+	connectionstate.ContextUser,
+	connectionstate.ConvertBool,
+)
+var propertyTransactionDeferrable = createConnectionProperty(
+	"transaction_deferrable",
+	"transaction_deferrable is a no-op on Spanner. It is defined in this driver for compatibility with PostgreSQL.",
+	false,
+	false,
+	nil,
+	connectionstate.ContextUser,
+	connectionstate.ConvertBool,
+)
 var propertyExcludeTxnFromChangeStreams = createConnectionProperty(
 	"exclude_txn_from_change_streams",
 	"exclude_txn_from_change_streams determines whether transactions on this connection should be excluded from "+

connectionstate/connection_state.go

@@ -131,7 +131,10 @@ func (cs *ConnectionState) SetValue(extension, name, value string, context Conte
 	return cs.setValue(extension, name, value, context, false)
 }
 
-func (cs *ConnectionState) SetLocalValue(extension, name, value string) error {
+func (cs *ConnectionState) SetLocalValue(extension, name, value string, isSetTransaction bool) error {
+	if isSetTransaction && !cs.inTransaction {
+		return status.Error(codes.FailedPrecondition, "SET TRANSACTION can only be used in transaction blocks")
+	}
 	return cs.setValue(extension, name, value, ContextUser, true)
 }
 

parser/simple_parser.go

@@ -290,6 +290,16 @@ func (p *simpleParser) eatKeywords(keywords []string) bool {
 	return true
 }
 
+// peekKeyword checks if the next keyword is the given keyword.
+// The position of the parser is not updated.
+func (p *simpleParser) peekKeyword(keyword string) bool {
+	pos := p.pos
+	defer func() {
+		p.pos = pos
+	}()
+	return p.eatKeyword(keyword)
+}
+
 // eatKeyword eats the given keyword at the current position of the parser if it exists
 // and returns true if the keyword was found. Otherwise, it returns false.
 func (p *simpleParser) eatKeyword(keyword string) bool {
@@ -323,8 +333,8 @@ func (p *simpleParser) readKeyword() string {
 		if isSpace(p.sql[p.pos]) {
 			break
 		}
-		// Only upper/lower-case letters are allowed in keywords.
-		if !((p.sql[p.pos] >= 'A' && p.sql[p.pos] <= 'Z') || (p.sql[p.pos] >= 'a' && p.sql[p.pos] <= 'z')) {
+		// Only upper/lower-case letters and underscores are allowed in keywords.
+		if !((p.sql[p.pos] >= 'A' && p.sql[p.pos] <= 'Z') || (p.sql[p.pos] >= 'a' && p.sql[p.pos] <= 'z')) && p.sql[p.pos] != '_' {
 			break
 		}
 	}

parser/statement_parser_test.go

@@ -2135,6 +2135,10 @@ func TestReadKeyword(t *testing.T) {
 			input: "Select from my_table",
 			want:  "Select",
 		},
+		{
+			input: "statement_tag",
+			want:  "statement_tag",
+		},
 	}
 	statementParser, err := NewStatementParser(databasepb.DatabaseDialect_GOOGLE_STANDARD_SQL, 1000)
 	if err != nil {
@@ -2404,6 +2408,36 @@ func TestCachedParamsAreImmutable(t *testing.T) {
 	}
 }
 
+func TestPeekKeyword(t *testing.T) {
+	t.Parallel()
+
+	parser, err := NewStatementParser(databasepb.DatabaseDialect_GOOGLE_STANDARD_SQL, 1000)
+	if err != nil {
+		t.Fatal(err)
+	}
+	sp := &simpleParser{sql: []byte("select * from foo"), statementParser: parser}
+	if !sp.peekKeyword("select") {
+		t.Fatal("peekKeyword should have returned true")
+	}
+	if g, w := sp.pos, 0; g != w {
+		t.Fatalf("position mismatch\n Got: %v\nWant: %v", g, w)
+	}
+
+	if !sp.eatKeyword("select") {
+		t.Fatal("eatKeyword should have returned true")
+	}
+	if !sp.eatToken('*') {
+		t.Fatal("eatToken should have returned true")
+	}
+	pos := sp.pos
+	if !sp.peekKeyword("from") {
+		t.Fatal("peekKeyword should have returned true")
+	}
+	if g, w := sp.pos, pos; g != w {
+		t.Fatalf("position mismatch\n Got: %v\nWant: %v", g, w)
+	}
+}
+
 func TestEatKeyword(t *testing.T) {
 	t.Parallel()
 

parser/statements.go

@@ -15,6 +15,8 @@
 package parser
 
 import (
+	"fmt"
+
 	"cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -142,11 +144,27 @@ func (s *ParsedShowStatement) parse(parser *StatementParser, query string) error
 
 // ParsedSetStatement is a statement of the form
 // SET [SESSION | LOCAL] [my_extension.]my_property {=|to} <value>
+//
+// It also covers statements of the form SET TRANSACTION. This is a
+// synonym for SET LOCAL, but is only supported for a specific set of
+// properties, and may only be executed before a transaction has been
+// activated. Examples include:
+// SET TRANSACTION READ ONLY
+// SET TRANSACTION ISOLATION LEVEL [SERIALIZABLE | REPEATABLE READ]
+//
+// One SET statement can set more than one property.
 type ParsedSetStatement struct {
-	query      string
-	Identifier Identifier
-	Literal    Literal
-	IsLocal    bool
+	query string
+	// Identifiers contains the properties that are being set. The number of elements in this slice
+	// must be equal to the number of Literals.
+	Identifiers []Identifier
+	// Literals contains the values that should be set for the properties.
+	Literals []Literal
+	// IsLocal indicates whether this is a SET LOCAL statement or not.
+	IsLocal bool
+	// IsTransaction indicates whether this is a SET TRANSACTION statement or not.
+	// IsTransaction automatically also implies IsLocal.
+	IsTransaction bool
 }
 
 func (s *ParsedSetStatement) Name() string {
@@ -165,10 +183,17 @@ func (s *ParsedSetStatement) parse(parser *StatementParser, query string) error
 		return status.Errorf(codes.InvalidArgument, "syntax error: expected SET")
 	}
 	isLocal := sp.eatKeyword("LOCAL")
-	if !isLocal && parser.Dialect == databasepb.DatabaseDialect_POSTGRESQL {
+	isTransaction := false
+	if !isLocal {
+		isTransaction = sp.eatKeyword("TRANSACTION")
+	}
+	if !isLocal && !isTransaction && parser.Dialect == databasepb.DatabaseDialect_POSTGRESQL {
 		// Just eat and ignore the SESSION keyword if it exists, as SESSION is the default.
 		_ = sp.eatKeyword("SESSION")
 	}
+	if isTransaction {
+		return s.parseSetTransaction(sp, query)
+	}
 	identifier, err := sp.eatIdentifier()
 	if err != nil {
 		return err
@@ -191,12 +216,93 @@ func (s *ParsedSetStatement) parse(parser *StatementParser, query string) error
 		return status.Errorf(codes.InvalidArgument, "unexpected tokens at position %d in %q", sp.pos, sp.sql)
 	}
 	s.query = query
-	s.Identifier = identifier
-	s.Literal = literalValue
+	s.Identifiers = []Identifier{identifier}
+	s.Literals = []Literal{literalValue}
 	s.IsLocal = isLocal
 	return nil
 }
 
+func (s *ParsedSetStatement) parseSetTransaction(sp *simpleParser, query string) error {
+	if !sp.hasMoreTokens() {
+		return status.Errorf(codes.InvalidArgument, "syntax error: missing TRANSACTION OPTION, expected one of ISOLATION LEVEL, READ WRITE, or READ ONLY")
+	}
+	s.query = query
+	s.IsLocal = true
+	s.IsTransaction = true
+
+	for {
+		if sp.peekKeyword("ISOLATION") {
+			if err := s.parseSetTransactionIsolationLevel(sp, query); err != nil {
+				return err
+			}
+		} else if sp.peekKeyword("READ") {
+			if err := s.parseSetTransactionMode(sp, query); err != nil {
+				return err
+			}
+		} else if sp.statementParser.Dialect == databasepb.DatabaseDialect_POSTGRESQL && (sp.peekKeyword("DEFERRABLE") || sp.peekKeyword("NOT")) {
+			// https://www.postgresql.org/docs/current/sql-set-transaction.html
+			if err := s.parseSetTransactionDeferrable(sp, query); err != nil {
+				return err
+			}
+		} else {
+			return status.Error(codes.InvalidArgument, "invalid TRANSACTION option, expected one of ISOLATION LEVEL, READ WRITE, or READ ONLY")
+		}
+		if !sp.hasMoreTokens() {
+			return nil
+		}
+		// Eat and ignore any commas separating the various options.
+		sp.eatToken(',')
+	}
+}
+
+func (s *ParsedSetStatement) parseSetTransactionIsolationLevel(sp *simpleParser, query string) error {
+	if !sp.eatKeywords([]string{"ISOLATION", "LEVEL"}) {
+		return status.Errorf(codes.InvalidArgument, "syntax error: expected ISOLATION LEVEL")
+	}
+	var value Literal
+	if sp.eatKeyword("SERIALIZABLE") {
+		value = Literal{Value: "serializable"}
+	} else if sp.eatKeywords([]string{"REPEATABLE", "READ"}) {
+		value = Literal{Value: "repeatable_read"}
+	} else {
+		return status.Errorf(codes.InvalidArgument, "syntax error: expected SERIALIZABLE OR REPETABLE READ")
+	}
+
+	s.Identifiers = append(s.Identifiers, Identifier{Parts: []string{"isolation_level"}})
+	s.Literals = append(s.Literals, value)
+	return nil
+}
+
+func (s *ParsedSetStatement) parseSetTransactionMode(sp *simpleParser, query string) error {
+	readOnly := false
+	if sp.eatKeywords([]string{"READ", "ONLY"}) {
+		readOnly = true
+	} else if sp.eatKeywords([]string{"READ", "WRITE"}) {
+		readOnly = false
+	} else {
+		return status.Errorf(codes.InvalidArgument, "syntax error: expected READ ONLY or READ WRITE")
+	}
+
+	s.Identifiers = append(s.Identifiers, Identifier{Parts: []string{"transaction_read_only"}})
+	s.Literals = append(s.Literals, Literal{Value: fmt.Sprintf("%v", readOnly)})
+	return nil
+}
+
+func (s *ParsedSetStatement) parseSetTransactionDeferrable(sp *simpleParser, query string) error {
+	deferrable := false
+	if sp.eatKeywords([]string{"NOT", "DEFERRABLE"}) {
+		deferrable = false
+	} else if sp.eatKeyword("DEFERRABLE") {
+		deferrable = true
+	} else {
+		return status.Errorf(codes.InvalidArgument, "syntax error: expected [NOT] DEFERRABLE")
+	}
+
+	s.Identifiers = append(s.Identifiers, Identifier{Parts: []string{"transaction_deferrable"}})
+	s.Literals = append(s.Literals, Literal{Value: fmt.Sprintf("%v", deferrable)})
+	return nil
+}
+
 // ParsedResetStatement is a statement of the form
 // RESET [my_extension.]my_property
 type ParsedResetStatement struct {
@@ -404,6 +510,7 @@ func (s *ParsedBeginStatement) parse(parser *StatementParser, query string) erro
 	// Parse a statement of the form
 	// GoogleSQL: BEGIN [TRANSACTION]
 	// PostgreSQL: {START | BEGIN} [{TRANSACTION | WORK}] (https://www.postgresql.org/docs/current/sql-begin.html)
+	// TODO: Support transaction modes in the BEGIN / START statement.
 	sp := &simpleParser{sql: []byte(query), statementParser: parser}
 	if sp.statementParser.Dialect == databasepb.DatabaseDialect_POSTGRESQL {
 		if !sp.eatKeyword("START") && !sp.eatKeyword("BEGIN") {