Repository: dbt-labs/dbt-core

Issue:
[CT-472] postgres: Retry n times on connection timeout
### Is there an existing feature request for this?

- [X] I have searched the existing issues

### Describe the Feature

Runs/builds/tests can create hundreds of independent database connections depending on the size of the project and a single connection timeout due to transient network connections, EC2 load (e.g. RDS case when connecting through EC2), or postgres load can cause an entire run to fail. Connection timeouts can be, and most often are, transient and will often succeed on a retry.

### Describe alternatives you've considered

_No response_

### Who will this benefit?

Anyone using the postgres connector.

### Are you interested in contributing this feature?

_No response_

### Anything else?

This would be similar to connect_retries on Snowflake.
See: https://github.com/dbt-labs/dbt-snowflake/pull/6/files



Discussion:
Thanks @barberscott! I think this will be the place for the change, for both Postgres and https://github.com/dbt-labs/dbt-redshift/issues/96, since they both use dbt-postgres's `psycopg2` connection.

There are related recommendations in #3303 around:
- Specific error codes/messages to capture as transient, and likely to succeed on retry
- A proposal for user-specified "retryable" errors, as another config in `profiles.yml`

At this point, given precedent in other adapters (Snowflake + Spark), I'm not opposed to a naive `retry_all` parameter. It makes sense to turn on in production (best-possible chance at query success), and leave off in development (quicker iteration on errors).

I sense we'd need to add both complementary configs (`connect_retries` + `connect_timeout`), with reasonable defaults. Perhaps we also want to add our attempt at `retry_on_database_errors`, based on the `psycopg2` errors we strongly believe to be transient.
Hey all, since this issue is significantly hindering us, I went ahead and modified the Postgres adapter to take a stab at fixing it.

