From 32468b4bf5824b2d2e7817e061d60f778cf8d386 Mon Sep 17 00:00:00 2001
From: AlexanderSabirov <61634896+AlexanderSabirov@users.noreply.github.com>
Date: Tue, 28 Jan 2025 11:22:20 +0300
Subject: [PATCH] Add support target_session_attrs

---
 dbt/adapters/postgres/connections.py |  5 +++++
 tests/unit/test_adapter.py           | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/dbt/adapters/postgres/connections.py b/dbt/adapters/postgres/connections.py
index e8f0abe5..d99cd508 100644
--- a/dbt/adapters/postgres/connections.py
+++ b/dbt/adapters/postgres/connections.py
@@ -35,6 +35,7 @@ class PostgresCredentials(Credentials):
     sslkey: Optional[str] = None
     sslrootcert: Optional[str] = None
     application_name: Optional[str] = "dbt"
+    target_session_attrs: Optional[str] = None
     retries: int = 1
 
     _ALIASES = {"dbname": "database", "pass": "password"}
@@ -63,6 +64,7 @@ def _connection_keys(self):
             "sslkey",
             "sslrootcert",
             "application_name",
+            "target_session_attrs",
             "retries",
         )
 
@@ -133,6 +135,9 @@ def open(cls, connection):
         if credentials.application_name:
             kwargs["application_name"] = credentials.application_name
 
+        if credentials.target_session_attrs:
+            kwargs["target_session_attrs"] = credentials.target_session_attrs
+
         def connect():
             handle = None
 
diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py
index 1c375ac0..3df99f20 100644
--- a/tests/unit/test_adapter.py
+++ b/tests/unit/test_adapter.py
@@ -308,6 +308,24 @@ def test_set_zero_keepalive(self, psycopg2):
             application_name="dbt",
         )
 
+    @mock.patch("dbt.adapters.postgres.connections.psycopg2")
+    def test_set_target_session_attrs(self, psycopg2):
+        self.config.credentials = self.config.credentials.replace(target_session_attrs="any")
+        connection = self.adapter.acquire_connection("dummy")
+
+        psycopg2.connect.assert_not_called()
+        connection.handle
+        psycopg2.connect.assert_called_once_with(
+            dbname="postgres",
+            user="root",
+            host="thishostshouldnotexist",
+            password="password",
+            port=5432,
+            connect_timeout=10,
+            application_name="dbt",
+            target_session_attrs="any",
+        )
+
     @mock.patch.object(PostgresAdapter, "execute_macro")
     @mock.patch.object(PostgresAdapter, "_get_catalog_relations")
     def test_get_catalog_various_schemas(self, mock_get_relations, mock_execute):