fix port to host assignments in connection strings

src/library-tests/Main.hs

@@ -442,16 +442,114 @@ main = hspec do
                 Left err -> expectationFailure ("URL roundtrip parse error: " <> Text.unpack err)
                 Right cs2 -> cs2 `shouldBe` cs
 
+        it "postgresql://a,b,c (default for all hosts)" do
+          let input = "postgresql://a,b,c"
+          case ConnectionString.parse input of
+            Left err -> expectationFailure (Text.unpack err)
+            Right cs ->
+              ConnectionString.toHosts cs
+                `shouldBe` [("a", Nothing), ("b", Nothing), ("c", Nothing)]
+
         it "host=host1,host2,host3 port=1,2,3" do
           let input = "host=host1,host2,host3 port=1,2,3"
           case ConnectionString.parse input of
             Left err -> expectationFailure ("Parse error: " <> Text.unpack err)
-            Right cs -> do
-              -- In keyword/value format, multiple hosts are separated by commas in the value
-              -- This is a special case that may not be supported yet
-              -- For now, just verify it parses
-              let hosts = ConnectionString.toHosts cs
-              length hosts `shouldSatisfy` (> 0)
+            Right cs ->
+              ConnectionString.toHosts cs
+                `shouldBe` [("host1", Just 1), ("host2", Just 2), ("host3", Just 3)]
+
+        it "host=a,b,c port=1,2,3 (don't depend on specific hostname pattern)" do
+          let input = "host=a,b,c port=1,2,3"
+          case ConnectionString.parse input of
+            Left err -> expectationFailure ("Parse error: " <> Text.unpack err)
+            Right cs ->
+              ConnectionString.toHosts cs
+                `shouldBe` [("a", Just 1), ("b", Just 2), ("c", Just 3)]
+
+        it "host=host1,host2 port=1,2" do
+          let input = "host=host1,host2 port=1,2"
+          case ConnectionString.parse input of
+            Left err -> expectationFailure ("Parse error: " <> Text.unpack err)
+            Right cs ->
+              ConnectionString.toHosts cs `shouldBe` [("host1", Just 1), ("host2", Just 2)]
+
+        it "host=host1,host2 port=1 (single port applies to all hosts)" do
+          let input = "host=host1,host2 port=1"
+          case ConnectionString.parse input of
+            Left err -> expectationFailure ("Parse error: " <> Text.unpack err)
+            Right cs ->
+              ConnectionString.toHosts cs `shouldBe` [("host1", Just 1), ("host2", Just 1)]
+
+        it "host=host1,host2,host3 port=1 (single port applies to all hosts)" do
+          let input = "host=host1,host2,host3 port=1"
+          case ConnectionString.parse input of
+            Left err -> expectationFailure ("Parse error: " <> Text.unpack err)
+            Right cs ->
+              ConnectionString.toHosts cs `shouldBe` [("host1", Just 1), ("host2", Just 1), ("host3", Just 1)]
+
+        it "host=a,b,c port=1 (don't depend on specific hostname pattern)" do
+          let input = "host=a,b,c port=1"
+          case ConnectionString.parse input of
+            Left err -> expectationFailure ("Parse error: " <> Text.unpack err)
+            Right cs ->
+              ConnectionString.toHosts cs `shouldBe` [("a", Just 1), ("b", Just 1), ("c", Just 1)]
+
+        it "rejects host=host1,host2,host3 port=1,2 (port count mismatch)" do
+          let input = "host=host1,host2,host3 port=1,2"
+          case ConnectionString.parse input of
+            Left err -> Text.unpack err `shouldContain` "could not match 2 port numbers to 3 hosts"
+            Right _ -> expectationFailure "Expected port count mismatch to be rejected"
+
+        it "rejects host=host1,host2 port=1,2,3 (port count mismatch)" do
+          let input = "host=host1,host2 port=1,2,3"
+          case ConnectionString.parse input of
+            Left err -> Text.unpack err `shouldContain` "could not match 3 port numbers to 2 hosts"
+            Right _ -> expectationFailure "Expected port count mismatch to be rejected"
+
+        it "rejects port=1,2 without host (port count mismatch)" do
+          let input = "port=1,2"
+          case ConnectionString.parse input of
+            Left err -> Text.unpack err `shouldContain` "could not match 2 port numbers to 1 hosts"
+            Right _ -> expectationFailure "Expected port count mismatch to be rejected"
+
+        it "host=a,b,c (default for all hosts)" do
+          let input = "host=a,b,c"
+          case ConnectionString.parse input of
+            Left err -> expectationFailure (Text.unpack err)
+            Right cs ->
+              ConnectionString.toHosts cs
+                `shouldBe` [("a", Nothing), ("b", Nothing), ("c", Nothing)]
+
+        it "host=a,b port=,2 (leading empty item selects default)" do
+          let input = "host=a,b port=,2"
+          case ConnectionString.parse input of
+            Left err -> expectationFailure (Text.unpack err)
+            Right cs ->
+              ConnectionString.toHosts cs `shouldBe` [("a", Nothing), ("b", Just 2)]
+
+        it "host=a,b,c port=1,,3 (middle empty item selects default)" do
+          let input = "host=a,b,c port=1,,3"
+          case ConnectionString.parse input of
+            Left err -> expectationFailure (Text.unpack err)
+            Right cs ->
+              ConnectionString.toHosts cs
+                `shouldBe` [("a", Just 1), ("b", Nothing), ("c", Just 3)]
+
+        it "host=a,b,c port=1,, (trailing empty items select default)" do
+          let input = "host=a,b,c port=1,,"
+          case ConnectionString.parse input of
+            Left err -> expectationFailure (Text.unpack err)
+            Right cs ->
+              ConnectionString.toHosts cs
+                `shouldBe` [("a", Just 1), ("b", Nothing), ("c", Nothing)]
+
+        it "host=a,b port='' (empty string selects default)" do
+          let input = "host=a,b port=''"
+          case ConnectionString.parse input of
+            Left err -> expectationFailure (Text.unpack err)
+            Right cs ->
+              ConnectionString.toHosts cs
+                `shouldBe` [("a", Nothing), ("b", Nothing)]
 
       describe "equivalence tests of the internal representation" do
         it "postgresql://host1:1,host2:2,host3:3/ is equivalent to host=host1,host2,host3 port=1,2,3" do
