Repository: GoogleCloudPlatform/pgadapter (java)

Issue:
Exit PGAdapter Sidecar when main container exits
Thank you in advance for your time!

I have a kubernetes job that's using pgadapter as a side car. When the main container finishes, the pod get stuck in `NotReady` because pg adpter is running.

Is there native way to handle this properly?

I've seen this stackoverflow post https://stackoverflow.com/questions/41679364/kubernetes-stop-cloudsql-proxy-sidecar-container-in-multi-container-pod-job but this looks a bit of a hack.  

Pull request:
feat: support SHUTDOWN SQL command
Adds support for a custom SQL command to shut down the server:

```
SHUTDOWN [SMART | FAST | IMMEDIATE]
```

This command only shuts down the server if PGAdapter has been started with the command line argument `--allow_shutdown_command`. This argument should only be used when PGAdapter runs in a trusted environment, for example as a side-car container. The SQL statement can be used to shut down PGAdapter after finishing a batch job, or when the client application knows that the cluster / pod / VM / ... it is running in is shutting down.

Fixes #2040

Patch:
diff --git a/README.md b/README.md
index 806f3c7b..cc65189b 100644
--- a/README.md
+++ b/README.md
@@ -271,6 +271,17 @@ The following list contains the most frequently used startup options for PGAdapt
     Use the -x switch to turn off the localhost check. This is required when running PGAdapter in a
     Docker container, as the connections from the host machine will not be seen as a connection from
     localhost in the container.