The basic premise was modifying [`PostgresCredentials`](https://github.com/dbt-labs/dbt-core/blob/8d39ef16b67fcc3c30e8a18751b4e64831ec9aaf/plugins/postgres/dbt/adapters/postgres/connections.py#L17=) to take in the new `connect_retries` and `retry_on_errors` parameters:

```python
@dataclass
class PostgresCredentials(Credentials):
    host: str
    user: str
    port: Port
    password: str  # on postgres the password is mandatory
    connect_retries: int = 3  # 3 seemed like a sensible default number of retries
    connect_timeout: int = 10
    role: Optional[str] = None
    search_path: Optional[str] = None
    keepalives_idle: int = 0  # 0 means to use the default value
    sslmode: Optional[str] = None
    sslcert: Optional[str] = None
    sslkey: Optional[str] = None
    sslrootcert: Optional[str] = None
    application_name: Optional[str] = "dbt"
    retry_on_errors: list[str] = [
        "08006", # ConnectionFailure
    ]
...
```

Users can set SQLSTATE codes that they wish to retry on. I'm uncertain whether taking codes is the right approach instead of exception class names (or allowing both).

Afterwards, I added an `except` clause for these errors in a connection loop added to [`PostgresConnectionManager.open`](https://github.com/dbt-labs/dbt-core/blob/8d39ef16b67fcc3c30e8a18751b4e64831ec9aaf/plugins/postgres/dbt/adapters/postgres/connections.py#L74=):

```python
    @classmethod
    def open(cls, connection):
        ...

        retryable_errors = tuple([psycopg2.errors.lookup(error_code) for error_code in credentials.retry_on_errors])

        attempt = 0
        while True:
            try:
                handle = psycopg2.connect(
                    dbname=credentials.database,
                    user=credentials.user,
                    host=credentials.host,
                    password=credentials.password,
                    port=credentials.port,
                    connect_timeout=credentials.connect_timeout,
                    **kwargs,
                )

                if credentials.role:
                    handle.cursor().execute("set role {}".format(credentials.role))

                connection.handle = handle
                connection.state = "open"

            except retryable_errors as e:
                logger.debug(
                    "Got a retryable error on attempt {} to open a postgres " "connection: '{}'".format(attempt, e)
                )

                attempt += 1
                if attempt < credentials.connect_retries:
                    continue

                logger.debug(
                    "Attempt number reached or exceeded {} retries when opening a postgres connection".format(credentials.connect_retries)
                )

                connection.handle = None
                connection.state = "fail"

                raise dbt.exceptions.FailedToConnectException(str(e))

            except psycopg2.Error as e:
                logger.debug(
                    "Got an unknown error on attempt {} to open a postgres " "connection: '{}'".format(attempt, e)
                )

                connection.handle = None
                connection.state = "fail"

                raise dbt.exceptions.FailedToConnectException(str(e))

        return connection
```

Some iteration over the code is most likely needed (as well as documentation). However, would this work for a PR? I'll be glad to add any necessary testing + documentation and contribute this.

Thanks.
@tomasfarias An approach in this vein makes a lot of sense to me! Thanks for sharing your code.

SQLState errors offer a concise way of specifying retryable errors (if a bit arcane). I think we need to establish a reasonable default set, so that 99% of users never need to override (or even know about) that config. We could also consider a brute force "retry all" option.

`pycopg2` doesn't offer us in-built `retry`. Would it better to look for another OSS `retry` alternative (e.g. [tenacity](https://github.com/jd/tenacity)), rather than rolling our own? It's not very complex code, but we do now have it in a few places ([snowflake](https://github.com/dbt-labs/dbt-snowflake/blob/80c400148b4847171cae425ee0587d8c55424e40/dbt/adapters/snowflake/connections.py#L261), [spark](https://github.com/dbt-labs/dbt-spark/blob/bc9fc0baaacce28706d0542391704d9ea971cee4/dbt/adapters/spark/connections.py#L337), [bigquery](https://github.com/dbt-labs/dbt-bigquery/blob/1b85fc50287ac9afd81fd92e715dffa5f0757379/dbt/adapters/bigquery/connections.py#L636) sorta) — and I think there's a lot of benefit in standardizing this in core / across adapters, both for users and for us as maintainers.

cc @nathaniel-may — I know retry is a thing you know / care a lot about, so I'd be curious to hear any thoughts you might have here.
Thanks for tagging me in @jtcohen6. This is one of those cases I would lean on "correctness above all else" rather than ask users to know enough about warehouse internals to list retryable error codes. Thankfully most dbt operations are retryable because they're idempotent in nature, but we do have a few places where we do stateful operations such as incremental models where we may need to be abundantly careful about retrying.

My preference is to write our own retry logic unless this library solves a particularly difficult problem for us. Personally, I don't see much value-add in a general purpose "retry library," but I could be convinced if I'm missing something.
@tomasfarias I'd be very supportive of a PR for this, if you're open to submitting one! My main recommendation would be, in keeping with Nate's suggestion — rather than ask users to list error codes that are retryable, let's try to develop a good set for starting out. I'm particularly interested in connection errors / timeouts.
I'll submit a PR for this week. I've also been looking into dbt-bigquery and dbt-snowflake to ensure the retry implementation can be re-used across adapters, and it's not just a dbt-postgres/redshif thing.

Thanks for the patience, I've been on vacation for the last couple of weeks.

Resolving patch:
diff --git a/core/dbt/adapters/base/connections.py b/core/dbt/adapters/base/connections.py
--- a/core/dbt/adapters/base/connections.py
+++ b/core/dbt/adapters/base/connections.py
@@ -1,10 +1,24 @@
 import abc
 import os
+from time import sleep
+import sys
 
 # multiprocessing.RLock is a function returning this type
 from multiprocessing.synchronize import RLock
 from threading import get_ident
-from typing import Dict, Tuple, Hashable, Optional, ContextManager, List
+from typing import (
+    Any,
+    Dict,
+    Tuple,
+    Hashable,
+    Optional,
+    ContextManager,
+    List,
+    Type,
+    Union,
+    Iterable,
+    Callable,
+)
 
 import agate
 
@@ -21,6 +35,7 @@
 from dbt.adapters.base.query_headers import (
     MacroQueryStringSetter,
 )
+from dbt.events import AdapterLogger
 from dbt.events.functions import fire_event
 from dbt.events.types import (
     NewConnection,
@@ -34,6 +49,9 @@
 )
 from dbt import flags
 
+SleepTime = Union[int, float]  # As taken by time.sleep.
+AdapterHandle = Any  # Adapter connection handle objects can be any class.
+
 
 class BaseConnectionManager(metaclass=abc.ABCMeta):
     """Methods to implement:
@@ -159,6 +177,94 @@ def set_connection_name(self, name: Optional[str] = None) -> Connection:
         conn.name = conn_name
         return conn
 
+    @classmethod
+    def retry_connection(
+        cls,
+        connection: Connection,
+        connect: Callable[[], AdapterHandle],
+        logger: AdapterLogger,
+        retryable_exceptions: Iterable[Type[Exception]],
+        retry_limit: int = 1,
+        retry_timeout: Union[Callable[[int], SleepTime], SleepTime] = 1,
+        _attempts: int = 0,
+    ) -> Connection:
+        """Given a Connection, set its handle by calling connect.
+
+        The calls to connect will be retried up to retry_limit times to deal with transient
+        connection errors. By default, one retry will be attempted if retryable_exceptions is set.
+
+        :param Connection connection: An instance of a Connection that needs a handle to be set,
+            usually when attempting to open it.
+        :param connect: A callable that returns the appropiate connection handle for a
+            given adapter. This callable will be retried retry_limit times if a subclass of any
+            Exception in retryable_exceptions is raised by connect.
+        :type connect: Callable[[], AdapterHandle]
+        :param AdapterLogger logger: A logger to emit messages on retry attempts or errors. When
+            handling expected errors, we call debug, and call warning on unexpected errors or when
+            all retry attempts have been exhausted.
+        :param retryable_exceptions: An iterable of exception classes that if raised by
+            connect should trigger a retry.
+        :type retryable_exceptions: Iterable[Type[Exception]]
+        :param int retry_limit: How many times to retry the call to connect. If this limit
+            is exceeded before a successful call, a FailedToConnectException will be raised.
+            Must be non-negative.
+        :param retry_timeout: Time to wait between attempts to connect. Can also take a
+            Callable that takes the number of attempts so far, beginning at 0, and returns an int
+            or float to be passed to time.sleep.
+        :type retry_timeout: Union[Callable[[int], SleepTime], SleepTime] = 1
+        :param int _attempts: Parameter used to keep track of the number of attempts in calling the
+            connect function across recursive calls. Passed as an argument to retry_timeout if it
+            is a Callable. This parameter should not be set by the initial caller.
+        :raises dbt.exceptions.FailedToConnectException: Upon exhausting all retry attempts without
+            successfully acquiring a handle.
+        :return: The given connection with its appropriate state and handle attributes set
+            depending on whether we successfully acquired a handle or not.
+        """
+        timeout = retry_timeout(_attempts) if callable(retry_timeout) else retry_timeout
+        if timeout < 0:
+            raise dbt.exceptions.FailedToConnectException(
+                "retry_timeout cannot be negative or return a negative time."
+            )
+
+        if retry_limit < 0 or retry_limit > sys.getrecursionlimit():
+            # This guard is not perfect others may add to the recursion limit (e.g. built-ins).
+            connection.handle = None
+            connection.state = ConnectionState.FAIL
+            raise dbt.exceptions.FailedToConnectException("retry_limit cannot be negative")
+
+        try:
+            connection.handle = connect()
+            connection.state = ConnectionState.OPEN
+            return connection
+
+        except tuple(retryable_exceptions) as e:
+            if retry_limit <= 0:
+                connection.handle = None
+                connection.state = ConnectionState.FAIL
+                raise dbt.exceptions.FailedToConnectException(str(e))
+
+            logger.debug(
+                f"Got a retryable error when attempting to open a {cls.TYPE} connection.\n"
+                f"{retry_limit} attempts remaining. Retrying in {timeout} seconds.\n"
+                f"Error:\n{e}"
+            )
+
+            sleep(timeout)
+            return cls.retry_connection(
+                connection=connection,
+                connect=connect,
+                logger=logger,
+                retry_limit=retry_limit - 1,
+                retry_timeout=retry_timeout,
+                retryable_exceptions=retryable_exceptions,
+                _attempts=_attempts + 1,
+            )
+
+        except Exception as e:
+            connection.handle = None
+            connection.state = ConnectionState.FAIL
+            raise dbt.exceptions.FailedToConnectException(str(e))
+
     @abc.abstractmethod
     def cancel_open(self) -> Optional[List[str]]:
         """Cancel all open connections on the adapter. (passable)"""
diff --git a/plugins/postgres/dbt/adapters/postgres/connections.py b/plugins/postgres/dbt/adapters/postgres/connections.py
--- a/plugins/postgres/dbt/adapters/postgres/connections.py
+++ b/plugins/postgres/dbt/adapters/postgres/connections.py
@@ -31,6 +31,7 @@ class PostgresCredentials(Credentials):
     sslkey: Optional[str] = None
     sslrootcert: Optional[str] = None
     application_name: Optional[str] = "dbt"
+    retries: int = 1
 
     _ALIASES = {"dbname": "database", "pass": "password"}
 
@@ -121,7 +122,7 @@ def open(cls, connection):
         if credentials.application_name:
             kwargs["application_name"] = credentials.application_name
 
-        try:
+        def connect():
             handle = psycopg2.connect(
                 dbname=credentials.database,
                 user=credentials.user,
@@ -131,23 +132,26 @@ def open(cls, connection):
                 connect_timeout=credentials.connect_timeout,
                 **kwargs,
             )
-
             if credentials.role:
                 handle.cursor().execute("set role {}".format(credentials.role))
-
-            connection.handle = handle
-            connection.state = "open"
-        except psycopg2.Error as e:
-            logger.debug(
-                "Got an error when attempting to open a postgres " "connection: '{}'".format(e)
-            )
-
-            connection.handle = None
-            connection.state = "fail"
-
-            raise dbt.exceptions.FailedToConnectException(str(e))
-
-        return connection
+            return handle
+
+        retryable_exceptions = [
+            # OperationalError is subclassed by all psycopg2 Connection Exceptions and it's raised
+            # by generic connection timeouts without an error code. This is a limitation of
+            # psycopg2 which doesn't provide subclasses for errors without a SQLSTATE error code.
+            # The limitation has been known for a while and there are no efforts to tackle it.
+            # See: https://github.com/psycopg/psycopg2/issues/682
+            psycopg2.errors.OperationalError,
+        ]
+
+        return cls.retry_connection(
+            connection,
+            connect=connect,
+            logger=logger,
+            retry_limit=credentials.retries,
+            retryable_exceptions=retryable_exceptions,
+        )
 
     def cancel(self, connection):
         connection_name = connection.name


Test patch:
diff --git a/test/unit/test_adapter_connection_manager.py b/test/unit/test_adapter_connection_manager.py
new file mode 100644
--- /dev/null
+++ b/test/unit/test_adapter_connection_manager.py
@@ -0,0 +1,494 @@
+import unittest
+from unittest import mock
+import sys
+
+import dbt.exceptions
+
+import psycopg2
+
+from dbt.contracts.connection import Connection
+from dbt.adapters.base import BaseConnectionManager
+from dbt.adapters.postgres import PostgresCredentials, PostgresConnectionManager
+from dbt.events import AdapterLogger
+
+
+class BaseConnectionManagerTest(unittest.TestCase):
+    def setUp(self):
+        self.postgres_credentials = PostgresCredentials(
+            host="localhost",
+            user="test-user",
+            port=1111,
+            password="test-password",
+            database="test-db",
+            schema="test-schema",
+        )
+        self.logger = AdapterLogger("test")
+        self.postgres_connection = Connection("postgres", None, self.postgres_credentials)
+
+    def test_retry_connection(self):
+        """Test a dummy handle is set on a connection on the first attempt.
+
+        This test uses a Connection populated with test PostgresCredentials values, and
+        expects the Connection.handle attribute to be set to True and it's state to
+        "open", after calling retry_connection.
+
+        Moreover, the attribute should be set in the first attempt as no exception would
+        be raised for retrying. A mock connect function is used to simulate a real connection
+        passing on the first attempt.
+        """
+        conn = self.postgres_connection
+        attempts = 0
+
+        def connect():
+            nonlocal attempts
+            attempts += 1
+            return True
+
+        conn = BaseConnectionManager.retry_connection(
+            conn,
+            connect,
+            self.logger,
+            retryable_exceptions=[],
+        )
+
+        assert conn.state == "open"
+        assert conn.handle is True
+        assert attempts == 1
+
+    def test_retry_connection_fails_unhandled(self):
+        """Test setting a handle fails upon raising a non-handled exception.
+
+        This test uses a Connection populated with test PostgresCredentials values, and
+        expects a ValueError to be raised by a mock connect function. As a
+        result:
+        * The Connection state should be "fail" and the handle None.
+        * The resulting attempt count should be 1 as we are not explicitly configured to handle a
+          ValueError.
+        * retry_connection should raise a FailedToConnectException with the Exception message.
+        """
+        conn = self.postgres_connection
+        attempts = 0
+
+        def connect():
+            nonlocal attempts
+            attempts += 1
+            raise ValueError("Something went horribly wrong")
+
+        with self.assertRaisesRegex(
+            dbt.exceptions.FailedToConnectException, "Something went horribly wrong"
+        ):
+
+            BaseConnectionManager.retry_connection(
+                conn,
+                connect,
+                self.logger,
+                retry_limit=1,
+                retry_timeout=lambda attempt: 0,
+                retryable_exceptions=(TypeError,),
+            )
+
+        assert conn.state == "fail"
+        assert conn.handle is None
+        assert attempts == 1
+
+    def test_retry_connection_fails_handled(self):
+        """Test setting a handle fails upon raising a handled exception.
+
+        This test uses a Connection populated with test PostgresCredentials values, and
+        expects a ValueError to be raised by a mock connect function.
+        As a result:
+        * The Connection state should be "fail" and the handle None.
+        * The resulting attempt count should be 2 as we are configured to handle a ValueError.
+        * retry_connection should raise a FailedToConnectException with the Exception message.
+        """
+        conn = self.postgres_connection
+        attempts = 0
+
+        def connect():
+            nonlocal attempts
+            attempts += 1
+            raise ValueError("Something went horribly wrong")
+
+        with self.assertRaisesRegex(
+            dbt.exceptions.FailedToConnectException, "Something went horribly wrong"
+        ):
+
+            BaseConnectionManager.retry_connection(
+                conn,
+                connect,
+                self.logger,
+                retry_timeout=0,
+                retryable_exceptions=(ValueError,),
+                retry_limit=1,
+            )
+
+        assert conn.state == "fail"
+        assert conn.handle is None
+
+    def test_retry_connection_passes_handled(self):
+        """Test setting a handle fails upon raising a handled exception.
+
+        This test uses a Connection populated with test PostgresCredentials values, and
+        expects a ValueError to be raised by a mock connect function only the first
+        time is called. Upon handling the exception once, connect should return.
+        As a result:
+        * The Connection state should be "open" and the handle True.
+        * The resulting attempt count should be 2 as we are configured to handle a ValueError.
+        """
+        conn = self.postgres_connection
+        is_handled = False
+        attempts = 0
+
+        def connect():
+            nonlocal is_handled
+            nonlocal attempts
+
+            attempts += 1
+
+            if is_handled:
+                return True
+
+            is_handled = True
+            raise ValueError("Something went horribly wrong")
+
+        conn = BaseConnectionManager.retry_connection(
+            conn,
+            connect,
+            self.logger,
+            retry_timeout=0,
+            retryable_exceptions=(ValueError,),
+            retry_limit=1,
+        )
+
+        assert conn.state == "open"
+        assert conn.handle is True
+        assert is_handled is True
+        assert attempts == 2
+
+    def test_retry_connection_attempts(self):
+        """Test setting a handle fails upon raising a handled exception multiple times.
+
+        This test uses a Connection populated with test PostgresCredentials values, and
+        expects a ValueError to be raised by a mock connect function. As a result:
+        * The Connection state should be "fail" and the handle None, as connect
+          never returns.
+        * The resulting attempt count should be 11 as we are configured to handle a ValueError.
+        * retry_connection should raise a FailedToConnectException with the Exception message.
+        """
+        conn = self.postgres_connection
+        attempts = 0
+
+        def connect():
+            nonlocal attempts
+            attempts += 1
+
+            raise ValueError("Something went horribly wrong")
+
+        with self.assertRaisesRegex(
+            dbt.exceptions.FailedToConnectException, "Something went horribly wrong"
+        ):
+            BaseConnectionManager.retry_connection(
+                conn,
+                connect,
+                self.logger,
+                retry_timeout=0,
+                retryable_exceptions=(ValueError,),
+                retry_limit=10,
+            )
+
+        assert conn.state == "fail"
+        assert conn.handle is None
+        assert attempts == 11
+
+    def test_retry_connection_fails_handling_all_exceptions(self):
+        """Test setting a handle fails after exhausting all attempts.
+
+        This test uses a Connection populated with test PostgresCredentials values, and
+        expects a TypeError to be raised by a mock connect function. As a result:
+        * The Connection state should be "fail" and the handle None, as connect
+          never returns.
+        * The resulting attempt count should be 11 as we are configured to handle all Exceptions.
+        * retry_connection should raise a FailedToConnectException with the Exception message.
+        """
+        conn = self.postgres_connection
+        attempts = 0
+
+        def connect():
+            nonlocal attempts
+            attempts += 1
+
+            raise TypeError("An unhandled thing went horribly wrong")
+
+        with self.assertRaisesRegex(
+            dbt.exceptions.FailedToConnectException, "An unhandled thing went horribly wrong"
+        ):
+            BaseConnectionManager.retry_connection(
+                conn,
+                connect,
+                self.logger,
+                retry_timeout=0,
+                retryable_exceptions=[Exception],
+                retry_limit=15,
+            )
+
+        assert conn.state == "fail"
+        assert conn.handle is None
+        assert attempts == 16
+
+    def test_retry_connection_passes_multiple_handled(self):
+        """Test setting a handle passes upon handling multiple exceptions.
+
+        This test uses a Connection populated with test PostgresCredentials values, and
+        expects a mock connect to raise a ValueError in the first invocation and a
+        TypeError in the second invocation. As a result:
+        * The Connection state should be "open" and the handle True, as connect
+          returns after both exceptions have been handled.
+        * The resulting attempt count should be 3.
+        """
+        conn = self.postgres_connection
+        is_value_err_handled = False
+        is_type_err_handled = False
+        attempts = 0
+
+        def connect():
+            nonlocal is_value_err_handled
+            nonlocal is_type_err_handled
+            nonlocal attempts
+
+            attempts += 1
+
+            if is_value_err_handled and is_type_err_handled:
+                return True
+            elif is_type_err_handled:
+                is_value_err_handled = True
+                raise ValueError("Something went horribly wrong")
+            else:
+                is_type_err_handled = True
+                raise TypeError("An unhandled thing went horribly wrong")
+
+        conn = BaseConnectionManager.retry_connection(
+            conn,
+            connect,
+            self.logger,
+            retry_timeout=0,
+            retryable_exceptions=(ValueError, TypeError),
+            retry_limit=2,
+        )
+
+        assert conn.state == "open"
+        assert conn.handle is True
+        assert is_type_err_handled is True
+        assert is_value_err_handled is True
+        assert attempts == 3
+
+    def test_retry_connection_passes_none_excluded(self):
+        """Test setting a handle passes upon handling multiple exceptions.
+
+        This test uses a Connection populated with test PostgresCredentials values, and
+        expects a mock connect to raise a ValueError in the first invocation and a
+        TypeError in the second invocation. As a result:
+        * The Connection state should be "open" and the handle True, as connect
+          returns after both exceptions have been handled.
+        * The resulting attempt count should be 3.
+        """
+        conn = self.postgres_connection
+        is_value_err_handled = False
+        is_type_err_handled = False
+        attempts = 0
+
+        def connect():
+            nonlocal is_value_err_handled
+            nonlocal is_type_err_handled
+            nonlocal attempts
+
+            attempts += 1
+
+            if is_value_err_handled and is_type_err_handled:
+                return True
+            elif is_type_err_handled:
+                is_value_err_handled = True
+                raise ValueError("Something went horribly wrong")
+            else:
+                is_type_err_handled = True
+                raise TypeError("An unhandled thing went horribly wrong")
+
+        conn = BaseConnectionManager.retry_connection(
+            conn,
+            connect,
+            self.logger,
+            retry_timeout=0,
+            retryable_exceptions=(ValueError, TypeError),
+            retry_limit=2,
+        )
+
+        assert conn.state == "open"
+        assert conn.handle is True
+        assert is_type_err_handled is True
+        assert is_value_err_handled is True
+        assert attempts == 3
+
+    def test_retry_connection_retry_limit(self):
+        """Test retry_connection raises an exception with a negative retry limit."""
+        conn = self.postgres_connection
+        attempts = 0
+
+        def connect():
+            nonlocal attempts
+            attempts += 1
+            return True
+
+        with self.assertRaisesRegex(
+            dbt.exceptions.FailedToConnectException, "retry_limit cannot be negative"
+        ):
+            BaseConnectionManager.retry_connection(
+                conn,
+                connect,
+                self.logger,
+                retry_timeout=0,
+                retryable_exceptions=(ValueError,),
+                retry_limit=-2,
+            )
+
+        assert conn.state == "fail"
+        assert conn.handle is None
+        assert attempts == 0
+
+    def test_retry_connection_retry_timeout(self):
+        """Test retry_connection raises an exception with a negative timeout."""
+        conn = self.postgres_connection
+        attempts = 0
+
+        def connect():
+            nonlocal attempts
+            attempts += 1
+            return True
+
+        for retry_timeout in [-10, -2.5, lambda _: -100, lambda _: -10.1]:
+            with self.assertRaisesRegex(
+                dbt.exceptions.FailedToConnectException,
+                "retry_timeout cannot be negative or return a negative time",
+            ):
+                BaseConnectionManager.retry_connection(
+                    conn,
+                    connect,
+                    self.logger,
+                    retry_timeout=-10,
+                    retryable_exceptions=(ValueError,),
+                    retry_limit=2,
+                )
+
+        assert conn.state == "init"
+        assert conn.handle is None
+        assert attempts == 0
+
+    def test_retry_connection_exceeds_recursion_limit(self):
+        """Test retry_connection raises an exception with retries that exceed recursion limit."""
+        conn = self.postgres_connection
+        attempts = 0
+
+        def connect():
+            nonlocal attempts
+            attempts += 1
+            return True
+
+        with self.assertRaisesRegex(
+            dbt.exceptions.FailedToConnectException,
+            "retry_limit cannot be negative",
+        ):
+            BaseConnectionManager.retry_connection(
+                conn,
+                connect,
+                self.logger,
+                retry_timeout=2,
+                retryable_exceptions=(ValueError,),
+                retry_limit=sys.getrecursionlimit() + 1,
+            )
+
+        assert conn.state == "fail"
+        assert conn.handle is None
+        assert attempts == 0
+
+    def test_retry_connection_with_exponential_backoff_timeout(self):
+        """Test retry_connection with an exponential backoff timeout.
+
+        We assert the provided exponential backoff function gets passed the right attempt number
+        and produces the expected timeouts.
+        """
+        conn = self.postgres_connection
+        attempts = 0
+        timeouts = []
+
+        def connect():
+            nonlocal attempts
+            attempts += 1
+
+            if attempts < 12:
+                raise ValueError("Keep trying!")
+            return True
+
+        def exp_backoff(n):
+            nonlocal timeouts
+            computed = 2**n
+            # We store the computed values to ensure they match the expected backoff...
+            timeouts.append((n, computed))
+            # but we return 0 as we don't want the test to go on forever.
+            return 0
+
+        conn = BaseConnectionManager.retry_connection(
+            conn,
+            connect,
+            self.logger,
+            retry_timeout=exp_backoff,
+            retryable_exceptions=(ValueError,),
+            retry_limit=12,
+        )
+
+        assert conn.state == "open"
+        assert conn.handle is True
+        assert attempts == 12
+        assert timeouts == [(n, 2**n) for n in range(12)]
+
+
+class PostgresConnectionManagerTest(unittest.TestCase):
+    def setUp(self):
+        self.credentials = PostgresCredentials(
+            host="localhost",
+            user="test-user",
+            port=1111,
+            password="test-password",
+            database="test-db",
+            schema="test-schema",
+            retries=2,
+        )
+        self.connection = Connection("postgres", None, self.credentials)
+
+    def test_open(self):
+        """Test opening a Postgres Connection with failures in the first 3 attempts.
+
+        This test uses a Connection populated with test PostgresCredentials values, and
+        expects a mock connect to raise a psycopg2.errors.ConnectionFailuer
+        in the first 3 invocations, after which the mock should return True. As a result:
+        * The Connection state should be "open" and the handle True, as connect
+          returns in the 4th attempt.
+        * The resulting attempt count should be 4.
+        """
+        conn = self.connection
+        attempt = 0
+
+        def connect(*args, **kwargs):
+            nonlocal attempt
+            attempt += 1
+
+            if attempt <= 2:
+                raise psycopg2.errors.ConnectionFailure("Connection has failed")
+
+            return True
+
+        with mock.patch("psycopg2.connect", wraps=connect) as mock_connect:
+            PostgresConnectionManager.open(conn)
+
+            assert mock_connect.call_count == 3
+
+        assert attempt == 3
+        assert conn.state == "open"
+        assert conn.handle is True