🐛 Gracefully handle tcp connection errors

CHANGELOG.md

@@ -1,5 +1,12 @@
 # CHANGELOG
 
+## Unreleased
+
+- Fixed a bug where Squirrel would panic if not able to establish a TCP
+  connection to the postgres server. Now it gracefully handles the error by
+  showing an appropriate error message.
+  ([Giacomo Cavalieri](https://github.com/giacomocavalieri))
+
 ## v1.3.0 - 2024-08-13
 
 - One can now use the `DATABASE_URL` env variable to specify a connection string

src/squirrel/internal/database/postgres.gleam

@@ -28,12 +28,13 @@ import gleam/string
 import squirrel/internal/database/postgres_protocol as pg
 import squirrel/internal/error.{
   type Error, type Pointer, type ValueIdentifierError, ByteIndex,
-  CannotParseQuery, PgCannotConnectUserDatabase, PgCannotDecodeReceivedMessage,
-  PgCannotDescribeQuery, PgCannotReceiveMessage, PgCannotSendMessage,
-  PgInvalidPassword, PgInvalidSha256ServerProof, PgPermissionDenied,
-  PgUnexpectedAuthMethodMessage, PgUnexpectedCleartextAuthMessage,
-  PgUnexpectedSha256AuthMessage, PgUnsupportedAuthentication, Pointer,
-  QueryHasInvalidColumn, QueryHasUnsupportedType,
+  CannotParseQuery, PgCannotDecodeReceivedMessage, PgCannotDescribeQuery,
+  PgCannotEstablishTcpConnection, PgCannotReceiveMessage, PgCannotSendMessage,
+  PgInvalidPassword, PgInvalidSha256ServerProof, PgInvalidUserDatabase,
+  PgPermissionDenied, PgUnexpectedAuthMethodMessage,
+  PgUnexpectedCleartextAuthMessage, PgUnexpectedSha256AuthMessage,
+  PgUnsupportedAuthentication, Pointer, QueryHasInvalidColumn,
+  QueryHasUnsupportedType,
 }
 import squirrel/internal/eval_extra
 import squirrel/internal/gleam
@@ -223,12 +224,17 @@ pub fn main(
   queries: List(UntypedQuery),
   connection: ConnectionOptions,
 ) -> Result(#(List(TypedQuery), List(Error)), Error) {
+  use db <- result.try(
+    pg.connect(connection.host, connection.port, connection.timeout)
+    |> result.map_error(PgCannotEstablishTcpConnection(
+      host: connection.host,
+      port: connection.port,
+      reason: _,
+    )),
+  )
+
   let context =
-    Context(
-      db: pg.connect(connection.host, connection.port, connection.timeout),
-      gleam_types: dict.new(),
-      column_nullability: dict.new(),
-    )
+    Context(db: db, gleam_types: dict.new(), column_nullability: dict.new())
 
   // Once the server has confirmed that it is ready to accept query requests we
   // can start gathering information about all the different queries.
@@ -289,7 +295,7 @@ fn authenticate(connection: ConnectionOptions) -> Db(Nil) {
     // In case there's a receive error while waiting for the server to be ready
     // we want to display a more helpful error message because the problem here
     // must be with an invalid username/database combination.
-    |> eval.replace_error(PgCannotConnectUserDatabase(
+    |> eval.replace_error(PgInvalidUserDatabase(
       user: connection.user,
       database: connection.database,
     )),

src/squirrel/internal/database/postgres_protocol.gleam

@@ -4,6 +4,8 @@
 ////
 //// I had to make a little change to the existing library:
 ////  - do not fail with an unexpected Command if the outer message is ok
+////  - change the connect function to take in a mug socket instead of asserting
+////    and creating one
 ////
 //// This library parses and generates packages for the PostgreSQL Binary Protocol
 //// It also provides a basic connection abstraction, but this hasn't been used
@@ -23,10 +25,9 @@ pub type Connection {
 }
 
 pub fn connect(host, port, timeout) {
-  let assert Ok(socket) =
-    mug.connect(mug.ConnectionOptions(host: host, port: port, timeout: timeout))
-
-  Connection(socket: socket, buffer: <<>>, timeout: timeout)
+  let options = mug.ConnectionOptions(host: host, port: port, timeout: timeout)
+  use socket <- result.try(mug.connect(options))
+  Ok(Connection(socket: socket, buffer: <<>>, timeout: timeout))
 }
 
 pub type StateInitial {

src/squirrel/internal/error.gleam

@@ -6,14 +6,19 @@ import gleam/regex
 import gleam/result
 import gleam/string
 import gleam_community/ansi
+import mug
 import simplifile
 
 pub type Error {
   // --- POSTGRES RELATED ERRORS -----------------------------------------------
+  /// When we can't establish a TCP connection to the server.
+  ///
+  PgCannotEstablishTcpConnection(host: String, port: Int, reason: mug.Error)
+
   /// When the server immediately closes the connection right after we try to
   /// connect using the given username and database.
   ///
-  PgCannotConnectUserDatabase(user: String, database: String)
+  PgInvalidUserDatabase(user: String, database: String)
 
   /// When we receive an unexpected message while waiting for the authentication
   /// methods allowed by the server.
@@ -163,7 +168,43 @@ pub fn to_doc(error: Error) -> Document {
   // that is actually printed and we do not have to make any effort to add and
   // print new errors.
   let printable_error = case error {
-    PgCannotConnectUserDatabase(user: user, database: database) ->
+    PgCannotEstablishTcpConnection(host: host, port: port, reason: reason) ->
+      printable_error("Cannot establish TCP connection")
+      |> add_paragraph(case reason {
+        mug.Econnrefused ->
+          "I couldn't connect to the database because "
+          <> style_inline_code(host)
+          <> " refused the connection to port "
+          <> int.to_string(port)
+          <> "."
+
+        mug.Closed ->
+          "I couldn't connect to the database because "
+          <> style_inline_code(host)
+          <> " closed the connection to port "
+          <> int.to_string(port)
+          <> "."
+
+        mug.Ehostunreach ->
+          "I couldn't connect to the database because "
+          <> style_inline_code(host)
+          <> " is unreachable."
+
+        mug.Timeout ->
+          "I couldn't connect to "
+          <> style_inline_code(host)
+          <> " at port "
+          <> int.to_string(port)
+          <> " because the connection timed out."
+
+        _ ->
+          "I couldn't connect to the database because I ran into the following
+problem while trying to establish a TCP connection to
+" <> style_inline_code(host) <> " at port " <> int.to_string(port) <> ":
+" <> string.inspect(reason)
+      })
+
+    PgInvalidUserDatabase(user: user, database: database) ->
       printable_error("Cannot connect")
       |> add_paragraph(
         "I couldn't connect to database "