Repository: GoogleCloudPlatform/pgadapter
Pull request #338 (2022-08-28 08:03:31)

feat: add COPY settings to session state
Adds COPY settings to the session state so these can be managed using
normal SET/SHOW session commands.

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/338?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#338](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/338?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (d5e079f) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/c3398882cb404dd55086f6177f07811a7d9767c6?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (c339888) will **decrease** coverage by `0.46%`.
> The diff coverage is `64.35%`.

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #338      +/-   ##
========================================================
- Coverage                 83.95%   83.49%   -0.47%     
- Complexity                 1593     1610      +17     
========================================================
  Files                       114      115       +1     
  Lines                      5417     5477      +60     
  Branches                    727      738      +11     
========================================================
+ Hits                       4548     4573      +25     
- Misses                      643      669      +26     
- Partials                    226      235       +9     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `83.49% <64.35%> (-0.47%)` | :arrow_down: |

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/338?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/338/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3Nlc3Npb24vU2Vzc2lvblN0YXRlLmphdmE=) | `90.16% <47.61%> (-8.88%)` | :arrow_down: |
| [.../cloud/spanner/pgadapter/session/CopySettings.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/338/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3Nlc3Npb24vQ29weVNldHRpbmdzLmphdmE=) | `55.76% <55.76%> (ø)` | |
| [.../cloud/spanner/pgadapter/utils/MutationWriter.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/338/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL011dGF0aW9uV3JpdGVyLmphdmE=) | `92.07% <91.66%> (-0.98%)` | :arrow_down: |
| [...panner/pgadapter/statements/BackendConnection.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/338/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQmFja2VuZENvbm5lY3Rpb24uamF2YQ==) | `90.69% <100.00%> (+0.02%)` | :arrow_up: |
| [...ud/spanner/pgadapter/statements/CopyStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/338/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQ29weVN0YXRlbWVudC5qYXZh) | `81.06% <100.00%> (+0.44%)` | :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
new file mode 100644
index 00000000..a2eacb44
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/session/CopySettings.java
@@ -0,0 +1,136 @@
+// Copyright 2022 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.cloud.spanner.pgadapter.session;
+
+import java.util.Map;
+
+public class CopySettings {
+  private static final int DEFAULT_PIPE_BUFFER_SIZE = 1 << 16;
+  private static final int DEFAULT_MAX_MUTATION_LIMIT = 20_000; // 20k mutations limit
+
+  private final SessionState sessionState;
+
+  static void initCopySettings(Map<String, PGSetting> settings) {
+    // COPY settings could originally be set using System properties. Use these system properties as
+    // the default if they have been set.
+    if (System.getProperty("copy_in_max_parallelism") != null) {
+      int maxParallelism = Integer.parseInt(System.getProperty("copy_in_max_parallelism"));
+      settings
+          .get("spanner.copy_max_parallelism")
+          .initSettingValue(String.valueOf(Math.max(maxParallelism, 1)));
+    }
+    if (System.getProperty("copy_in_insert_or_update") != null) {
+      settings
+          .get("spanner.copy_upsert")
+          .initSettingValue(System.getProperty("copy_in_insert_or_update"));
+    }
+    if (System.getProperty("copy_in_mutation_limit") != null) {
+      int mutationLimit = Integer.parseInt(System.getProperty("copy_in_mutation_limit"));
+      settings
+          .get("spanner.copy_max_atomic_mutations")
+          .initSettingValue(String.valueOf(Math.max(mutationLimit, 1)));
+      settings
+          .get("spanner.copy_batch_size")
+          .initSettingValue(String.valueOf(Math.max(mutationLimit, 1)));
+    }
+    if (System.getProperty("copy_in_commit_limit") != null) {
+      int commitSizeLimit = Integer.parseInt(System.getProperty("copy_in_commit_limit"));
+      settings
+          .get("spanner.copy_max_atomic_commit_size")
+          .initSettingValue(String.valueOf(commitSizeLimit));
+      settings
+          .get("spanner.copy_max_non_atomic_commit_size")
+          .initSettingValue(String.valueOf(commitSizeLimit));
+    }
+    if (System.getProperty("copy_in_commit_limit_multiplier_factor") != null) {
+      float commitLimitMultiplierFactor =
+          Float.parseFloat(System.getProperty("copy_in_commit_limit_multiplier_factor"));
+      settings
+          .get("spanner.copy_commit_size_multiplier_factor")
+          .initSettingValue(String.valueOf(Math.max(commitLimitMultiplierFactor, 1.0f)));
+    }
+    if (System.getProperty("copy_in_pipe_buffer_size") != null) {
+      int pipeBufferSize =
+          Integer.parseInt(
+              System.getProperty(
+                  "copy_in_pipe_buffer_size", String.valueOf(DEFAULT_PIPE_BUFFER_SIZE)));
+      settings
+          .get("spanner.copy_pipe_buffer_size")
+          .initSettingValue(String.valueOf(Math.max(pipeBufferSize, 1024)));
+    }
+  }
+
+  public CopySettings(SessionState sessionState) {
+    this.sessionState = sessionState;
+  }
+
+  /** Returns the maximum number of parallel transactions for a single COPY operation. */
+  public int getMaxParallelism() {
+    return sessionState.getIntegerSetting("spanner", "copy_max_parallelism", 128);
+  }
+
+  /** Returns the commit timeout for COPY operations in seconds. */
+  public int getCommitTimeoutSeconds() {
+    return sessionState.getIntegerSetting("spanner", "copy_commit_timeout", 300);
+  }
+
+  /** Returns the batch size to use for non-atomic COPY operations. */
+  public int getNonAtomicBatchSize() {
+    return sessionState.getIntegerSetting("spanner", "copy_batch_size", 5000);
+  }
+
+  /** Returns the maximum number of mutations in a single commit request. */
+  public int getMaxAtomicMutationsLimit() {
+    return sessionState.getIntegerSetting("spanner", "copy_max_atomic_mutations", 20_000);
+  }
+
+  /** Returns the maximum number of bytes in a single commit request. */
+  public int getMaxAtomicCommitSize() {
+    return sessionState.getIntegerSetting("spanner", "copy_max_atomic_commit_size", 100_000_000);
+  }
+
+  /**
+   * Returns the maximum number of bytes in a single commit request for non-atomic COPY operations.
+   */
+  public int getMaxNonAtomicCommitSize() {
+    return sessionState.getIntegerSetting("spanner", "copy_max_non_atomic_commit_size", 5_000_000);
+  }
+
+  /**
+   * Returns the multiplier that will be applied to the calculated commit size to ensure that a
+   * commit request in a non-atomic COPY operation does not exceed the fixed commit size limit of
+   * Cloud Spanner.
+   */
+  public float getCommitSizeMultiplier() {
+    return sessionState.getFloatSetting("spanner", "copy_commit_size_multiplier_factor", 2.0f);
+  }
+
+  /** Returns the buffer size to use for incoming COPY data messages. */
+  public int getPipeBufferSize() {
+    return sessionState.getIntegerSetting("spanner", "copy_pipe_buffer_size", 1 << 16);
+  }
+
+  /**
+   * Returns whether COPY operations should use upsert instead of insert.
+   *
+   * <p>COPY will INSERT records by default. This is consistent with how COPY on PostgreSQL works.
+   * This option allows PGAdapter to use InsertOrUpdate instead. This can be slightly more efficient
+   * for bulk uploading, and it makes it easier to retry a failed non-atomic batch that might have
+   * already uploaded some but not all data.
+   */
+  public boolean isCopyUpsert() {
+    return sessionState.getBoolSetting("spanner", "copy_upsert", false);
+  }
+}
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 82172086..a425a4b3 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
@@ -14,6 +14,8 @@
 
 package com.google.cloud.spanner.pgadapter.session;
 
