From 2ca56d251e95802d17034eea495b78dc7b670992 Mon Sep 17 00:00:00 2001
From: Victor Hugo Borja <vborja@apache.org>
Date: Sat, 26 Jun 2021 03:14:14 -0500
Subject: [PATCH] Support sending SET commands.

Previous to this change, if you tried to issue a `SET` command,
finagle-postgres would raise an exception telling no completion command
is known for tag `SET`. This patch implements such a response from
server.

`SET` commands are needed for sending runtime settings to postgres
or when trying to set the current transaction isolation level.

See:

https://github.com/getquill/quill/issues/29#issuecomment-560162692
---
 .../finagle/postgres/connection/ConnectionStateMachine.scala | 2 ++
 .../finagle/postgres/messages/BackendMessageParser.scala     | 1 +
 .../twitter/finagle/postgres/messages/BackendMessages.scala  | 2 ++
 .../finagle/postgres/integration/IntegrationSpec.scala       | 5 +++++
 4 files changed, 10 insertions(+)

diff --git a/src/main/scala/com/twitter/finagle/postgres/connection/ConnectionStateMachine.scala b/src/main/scala/com/twitter/finagle/postgres/connection/ConnectionStateMachine.scala
index f797833e..2a837c87 100644
--- a/src/main/scala/com/twitter/finagle/postgres/connection/ConnectionStateMachine.scala
+++ b/src/main/scala/com/twitter/finagle/postgres/connection/ConnectionStateMachine.scala
@@ -99,6 +99,7 @@ class ConnectionStateMachine(state: State = AuthenticationRequired, val id: Int)
     case (CommandComplete(DiscardAll), SimpleQuery) => (None, EmitOnReadyForQuery(CommandCompleteResponse(1)))
     case (CommandComplete(Begin | Savepoint | Release | RollBack | Commit), SimpleQuery) =>
       (None, EmitOnReadyForQuery(CommandCompleteResponse(1)))
+    case (CommandComplete(Set), SimpleQuery) => (None, EmitOnReadyForQuery(CommandCompleteResponse(1)))
     case (CommandComplete(Do), SimpleQuery) => (None, EmitOnReadyForQuery(CommandCompleteResponse(1)))
 
     case (RowDescription(fields), SimpleQuery) =>
@@ -123,6 +124,7 @@ class ConnectionStateMachine(state: State = AuthenticationRequired, val id: Int)
     case (CommandComplete(Update(count)), ExecutePreparedStatement) => (Some(CommandCompleteResponse(count)), Connected)
     case (CommandComplete(Delete(count)), ExecutePreparedStatement) => (Some(CommandCompleteResponse(count)), Connected)
     case (CommandComplete(Begin), ExecutePreparedStatement) => (Some(CommandCompleteResponse(1)), Connected)
+    case (CommandComplete(Set), ExecutePreparedStatement) => (Some(CommandCompleteResponse(1)), Connected)
     case (CommandComplete(Savepoint), ExecutePreparedStatement) => (Some(CommandCompleteResponse(1)), Connected)
     case (CommandComplete(RollBack), ExecutePreparedStatement) => (Some(CommandCompleteResponse(1)), Connected)
     case (CommandComplete(Commit), ExecutePreparedStatement) => (Some(CommandCompleteResponse(1)), Connected)
diff --git a/src/main/scala/com/twitter/finagle/postgres/messages/BackendMessageParser.scala b/src/main/scala/com/twitter/finagle/postgres/messages/BackendMessageParser.scala
index 2625ec55..8ae05287 100644
--- a/src/main/scala/com/twitter/finagle/postgres/messages/BackendMessageParser.scala
+++ b/src/main/scala/com/twitter/finagle/postgres/messages/BackendMessageParser.scala
@@ -238,6 +238,7 @@ class BackendMessageParser {
           case "DELETE" => Delete(parts(1).toInt)
           case "UPDATE" => Update(parts(1).toInt)
           case "BEGIN"  => Begin
+          case "SET"    => Set
           case "SAVEPOINT" => Savepoint
           case "RELEASE" => Release
           case "ROLLBACK"  => RollBack
diff --git a/src/main/scala/com/twitter/finagle/postgres/messages/BackendMessages.scala b/src/main/scala/com/twitter/finagle/postgres/messages/BackendMessages.scala
index 06d67fe2..a453b1cf 100644
--- a/src/main/scala/com/twitter/finagle/postgres/messages/BackendMessages.scala
+++ b/src/main/scala/com/twitter/finagle/postgres/messages/BackendMessages.scala
@@ -75,6 +75,8 @@ case class Select(count: Int) extends CommandCompleteStatus
 
 case object Begin extends CommandCompleteStatus
 
+case object Set extends CommandCompleteStatus
+
 case object Savepoint extends CommandCompleteStatus
 
 case object Release extends CommandCompleteStatus
diff --git a/src/test/scala/com/twitter/finagle/postgres/integration/IntegrationSpec.scala b/src/test/scala/com/twitter/finagle/postgres/integration/IntegrationSpec.scala
index 6f68f816..96361501 100644
--- a/src/test/scala/com/twitter/finagle/postgres/integration/IntegrationSpec.scala
+++ b/src/test/scala/com/twitter/finagle/postgres/integration/IntegrationSpec.scala
@@ -380,6 +380,11 @@ class IntegrationSpec extends Spec {
         Await.result(result)
       }
 
+      "support sending SET commands" in {
+        val client = getClient
+        val result = client.query("SET TIME ZONE DEFAULT;")
+        Await.result(result)
+      }
 
       "throw a ServerError" when {
         "query has error" in {