Repository: GoogleCloudPlatform/pgadapter
Pull request #426 (2022-10-24 11:33:40)

feat: add setting for copy_commit_priority
Adds the setting 'spanner.copy_commit_priority' that can be used to set the Rpc priority for Commit requests for COPY operations. The default is MEDIUM to be consistent with Dataflow.

Fixes b/254330479

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/426?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#426](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/426?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (835dd87) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/39bda98787cbb0c2bf38e8f22e8922e900737b1b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (39bda98) will **increase** coverage by `0.00%`.
> The diff coverage is `88.00%`.

```diff
@@                  Coverage Diff                  @@
##             postgresql-dialect     #426   +/-   ##
=====================================================
  Coverage                 86.29%   86.29%           
- Complexity                 1908     1913    +5     
=====================================================
  Files                       122      122           
  Lines                      6187     6209   +22     
  Branches                    845      848    +3     
=====================================================
+ Hits                       5339     5358   +19     
- Misses                      605      607    +2     
- Partials                    243      244    +1     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `86.29% <88.00%> (+<0.01%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/426?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [.../cloud/spanner/pgadapter/session/SessionState.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/426/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3Nlc3Npb24vU2Vzc2lvblN0YXRlLmphdmE=) | `95.02% <50.00%> (-0.51%)` | :arrow_down: |
| [.../cloud/spanner/pgadapter/session/CopySettings.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/426/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3Nlc3Npb24vQ29weVNldHRpbmdzLmphdmE=) | `96.55% <66.66%> (-3.45%)` | :arrow_down: |
| [...gle/cloud/spanner/pgadapter/session/PGSetting.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/426/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3Nlc3Npb24vUEdTZXR0aW5nLmphdmE=) | `93.54% <100.00%> (+0.44%)` | :arrow_up: |
| [.../cloud/spanner/pgadapter/utils/MutationWriter.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/426/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL011dGF0aW9uV3JpdGVyLmphdmE=) | `90.33% <100.00%> (+0.09%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform)


Patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/session/CopySettings.java b/src/main/java/com/google/cloud/spanner/pgadapter/session/CopySettings.java
index a2eacb44..3c9de5ce 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/session/CopySettings.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/session/CopySettings.java
@@ -14,6 +14,8 @@
 
 package com.google.cloud.spanner.pgadapter.session;
 
+import com.google.cloud.spanner.Options.RpcPriority;
+import java.util.Locale;
 import java.util.Map;
 
 public class CopySettings {
@@ -86,6 +88,19 @@ public class CopySettings {
     return sessionState.getIntegerSetting("spanner", "copy_commit_timeout", 300);
   }
 
+  /** Returns the request priority for commits executed by COPY operations. */
+  public RpcPriority getCommitPriority() {
+    String setting =
+        sessionState
+            .getStringSetting("spanner", "copy_commit_priority", "medium")
+            .toUpperCase(Locale.ENGLISH);
+    try {
+      return RpcPriority.valueOf(setting);
+    } catch (IllegalArgumentException e) {
+      return RpcPriority.MEDIUM;
+    }
+  }
+
   /** Returns the batch size to use for non-atomic COPY operations. */
   public int getNonAtomicBatchSize() {
     return sessionState.getIntegerSetting("spanner", "copy_batch_size", 5000);
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/session/PGSetting.java b/src/main/java/com/google/cloud/spanner/pgadapter/session/PGSetting.java
index 69218bc2..63300cbf 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/session/PGSetting.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/session/PGSetting.java
@@ -22,10 +22,11 @@ import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.pgadapter.error.PGException;
 import com.google.cloud.spanner.pgadapter.parsers.BooleanParser;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Iterables;
 import java.util.Arrays;
+import java.util.Locale;
 import java.util.Objects;
 import java.util.Scanner;
 import java.util.stream.Collectors;
@@ -96,6 +97,7 @@ public class PGSetting {
   private final String minVal;
   private final String maxVal;
   private final String[] enumVals;
+  private final ImmutableList upperCaseEnumVals;
 
   private String setting;
   private String unit;
@@ -175,6 +177,7 @@ public class PGSetting {
     this.minVal = null;
     this.maxVal = null;
     this.enumVals = null;
+    this.upperCaseEnumVals = null;
   }
 
   @VisibleForTesting
@@ -207,6 +210,13 @@ public class PGSetting {
     this.minVal = minVal;
     this.maxVal = maxVal;
     this.enumVals = enumVals;
+    this.upperCaseEnumVals =
+        enumVals == null
+            ? null
+            : ImmutableList.copyOf(
+                Arrays.stream(enumVals)
+                    .map(v -> v.toUpperCase(Locale.ENGLISH))
+                    .collect(Collectors.toList()));
     this.setting = setting;
     this.unit = unit;
     this.source = source;
@@ -431,8 +441,10 @@ public class PGSetting {
       } catch (NumberFormatException exception) {
         throw invalidValueError(getCasePreservingKey(), value);
       }
-    } else if (enumVals != null && !Iterables.contains(Arrays.asList(this.enumVals), value)) {
-      throw invalidValueError(getCasePreservingKey(), value);
+    } else if (upperCaseEnumVals != null
+        && !upperCaseEnumVals.contains(
+            MoreObjects.firstNonNull(value, "").toUpperCase(Locale.ENGLISH))) {
+      throw invalidEnumError(getCasePreservingKey(), value, enumVals);
     }
   }
 
@@ -448,6 +460,13 @@ public class PGSetting {
         String.format("invalid value for parameter \"%s\": \"%s\"", key, value));
   }
 
+  static PGException invalidEnumError(String key, String value, String[] enumVals) {
+    return PGException.newBuilder(
+            String.format("invalid value for parameter \"%s\": \"%s\"", key, value))
+        .setHints(String.format("Available values: %s.", String.join(", ", enumVals)))
+        .build();
+  }
+
   public String getCategory() {
     return category;
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java b/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
index 6825e135..060f564c 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
@@ -271,6 +271,11 @@ public class SessionState {
     return null;
   }
 
+  String getStringSetting(String extension, String name, String defaultValue) {
+    PGSetting setting = internalGet(toKey(extension, name), false);
+    return setting != null ? setting.getSetting() : defaultValue;
+  }
+
   boolean getBoolSetting(String extension, String name, boolean defaultValue) {
     PGSetting setting = internalGet(toKey(extension, name), false);
     if (setting != null) {
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java
index cbd660ab..df6d4089 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java
@@ -25,6 +25,7 @@ import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.Mutation.WriteBuilder;
+import com.google.cloud.spanner.Options;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.SpannerOptions;
@@ -383,7 +384,10 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
                                       Duration.ofSeconds(copySettings.getCommitTimeoutSeconds()));
                             }
                           });
-              context.run(() -> dbClient.write(immutableMutations));
+              context.run(
+                  () ->
+                      dbClient.writeWithOptions(
+                          immutableMutations, Options.priority(copySettings.getCommitPriority())));
               return null;
             });
     Futures.addCallback(
diff --git a/src/main/resources/com/google/cloud/spanner/pgadapter/session/pg_settings.txt b/src/main/resources/com/google/cloud/spanner/pgadapter/session/pg_settings.txt
index fb68fcad..50b31ef4 100644
--- a/src/main/resources/com/google/cloud/spanner/pgadapter/session/pg_settings.txt
+++ b/src/main/resources/com/google/cloud/spanner/pgadapter/session/pg_settings.txt
@@ -346,6 +346,7 @@ zero_damaged_pages	off	\N	Developer Options	Continues processing past damaged pa
 spanner.copy_batch_size	5000	\N	COPY / Batch size for non-atomic COPY operations	The number of mutations to include in one batch when executing a non-atomic COPY operation.	\N	user	integer	default	\N	\N	\N	5000	5000	\N	\N	f
 spanner.copy_max_parallelism	128	\N	COPY / Max concurrent transactions for a non-atomic COPY operation	The maximum number of concurrent transactions for a non-atomic COPY operation.	\N	user	integer	default	\N	\N	\N	128	128	\N	\N	f
 spanner.copy_commit_timeout	300	\N	COPY / Timeout in seconds for commits for COPY operations	Timeout in seconds for commit requests for COPY operations.	\N	user	integer	default	\N	\N	\N	300	300	\N	\N	f
+spanner.copy_commit_priority	medium	\N	COPY / RPC priority for commits for COPY operations	RPC priority for commit requests for COPY operations.	\N	user	enum	default	\N	\N	{low,medium,high}	medium	medium	\N	\N	f
 spanner.copy_upsert	off	\N	COPY / Use Upsert instead of Insert for COPY	Use insert-or-update instead of insert for COPY operations.	\N	user	bool	default	\N	\N	\N	off	off	\N	\N	f
 spanner.copy_max_atomic_mutations	20000	\N	COPY / Max number of mutations for atomic COPY operations	The maximum number of mutations in an atomic COPY operation.	\N	internal	integer	default	\N	\N	\N	20000	20000	\N	\N	f
 spanner.copy_max_atomic_commit_size	100000000	\N	COPY / Max number of bytes in an atomic COPY operation	The maximum number of bytes in an atomic COPY operation.	\N	internal	integer	default	\N	\N	\N	100000000	100000000	\N	\N	f


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
index fc178265..773c2c7d 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
@@ -40,6 +40,7 @@ import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.Mutation;
 import com.google.spanner.v1.Mutation.OperationCase;
+import com.google.spanner.v1.RequestOptions.Priority;
 import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.RollbackRequest;
 import com.google.spanner.v1.StructType;
@@ -138,6 +139,7 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
     assertEquals(1, commitRequests.size());
     CommitRequest commitRequest = commitRequests.get(0);
     assertEquals(1, commitRequest.getMutationsCount());
+    assertEquals(Priority.PRIORITY_MEDIUM, commitRequest.getRequestOptions().getPriority());
 
     Mutation mutation = commitRequest.getMutations(0);
     assertEquals(OperationCase.INSERT, mutation.getOperationCase());
@@ -145,6 +147,22 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
     assertEquals(3, mutation.getInsert().getColumnsCount());
   }
 
+  @Test
+  public void testCopyInWithPriority() throws SQLException, IOException {
+    setupCopyInformationSchemaResults();
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      connection.createStatement().execute("set spanner.copy_commit_priority=low");
+      CopyManager copyManager = new CopyManager(connection.unwrap(BaseConnection.class));
+      copyManager.copyIn("COPY users FROM STDIN;", new StringReader("5\t5\t5\n"));
+    }
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+    CommitRequest commitRequest = commitRequests.get(0);
+    assertEquals(Priority.PRIORITY_LOW, commitRequest.getRequestOptions().getPriority());
+  }
+
   @Test
   public void testCopyInWithColumnNames() throws SQLException, IOException {
     setupCopyInformationSchemaResults();
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
index 73cfaaff..37c20e54 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -2359,7 +2359,7 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
           }
           count++;
         }
-        assertEquals(357, count);
+        assertEquals(358, count);
       }
     }
   }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java
index caf50701..c0cd0ca9 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java
@@ -25,8 +25,10 @@ import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.pgadapter.error.PGException;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.DdlTransactionMode;
 import com.google.cloud.spanner.pgadapter.statements.PgCatalog;
@@ -158,7 +160,7 @@ public class SessionStateTest {
   public void testGetAll() {
     SessionState state = new SessionState(mock(OptionsMetadata.class));
     List<PGSetting> allSettings = state.getAll();
-    assertEquals(356, allSettings.size());
+    assertEquals(357, allSettings.size());
   }
 
   @Test
@@ -196,7 +198,7 @@ public class SessionStateTest {
     state.setLocal("spanner", "custom_local_setting", "value2");
 
     List<PGSetting> allSettings = state.getAll();
-    assertEquals(358, allSettings.size());
+    assertEquals(359, allSettings.size());
 
     PGSetting applicationName =
         allSettings.stream()
@@ -344,11 +346,11 @@ public class SessionStateTest {
   @Test
   public void testSetInvalidEnumValue() {
     SessionState state = new SessionState(mock(OptionsMetadata.class));
-    SpannerException exception =
-        assertThrows(SpannerException.class, () -> state.set(null, "bytea_output", "random_value"));
+    PGException exception =
+        assertThrows(PGException.class, () -> state.set(null, "bytea_output", "random_value"));
     assertEquals(
-        "INVALID_ARGUMENT: invalid value for parameter \"bytea_output\": \"random_value\"",
-        exception.getMessage());
+        "invalid value for parameter \"bytea_output\": \"random_value\"", exception.getMessage());
+    assertEquals("Available values: escape, hex.", exception.getHints());
   }
 
   @Test
@@ -713,6 +715,26 @@ public class SessionStateTest {
     assertFalse(state.isReplacePgCatalogTables());
   }
 
+  @Test
+  public void testCopyCommitPriority() {
+    SessionState state = new SessionState(ImmutableMap.of(), mock(OptionsMetadata.class));
+    CopySettings copySettings = new CopySettings(state);
+
+    assertEquals(RpcPriority.MEDIUM, copySettings.getCommitPriority());
+
+    state.set("spanner", "copy_commit_priority", "high");
+    assertEquals(RpcPriority.HIGH, copySettings.getCommitPriority());
+    state.set("spanner", "copy_commit_priority", "Low");
+    assertEquals(RpcPriority.LOW, copySettings.getCommitPriority());
+
+    PGException exception =
+        assertThrows(PGException.class, () -> state.set("spanner", "copy_commit_priority", "foo"));
+    assertEquals(
+        "invalid value for parameter \"spanner.copy_commit_priority\": \"foo\"",
+        exception.getMessage());
+    assertEquals("Available values: low, medium, high.", exception.getHints());
+  }
+
   @Test
   public void testDdlTransactionMode_noDefault() {
     SessionState state = new SessionState(ImmutableMap.of(), mock(OptionsMetadata.class));
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java
index a9694826..e7e9c0c5 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java
@@ -19,7 +19,9 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyIterable;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
@@ -151,7 +153,7 @@ public class MutationWriterTest {
                 .set("name")
                 .to("Two")
                 .build());
-    verify(databaseClient).write(expectedMutations);
+    verify(databaseClient).writeWithOptions(eq(expectedMutations), any());
   }
 
   @Test
@@ -238,7 +240,7 @@ public class MutationWriterTest {
       StatementResult updateCount = mutationWriter.call();
 
       assertEquals(5L, updateCount.getUpdateCount().longValue());
-      verify(databaseClient, times(3)).write(anyIterable());
+      verify(databaseClient, times(3)).writeWithOptions(anyIterable(), any());
     } finally {
       System.getProperties().remove("copy_in_mutation_limit");
     }
@@ -329,7 +331,7 @@ public class MutationWriterTest {
       // 4. The second batch contains 28 bytes. (3 - 'Three')
       // 5. The third batch contains 24 bytes. (4 - 'Four')
       // 6. the fourth batch contains 24 bytes. (5 - 'Five')
-      verify(databaseClient, times(4)).write(anyIterable());
+      verify(databaseClient, times(4)).writeWithOptions(anyIterable(), any());
     } finally {
       System.getProperties().remove("copy_in_commit_limit");
     }
@@ -399,7 +401,7 @@ public class MutationWriterTest {
                 .set("name")
                 .to("Five")
                 .build());
-    verify(databaseClient).write(expectedMutations);
+    verify(databaseClient).writeWithOptions(eq(expectedMutations), any());
 
     executor.shutdown();
   }