feat: add support for QueryOptions (#76)

* feat: add support for QueryOptions

* fix: review comments

* deps: temporarily update because of build and test errors

* deps: update to released versions

* fix: use grpc 1.27.2

* fix: fix invalid query hint in IT

Co-authored-by: skuruppu <[email protected]>

Author: olavloite

pom.xml

@@ -58,7 +58,7 @@
   </licenses>
   <properties>
     <site.installationModule>google-cloud-spanner-jdbc</site.installationModule>
-    <grpc.version>1.28.0</grpc.version>
+    <grpc.version>1.27.2</grpc.version>
     <api-client.version>1.30.9</api-client.version>
     <google.core.version>1.93.3</google.core.version>
     <gax.version>1.54.0</gax.version>

src/main/java/com/google/cloud/spanner/jdbc/ClientSideStatementValueConverters.java

@@ -213,6 +213,20 @@ public AutocommitDmlMode convert(String value) {
     }
   }
 
+  static class StringValueConverter implements ClientSideStatementValueConverter<String> {
+    public StringValueConverter(String allowedValues) {}
+
+    @Override
+    public Class<String> getParameterClass() {
+      return String.class;
+    }
+
+    @Override
+    public String convert(String value) {
+      return value;
+    }
+  }
+
   /** Converter for converting string values to {@link TransactionMode} values. */
   static class TransactionModeConverter
       implements ClientSideStatementValueConverter<TransactionMode> {

src/main/java/com/google/cloud/spanner/jdbc/Connection.java

@@ -78,6 +78,11 @@
  *   <li><code>
  *       SET READ_ONLY_STALENESS='STRONG' | 'MIN_READ_TIMESTAMP &lt;timestamp&gt;' | 'READ_TIMESTAMP &lt;timestamp&gt;' | 'MAX_STALENESS &lt;int64&gt;s|ms|mus|ns' | 'EXACT_STALENESS (&lt;int64&gt;s|ms|mus|ns)'
  *       </code>: Sets the value of <code>READ_ONLY_STALENESS</code> for this connection.
+ *   <li><code>SHOW OPTIMIZER_VERSION</code>: Returns the current value of <code>
+ *       OPTIMIZER_VERSION</code> of this connection as a {@link ResultSet}
+ *   <li><code>
+ *       SET OPTIMIZER_VERSION='&lt;version&gt;' | 'LATEST'
+ *       </code>: Sets the value of <code>OPTIMIZER_VERSION</code> for this connection.
  *   <li><code>BEGIN [TRANSACTION]</code>: Begins a new transaction. This statement is optional when
  *       the connection is not in autocommit mode, as a new transaction will automatically be
  *       started when a query or update statement is issued. In autocommit mode, this statement will
@@ -384,6 +389,24 @@ interface Connection extends AutoCloseable {
    */
   TimestampBound getReadOnlyStaleness();
 
+  /**
+   * Sets the query optimizer version to use for this connection.
+   *
+   * @param optimizerVersion The query optimizer version to use. Must be a valid optimizer version
+   *     number, the string <code>LATEST</code> or an empty string. The empty string will instruct
+   *     the connection to use the optimizer version that is defined in the environment variable
+   *     <code>SPANNER_OPTIMIZER_VERSION</code>. If no value is specified in the environment
+   *     variable, the default query optimizer of Cloud Spanner is used.
+   */
+  void setOptimizerVersion(String optimizerVersion);
+
+  /**
+   * Gets the current query optimizer version of this connection.
+   *
+   * @return The query optimizer version that is currently used by this connection.
+   */
+  String getOptimizerVersion();
+
   /**
    * Commits the current transaction of this connection. All mutations that have been buffered
    * during the current transaction will be written to the database.

src/main/java/com/google/cloud/spanner/jdbc/ConnectionImpl.java

@@ -35,6 +35,7 @@
 import com.google.cloud.spanner.jdbc.UnitOfWork.UnitOfWorkState;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Iterator;
@@ -192,6 +193,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
   private final List<TransactionRetryListener> transactionRetryListeners = new ArrayList<>();
   private AutocommitDmlMode autocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL;
   private TimestampBound readOnlyStaleness = TimestampBound.strong();
+  private QueryOptions queryOptions = QueryOptions.getDefaultInstance();
 
   /** Create a connection and register it in the SpannerPool. */
   ConnectionImpl(ConnectionOptions options) {
@@ -204,6 +206,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
     this.retryAbortsInternally = options.isRetryAbortsInternally();
     this.readOnly = options.isReadOnly();
     this.autocommit = options.isAutocommit();
+    this.queryOptions = this.queryOptions.toBuilder().mergeFrom(options.getQueryOptions()).build();
     this.ddlClient = createDdlClient();
     setDefaultTransactionOptions();
   }
@@ -389,6 +392,19 @@ public TimestampBound getReadOnlyStaleness() {
     return this.readOnlyStaleness;
   }
 
+  @Override
+  public void setOptimizerVersion(String optimizerVersion) {
+    Preconditions.checkNotNull(optimizerVersion);
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    this.queryOptions = queryOptions.toBuilder().setOptimizerVersion(optimizerVersion).build();
+  }
+
+  @Override
+  public String getOptimizerVersion() {
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    return this.queryOptions.getOptimizerVersion();
+  }
+
   @Override
   public void setStatementTimeout(long timeout, TimeUnit unit) {
     Preconditions.checkArgument(timeout > 0L, "Zero or negative timeout values are not allowed");
@@ -639,7 +655,7 @@ private void endCurrentTransaction(EndTransactionMethod endTransactionMethod) {
   public StatementResult execute(Statement statement) {
     Preconditions.checkNotNull(statement);
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
-    ParsedStatement parsedStatement = parser.parse(statement);
+    ParsedStatement parsedStatement = parser.parse(statement, this.queryOptions);
     switch (parsedStatement.getType()) {
       case CLIENT_SIDE:
         return parsedStatement
@@ -680,7 +696,7 @@ private ResultSet parseAndExecuteQuery(
     Preconditions.checkNotNull(query);
     Preconditions.checkNotNull(analyzeMode);
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
-    ParsedStatement parsedStatement = parser.parse(query);
+    ParsedStatement parsedStatement = parser.parse(query, this.queryOptions);
     if (parsedStatement.isQuery()) {
       switch (parsedStatement.getType()) {
         case CLIENT_SIDE:
@@ -809,7 +825,8 @@ private long[] internalExecuteBatchUpdate(final List<ParsedStatement> updates) {
    * Returns the current {@link UnitOfWork} of this connection, or creates a new one based on the
    * current transaction settings of the connection and returns that.
    */
-  private UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() {
+  @VisibleForTesting
+  UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() {
     if (this.currentUnitOfWork == null || !this.currentUnitOfWork.isActive()) {
       this.currentUnitOfWork = createNewUnitOfWork();
     }

src/main/java/com/google/cloud/spanner/jdbc/ConnectionOptions.java

@@ -31,6 +31,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Sets;
+import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -144,6 +145,7 @@ public String[] getValidValues() {
   private static final String DEFAULT_OAUTH_TOKEN = null;
   private static final String DEFAULT_NUM_CHANNELS = null;
   private static final String DEFAULT_USER_AGENT = null;
+  private static final String DEFAULT_OPTIMIZER_VERSION = "";
 
   private static final String PLAIN_TEXT_PROTOCOL = "http:";
   private static final String HOST_PROTOCOL = "https:";
@@ -166,6 +168,8 @@ public String[] getValidValues() {
   public static final String NUM_CHANNELS_PROPERTY_NAME = "numChannels";
   /** Custom user agent string is only for other Google libraries. */
   private static final String USER_AGENT_PROPERTY_NAME = "userAgent";
+  /** Query optimizer version to use for a connection. */
+  private static final String OPTIMIZER_VERSION_PROPERTY_NAME = "optimizerVersion";
 
   /** All valid connection properties. */
   public static final Set<ConnectionProperty> VALID_PROPERTIES =
@@ -183,7 +187,8 @@ public String[] getValidValues() {
                   ConnectionProperty.createStringProperty(NUM_CHANNELS_PROPERTY_NAME, ""),
                   ConnectionProperty.createBooleanProperty(
                       USE_PLAIN_TEXT_PROPERTY_NAME, "", DEFAULT_USE_PLAIN_TEXT),
-                  ConnectionProperty.createStringProperty(USER_AGENT_PROPERTY_NAME, ""))));
+                  ConnectionProperty.createStringProperty(USER_AGENT_PROPERTY_NAME, ""),
+                  ConnectionProperty.createStringProperty(OPTIMIZER_VERSION_PROPERTY_NAME, ""))));
 
   private static final Set<ConnectionProperty> INTERNAL_PROPERTIES =
       Collections.unmodifiableSet(
@@ -281,6 +286,7 @@ private boolean isValidUri(String uri) {
      *       false.
      *   <li>retryAbortsInternally (boolean): Sets the initial retryAbortsInternally mode for the
      *       connection. Default is true.
+     *   <li>optimizerVersion (string): Sets the query optimizer version to use for the connection.
      * </ul>
      *
      * @param uri The URI of the Spanner database to connect to.
@@ -373,6 +379,7 @@ public static Builder newBuilder() {
   private final Credentials credentials;
   private final Integer numChannels;
   private final String userAgent;
+  private final QueryOptions queryOptions;
 
   private final boolean autocommit;
   private final boolean readOnly;
@@ -397,6 +404,9 @@ private ConnectionOptions(Builder builder) {
 
     this.usePlainText = parseUsePlainText(this.uri);
     this.userAgent = parseUserAgent(this.uri);
+    QueryOptions.Builder queryOptionsBuilder = QueryOptions.newBuilder();
+    queryOptionsBuilder.setOptimizerVersion(parseOptimizerVersion(this.uri));
+    this.queryOptions = queryOptionsBuilder.build();
 
     this.host =
         matcher.group(Builder.HOST_GROUP) == null
@@ -501,6 +511,12 @@ static String parseUserAgent(String uri) {
     return value != null ? value : DEFAULT_USER_AGENT;
   }
 
+  @VisibleForTesting
+  static String parseOptimizerVersion(String uri) {
+    String value = parseUriProperty(uri, OPTIMIZER_VERSION_PROPERTY_NAME);
+    return value != null ? value : DEFAULT_OPTIMIZER_VERSION;
+  }
+
   @VisibleForTesting
   static String parseUriProperty(String uri, String property) {
     Pattern pattern = Pattern.compile(String.format("(?is)(?:;|\\?)%s=(.*?)(?:;|$)", property));
@@ -633,6 +649,11 @@ String getUserAgent() {
     return userAgent;
   }
 
+  /** The {@link QueryOptions} to use for the connection. */
+  QueryOptions getQueryOptions() {
+    return queryOptions;
+  }
+
   /** Interceptors that should be executed after each statement */
   List<StatementExecutionInterceptor> getStatementExecutionInterceptors() {
     return statementExecutionInterceptors;

src/main/java/com/google/cloud/spanner/jdbc/ConnectionStatementExecutor.java

@@ -60,6 +60,10 @@ interface ConnectionStatementExecutor {
 
   StatementResult statementShowReadOnlyStaleness();
 
+  StatementResult statementSetOptimizerVersion(String optimizerVersion);
+
+  StatementResult statementShowOptimizerVersion();
+
   StatementResult statementBeginTransaction();
 
   StatementResult statementCommit();

src/main/java/com/google/cloud/spanner/jdbc/ConnectionStatementExecutorImpl.java

@@ -23,6 +23,7 @@
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.RUN_BATCH;
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SET_AUTOCOMMIT;
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SET_AUTOCOMMIT_DML_MODE;
+import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SET_OPTIMIZER_VERSION;
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SET_READONLY;
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SET_READ_ONLY_STALENESS;
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SET_RETRY_ABORTS_INTERNALLY;
@@ -31,6 +32,7 @@
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SHOW_AUTOCOMMIT;
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SHOW_AUTOCOMMIT_DML_MODE;
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SHOW_COMMIT_TIMESTAMP;
+import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SHOW_OPTIMIZER_VERSION;
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SHOW_READONLY;
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SHOW_READ_ONLY_STALENESS;
 import static com.google.cloud.spanner.jdbc.StatementResult.ClientSideStatementType.SHOW_READ_TIMESTAMP;
@@ -183,6 +185,18 @@ public StatementResult statementShowReadOnlyStaleness() {
         SHOW_READ_ONLY_STALENESS);
   }
 
+  @Override
+  public StatementResult statementSetOptimizerVersion(String optimizerVersion) {
+    getConnection().setOptimizerVersion(optimizerVersion);
+    return noResult(SET_OPTIMIZER_VERSION);
+  }
+
+  @Override
+  public StatementResult statementShowOptimizerVersion() {
+    return resultSet(
+        "OPTIMIZER_VERSION", getConnection().getOptimizerVersion(), SHOW_OPTIMIZER_VERSION);
+  }
+
   @Override
   public StatementResult statementBeginTransaction() {
     getConnection().beginTransaction();

src/main/java/com/google/cloud/spanner/jdbc/JdbcDriver.java

@@ -86,6 +86,7 @@
  *       connection. Default is true. @see {@link
  *       com.google.cloud.spanner.jdbc.CloudSpannerJdbcConnection#setRetryAbortsInternally(boolean)}
  *       for more information.
+ *   <li>optimizerVersion (string): The query optimizer version to use for the connection. The value must be either a valid version number or <code>LATEST</code>. If no value is specified, the query optimizer version specified in the environment variable <code>SPANNER_OPTIMIZER_VERSION<code> will be used. If no query optimizer version is specified in the connection URL or in the environment variable, the default query optimizer version of Cloud Spanner will be used.
  * </ul>
  */
 public class JdbcDriver implements Driver {

src/main/java/com/google/cloud/spanner/jdbc/StatementParser.java

@@ -20,10 +20,13 @@
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.jdbc.ClientSideStatementImpl.CompileException;
+import com.google.cloud.spanner.jdbc.StatementParser.ParsedStatement;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableSet;
+import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import java.util.Collections;
+import java.util.Objects;
 import java.util.Set;
 
 /**
@@ -66,8 +69,10 @@ private static ParsedStatement ddl(Statement statement, String sqlWithoutComment
       return new ParsedStatement(StatementType.DDL, statement, sqlWithoutComments);
     }
 
-    private static ParsedStatement query(Statement statement, String sqlWithoutComments) {
-      return new ParsedStatement(StatementType.QUERY, statement, sqlWithoutComments);
+    private static ParsedStatement query(
+        Statement statement, String sqlWithoutComments, QueryOptions defaultQueryOptions) {
+      return new ParsedStatement(
+          StatementType.QUERY, statement, sqlWithoutComments, defaultQueryOptions);
     }
 
     private static ParsedStatement update(Statement statement, String sqlWithoutComments) {
@@ -91,14 +96,40 @@ private ParsedStatement(
     }
 
     private ParsedStatement(StatementType type, Statement statement, String sqlWithoutComments) {
+      this(type, statement, sqlWithoutComments, null);
+    }
+
+    private ParsedStatement(
+        StatementType type,
+        Statement statement,
+        String sqlWithoutComments,
+        QueryOptions defaultQueryOptions) {
       Preconditions.checkNotNull(type);
       Preconditions.checkNotNull(statement);
       this.type = type;
       this.clientSideStatement = null;
-      this.statement = statement;
+      this.statement = mergeQueryOptions(statement, defaultQueryOptions);
       this.sqlWithoutComments = sqlWithoutComments;
     }
 
+    @Override
+    public int hashCode() {
+      return Objects.hash(
+          this.type, this.clientSideStatement, this.statement, this.sqlWithoutComments);
+    }
+
+    @Override
+    public boolean equals(Object other) {
+      if (!(other instanceof ParsedStatement)) {
+        return false;
+      }
+      ParsedStatement o = (ParsedStatement) other;
+      return Objects.equals(this.type, o.type)
+          && Objects.equals(this.clientSideStatement, o.clientSideStatement)
+          && Objects.equals(this.statement, o.statement)
+          && Objects.equals(this.sqlWithoutComments, o.sqlWithoutComments);
+    }
+
     StatementType getType() {
       return type;
     }
@@ -148,6 +179,26 @@ Statement getStatement() {
       return statement;
     }
 
+    /**
+     * Merges the {@link QueryOptions} of the {@link Statement} with the current {@link
+     * QueryOptions} of this connection. The {@link QueryOptions} that are already present on the
+     * statement take precedence above the connection {@link QueryOptions}.
+     */
+    Statement mergeQueryOptions(Statement statement, QueryOptions defaultQueryOptions) {
+      if (defaultQueryOptions == null
+          || defaultQueryOptions.equals(QueryOptions.getDefaultInstance())) {
+        return statement;
+      }
+      if (statement.getQueryOptions() == null) {
+        return statement.toBuilder().withQueryOptions(defaultQueryOptions).build();
+      }
+      return statement
+          .toBuilder()
+          .withQueryOptions(
+              defaultQueryOptions.toBuilder().mergeFrom(statement.getQueryOptions()).build())
+          .build();
+    }
+
     String getSqlWithoutComments() {
       return sqlWithoutComments;
     }
@@ -183,12 +234,16 @@ private StatementParser() {
    * @return the parsed and categorized statement.
    */
   ParsedStatement parse(Statement statement) {
+    return parse(statement, null);
+  }
+
+  ParsedStatement parse(Statement statement, QueryOptions defaultQueryOptions) {
     String sql = removeCommentsAndTrim(statement.getSql());
     ClientSideStatementImpl client = parseClientSideStatement(sql);
     if (client != null) {
       return ParsedStatement.clientSideStatement(client, statement, sql);
     } else if (isQuery(sql)) {
-      return ParsedStatement.query(statement, sql);
+      return ParsedStatement.query(statement, sql, defaultQueryOptions);
     } else if (isUpdateStatement(sql)) {
       return ParsedStatement.update(statement, sql);
     } else if (isDdlStatement(sql)) {