+
+--allow_shutdown_command
+  * Enables the use of the custom SQL command `SHUTDOWN [SMART | FAST | IMMEDIATE]`. This command can
+    be used to shut down PGAdapter by just sending it a SQL statement. This option should only be enabled
+    when PGAdapter runs in a trusted environment, for example as a side-car container. The default
+    shutdown mode is `FAST`, which terminates all existing connections and then shuts down PGAdapter.
+    Use shutdown mode `SMART` to instruct PGAdapter to wait until all existing connections have been
+    terminated by the client before shutting down. See also https://www.postgresql.org/docs/current/server-shutdown.html
+    for more information about shutdown modes.
+  * Note that `SHUTDOWN [SMART | FAST | IMMEDIATE]` only works on PGAdapter. This command is not
+    supported by PostgreSQL.
 ```
 
 * See [command line arguments](docs/command_line_arguments.md) for a list of all supported arguments.
diff --git a/benchmarks/tpcc/src/main/java/com/google/cloud/pgadapter/tpcc/BenchmarkApplication.java b/benchmarks/tpcc/src/main/java/com/google/cloud/pgadapter/tpcc/BenchmarkApplication.java
index 1f2f704e..996e859f 100644
--- a/benchmarks/tpcc/src/main/java/com/google/cloud/pgadapter/tpcc/BenchmarkApplication.java
+++ b/benchmarks/tpcc/src/main/java/com/google/cloud/pgadapter/tpcc/BenchmarkApplication.java
@@ -245,7 +245,6 @@ public class BenchmarkApplication implements CommandLineRunner {
 
   static Attributes createMetricAttributes(SpannerConfiguration spannerConfiguration) {
     AttributesBuilder attributesBuilder = Attributes.builder();
-    attributesBuilder.put("connection_id", UUID.randomUUID().toString());
     attributesBuilder.put("database", spannerConfiguration.getDatabase());
     attributesBuilder.put("instance_id", spannerConfiguration.getInstance());
     attributesBuilder.put("project_id", spannerConfiguration.getProject());
diff --git a/docs/command_line_arguments.md b/docs/command_line_arguments.md
index 40d5048a..cc1784ca 100644
--- a/docs/command_line_arguments.md
+++ b/docs/command_line_arguments.md
@@ -70,6 +70,17 @@
     Docker container, as the connections from the host machine will not be seen as a connection from
     localhost in the container.
 
+--allow_shutdown_command
+  * Enables the use of the custom SQL command `SHUTDOWN [SMART | FAST | IMMEDIATE]`. This command can
+    be used to shut down PGAdapter by just sending it a SQL statement. This option should only be enabled
+    when PGAdapter runs in a trusted environment, for example as a side-car container. The default
+    shutdown mode is `FAST`, which terminates all existing connections and then shuts down PGAdapter.
+    Use shutdown mode `SMART` to instruct PGAdapter to wait until all existing connections have been
+    terminated by the client before shutting down. See also https://www.postgresql.org/docs/current/server-shutdown.html
+    for more information about shutdown modes.
+  * Note that `SHUTDOWN [SMART | FAST | IMMEDIATE]` only works on PGAdapter. This command is not
+    supported by PostgreSQL.
+
 -r
   * PGAdapter internally uses the core engine of the Cloud Spanner JDBC driver. This option specifies
     additional properties that will be used with that internal JDBC connection. All connection properties
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java b/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
index 72166cd4..f38ee4c9 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
@@ -114,6 +114,8 @@ public class ProxyServer extends AbstractApiService {
     IMMEDIATE,
   }
 
+  private final AtomicReference<ShutdownHandler> shutdownHandler = new AtomicReference<>();
+
   private final AtomicReference<ShutdownMode> shutdownMode = new AtomicReference<>();
 
   private final AtomicReference<CountDownLatch> allHandlersTerminatedLatch =
@@ -322,6 +324,13 @@ public class ProxyServer extends AbstractApiService {
     awaitTerminated();
   }
 
+  public synchronized ShutdownHandler getOrCreateShutdownHandler() {
+    if (this.shutdownHandler.get() == null) {
+      this.shutdownHandler.set(ShutdownHandler.createForServer(this));
+    }
+    return this.shutdownHandler.get();
+  }
+
   void setShutdownMode(ShutdownMode shutdownMode) {
     logger.log(Level.INFO, "Setting shutdown mode to {0}", shutdownMode);
     this.shutdownMode.set(shutdownMode);
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/Server.java b/src/main/java/com/google/cloud/spanner/pgadapter/Server.java
index 958a9a3b..361d6c39 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/Server.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/Server.java
@@ -62,7 +62,7 @@ import org.codehaus.mojo.animal_sniffer.IgnoreJRERequirement;
 public class Server {
   private static final Logger logger = Logger.getLogger(Server.class.getName());
 
-  private static ShutdownHandler shutdownHandler;
+  private static volatile ShutdownHandler shutdownHandler;
 
   /**
    * Main method for running a Spanner PostgreSQL Adapter {@link Server} as a stand-alone
@@ -78,7 +78,7 @@ public class Server {
 
       // Create a shutdown handler and register signal handlers for the signals that should
       // terminate the server.
-      Server.shutdownHandler = ShutdownHandler.createForServer(proxyServer);
+      Server.shutdownHandler = proxyServer.getOrCreateShutdownHandler();
       registerSignalHandlers();
     } catch (Exception e) {
       printError(e, System.err, System.out);
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/ShutdownHandler.java b/src/main/java/com/google/cloud/spanner/pgadapter/ShutdownHandler.java
index 4a991404..80037820 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ShutdownHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ShutdownHandler.java
@@ -14,6 +14,7 @@
 
 package com.google.cloud.spanner.pgadapter;
 
+import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.pgadapter.ProxyServer.ShutdownMode;
 import com.google.common.base.Preconditions;
 import java.util.concurrent.atomic.AtomicReference;
@@ -28,6 +29,7 @@ public class ShutdownHandler {
   private final Thread shutdownThread;
   private final AtomicReference<ShutdownMode> shutdownMode = new AtomicReference<>();
 
+  @InternalApi
   static ShutdownHandler createForServer(@Nonnull ProxyServer proxyServer) {
     return new ShutdownHandler(proxyServer);
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
index 1942510e..1249305d 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
@@ -100,6 +100,7 @@ public class OptionsMetadata {
     private String unixDomainSocketDirectory;
     private boolean autoConfigEmulator;
     private boolean logGrpcMessages;
+    private boolean allowShutdownStatement;
     private boolean debugMode;
     private String endpoint;
     private boolean usePlainText;
@@ -369,6 +370,15 @@ public class OptionsMetadata {
       return this;
     }
 
+    /**
+     * Enables the use of the SHUTDOWN [SMART | FAST | IMMEDIATE] statement to shutdown the proxy
+     * server.
+     */
+    public Builder setAllowShutdownStatement(boolean allowShutdownStatement) {
+      this.allowShutdownStatement = allowShutdownStatement;
+      return this;
+    }
+
     Builder enableDebugMode() {
       this.debugMode = true;
       return this;
@@ -486,6 +496,9 @@ public class OptionsMetadata {
       if (logGrpcMessages) {
         addOption(args, OPTION_LOG_GRPC_MESSAGES);
       }
+      if (allowShutdownStatement) {
+        addOption(args, OPTION_ALLOW_SHUTDOWN_STATEMENT);
+      }
       if (debugMode) {
         addOption(args, OPTION_INTERNAL_DEBUG_MODE);
         addOption(args, OPTION_SKIP_INTERNAL_DEBUG_MODE_WARNING);
@@ -604,6 +617,7 @@ public class OptionsMetadata {
   private static final String OPTION_DEBUG_MODE = "debug";
   private static final String OPTION_LEGACY_LOGGING = "legacy_logging";
   private static final String OPTION_LOG_GRPC_MESSAGES = "log_grpc_messages";
+  private static final String OPTION_ALLOW_SHUTDOWN_STATEMENT = "allow_shutdown_statement";
 
   private final Map<String, String> environment;
   private final String osName;
@@ -635,6 +649,7 @@ public class OptionsMetadata {
   private final boolean debugMode;
   private final Duration startupTimeout;
   private final boolean logGrpcMessages;
+  private final boolean allowShutdownStatement;
 
   /**
    * Creates a new instance of {@link OptionsMetadata} from the given arguments.
@@ -728,6 +743,7 @@ public class OptionsMetadata {
     this.serverVersion = commandLine.getOptionValue(OPTION_SERVER_VERSION, DEFAULT_SERVER_VERSION);
     this.debugMode = commandLine.hasOption(OPTION_INTERNAL_DEBUG_MODE);
     this.logGrpcMessages = commandLine.hasOption(OPTION_LOG_GRPC_MESSAGES);
+    this.allowShutdownStatement = commandLine.hasOption(OPTION_ALLOW_SHUTDOWN_STATEMENT);
     this.startupTimeout = startupTimeout;
   }
 
@@ -802,6 +818,7 @@ public class OptionsMetadata {
     this.serverVersion = DEFAULT_SERVER_VERSION;
     this.debugMode = false;
     this.logGrpcMessages = false;
+    this.allowShutdownStatement = false;
     this.startupTimeout = DEFAULT_STARTUP_TIMEOUT;
   }
 
@@ -1292,6 +1309,11 @@ public class OptionsMetadata {
         "Logs all gRPC messages that are sent by PGAdapter to Spanner to the default log handler (stdout).\n"
             + "This option should only be enabled to debug problems and/or to determine exactly which gRPC messages\n"
             + "are being sent to Spanner. Enabling this in production will cause a large number of messages to be logged.");
+    options.addOption(
+        OPTION_ALLOW_SHUTDOWN_STATEMENT,
+        "allow-shutdown-statement",
+        false,
+        "Allows the proxy server to be shutdown by executing the SHUTDOWN [SMART | FAST | IMMEDIATE] statement.");
     options.addOption(
         OPTION_INTERNAL_DEBUG_MODE,
         "internal-debug-mode",
@@ -1392,6 +1414,10 @@ public class OptionsMetadata {
     return this.logGrpcMessages;
   }
 
+  public boolean isAllowShutdownStatement() {
+    return this.allowShutdownStatement;
+  }
+
   public boolean isDebugMode() {
     return this.debugMode;
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java
index 3c74fce1..1cef475e 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java
@@ -66,9 +66,7 @@ public class ExtendedQueryProtocolHandler {
                 .getServer()
                 .getTracer(ConnectionHandler.class.getName(), getVersion()),
             connectionHandler.getServer().getMetrics(),
-            createMetricAttributes(
-                connectionHandler.getDatabaseId(),
-                connectionHandler.getTraceConnectionId().toString()),
+            createMetricAttributes(connectionHandler.getDatabaseId()),
             connectionHandler.getTraceConnectionId().toString(),
             connectionHandler::closeAllPortals,
             connectionHandler.getDatabaseId(),
@@ -97,9 +95,8 @@ public class ExtendedQueryProtocolHandler {
   }
 
   @VisibleForTesting
-  static Attributes createMetricAttributes(DatabaseId databaseId, String connectionId) {
+  static Attributes createMetricAttributes(DatabaseId databaseId) {
     AttributesBuilder attributesBuilder = Attributes.builder();
-    attributesBuilder.put("pgadapter.connection_id", connectionId);
     attributesBuilder.put("database", databaseId.getDatabase());
     attributesBuilder.put("instance_id", databaseId.getInstanceId().getInstance());
     attributesBuilder.put("project_id", databaseId.getInstanceId().getProject());
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/ShutdownStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/ShutdownStatement.java
new file mode 100644
index 00000000..329b3e25
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/ShutdownStatement.java
@@ -0,0 +1,156 @@
+// Copyright 2024 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.statements;
+
+import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
+import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
+import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler;
+import com.google.cloud.spanner.pgadapter.ProxyServer;
+import com.google.cloud.spanner.pgadapter.ProxyServer.ShutdownMode;
+import com.google.cloud.spanner.pgadapter.ShutdownHandler;
+import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
+import com.google.cloud.spanner.pgadapter.error.SQLState;
+import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection.NoResult;
+import com.google.cloud.spanner.pgadapter.statements.ShutdownStatement.ParsedShutdownStatement.Builder;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.Futures;
+import java.util.List;
+import java.util.concurrent.Future;
+
+/**
+ * <code>SHUTDOWN [FAST | SMART | IMMEDIATE]</code>
+ *
+ * <p>Shuts down the proxy server. This statement is a PGAdapter-only statement, and is not
+ * supported on normal PostgreSQL.
+ */
+public class ShutdownStatement extends IntermediatePortalStatement {
+  static final class ParsedShutdownStatement {
+    static final class Builder {
+      ShutdownMode shutdownMode = ShutdownMode.FAST;
+    }
+
+    /** The server that will be shut down. */
+    final ProxyServer server;
+
+    final ShutdownMode shutdownMode;
+
+    private ParsedShutdownStatement(Builder builder, ProxyServer server) {
+      this.server = server;
+      this.shutdownMode = builder.shutdownMode;
+    }
+  }
+
+  private final ParsedShutdownStatement parsedShutdownStatement;
+
+  private ShutdownHandler shutdownHandler;
+
+  public ShutdownStatement(
+      ConnectionHandler connectionHandler,
+      OptionsMetadata options,
+      String name,
+      ParsedStatement parsedStatement,
+      Statement originalStatement) {
+    super(
+        name,
+        new IntermediatePreparedStatement(
+            connectionHandler,
+            options,
+            name,
+            NO_PARAMETER_TYPES,
+            parsedStatement,
+            originalStatement),
+        NO_PARAMS,
+        ImmutableList.of(),
+        ImmutableList.of());
+    this.parsedShutdownStatement = parse(originalStatement.getSql(), connectionHandler.getServer());
+  }
+
+  @Override
+  public String getCommandTag() {
+    return "SHUTDOWN";
+  }
+
+  @Override
+  public StatementType getStatementType() {
+    return StatementType.UNKNOWN;
+  }
+
+  @Override
+  public void close() throws Exception {
+    super.close();
+    if (this.executed && this.shutdownHandler != null) {
+      this.shutdownHandler.shutdown(this.parsedShutdownStatement.shutdownMode);
+    }
+  }
+
+  @Override
+  public void executeAsync(BackendConnection backendConnection) {
+    if (!this.executed) {
+      this.executed = true;
+      if (options.isAllowShutdownStatement()) {
+        this.shutdownHandler = this.parsedShutdownStatement.server.getOrCreateShutdownHandler();
+        setFutureStatementResult(Futures.immediateFuture(new NoResult(getCommandTag())));
+      } else {
+        setFutureStatementResult(
+            Futures.immediateFailedFuture(
+                PGExceptionFactory.newPGException(
+                    "SHUTDOWN [SMART | FAST | IMMEDIATE] statement is not enabled for this server. Start PGAdapter with --allow_shutdown_statement to enable the use of the SHUTDOWN statement.",
+                    SQLState.FeatureNotSupported)));
+      }
+    }
+  }
+
+  @Override
+  public Future<StatementResult> describeAsync(BackendConnection backendConnection) {
+    // Return null to indicate that this SHUTDOWN statement does not return any
+    // RowDescriptionResponse.
+    return Futures.immediateFuture(null);
+  }
+
+  @Override
+  public IntermediatePortalStatement createPortal(
+      String name,
+      byte[][] parameters,
+      List<Short> parameterFormatCodes,
+      List<Short> resultFormatCodes) {
+    // SHUTDOWN does not support binding any parameters, so we just return the same statement.
+    return this;
+  }
+
+  static ParsedShutdownStatement parse(String sql, ProxyServer server) {
+    Preconditions.checkNotNull(sql);
+    Preconditions.checkNotNull(server);
+
+    SimpleParser parser = new SimpleParser(sql);
+    ParsedShutdownStatement.Builder builder = new Builder();
+    if (!parser.eatKeyword("shutdown")) {
+      throw PGExceptionFactory.newPGException(
+          "not a valid SHUTDOWN statement: " + sql, SQLState.SyntaxError);
+    }
+    if (parser.eatKeyword("fast")) {
+      builder.shutdownMode = ShutdownMode.FAST;
+    } else if (parser.eatKeyword("smart")) {
+      builder.shutdownMode = ShutdownMode.SMART;
+    } else if (parser.eatKeyword("immediate")) {
+      builder.shutdownMode = ShutdownMode.IMMEDIATE;
+    }
+    parser.throwIfHasMoreTokens();
+    return new ParsedShutdownStatement(builder, server);
+  }
+}
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
index b01fc53f..fcdc30a9 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
@@ -30,6 +30,7 @@ import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.SAVEP
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.SELECT_CURRENT_SETTING;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.SELECT_SET_CONFIG;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.SHOW_DATABASE_DDL;
+import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.SHUTDOWN;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.TRUNCATE;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.VACUUM;
 
@@ -58,6 +59,7 @@ import com.google.cloud.spanner.pgadapter.statements.SavepointStatement;
 import com.google.cloud.spanner.pgadapter.statements.SelectCurrentSettingStatement;
 import com.google.cloud.spanner.pgadapter.statements.SelectSetConfigStatement;
 import com.google.cloud.spanner.pgadapter.statements.ShowDatabaseDdlStatement;
+import com.google.cloud.spanner.pgadapter.statements.ShutdownStatement;
 import com.google.cloud.spanner.pgadapter.statements.TruncateStatement;
 import com.google.cloud.spanner.pgadapter.statements.VacuumStatement;
 import com.google.cloud.spanner.pgadapter.wireoutput.ParseCompleteResponse;
@@ -244,6 +246,13 @@ public class ParseMessage extends AbstractQueryProtocolMessage {
             name,
             parsedStatement,
             originalStatement);
+      } else if (isCommand(SHUTDOWN, originalStatement.getSql())) {
+        return new ShutdownStatement(
+            connectionHandler,
+            connectionHandler.getServer().getOptions(),
+            name,
+            parsedStatement,
+            originalStatement);
       } else {
         return new IntermediatePreparedStatement(
             connectionHandler,
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java
index 4f417793..02bd901c 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java
@@ -41,6 +41,7 @@ public class QueryMessage extends ControlMessage {
   public static final String FETCH = "FETCH";
   public static final String MOVE = "MOVE";
   public static final String CLOSE = "CLOSE";
+  public static final String SHUTDOWN = "SHUTDOWN";
   public static final ImmutableList<String> SHOW_DATABASE_DDL =
       ImmutableList.of("SHOW", "DATABASE", "DDL");
   public static final ImmutableList<String> SELECT_CURRENT_SETTING =


Test patch:
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 61e783fc..f4f8eb86 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -5596,6 +5596,19 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
         exception.getMessage().contains("The server does not support GSS Encryption"));
   }
 
+  @Test
+  public void testShutdown_failsByDefault() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(createUrl());
+        java.sql.Statement statement = connection.createStatement()) {
+      PSQLException exception =
+          assertThrows(PSQLException.class, () -> statement.execute("shutdown"));
+      assertEquals(
+          "ERROR: SHUTDOWN [SMART | FAST | IMMEDIATE] statement is not enabled for this server. "
+              + "Start PGAdapter with --allow_shutdown_statement to enable the use of the SHUTDOWN statement.",
+          exception.getMessage());
+    }
+  }
+
   private void verifySelect1(java.sql.Statement statement) throws SQLException {
     try (ResultSet resultSet = statement.executeQuery("SELECT 1")) {
       assertTrue(resultSet.next());
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ShutdownModeMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ShutdownModeMockServerTest.java
index e7d3d61b..ff3e5ca8 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ShutdownModeMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ShutdownModeMockServerTest.java
@@ -32,15 +32,26 @@ import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
 import java.util.concurrent.TimeUnit;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
 
-@RunWith(JUnit4.class)
+@RunWith(Parameterized.class)
 public class ShutdownModeMockServerTest extends AbstractMockServerTest {
+
+  @Parameter public boolean useSqlStatement;
+
+  @Parameters(name = "useSqlStatement = {0}")
+  public static Object[] data() {
+    return new Object[] {false, true};
+  }
+
   private ProxyServer proxyServer;
 
   @Before
@@ -51,7 +62,8 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
         .enableDebugMode()
         .setUsePlainText()
         .setEndpoint(String.format("localhost:%d", spannerServer.getPort()))
-        .setCredentials(NoCredentials.getInstance());
+        .setCredentials(NoCredentials.getInstance())
+        .setAllowShutdownStatement(true);
     proxyServer = new ProxyServer(builder.build(), OpenTelemetry.noop());
     proxyServer.startServer();
   }
@@ -71,7 +83,7 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
   public void testSmartShutdown() throws SQLException {
     String sql = "SELECT 1";
 
-    ShutdownHandler shutdownHandler = ShutdownHandler.createForServer(proxyServer);
+    ShutdownHandler shutdownHandler = createShutdownHandler();
     // Verify that the server is running.
     assertEquals(State.RUNNING, proxyServer.state());
 
@@ -83,7 +95,11 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
         assertFalse(resultSet.next());
       }
       // Initiate a smart shutdown.
-      shutdownHandler.shutdown(ShutdownMode.SMART);
+      if (useSqlStatement) {
+        connection.createStatement().execute("shutdown smart");
+      } else {
+        shutdownHandler.shutdown(ShutdownMode.SMART);
+      }
       // Wait until the server port is no longer open.
       Stopwatch stopwatch = Stopwatch.createStarted();
       while (isPortListening(proxyServer.getLocalPort())
@@ -127,7 +143,7 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
   public void testSmartShutdownWithoutConnections() throws SQLException {
     String sql = "SELECT 1";
 
-    ShutdownHandler shutdownHandler = ShutdownHandler.createForServer(proxyServer);
+    ShutdownHandler shutdownHandler = createShutdownHandler();
     // Verify that the server is running.
     assertEquals(State.RUNNING, proxyServer.state());
 
@@ -140,7 +156,14 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
       }
     }
     // Initiate a smart shutdown.
-    shutdownHandler.shutdown(ShutdownMode.SMART);
+    if (useSqlStatement) {
+      try (Connection connection = DriverManager.getConnection(createUrl());
+          Statement statement = connection.createStatement()) {
+        statement.execute("shutdown smart");
+      }
+    } else {
+      shutdownHandler.shutdown(ShutdownMode.SMART);
+    }
     // Wait until the server port is no longer open.
     Stopwatch stopwatch = Stopwatch.createStarted();
     while (isPortListening(proxyServer.getLocalPort())
@@ -169,7 +192,7 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
   public void testFastShutdown() throws SQLException {
     String sql = "SELECT 1";
 
-    ShutdownHandler shutdownHandler = ShutdownHandler.createForServer(proxyServer);
+    ShutdownHandler shutdownHandler = createShutdownHandler();
     // Verify that the server is running.
     assertEquals(State.RUNNING, proxyServer.state());
 
@@ -181,7 +204,11 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
         assertFalse(resultSet.next());
       }
       // Initiate a fast shutdown.
-      shutdownHandler.shutdown(ShutdownMode.FAST);
+      if (useSqlStatement) {
+        connection.createStatement().execute("shutdown fast");
+      } else {
+        shutdownHandler.shutdown(ShutdownMode.FAST);
+      }
       // Verify that the server stops.
       Stopwatch stopwatch = Stopwatch.createStarted();
       while (proxyServer.state() != State.TERMINATED
@@ -209,7 +236,7 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
   public void testSmartShutdownFollowedByFastShutdown() throws SQLException {
     String sql = "SELECT 1";
 
-    ShutdownHandler shutdownHandler = ShutdownHandler.createForServer(proxyServer);
+    ShutdownHandler shutdownHandler = createShutdownHandler();
     // Verify that the server is running.
     assertEquals(State.RUNNING, proxyServer.state());
 
@@ -221,7 +248,11 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
         assertFalse(resultSet.next());
       }
       // Initiate a smart shutdown.
-      shutdownHandler.shutdown(ShutdownMode.SMART);
+      if (useSqlStatement) {
+        connection.createStatement().execute("shutdown smart");
+      } else {
+        shutdownHandler.shutdown(ShutdownMode.SMART);
+      }
       // Wait until the server port is no longer open.
       Stopwatch stopwatch = Stopwatch.createStarted();
       while (isPortListening(proxyServer.getLocalPort())
@@ -248,7 +279,11 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
       assertEquals(State.STOPPING, proxyServer.state());
 
       // Now force the server to stop by initiating a fast shutdown.
-      shutdownHandler.shutdown(ShutdownMode.FAST);
+      if (useSqlStatement) {
+        connection.createStatement().execute("shutdown fast");
+      } else {
+        shutdownHandler.shutdown(ShutdownMode.FAST);
+      }
       stopwatch = Stopwatch.createStarted();
       while (proxyServer.state() != State.TERMINATED
           && stopwatch.elapsed(TimeUnit.MILLISECONDS) < 2000) {
@@ -263,6 +298,10 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
     }
   }
 
+  ShutdownHandler createShutdownHandler() {
+    return useSqlStatement ? null : ShutdownHandler.createForServer(this.proxyServer);
+  }
+
   boolean isPortListening(int port) {
     try (Socket ignored = new Socket("localhost", port)) {
       return true;
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
index 58a06cd4..7edfa78a 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
@@ -87,7 +87,7 @@ public class BackendConnectionTest {
   private static final Runnable DO_NOTHING = () -> {};
   private static final DatabaseId DATABASE_ID = DatabaseId.of("p", "i", "d");
   private static final Attributes METRIC_ATTRIBUTES =
-      ExtendedQueryProtocolHandler.createMetricAttributes(DATABASE_ID, "test-id");
+      ExtendedQueryProtocolHandler.createMetricAttributes(DATABASE_ID);
 
   @Test
   public void testExtractDdlUpdateCounts() {
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/ShutdownStatementTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/ShutdownStatementTest.java
new file mode 100644
index 00000000..f3369114
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/ShutdownStatementTest.java
@@ -0,0 +1,163 @@
+// 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.statements;
+
+import static com.google.cloud.spanner.pgadapter.statements.ShutdownStatement.parse;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.connection.AbstractStatementParser;
+import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler;
+import com.google.cloud.spanner.pgadapter.ProxyServer;
+import com.google.cloud.spanner.pgadapter.ProxyServer.ShutdownMode;
+import com.google.cloud.spanner.pgadapter.ShutdownHandler;
+import com.google.cloud.spanner.pgadapter.error.PGException;
+import com.google.cloud.spanner.pgadapter.error.SQLState;
+import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
+import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement.ResultNotReadyBehavior;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ShutdownStatementTest {
+
+  @Test
+  public void testParseShutdown() {
+    ProxyServer server = mock(ProxyServer.class);
+
+    assertEquals(ShutdownMode.FAST, parse("shutdown", server).shutdownMode);
+    assertEquals(ShutdownMode.FAST, parse("Shutdown", server).shutdownMode);
+    assertEquals(ShutdownMode.FAST, parse("SHUTDOWN", server).shutdownMode);
+    assertEquals(ShutdownMode.FAST, parse("\t shutdown \n", server).shutdownMode);
+    assertEquals(ShutdownMode.IMMEDIATE, parse("shutdown immediate", server).shutdownMode);
+    assertEquals(ShutdownMode.SMART, parse("shutdown smart", server).shutdownMode);
+    assertEquals(ShutdownMode.FAST, parse("shutdown fast", server).shutdownMode);
+    assertEquals(ShutdownMode.FAST, parse("shutdown\tfast", server).shutdownMode);
+    assertEquals(ShutdownMode.FAST, parse("shutdown\nfast", server).shutdownMode);
+    assertEquals(ShutdownMode.FAST, parse("shutdown\n\t   fast", server).shutdownMode);
+
+    PGException exception;
+    exception = assertThrows(PGException.class, () -> parse("shutdown fast immediate", server));
+    assertEquals(SQLState.SyntaxError, exception.getSQLState());
+    exception = assertThrows(PGException.class, () -> parse("shutdown foo", server));
+    assertEquals(SQLState.SyntaxError, exception.getSQLState());
+    exception = assertThrows(PGException.class, () -> parse("shutdown foo fast", server));
+    assertEquals(SQLState.SyntaxError, exception.getSQLState());
+    exception = assertThrows(PGException.class, () -> parse("shutdown fast,", server));
+    assertEquals(SQLState.SyntaxError, exception.getSQLState());
+    exception = assertThrows(PGException.class, () -> parse("fast shutdown", server));
+    assertEquals(SQLState.SyntaxError, exception.getSQLState());
+  }
+
+  @Test
+  public void testShutdownStatement_usesDefaultShutdownMode() throws Exception {
+    OptionsMetadata options = mock(OptionsMetadata.class);
+    when(options.isAllowShutdownStatement()).thenReturn(true);
+    ConnectionHandler connectionHandler = mock(ConnectionHandler.class);
+    ProxyServer server = mock(ProxyServer.class);
+    when(connectionHandler.getServer()).thenReturn(server);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(mock(ConnectionMetadata.class));
+    ShutdownHandler shutdownHandler = mock(ShutdownHandler.class);
+    when(server.getOrCreateShutdownHandler()).thenReturn(shutdownHandler);
+
+    String sql = "shutdown";
+    Statement originalStatement = Statement.of(sql);
+    ParsedStatement parsedStatement =
+        AbstractStatementParser.getInstance(Dialect.POSTGRESQL).parse(originalStatement);
+    ShutdownStatement statement =
+        new ShutdownStatement(
+            connectionHandler, options, /* name= */ "", parsedStatement, originalStatement);
+
+    statement.executeAsync(mock(BackendConnection.class));
+    statement.close();
+
+    // PGAdapter uses FAST shutdown mode by default.
+    // We need to use a timeout, as the shutdown handler runs async. The wait time is normally much
+    // less than the 1000 milliseconds that is used as the timeout value.
+    verify(shutdownHandler, timeout(1000L)).shutdown(ShutdownMode.FAST);
+  }
+
+  @Test
+  public void testShutdownStatement_usesParsedShutdownMode() throws Exception {
+    OptionsMetadata options = mock(OptionsMetadata.class);
+    when(options.isAllowShutdownStatement()).thenReturn(true);
+    ConnectionHandler connectionHandler = mock(ConnectionHandler.class);
+    ProxyServer server = mock(ProxyServer.class);
+    when(connectionHandler.getServer()).thenReturn(server);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(mock(ConnectionMetadata.class));
+    ShutdownHandler shutdownHandler = mock(ShutdownHandler.class);
+    when(server.getOrCreateShutdownHandler()).thenReturn(shutdownHandler);
+
+    String sql = "shutdown smart";
+    Statement originalStatement = Statement.of(sql);
+    ParsedStatement parsedStatement =
+        AbstractStatementParser.getInstance(Dialect.POSTGRESQL).parse(originalStatement);
+    ShutdownStatement statement =
+        new ShutdownStatement(
+            connectionHandler, options, /* name= */ "", parsedStatement, originalStatement);
+
+    statement.executeAsync(mock(BackendConnection.class));
+    statement.close();
+
+    // The shutdown mode from the SQL statement should be used.
+    // We need to use a timeout, as the shutdown handler runs async. The wait time is normally much
+    // less than the 1000 milliseconds that is used as the timeout value.
+    verify(shutdownHandler, timeout(1000L)).shutdown(ShutdownMode.SMART);
+  }
+
+  @Test
+  public void testShutdownStatement_returnsErrorIfNotEnabled() throws Exception {
+    OptionsMetadata options = mock(OptionsMetadata.class);
+    when(options.isAllowShutdownStatement()).thenReturn(false);
+    ConnectionHandler connectionHandler = mock(ConnectionHandler.class);
+    ProxyServer server = mock(ProxyServer.class);
+    when(connectionHandler.getServer()).thenReturn(server);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(mock(ConnectionMetadata.class));
+
+    String sql = "shutdown";
+    Statement originalStatement = Statement.of(sql);
+    ParsedStatement parsedStatement =
+        AbstractStatementParser.getInstance(Dialect.POSTGRESQL).parse(originalStatement);
+    ShutdownStatement statement =
+        new ShutdownStatement(
+            connectionHandler, options, /* name= */ "", parsedStatement, originalStatement);
+
+    // Execute the statement and then get the result.
+    statement.executeAsync(mock(BackendConnection.class));
+    statement.close();
+    // The result should be an exception.
+    statement.initFutureResult(ResultNotReadyBehavior.FAIL);
+    assertNotNull(statement.getException());
+    assertEquals(
+        "SHUTDOWN [SMART | FAST | IMMEDIATE] statement is not enabled for this server. "
+            + "Start PGAdapter with --allow_shutdown_statement to enable the use of the SHUTDOWN statement.",
+        statement.getException().getMessage());
+
+    // The server should not be shut down.
+    verify(server, never()).stopServer(any(ShutdownMode.class));
+  }
+}
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 9cac0270..a5513ef5 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
@@ -97,7 +97,7 @@ public class StatementTest {
   private static final Runnable DO_NOTHING = () -> {};
   private static final DatabaseId DATABASE_ID = DatabaseId.of("p", "i", "d");
   private static final Attributes METRIC_ATTRIBUTES =
-      ExtendedQueryProtocolHandler.createMetricAttributes(DATABASE_ID, "test-connection");
+      ExtendedQueryProtocolHandler.createMetricAttributes(DATABASE_ID);
 
   private static ParsedStatement parse(String sql) {
     return PARSER.parse(Statement.of(sql));
@@ -521,7 +521,7 @@ public class StatementTest {
         new BackendConnection(
             NOOP_OTEL,
             NOOP_OTEL_METER,
-            ExtendedQueryProtocolHandler.createMetricAttributes(databaseId, "test-connection"),
+            ExtendedQueryProtocolHandler.createMetricAttributes(databaseId),
             UUID.randomUUID().toString(),
             DO_NOTHING,
             databaseId,
@@ -629,7 +629,7 @@ public class StatementTest {
         new BackendConnection(
             NOOP_OTEL,
             NOOP_OTEL_METER,
-            ExtendedQueryProtocolHandler.createMetricAttributes(databaseId, "test-connection"),
+            ExtendedQueryProtocolHandler.createMetricAttributes(databaseId),
             UUID.randomUUID().toString(),
             DO_NOTHING,
             databaseId,
@@ -688,7 +688,7 @@ public class StatementTest {
         new BackendConnection(
             NOOP_OTEL,
             NOOP_OTEL_METER,
-            ExtendedQueryProtocolHandler.createMetricAttributes(databaseId, "test-connection"),
+            ExtendedQueryProtocolHandler.createMetricAttributes(databaseId),
             UUID.randomUUID().toString(),
             DO_NOTHING,
             databaseId,