Merge branch 'main' into new-mode-support

Author: nareshz

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractMultiplexedSessionDatabaseClient.java

@@ -26,11 +26,6 @@
  */
 abstract class AbstractMultiplexedSessionDatabaseClient implements DatabaseClient {
 
-  @Override
-  public Dialect getDialect() {
-    throw new UnsupportedOperationException();
-  }
-
   @Override
   public String getDatabaseRole() {
     throw new UnsupportedOperationException();

google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java

@@ -132,6 +132,10 @@ private boolean canUseMultiplexedSessionsForPartitionedOps() {
 
   @Override
   public Dialect getDialect() {
+    MultiplexedSessionDatabaseClient client = getMultiplexedSessionDatabaseClient();
+    if (client != null) {
+      return client.getDialect();
+    }
     return pool.getDialect();
   }
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java

@@ -277,6 +277,13 @@ public void onSessionReady(SessionImpl session) {
                     .getSkipVerifyBeginTransactionForMuxRW()) {
               verifyBeginTransactionWithRWOnMultiplexedSessionAsync(session.getName());
             }
+            if (sessionClient
+                .getSpanner()
+                .getOptions()
+                .getSessionPoolOptions()
+                .isAutoDetectDialect()) {
+              MAINTAINER_SERVICE.submit(() -> getDialect());
+            }
           }
 
           @Override
@@ -513,6 +520,30 @@ private int getSingleUseChannelHint() {
     }
   }
 
+  private final AbstractLazyInitializer<Dialect> dialectSupplier =
+      new AbstractLazyInitializer<Dialect>() {
+        @Override
+        protected Dialect initialize() {
+          try (ResultSet dialectResultSet =
+              singleUse().executeQuery(SessionPool.DETERMINE_DIALECT_STATEMENT)) {
+            if (dialectResultSet.next()) {
+              return Dialect.fromName(dialectResultSet.getString(0));
+            }
+          }
+          // This should not really happen, but it is the safest fallback value.
+          return Dialect.GOOGLE_STANDARD_SQL;
+        }
+      };
+
+  @Override
+  public Dialect getDialect() {
+    try {
+      return dialectSupplier.get();
+    } catch (Exception exception) {
+      throw SpannerExceptionFactory.asSpannerException(exception);
+    }
+  }
+
   @Override
   public Timestamp write(Iterable<Mutation> mutations) throws SpannerException {
     return createMultiplexedSessionTransaction(/* singleUse = */ false).write(mutations);

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java

@@ -2256,13 +2256,9 @@ enum Position {
   @VisibleForTesting
   static final Statement DETERMINE_DIALECT_STATEMENT =
       Statement.newBuilder(
-              "SELECT 'POSTGRESQL' AS DIALECT\n"
-                  + "FROM INFORMATION_SCHEMA.SCHEMATA\n"
-                  + "WHERE SCHEMA_NAME='information_schema'\n"
-                  + "UNION ALL\n"
-                  + "SELECT 'GOOGLE_STANDARD_SQL' AS DIALECT\n"
-                  + "FROM INFORMATION_SCHEMA.SCHEMATA\n"
-                  + "WHERE SCHEMA_NAME='INFORMATION_SCHEMA' AND CATALOG_NAME=''")
+              "select option_value "
+                  + "from information_schema.database_options "
+                  + "where option_name='database_dialect'")
           .build();
 
   private final SessionPoolOptions options;
@@ -3211,7 +3207,9 @@ public void onSessionReady(SessionImpl session) {
           if (allSessions.size() >= minSessions) {
             waitOnMinSessionsLatch.countDown();
           }
-          if (options.isAutoDetectDialect() && !detectDialectStarted) {
+          if (options.isAutoDetectDialect()
+              && !detectDialectStarted
+              && !options.getUseMultiplexedSession()) {
             // Get the dialect of the underlying database if that has not yet been done. Note that
             // this method will release the session into the pool once it is done.
             detectDialectStarted = true;

google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -3998,7 +3998,7 @@ public void testCreateSessionsFailure_shouldNotPropagateToCloseMethod() {
     try {
       // Simulate session creation failures on the backend.
       mockSpanner.setCreateSessionExecutionTime(
-          SimulatedExecutionTime.ofStickyException(Status.RESOURCE_EXHAUSTED.asRuntimeException()));
+          SimulatedExecutionTime.ofStickyException(Status.PERMISSION_DENIED.asRuntimeException()));
       // This will not cause any failure as getting a session from the pool is guaranteed to be
       // non-blocking, and any exceptions will be delayed until actual query execution.
       mockSpanner.freeze();
@@ -4007,8 +4007,8 @@ public void testCreateSessionsFailure_shouldNotPropagateToCloseMethod() {
               DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
       try (ResultSet rs = client.singleUse().executeQuery(SELECT1)) {
         mockSpanner.unfreeze();
-        SpannerException e = assertThrows(SpannerException.class, rs::next);
-        assertThat(e.getErrorCode()).isEqualTo(ErrorCode.RESOURCE_EXHAUSTED);
+        SpannerException exception = assertThrows(SpannerException.class, rs::next);
+        assertEquals(ErrorCode.PERMISSION_DENIED, exception.getErrorCode());
       }
     } finally {
       mockSpanner.setCreateSessionExecutionTime(SimulatedExecutionTime.none());
@@ -4547,6 +4547,8 @@ public void testGetDialectPostgreSQLPreloaded() {
   public void testGetDialect_FailsDirectlyIfDatabaseNotFound() {
     mockSpanner.setBatchCreateSessionsExecutionTime(
         SimulatedExecutionTime.stickyDatabaseNotFoundException("invalid-database"));
+    mockSpanner.setCreateSessionExecutionTime(
+        SimulatedExecutionTime.stickyDatabaseNotFoundException("invalid-database"));
     DatabaseClient client =
         spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
 
@@ -4563,6 +4565,8 @@ public void testGetDialect_FailsDirectlyIfDatabaseNotFound() {
   public void testGetDialectDefaultPreloaded_FailsDirectlyIfDatabaseNotFound() {
     mockSpanner.setBatchCreateSessionsExecutionTime(
         SimulatedExecutionTime.stickyDatabaseNotFoundException("invalid-database"));
+    mockSpanner.setCreateSessionExecutionTime(
+        SimulatedExecutionTime.stickyDatabaseNotFoundException("invalid-database"));
     try (Spanner spanner =
         this.spanner
             .getOptions()

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java

@@ -16,6 +16,7 @@
 
 package com.google.cloud.spanner.connection;
 
+import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ForceCloseSpannerFunction;
 import com.google.cloud.spanner.MockSpannerServiceImpl;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
@@ -198,6 +199,8 @@ public void getOperation(
     mockSpanner.putStatementResult(
         StatementResult.query(SELECT_RANDOM_STATEMENT, RANDOM_RESULT_SET));
     mockSpanner.putStatementResult(StatementResult.query(SELECT1_STATEMENT, SELECT1_RESULTSET));
+    mockSpanner.putStatementResult(
+        StatementResult.detectDialectResult(Dialect.GOOGLE_STANDARD_SQL));
 
     futureParentHandlers = Logger.getLogger(AbstractFuture.class.getName()).getUseParentHandlers();
     exceptionRunnableParentHandlers =

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionTest.java

@@ -674,6 +674,8 @@ public void testPostgreSQLGetDialect() {
     public void testGetDialect_DatabaseNotFound() throws Exception {
       mockSpanner.setBatchCreateSessionsExecutionTime(
           SimulatedExecutionTime.stickyDatabaseNotFoundException("invalid-database"));
+      mockSpanner.setCreateSessionExecutionTime(
+          SimulatedExecutionTime.stickyDatabaseNotFoundException("invalid-database"));
       try (Connection connection = createConnection()) {
         SpannerException exception = assertThrows(SpannerException.class, connection::getDialect);
         assertEquals(ErrorCode.NOT_FOUND, exception.getErrorCode());