+import static com.google.cloud.spanner.pgadapter.session.CopySettings.initCopySettings;
+
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.SpannerException;
@@ -24,6 +26,8 @@ import com.google.cloud.spanner.pgadapter.parsers.BooleanParser;
 import com.google.cloud.spanner.pgadapter.session.PGSetting.Context;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 import java.util.ArrayList;
@@ -83,22 +87,29 @@ public class SessionState {
   private Map<String, PGSetting> localSettings;
 
   public SessionState(OptionsMetadata options) {
-    this(SERVER_SETTINGS, options);
+    this(ImmutableMap.of(), options);
   }
 
   @VisibleForTesting
-  SessionState(Map<String, PGSetting> serverSettings, OptionsMetadata options) {
-    this.settings = new HashMap<>(serverSettings.size());
-    for (Entry<String, PGSetting> entry : serverSettings.entrySet()) {
+  SessionState(Map<String, PGSetting> extraServerSettings, OptionsMetadata options) {
+    Preconditions.checkNotNull(extraServerSettings);
+    Preconditions.checkNotNull(options);
+    this.settings = new HashMap<>(SERVER_SETTINGS.size() + extraServerSettings.size());
+    for (Entry<String, PGSetting> entry : SERVER_SETTINGS.entrySet()) {
+      this.settings.put(entry.getKey(), entry.getValue().copy());
+    }
+    for (Entry<String, PGSetting> entry : extraServerSettings.entrySet()) {
       this.settings.put(entry.getKey(), entry.getValue().copy());
     }
-    initSettingValue("server_version", options.getServerVersion());
-    initSettingValue("server_version_num", options.getServerVersionNum());
+    this.settings.get("server_version").initSettingValue(options.getServerVersion());
+    this.settings.get("server_version_num").initSettingValue(options.getServerVersionNum());
     initSettingValue(
         "spanner.ddl_transaction_mode",
         MoreObjects.firstNonNull(options.getDdlTransactionMode(), DdlTransactionMode.Batch).name());
     initSettingValue(
         "spanner.replace_pg_catalog_tables", Boolean.toString(options.replacePgCatalogTables()));
+
+    initCopySettings(this.settings);
   }
 
   void initSettingValue(String key, String value) {
@@ -257,6 +268,39 @@ public class SessionState {
     return null;
   }
 
+  boolean getBoolSetting(String extension, String name, boolean defaultValue) {
+    PGSetting setting = internalGet(toKey(extension, name), false);
+    if (setting != null) {
+      try {
+        return BooleanParser.toBoolean(setting.getSetting());
+      } catch (Exception ignore) {
+      }
+    }
+    return defaultValue;
+  }
+
+  int getIntegerSetting(String extension, String name, int defaultValue) {
+    PGSetting setting = internalGet(toKey(extension, name), false);
+    if (setting != null) {
+      try {
+        return Integer.parseInt(setting.getSetting());
+      } catch (Exception ignore) {
+      }
+    }
+    return defaultValue;
+  }
+
+  float getFloatSetting(String extension, String name, float defaultValue) {
+    PGSetting setting = internalGet(toKey(extension, name), false);
+    if (setting != null) {
+      try {
+        return Float.parseFloat(setting.getSetting());
+      } catch (Exception ignore) {
+      }
+    }
+    return defaultValue;
+  }
+
   /** Returns all settings and their current values. */
   public List<PGSetting> getAll() {
     List<PGSetting> result =
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java
index 2ccd6fc6..a643b184 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java
@@ -323,6 +323,10 @@ public class CopyStatement extends IntermediatePortalStatement {
       setParserFormat(this.options);
       mutationWriter =
           new MutationWriter(
+              connectionHandler
+                  .getExtendedQueryProtocolHandler()
+                  .getBackendConnection()
+                  .getSessionState(),
               getTransactionMode(),
               connection,
               options.getTableName(),
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 33e9c6d6..2d93e3d8 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
@@ -32,6 +32,8 @@ import com.google.cloud.spanner.Value;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.pgadapter.parsers.copy.CopyTreeParser.CopyOptions;
+import com.google.cloud.spanner.pgadapter.session.CopySettings;
+import com.google.cloud.spanner.pgadapter.session.SessionState;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.UpdateCount;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
@@ -85,45 +87,15 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
 
   private static final Logger logger = Logger.getLogger(MutationWriter.class.getName());
 
-  private static final int DEFAULT_MUTATION_LIMIT = 20_000; // 20k mutation count limit
-  private static final int DEFAULT_COMMIT_LIMIT =
-      100_000_000; // 100MB mutation API commit size limit
-  /**
-   * The factor that the size of the incoming payload is multiplied with to estimate whether the
-   * current commit size will exceed the maximum commit size. The factor is chosen conservatively to
-   * ensure that a COPY operation of a large file does not unnecessarily break because of an
-   * exceeded commit size.
-   */
-  private static final float DEFAULT_COMMIT_LIMIT_MULTIPLIER_FACTOR = 2.0f;
-
-  private static final int DEFAULT_COMMIT_TIMEOUT_SECONDS = 300;
-  private static final int DEFAULT_MAX_PARALLELISM = 128;
-  private static final int DEFAULT_PIPE_BUFFER_SIZE = 1 << 16;
-
-  /**
-   * COPY will INSERT records by default. This is consistent with how COPY on PostgreSQL works. This
-   * option allows PGAdapter to use InsertOrUpdate instead. This can be slightly more efficient for
-   * bulk uploading, and it makes it easier to retry a failed non-atomic batch that might have
-   * already uploaded some but not all data.
-   */
-  private final boolean insertOrUpdate =
-      Boolean.parseBoolean(System.getProperty("copy_in_insert_or_update", "false"));
-
-  private final int commitSizeLimit =
-      Integer.parseInt(
-          System.getProperty("copy_in_commit_limit", String.valueOf(DEFAULT_COMMIT_LIMIT)));
-  private final int commitTimeoutSeconds =
-      Integer.parseInt(
-          System.getProperty(
-              "copy_in_commit_timeout_seconds", String.valueOf(DEFAULT_COMMIT_TIMEOUT_SECONDS)));
   private final CopyTransactionMode transactionMode;
   private long rowCount;
   private final Connection connection;
   private final String tableName;
   private final Map<String, Type> tableColumns;
-  private final int maxBatchSize;
+  private final int maxAtomicBatchSize;
+  private final int nonAtomicBatchSize;
   private final long commitSizeLimitForBatching;
-  private final int maxParallelism;
+  private final CopySettings copySettings;
   private final CopyInParser parser;
   private final PipedOutputStream payload = new PipedOutputStream();
   private final AtomicBoolean commit = new AtomicBoolean(false);
@@ -138,6 +110,7 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
   private SpannerException exception;
 
   public MutationWriter(
+      SessionState sessionState,
       CopyTransactionMode transactionMode,
       Connection connection,
       String tableName,
@@ -151,35 +124,19 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
     this.connection = connection;
     this.tableName = tableName;
     this.tableColumns = tableColumns;
-    int mutationLimit =
-        Math.max(
-            Integer.parseInt(
-                System.getProperty(
-                    "copy_in_mutation_limit", String.valueOf(DEFAULT_MUTATION_LIMIT))),
-            1);
-    this.maxBatchSize = Math.max(mutationLimit / (tableColumns.size() + indexedColumnsCount), 1);
-    float commitLimitMultiplierFactor =
-        Math.max(
-            Float.parseFloat(
-                System.getProperty(
-                    "copy_in_commit_limit_multiplier_factor",
-                    String.valueOf(DEFAULT_COMMIT_LIMIT_MULTIPLIER_FACTOR))),
-            1.0f);
+    this.copySettings = new CopySettings(sessionState);
+    int atomicMutationLimit = copySettings.getMaxAtomicMutationsLimit();
+    this.maxAtomicBatchSize =
+        Math.max(atomicMutationLimit / (tableColumns.size() + indexedColumnsCount), 1);
+    int nonAtomicMutations = copySettings.getNonAtomicBatchSize();
+    this.nonAtomicBatchSize =
+        Math.max(nonAtomicMutations / (tableColumns.size() + indexedColumnsCount), 1);
     this.commitSizeLimitForBatching =
-        Math.round((float) commitSizeLimit / commitLimitMultiplierFactor);
-    this.maxParallelism =
-        Math.max(
-            Integer.parseInt(
-                System.getProperty(
-                    "copy_in_max_parallelism", String.valueOf(DEFAULT_MAX_PARALLELISM))),
-            1);
-    int pipeBufferSize =
-        Math.max(
-            Integer.parseInt(
-                System.getProperty(
-                    "copy_in_pipe_buffer_size", String.valueOf(DEFAULT_PIPE_BUFFER_SIZE))),
-            1024);
-    this.parser = CopyInParser.create(copyFormat, format, payload, pipeBufferSize, hasHeader);
+        Math.round(
+            (float) copySettings.getMaxAtomicCommitSize() / copySettings.getCommitSizeMultiplier());
+    this.parser =
+        CopyInParser.create(
+            copyFormat, format, payload, copySettings.getPipeBufferSize(), hasHeader);
   }
 
   /** @return number of rows copied into Spanner */
@@ -240,7 +197,7 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
     // pressure in our entire pipeline that consists of:
     // Client app (psql) -> CopyData message -> CSVParser -> Transaction.
     LinkedBlockingDeque<ApiFuture<Void>> activeCommitFutures =
-        new LinkedBlockingDeque<>(maxParallelism);
+        new LinkedBlockingDeque<>(copySettings.getMaxParallelism());
     // This list holds all transactions that we have started. We will wait on this entire list
     // before finishing, to ensure that all data has been written before we signal that we are done.
     List<ApiFuture<Void>> allCommitFutures = new ArrayList<>();
@@ -278,28 +235,28 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
         } else {
           mutations.add(mutation);
           currentBufferByteSize += mutationSize;
-          if (mutations.size() > maxBatchSize) {
+          if (mutations.size() > maxAtomicBatchSize) {
             throw SpannerExceptionFactory.newSpannerException(
                 ErrorCode.FAILED_PRECONDITION,
                 "Record count: "
                     + mutations.size()
                     + " has exceeded the limit: "
-                    + maxBatchSize
+                    + maxAtomicBatchSize
                     + ".\n\nThe number of mutations per record is equal to the number of columns in the record "
                     + "plus the number of indexed columns in the record. The maximum number of mutations "
                     + "in one transaction is "
-                    + DEFAULT_MUTATION_LIMIT
+                    + copySettings.getMaxAtomicMutationsLimit()
                     + ".\n\nExecute `SET SPANNER.AUTOCOMMIT_DML_MODE='PARTITIONED_NON_ATOMIC'` before executing a large COPY operation "
                     + "to instruct PGAdapter to automatically break large transactions into multiple smaller. "
                     + "This will make the COPY operation non-atomic.\n\n");
           }
-          if (currentBufferByteSize > commitSizeLimit) {
+          if (currentBufferByteSize > copySettings.getMaxAtomicCommitSize()) {
             throw SpannerExceptionFactory.newSpannerException(
                 ErrorCode.FAILED_PRECONDITION,
                 "Commit size: "
                     + currentBufferByteSize
                     + " has exceeded the limit: "
-                    + commitSizeLimit
+                    + copySettings.getMaxAtomicCommitSize()
                     + ".\n\nExecute `SET SPANNER.AUTOCOMMIT_DML_MODE='PARTITIONED_NON_ATOMIC'` before executing a large COPY operation "
                     + "to instruct PGAdapter to automatically break large transactions into multiple smaller. "
                     + "This will make the COPY operation non-atomic.\n\n");
@@ -363,8 +320,11 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
       long currentBufferByteSize,
       int mutationSize)
       throws Exception {
-    // Flush before adding to the buffer if the commit size would be exceeded.
-    if (!mutations.isEmpty() && currentBufferByteSize + mutationSize > commitSizeLimitForBatching) {
+    // Flush before adding to the buffer if either of the commit size limits would be exceeded.
+    long estimatedNextSize = currentBufferByteSize + mutationSize;
+    if (!mutations.isEmpty()
+        && (estimatedNextSize > commitSizeLimitForBatching
+            || estimatedNextSize > copySettings.getMaxNonAtomicCommitSize())) {
       allCommitFutures.add(writeToSpannerAsync(activeCommitFutures, mutations));
       mutations.clear();
       mutations.add(mutation);
@@ -372,7 +332,7 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
     }
 
     mutations.add(mutation);
-    if (mutations.size() == maxBatchSize) {
+    if (mutations.size() == nonAtomicBatchSize) {
       allCommitFutures.add(writeToSpannerAsync(activeCommitFutures, mutations));
       mutations.clear();
       return 0L; // Buffer is empty, so the batch size in bytes is now back to zero.
@@ -400,7 +360,8 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
                       .withValue(
                           SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY,
                           SpannerCallContextTimeoutConfigurator.create()
-                              .withCommitTimeout(Duration.ofSeconds(100L)));
+                              .withCommitTimeout(
+                                  Duration.ofSeconds(copySettings.getCommitTimeoutSeconds())));
               context.run(() -> dbClient.write(immutableMutations));
               return null;
             });
@@ -513,7 +474,7 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
     // useful for importing large datasets using PartitionedNonAtomic mode. If an import attempt
     // fails halfway, it can easily be retried with InsertOrUpdate as it will just overwrite
     // existing records instead of failing on a UniqueKeyConstraint violation.
-    if (this.insertOrUpdate) {
+    if (copySettings.isCopyUpsert()) {
       builder = Mutation.newInsertOrUpdateBuilder(this.tableName);
     } else {
       builder = Mutation.newInsertBuilder(this.tableName);
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 965358a4..fb68fcad 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
@@ -343,5 +343,14 @@ work_mem	4096	kB	Resource Usage / Memory	Sets the maximum memory to be used for
 xmlbinary	base64	\N	Client Connection Defaults / Statement Behavior	Sets how binary values are to be encoded in XML.	\N	user	enum	default	\N	\N	{base64,hex}	base64	base64	\N	\N	f
 xmloption	content	\N	Client Connection Defaults / Statement Behavior	Sets whether XML data in implicit parsing and serialization operations is to be considered as documents or content fragments.	\N	user	enum	default	\N	\N	{content,document}	content	content	\N	\N	f
 zero_damaged_pages	off	\N	Developer Options	Continues processing past damaged page headers.	Detection of a damaged page header normally causes PostgreSQL to report an error, aborting the current transaction. Setting zero_damaged_pages to true causes the system to instead report a warning, zero out the damaged page, and continue processing. This behavior will destroy data, namely all the rows on the damaged page.	superuser	bool	default	\N	\N	\N	off	off	\N	\N	f
+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_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
+spanner.copy_max_non_atomic_commit_size	5000000	\N	COPY / The max number of bytes per commit in a non-atomic COPY operation	The max number of bytes per commit in a non-atomic COPY operation.	\N	user	integer	default	\N	\N	\N	5000000	5000000	\N	\N	f
+spanner.copy_commit_size_multiplier_factor	2.0	\N	COPY / Factor for estimating COPY commit size	The factor that the size of the incoming payload is multiplied with to estimate whether the current commit size will exceed the maximum commit size.	\N	user	real	default	\N	\N	\N	2.0	2.0	\N	\N	f
+spanner.copy_pipe_buffer_size	65536	\N	COPY / Buffer size for incoming COPY data messages	Internal buffer size for incoming COPY data messages.	\N	internal	integer	default	\N	\N	\N	100000000	100000000	\N	\N	f
 spanner.ddl_transaction_mode	Batch	\N	PGAdapter Options	DDL transaction mode.	\N	user	enum	default	\N	\N	{Single,Batch,AutocommitImplicitTransaction,AutocommitExplicitTransaction}	Batch	Batch	\N	\N	f
 spanner.replace_pg_catalog_tables	on	\N	PGAdapter Options	Automatically replace pg_catalog table references with Common Table Expressions.	\N	user	bool	default	\N	\N	\N	on	on	\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 48aaf0e3..085e184f 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
@@ -436,8 +436,8 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
   public void testCopyInExceedsCommitSizeLimit_FailsInAtomicMode() throws SQLException {
     setupCopyInformationSchemaResults();
 
+    System.setProperty("copy_in_commit_limit", "10");
     try (Connection connection = DriverManager.getConnection(createUrl())) {
-      System.setProperty("copy_in_commit_limit", "10");
       CopyManager copyManager = new CopyManager(connection.unwrap(BaseConnection.class));
       SQLException exception =
           assertThrows(
@@ -460,8 +460,8 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
   public void testCopyInExceedsCommitSizeLimit_BatchesInNonAtomicMode()
       throws SQLException, IOException {
     setupCopyInformationSchemaResults();
+    System.setProperty("copy_in_commit_limit", "10");
     try (Connection connection = DriverManager.getConnection(createUrl())) {
-      System.setProperty("copy_in_commit_limit", "10");
       connection
           .createStatement()
           .execute("set spanner.autocommit_dml_mode='partitioned_non_atomic'");
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 4e5ff5cf..72d80494 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -2057,7 +2057,7 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
           }
           count++;
         }
-        assertEquals(347, count);
+        assertEquals(356, count);
       }
     }
   }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/session/CopySettingsTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/session/CopySettingsTest.java
new file mode 100644
index 00000000..3c4b598d
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/session/CopySettingsTest.java
@@ -0,0 +1,149 @@
+// Copyright 2022 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.cloud.spanner.pgadapter.session;
+
+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.Mockito.mock;
+
+import com.google.cloud.spanner.SpannerException;
+import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class CopySettingsTest {
+
+  @Test
+  public void testMaxParallelism() {
+    try {
+      System.setProperty("copy_in_max_parallelism", "15");
+      SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+      CopySettings copySettings = new CopySettings(sessionState);
+      assertEquals(15, copySettings.getMaxParallelism());
+    } finally {
+      System.clearProperty("copy_in_max_parallelism");
+    }
+    SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+    CopySettings copySettings = new CopySettings(sessionState);
+    assertEquals(128, copySettings.getMaxParallelism());
+
+    sessionState.set("spanner", "copy_max_parallelism", "256");
+    assertEquals(copySettings.getMaxParallelism(), 256);
+  }
+
+  @Test
+  public void testUpsert() {
+    try {
+      System.setProperty("copy_in_insert_or_update", "true");
+      SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+      CopySettings copySettings = new CopySettings(sessionState);
+      assertTrue(copySettings.isCopyUpsert());
+    } finally {
+      System.clearProperty("copy_in_insert_or_update");
+    }
+    SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+    CopySettings copySettings = new CopySettings(sessionState);
+    assertFalse(copySettings.isCopyUpsert());
+
+    sessionState.set("spanner", "copy_upsert", "true");
+    assertTrue(copySettings.isCopyUpsert());
+  }
+
+  @Test
+  public void testMutationLimit() {
+    try {
+      System.setProperty("copy_in_mutation_limit", "100");
+      SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+      CopySettings copySettings = new CopySettings(sessionState);
+      assertEquals(100, copySettings.getMaxAtomicMutationsLimit());
+      assertEquals(100, copySettings.getNonAtomicBatchSize());
+    } finally {
+      System.clearProperty("copy_in_mutation_limit");
+    }
+    SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+    CopySettings copySettings = new CopySettings(sessionState);
+    assertEquals(20_000, copySettings.getMaxAtomicMutationsLimit());
+    assertEquals(5000, copySettings.getNonAtomicBatchSize());
+
+    assertThrows(
+        SpannerException.class,
+        () -> sessionState.set("spanner", "copy_max_atomic_mutations", "1000"));
+    sessionState.set("spanner", "copy_batch_size", "1000");
+    assertEquals(1000, copySettings.getNonAtomicBatchSize());
+  }
+
+  @Test
+  public void testCommitSizeLimit() {
+    try {
+      System.setProperty("copy_in_commit_limit", "1000");
+      SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+      CopySettings copySettings = new CopySettings(sessionState);
+      assertEquals(1000, copySettings.getMaxAtomicCommitSize());
+      assertEquals(1000, copySettings.getMaxNonAtomicCommitSize());
+    } finally {
+      System.clearProperty("copy_in_commit_limit");
+    }
+    SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+    CopySettings copySettings = new CopySettings(sessionState);
+    assertEquals(100_000_000, copySettings.getMaxAtomicCommitSize());
+    assertEquals(5_000_000, copySettings.getMaxNonAtomicCommitSize());
+
+    assertThrows(
+        SpannerException.class,
+        () -> sessionState.set("spanner", "copy_max_atomic_commit_size", "1000"));
+    sessionState.set("spanner", "copy_max_non_atomic_commit_size", "1000");
+    assertEquals(1000, copySettings.getMaxNonAtomicCommitSize());
+  }
+
+  @Test
+  public void testCommitLimitMultiplierFactor() {
+    try {
+      System.setProperty("copy_in_commit_limit_multiplier_factor", "1.5");
+      SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+      CopySettings copySettings = new CopySettings(sessionState);
+      assertEquals(1.5f, copySettings.getCommitSizeMultiplier(), 0.0f);
+    } finally {
+      System.clearProperty("copy_in_commit_limit_multiplier_factor");
+    }
+    SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+    CopySettings copySettings = new CopySettings(sessionState);
+    assertEquals(2.0f, copySettings.getCommitSizeMultiplier(), 0.0f);
+
+    sessionState.set("spanner", "copy_commit_size_multiplier_factor", "3.0");
+    assertEquals(3.0f, copySettings.getCommitSizeMultiplier(), 0.0f);
+  }
+
+  @Test
+  public void testPipeBufferSize() {
+    try {
+      System.setProperty("copy_in_pipe_buffer_size", "5000");
+      SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+      CopySettings copySettings = new CopySettings(sessionState);
+      assertEquals(5000, copySettings.getPipeBufferSize());
+    } finally {
+      System.clearProperty("copy_in_pipe_buffer_size");
+    }
+    SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+    CopySettings copySettings = new CopySettings(sessionState);
+    assertEquals(1 << 16, copySettings.getPipeBufferSize());
+
+    assertThrows(
+        SpannerException.class, () -> sessionState.set("spanner", "copy_pipe_buffer_size", "1000"));
+  }
+}
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 4bd15f02..d8dd98bf 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
@@ -156,7 +156,7 @@ public class SessionStateTest {
   public void testGetAll() {
     SessionState state = new SessionState(mock(OptionsMetadata.class));
     List<PGSetting> allSettings = state.getAll();
-    assertEquals(347, allSettings.size());
+    assertEquals(356, allSettings.size());
   }
 
   @Test
@@ -194,7 +194,7 @@ public class SessionStateTest {
     state.setLocal("spanner", "custom_local_setting", "value2");
 
     List<PGSetting> allSettings = state.getAll();
-    assertEquals(349, allSettings.size());
+    assertEquals(358, allSettings.size());
 
     PGSetting applicationName =
         allSettings.stream()
@@ -548,6 +548,57 @@ public class SessionStateTest {
         withSessionState.getSql());
   }
 
+  @Test
+  public void testGetBoolSetting() {
+    SessionState state = new SessionState(mock(OptionsMetadata.class));
+    assertFalse(state.getBoolSetting("spanner", "unknown_setting", false));
+    assertTrue(state.getBoolSetting("spanner", "unknown_setting", true));
+
+    state.set("spanner", "custom_setting", "on");
+    assertTrue(state.getBoolSetting("spanner", "custom_setting", false));
+    state.set("spanner", "custom_setting", "off");
+    assertFalse(state.getBoolSetting("spanner", "custom_setting", true));
+
+    state.set("spanner", "custom_setting", "foo");
+    assertFalse(state.getBoolSetting("spanner", "custom_setting", false));
+    state.set("spanner", "custom_setting", "foo");
+    assertTrue(state.getBoolSetting("spanner", "custom_setting", true));
+  }
+
+  @Test
+  public void testGetIntegerSetting() {
+    SessionState state = new SessionState(mock(OptionsMetadata.class));
+    assertEquals(100, state.getIntegerSetting("spanner", "unknown_setting", 100));
+    assertEquals(0, state.getIntegerSetting("spanner", "unknown_setting", 0));
+
+    state.set("spanner", "custom_setting", "200");
+    assertEquals(200, state.getIntegerSetting("spanner", "custom_setting", 100));
+    state.set("spanner", "custom_setting", "-200");
+    assertEquals(-200, state.getIntegerSetting("spanner", "custom_setting", 100));
+
+    state.set("spanner", "custom_setting", "foo");
+    assertEquals(100, state.getIntegerSetting("spanner", "custom_setting", 100));
+    state.set("spanner", "custom_setting", "bar");
+    assertEquals(0, state.getIntegerSetting("spanner", "custom_setting", 0));
+  }
+
+  @Test
+  public void testGetFloatSetting() {
+    SessionState state = new SessionState(mock(OptionsMetadata.class));
+    assertEquals(1.1f, state.getFloatSetting("spanner", "unknown_setting", 1.1f), 0.0f);
+    assertEquals(0.0f, state.getFloatSetting("spanner", "unknown_setting", 0.0f), 0.0f);
+
+    state.set("spanner", "custom_setting", "2");
+    assertEquals(2f, state.getFloatSetting("spanner", "custom_setting", 1f), 0.0f);
+    state.set("spanner", "custom_setting", "-2.9");
+    assertEquals(-2.9f, state.getFloatSetting("spanner", "custom_setting", 1.5f), 0.0f);
+
+    state.set("spanner", "custom_setting", "foo");
+    assertEquals(4.4f, state.getFloatSetting("spanner", "custom_setting", 4.4f), 0.0f);
+    state.set("spanner", "custom_setting", "bar");
+    assertEquals(0f, state.getFloatSetting("spanner", "custom_setting", 0f), 0.0f);
+  }
+
   @Test
   public void testTryGet() {
     assertNull(
@@ -568,8 +619,8 @@ public class SessionStateTest {
 
   @Test
   public void testIsReplacePgCatalogTables_noDefault() {
-    SessionState state = new SessionState(ImmutableMap.of(), mock(OptionsMetadata.class));
-    assertTrue(state.isReplacePgCatalogTables());
+    SessionState state = new SessionState(mock(OptionsMetadata.class));
+    assertFalse(state.isReplacePgCatalogTables());
   }
 
   @Test
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java
index 7806580d..333963f4 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java
@@ -45,6 +45,7 @@ import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.ProxyServer;
 import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import com.google.cloud.spanner.pgadapter.session.SessionState;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage;
@@ -437,6 +438,13 @@ public class StatementTest {
   public void testCopyInvalidBuildMutation() throws Exception {
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+    ExtendedQueryProtocolHandler extendedQueryProtocolHandler =
+        mock(ExtendedQueryProtocolHandler.class);
+    when(extendedQueryProtocolHandler.getBackendConnection()).thenReturn(backendConnection);
+    when(connectionHandler.getExtendedQueryProtocolHandler())
+        .thenReturn(extendedQueryProtocolHandler);
+    SessionState sessionState = new SessionState(options);
+    when(backendConnection.getSessionState()).thenReturn(sessionState);
     setupQueryInformationSchemaResults();
 
     String sql = "COPY keyvalue FROM STDIN;";
@@ -514,6 +522,13 @@ public class StatementTest {
   @Test
   public void testCopyBatchSizeLimit() throws Exception {
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
+    ExtendedQueryProtocolHandler extendedQueryProtocolHandler =
+        mock(ExtendedQueryProtocolHandler.class);
+    when(extendedQueryProtocolHandler.getBackendConnection()).thenReturn(backendConnection);
+    when(connectionHandler.getExtendedQueryProtocolHandler())
+        .thenReturn(extendedQueryProtocolHandler);
+    SessionState sessionState = new SessionState(options);
+    when(backendConnection.getSessionState()).thenReturn(sessionState);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     setupQueryInformationSchemaResults();
     BackendConnection backendConnection =
@@ -556,6 +571,13 @@ public class StatementTest {
   @Test
   public void testCopyDataRowLengthMismatchLimit() throws Exception {
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+    ExtendedQueryProtocolHandler extendedQueryProtocolHandler =
+        mock(ExtendedQueryProtocolHandler.class);
+    when(extendedQueryProtocolHandler.getBackendConnection()).thenReturn(backendConnection);
+    when(connectionHandler.getExtendedQueryProtocolHandler())
+        .thenReturn(extendedQueryProtocolHandler);
+    SessionState sessionState = new SessionState(options);
+    when(backendConnection.getSessionState()).thenReturn(sessionState);
     setupQueryInformationSchemaResults();
     BackendConnection backendConnection =
         new BackendConnection(
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 e6512c23..cbebf9e0 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
@@ -36,7 +36,9 @@ import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.Value;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.parsers.copy.CopyTreeParser.CopyOptions.Format;
+import com.google.cloud.spanner.pgadapter.session.SessionState;
 import com.google.cloud.spanner.pgadapter.utils.CsvCopyParser.CsvCopyRecord;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter.CopyTransactionMode;
 import com.google.common.collect.ImmutableList;
@@ -117,11 +119,13 @@ public class MutationWriterTest {
             .builder()
             .setHeader(tableColumns.keySet().toArray(new String[0]))
             .build();
+    SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
     Connection connection = mock(Connection.class);
     DatabaseClient databaseClient = mock(DatabaseClient.class);
     when(connection.getDatabaseClient()).thenReturn(databaseClient);
     MutationWriter mutationWriter =
         new MutationWriter(
+            sessionState,
             CopyTransactionMode.ImplicitAtomic,
             connection,
             "numbers",
@@ -161,9 +165,11 @@ public class MutationWriterTest {
               .builder()
               .setHeader(tableColumns.keySet().toArray(new String[0]))
               .build();
+      SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
       Connection connection = mock(Connection.class);
       MutationWriter mutationWriter =
           new MutationWriter(
+              sessionState,
               CopyTransactionMode.ImplicitAtomic,
               connection,
               "numbers",
@@ -180,7 +186,7 @@ public class MutationWriterTest {
       assertEquals(
           "FAILED_PRECONDITION: Record count: 2 has exceeded the limit: 1.\n"
               + "\n"
-              + "The number of mutations per record is equal to the number of columns in the record plus the number of indexed columns in the record. The maximum number of mutations in one transaction is 20000.\n"
+              + "The number of mutations per record is equal to the number of columns in the record plus the number of indexed columns in the record. The maximum number of mutations in one transaction is 1.\n"
               + "\n"
               + "Execute `SET SPANNER.AUTOCOMMIT_DML_MODE='PARTITIONED_NON_ATOMIC'` before executing a large COPY operation to instruct PGAdapter to automatically break large transactions into multiple smaller. This will make the COPY operation non-atomic.\n\n",
           exception.getMessage());
@@ -203,12 +209,17 @@ public class MutationWriterTest {
               .builder()
               .setHeader(tableColumns.keySet().toArray(new String[0]))
               .build();
+      SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
+      // 6 == 2 mutations per batch, as we have 2 columns + 1 indexed column.
+      sessionState.set("spanner", "copy_batch_size", "6");
+      sessionState.commit();
       Connection connection = mock(Connection.class);
       DatabaseClient databaseClient = mock(DatabaseClient.class);
       when(connection.getDatabaseClient()).thenReturn(databaseClient);
 
       MutationWriter mutationWriter =
           new MutationWriter(
+              sessionState,
               CopyTransactionMode.ImplicitNonAtomic,
               connection,
               "numbers",
@@ -244,9 +255,11 @@ public class MutationWriterTest {
               .builder()
               .setHeader(tableColumns.keySet().toArray(new String[0]))
               .build();
+      SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
       Connection connection = mock(Connection.class);
       MutationWriter mutationWriter =
           new MutationWriter(
+              sessionState,
               CopyTransactionMode.ImplicitAtomic,
               connection,
               "numbers",
@@ -283,12 +296,14 @@ public class MutationWriterTest {
               .builder()
               .setHeader(tableColumns.keySet().toArray(new String[0]))
               .build();
+      SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
       Connection connection = mock(Connection.class);
       DatabaseClient databaseClient = mock(DatabaseClient.class);
       when(connection.getDatabaseClient()).thenReturn(databaseClient);
 
       MutationWriter mutationWriter =
           new MutationWriter(
+              sessionState,
               CopyTransactionMode.ImplicitNonAtomic,
               connection,
               "numbers",
@@ -328,12 +343,14 @@ public class MutationWriterTest {
             .builder()
             .setHeader(tableColumns.keySet().toArray(new String[0]))
             .build();
+    SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
     Connection connection = mock(Connection.class);
     DatabaseClient databaseClient = mock(DatabaseClient.class);
     when(connection.getDatabaseClient()).thenReturn(databaseClient);
 
     MutationWriter mutationWriter =
         new MutationWriter(
+            sessionState,
             CopyTransactionMode.ImplicitNonAtomic,
             connection,
             "numbers",
@@ -509,6 +526,7 @@ public class MutationWriterTest {
   @Test
   public void testBuildMutationNulls() throws IOException {
     // TODO(b/237831798): Add support for ARRAY in COPY operations.
+    SessionState sessionState = new SessionState(mock(OptionsMetadata.class));
     Connection connection = mock(Connection.class);
     for (Type type :
         new Type[] {
@@ -524,6 +542,7 @@ public class MutationWriterTest {
         }) {
       MutationWriter mutationWriter =
           new MutationWriter(
+              sessionState,
               CopyTransactionMode.ImplicitAtomic,
               connection,
               "my_table",
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
index eab676ee..8a4b2fa1 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
@@ -1419,6 +1419,13 @@ public class ProtocolTest {
   @Test
   public void testCopyFromFilePipe() throws Exception {
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+    ExtendedQueryProtocolHandler extendedQueryProtocolHandler =
+        mock(ExtendedQueryProtocolHandler.class);
+    when(extendedQueryProtocolHandler.getBackendConnection()).thenReturn(backendConnection);
+    when(connectionHandler.getExtendedQueryProtocolHandler())
+        .thenReturn(extendedQueryProtocolHandler);
+    SessionState sessionState = new SessionState(options);
+    when(backendConnection.getSessionState()).thenReturn(sessionState);
     setupQueryInformationSchemaResults();
 
     byte[] payload = Files.readAllBytes(Paths.get("./src/test/resources/small-file-test.txt"));
@@ -1533,7 +1540,6 @@ public class ProtocolTest {
     when(sessionState.get(null, "server_version")).thenReturn(serverVersionSetting);
     when(backendConnection.getSessionState()).thenReturn(sessionState);
     when(server.getOptions()).thenReturn(options);
-    when(options.getServerVersion()).thenReturn("13.4");
     when(options.shouldAuthenticate()).thenReturn(false);
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);
     when(connectionMetadata.getOutputStream()).thenReturn(outputStream);