From 2bae001c3fe50a41eb34204a0b08534d98f9e4e7 Mon Sep 17 00:00:00 2001
From: steve-chavez <stevechavezast@gmail.com>
Date: Wed, 1 May 2024 21:10:13 -0500
Subject: [PATCH] wip: force read-write for listener connection

---
 src/PostgREST/AppState.hs |  4 +++-
 src/PostgREST/Config.hs   | 26 ++++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/PostgREST/AppState.hs b/src/PostgREST/AppState.hs
index 844a99cb029..d975d8a1fa2 100644
--- a/src/PostgREST/AppState.hs
+++ b/src/PostgREST/AppState.hs
@@ -62,6 +62,7 @@ import Data.Time.Clock    (UTCTime, getCurrentTime)
 
 import PostgREST.Config                  (AppConfig (..),
                                           addFallbackAppName,
+                                          addTargetSessionAttrs,
                                           readAppConfig)
 import PostgREST.Config.Database         (queryDbSettings,
                                           queryPgVersion,
@@ -499,7 +500,8 @@ listener appState@AppState{stateObserver=observer} conf@AppConfig{..} = do
 
   -- forkFinally allows to detect if the thread dies
   void . flip forkFinally (handleFinally dbChannel configDbPoolAutomaticRecovery) $ do
-    dbOrError <- acquire $ toUtf8 (addFallbackAppName prettyVersion configDbUri)
+    let connString = addTargetSessionAttrs $ addFallbackAppName prettyVersion configDbUri
+    dbOrError <- acquire $ toUtf8 connString
     case dbOrError of
       Right db -> do
         observer $ DBListenerStart dbChannel
diff --git a/src/PostgREST/Config.hs b/src/PostgREST/Config.hs
index e13cbf49976..9ef339d17be 100644
--- a/src/PostgREST/Config.hs
+++ b/src/PostgREST/Config.hs
@@ -25,6 +25,7 @@ module PostgREST.Config
   , toURI
   , parseSecret
   , addFallbackAppName
+  , addTargetSessionAttrs
   ) where
 
 import qualified Crypto.JOSE.Types      as JOSE
@@ -518,3 +519,28 @@ addFallbackAppName version dbUri = dbUri <>
     uriFmt = pKeyWord <> toS (escapeURIString isUnescapedInURIComponent $ toS pgrstVer)
     pKeyWord = "fallback_application_name="
     pgrstVer = "PostgREST " <> T.decodeUtf8 version
+
+
+-- | Adds `target_session_attrs=read-write` to the connection string. This allows using PostgREST listener when multiple hosts are specified in the connection string.
+--
+-- >>> addTargetSessionAttrs "postgres:///postgres?host=/dir/0kN/socket_replica_24378,/dir/0kN/socket"
+-- "postgres:///postgres?host=/dir/0kN/socket_replica_24378,/dir/0kN/socket&target_session_attrs=read-write"
+--
+-- >>> addTargetSessionAttrs "postgresql://host1:123,host2:456/somedb"
+-- "postgresql://host1:123,host2:456/somedb?target_session_attrs=read-write"
+--
+-- >>> addTargetSessionAttrs "postgresql://host1:123,host2:456/somedb?fallback_application_name=foo"
+-- "postgresql://host1:123,host2:456/somedb?fallback_application_name=foo&target_session_attrs=read-write"
+--
+-- >>> addTargetSessionAttrs "invalid_uri1=val1 invalid_uri2=val2"
+-- "invalid_uri1=val1 invalid_uri2=val2"
+addTargetSessionAttrs :: Text -> Text
+addTargetSessionAttrs dbUri = dbUri <>
+  case uriQuery <$> parseURI (toS dbUri) of
+    -- Don't do anything to key=val connection strings or invalid URIs
+    Nothing  -> mempty
+    Just ""  -> "?" <> part
+    Just "?" -> part
+    _        -> "&" <> part
+  where
+    part = "target_session_attrs=read-write"