Fix tests

Author: sakthivelmanii

google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/ClientSideStatements.json

@@ -754,16 +754,19 @@
 	  }
 	},
   {
-    "name": "SET LOCAL DML_BATCH_UPDATE_COUNT = <INT64>",
+    "name": "SET [LOCAL] DML_BATCH_UPDATE_COUNT = <INT64>",
     "executorName": "ClientSideStatementSetExecutor",
     "resultType": "NO_RESULT",
     "statementType": "SET_DML_BATCH_UPDATE_COUNT",
-    "regex": "(?is)\\A\\s*set\\s+((?:session|local)\\s+)dml_batch_update_count\\s*(?:=)\\s*(.*)\\z",
+    "regex": "(?is)\\A\\s*set\\s+((?:session|local)\\s+)?dml_batch_update_count\\s*(?:=)\\s*(.*)\\z",
     "method": "statementSetDmlBatchUpdateCount",
     "exampleStatements": [
       "set local dml_batch_update_count = 0",
-      "set local dml_batch_update_count = 100"
+      "set local dml_batch_update_count = 100",
+      "set dml_batch_update_count = 1",
+      "set dml_batch_update_count = 100"
     ],
+    "examplePrerequisiteStatements": ["set readonly = false", "set autocommit = false"],
     "setStatement": {
       "propertyName": "DML_BATCH_UPDATE_COUNT",
       "separator": "(?:=|\\s+TO\\s+)",

google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json

@@ -967,7 +967,7 @@
       }
     },
     {
-      "name": "SET LOCAL SPANNER.DML_BATCH_UPDATE_COUNT =|TO <bigint>",
+      "name": "SET [LOCAL] SPANNER.DML_BATCH_UPDATE_COUNT =|TO <bigint>",
       "executorName": "ClientSideStatementSetExecutor",
       "resultType": "NO_RESULT",
       "statementType": "SET_DML_BATCH_UPDATE_COUNT",
@@ -977,8 +977,10 @@
         "set local spanner.dml_batch_update_count = 0",
         "set local spanner.dml_batch_update_count = 100",
         "set local spanner.dml_batch_update_count to 1",
-        "set local spanner.dml_batch_update_count to 1"
+        "set spanner.dml_batch_update_count to 1",
+        "set spanner.dml_batch_update_count = 1"
       ],
+      "examplePrerequisiteStatements": ["set spanner.readonly = false", "set autocommit = false"],
       "setStatement": {
         "propertyName": "SPANNER.DML_BATCH_UPDATE_COUNT",
         "separator": "(?:=|\\s+TO\\s+)",

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

@@ -341,29 +341,38 @@ public void testAutocommitRunBatch() {
 
   @Test
   public void testDmlBatchUpdateCount() {
-    SpannerPool.closeSpannerPool();
-    try {
-      try (Connection connection = createConnection()) {
-        connection.execute(Statement.of("set local spanner.dml_batch_update_count = 1"));
-        connection.execute(Statement.of("START BATCH DML"));
-        List<Statement> statements = Arrays.asList(INSERT_STATEMENT, INSERT_STATEMENT);
-        long[] updateCounts = connection.executeBatchUpdate(statements);
-        assertThat(updateCounts).asList().containsExactly(1L, 1L);
-        connection.execute(Statement.of("RUN BATCH"));
-        connection.commit();
-
-        connection.execute(Statement.of("START BATCH DML"));
-        statements = Arrays.asList(INSERT_STATEMENT, INSERT_STATEMENT);
-        updateCounts = connection.executeBatchUpdate(statements);
-        assertThat(updateCounts).asList().containsExactly(-1L, -1L);
-        connection.execute(Statement.of("RUN BATCH"));
-        connection.commit();
-      }
-    } finally {
-      SpannerPool.closeSpannerPool();
-      mockSpanner.putStatementResult(
-          MockSpannerServiceImpl.StatementResult.detectDialectResult(Dialect.GOOGLE_STANDARD_SQL));
-    }
+    Arrays.asList(Dialect.POSTGRESQL, Dialect.GOOGLE_STANDARD_SQL)
+        .forEach(
+            dialect -> {
+              String prefix = dialect == Dialect.POSTGRESQL ? "spanner." : "";
+              mockSpanner.putStatementResult(
+                  MockSpannerServiceImpl.StatementResult.detectDialectResult(dialect));
+              SpannerPool.closeSpannerPool();
+              try {
+                try (Connection connection = createConnection()) {
+                  connection.execute(
+                      Statement.of("set local " + prefix + "dml_batch_update_count = 1"));
+                  connection.execute(Statement.of("START BATCH DML"));
+                  List<Statement> statements = Arrays.asList(INSERT_STATEMENT, INSERT_STATEMENT);
+                  long[] updateCounts = connection.executeBatchUpdate(statements);
+                  assertThat(updateCounts).asList().containsExactly(1L, 1L);
+                  connection.execute(Statement.of("RUN BATCH"));
+                  connection.commit();
+
+                  connection.execute(Statement.of("START BATCH DML"));
+                  statements = Arrays.asList(INSERT_STATEMENT, INSERT_STATEMENT);
+                  updateCounts = connection.executeBatchUpdate(statements);
+                  assertThat(updateCounts).asList().containsExactly(-1L, -1L);
+                  connection.execute(Statement.of("RUN BATCH"));
+                  connection.commit();
+                }
+              } finally {
+                SpannerPool.closeSpannerPool();
+                mockSpanner.putStatementResult(
+                    MockSpannerServiceImpl.StatementResult.detectDialectResult(
+                        Dialect.GOOGLE_STANDARD_SQL));
+              }
+            });
   }
 
   @Test