From 4dbb033ed4bec22afb330c13a02a308a82166cbd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Mon, 27 May 2024 15:11:14 +0200
Subject: [PATCH] chore: discard all should reset all session state

---
 .../pgadapter/statements/DiscardStatement.java   |  1 +
 .../pgadapter/JdbcSimpleModeMockServerTest.java  | 16 ++++++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/DiscardStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/DiscardStatement.java
index 09b1d9ab0..3a95dddc1 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/DiscardStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/DiscardStatement.java
@@ -94,6 +94,7 @@ public void executeAsync(BackendConnection backendConnection) {
                         SQLState.ActiveSqlTransaction)));
             return;
           }
+          backendConnection.getSessionState().resetAll();
           connectionHandler.closeAllStatements();
           break;
         case PLANS:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java
index f1620f7e6..aa8401ddd 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java
@@ -597,6 +597,22 @@ public void testDiscard() throws SQLException {
           SQLState.InvalidSqlStatementName.toString(),
           exception.getServerErrorMessage().getSQLState());
 
+      // Verify that DISCARD ALL resets all session state.
+      connection.createStatement().execute("set spanner.copy_upsert=true");
+      try (ResultSet resultSet =
+          connection.createStatement().executeQuery("show spanner.copy_upsert")) {
+        assertTrue(resultSet.next());
+        assertTrue(resultSet.getBoolean(1));
+        assertFalse(resultSet.next());
+      }
+      connection.createStatement().execute("discard all");
+      try (ResultSet resultSet =
+          connection.createStatement().executeQuery("show spanner.copy_upsert")) {
+        assertTrue(resultSet.next());
+        assertFalse(resultSet.getBoolean(1));
+        assertFalse(resultSet.next());
+      }
+
       // Verify that 'discard all' is not accepted in a transaction block.
       connection.setAutoCommit(false);
       exception =