@@ -460,19 +558,22 @@ main = hspec do
           case (ConnectionString.parse url, ConnectionString.parse kv) of
             (Right cs1, Right cs2) -> do
               -- They should represent the same connection
-              -- At minimum, they should have the same number of hosts
-              length (ConnectionString.toHosts cs1) `shouldBe` length (ConnectionString.toHosts cs2)
+              let hosts1 = ConnectionString.toHosts cs1
+                  hosts2 = ConnectionString.toHosts cs2
+              hosts1 `shouldBe` [("host1", Just 1), ("host2", Just 2), ("host3", Just 3)]
+              hosts1 `shouldBe` hosts2
             (Left err, _) -> expectationFailure ("URL parse error: " <> Text.unpack err)
             (_, Left err) -> expectationFailure ("KV parse error: " <> Text.unpack err)
 
-        it "postgresql://host1:1,host2:2,host3/ is equivalent to host=host1,host2,host3 port=1,2" do
+        it "postgresql://host1:1,host2:2,host3/ is equivalent to host=host1,host2,host3 port=1,2," do
           let url = "postgresql://host1:1,host2:2,host3/"
-              kv = "host=host1,host2,host3 port=1,2"
+              kv = "host=host1,host2,host3 port=1,2,"
           case (ConnectionString.parse url, ConnectionString.parse kv) of
             (Right cs1, Right cs2) -> do
               -- They should represent the same connection
               let hosts1 = ConnectionString.toHosts cs1
                   hosts2 = ConnectionString.toHosts cs2
+              hosts1 `shouldBe` [("host1", Just 1), ("host2", Just 2), ("host3", Nothing)]
               hosts1 `shouldBe` hosts2
             (Left err, _) -> expectationFailure ("URL parse error: " <> Text.unpack err)
             (_, Left err) -> expectationFailure ("KV parse error: " <> Text.unpack err)

src/library/PostgresqlConnectionString/Parsers.hs

@@ -3,7 +3,6 @@
 module PostgresqlConnectionString.Parsers where
 
 import qualified Data.CharSet as CharSet
-import qualified Data.List as List
 import qualified Data.Map.Strict as Map
 import qualified Data.Text as Text
 import qualified PercentEncoding
@@ -78,10 +77,10 @@ getUriConnectionString = do
 
 getKeyValueConnectionString :: P ConnectionString
 getKeyValueConnectionString =
-  fromKeyValueParams <$> getKeyValueParams
+  fromKeyValueParams =<< getKeyValueParams
 
-fromKeyValueParams :: Map.Map Text Text -> ConnectionString
-fromKeyValueParams params =
+fromKeyValueParams :: Map.Map Text Text -> P ConnectionString
+fromKeyValueParams params = do
   -- Extract known connection parameters
   let user = Map.lookup "user" params
       password = Map.lookup "password" params
@@ -98,21 +97,39 @@ fromKeyValueParams params =
         )
           params
 
-      -- Parse hosts if present - handle comma-separated hosts and ports
-      hosts = case hostText of
+  -- Validate the port count against the host count. A single port applies to
+  -- all hosts; any other count must match exactly (mirrors libpq's "could not
+  -- match N port numbers to M hosts" error).
+  let hostCount = maybe 1 (length . Text.splitOn ",") hostText
+  case portText of
+    Just portValue ->
+      let portCount = length (Text.splitOn "," portValue)
+       in unless (portCount == 1 || portCount == hostCount) $
+            fail
+              ( "could not match "
+                  <> show portCount
+                  <> " port numbers to "
+                  <> show hostCount
+                  <> " hosts"
+              )
+    Nothing -> pure ()
+
+  -- Parse hosts if present - handle comma-separated hosts and ports
+  let hosts = case hostText of
         Nothing -> []
         Just h ->
           let hostList = Text.splitOn "," h
               portList = maybe [] (Text.splitOn ",") portText
-              -- Pair up hosts with ports, padding with Nothing if needed
-              -- Exception: if just one port is provided, use it for all hosts
-              portList'
-                | length portList == 1 = repeat (fmap fst (List.uncons portList))
-                | otherwise = (map Just portList) ++ repeat Nothing
-              pairs = zipWith (\host mPort -> (host, mPort)) hostList portList'
-           in map (\(host, mPortText) -> Host host (mPortText >>= parsePort)) pairs
+              -- A single port applies to all hosts; otherwise the lists match
+              -- in length. An absent port uses the default port for every
+              -- host.
+              ports = case portList of
+                [] -> repeat Nothing
+                [single] -> repeat (Just single)
+                list -> map Just list
+           in zipWith (\host mPortText -> Host host (mPortText >>= parsePort)) hostList ports
 
-  in ConnectionString user password hosts dbname remainingParams
+  pure (ConnectionString user password hosts dbname remainingParams)
   where
     parsePort :: Text -> Maybe Word16
     parsePort t = case reads (Text.unpack t) of