From 58c229a3f7371f159a0d425c4dcfb9fe76285555 Mon Sep 17 00:00:00 2001 From: Andrei Date: Fri, 22 Nov 2024 11:24:42 +0000 Subject: [PATCH] black format --- dbt/adapters/duckdb/plugins/postgres.py | 34 ++++++++------- tests/functional/plugins/test_postgres.py | 51 ++++++++++++++++------- 2 files changed, 54 insertions(+), 31 deletions(-) diff --git a/dbt/adapters/duckdb/plugins/postgres.py b/dbt/adapters/duckdb/plugins/postgres.py index ffecbdd3..e891ecc7 100644 --- a/dbt/adapters/duckdb/plugins/postgres.py +++ b/dbt/adapters/duckdb/plugins/postgres.py @@ -1,11 +1,14 @@ -from typing import Any, Dict, Optional, List, Tuple +from typing import Any, Dict, List, Optional, Tuple + from duckdb import DuckDBPyConnection + from dbt.adapters.events.logging import AdapterLogger from . import BasePlugin PG_EXT = "postgres" + class Plugin(BasePlugin): logger = AdapterLogger("DuckDB_PostgresPlugin") @@ -15,9 +18,7 @@ def __init__(self, name: str, plugin_config: Dict[str, Any]): """ super().__init__(name, plugin_config) self.logger.debug( - "Plugin __init__ called with name: %s and config: %s", - name, - plugin_config + "Plugin __init__ called with name: %s and config: %s", name, plugin_config ) self.initialize(plugin_config) @@ -32,16 +33,18 @@ def initialize(self, config: Dict[str, Any]): self.logger.error( "Initialization failed: 'dsn' is a required argument for the postgres plugin!" ) - raise ValueError( - "'dsn' is a required argument for the postgres plugin!" - ) + raise ValueError("'dsn' is a required argument for the postgres plugin!") self._pg_schema: Optional[str] = config.get("pg_schema") # Can be None self._duckdb_alias: str = config.get("duckdb_alias", "postgres_db") self._read_only: bool = config.get("read_only", False) self._secret: Optional[str] = config.get("secret") - self._attach_options: Dict[str, Any] = config.get("attach_options", {}) # Additional ATTACH options - self._settings: Dict[str, Any] = config.get("settings", {}) # Extension settings via SET commands + self._attach_options: Dict[str, Any] = config.get( + "attach_options", {} + ) # Additional ATTACH options + self._settings: Dict[str, Any] = config.get( + "settings", {} + ) # Extension settings via SET commands self.logger.info( "PostgreSQL plugin initialized with dsn='%s', pg_schema='%s', " @@ -50,7 +53,7 @@ def initialize(self, config: Dict[str, Any]): self._pg_schema, self._duckdb_alias, self._read_only, - self._secret + self._secret, ) def configure_connection(self, conn: DuckDBPyConnection): @@ -72,8 +75,7 @@ def configure_connection(self, conn: DuckDBPyConnection): try: conn.execute(attach_stmt) self.logger.info( - "Successfully attached PostgreSQL database with DSN: %s", - self._dsn + "Successfully attached PostgreSQL database with DSN: %s", self._dsn ) except Exception as e: self.logger.error("Failed to attach PostgreSQL database: %s", e) @@ -88,7 +90,7 @@ def _set_extension_settings(self, conn: DuckDBPyConnection): if isinstance(value, str): value = f"'{value}'" elif isinstance(value, bool): - value = 'true' if value else 'false' + value = "true" if value else "false" set_stmt = f"SET {setting} = {value};" self.logger.debug("Setting extension option: %s", set_stmt) try: @@ -112,7 +114,7 @@ def _build_attach_statement(self) -> str: # Additional attach options for k, v in self._attach_options.items(): if isinstance(v, bool): - v = 'true' if v else 'false' + v = "true" if v else "false" elif isinstance(v, str): v = f"'{v}'" attach_options.append((k.upper(), v)) @@ -125,5 +127,7 @@ def _build_attach_statement(self) -> str: f"{k} {v}" if v is not None else k for k, v in attach_options ) - attach_stmt = f"ATTACH '{self._dsn}' AS {self._duckdb_alias} ({attach_options_str});" + attach_stmt = ( + f"ATTACH '{self._dsn}' AS {self._duckdb_alias} ({attach_options_str});" + ) return attach_stmt diff --git a/tests/functional/plugins/test_postgres.py b/tests/functional/plugins/test_postgres.py index 870978bf..04ef3ee6 100644 --- a/tests/functional/plugins/test_postgres.py +++ b/tests/functional/plugins/test_postgres.py @@ -1,13 +1,17 @@ import os + import pytest from testcontainers.postgres import PostgresContainer + from dbt.adapters.duckdb.plugins.postgres import Plugin from dbt.adapters.events.logging import AdapterLogger from dbt.tests.util import run_dbt, run_sql_with_adapter logger = AdapterLogger("DuckDB_PostgresPlugin") -@pytest.mark.skip(reason=""" + +@pytest.mark.skip( + reason=""" Skipping this b/c it requires running a properly setup Postgres server when testing it locally and also b/c I think there is something wrong with profiles_config_update since it can't be used in multiple @@ -15,7 +19,8 @@ Exercise locally with: pytest --profile=file tests/functional/plugins/test_postgres.py there are networking issues testing this in a devcontainer, more here https://github.com/testcontainers/testcontainers-python/issues/538 -the test works outside the devcontainer though""") +the test works outside the devcontainer though""" +) @pytest.mark.usefixtures("postgres_container") class TestPostgresPluginWithDbt: @pytest.fixture(scope="class") @@ -48,7 +53,8 @@ def postgres_container(self): @pytest.fixture(scope="class", autouse=True) def attach_postgres_db(self, project, postgres_container): import time - from psycopg import connect, OperationalError + + from psycopg import OperationalError, connect db_host = os.getenv("DB_HOST") db_port = os.getenv("DB_PORT") @@ -61,7 +67,11 @@ def attach_postgres_db(self, project, postgres_container): while retries > 0: try: with connect( - host=db_host, port=db_port, user=db_username, password=db_password, dbname=db_name + host=db_host, + port=db_port, + user=db_username, + password=db_password, + dbname=db_name, ) as conn: logger.info("PostgreSQL container is ready.") break @@ -154,7 +164,7 @@ def test_postgres_plugin_write(self, project): res = run_sql_with_adapter( project.adapter, "SELECT i FROM postgres_db.public.foo WHERE id = 4;", - fetch="one" + fetch="one", ) assert res[0] == 4, "The value of 'i' should be 4." @@ -194,9 +204,11 @@ def test_postgres_plugin_read_only(self, project): id INTEGER PRIMARY KEY, value TEXT ); - """ + """, ) - assert "read-only" in str(exc_info.value).lower(), "Write operation should fail in read-only mode." + assert ( + "read-only" in str(exc_info.value).lower() + ), "Write operation should fail in read-only mode." def test_postgres_plugin_extension_settings(self, project): """ @@ -209,7 +221,7 @@ def test_postgres_plugin_extension_settings(self, project): res = run_sql_with_adapter( project.adapter, "SELECT current_setting('pg_use_binary_copy');", - fetch="one" + fetch="one", ) assert res[0] == True, "The setting 'pg_use_binary_copy' should be 'true'." @@ -219,7 +231,9 @@ def test_postgres_plugin_attach_options(self, project): """ # Detach any existing databases run_sql_with_adapter(project.adapter, "DETACH DATABASE IF EXISTS postgres_db;") - run_sql_with_adapter(project.adapter, "DETACH DATABASE IF EXISTS postgres_db_custom;") + run_sql_with_adapter( + project.adapter, "DETACH DATABASE IF EXISTS postgres_db_custom;" + ) # Re-attach with a custom alias and additional settings db_host = os.getenv("DB_HOST") @@ -255,7 +269,9 @@ def test_postgres_plugin_attach_options(self, project): # Verify that the database is attached with the custom alias result = run_sql_with_adapter(project.adapter, "SHOW DATABASES;", fetch="all") aliases = [row[0] for row in result] - assert "postgres_db_custom" in aliases, "Database 'postgres_db_custom' should be attached." + assert ( + "postgres_db_custom" in aliases + ), "Database 'postgres_db_custom' should be attached." # Detach the database after validation run_sql_with_adapter(project.adapter, "DETACH DATABASE postgres_db_custom;") @@ -310,7 +326,7 @@ def test_postgres_plugin_pg_array_as_varchar(self, project): res = run_sql_with_adapter( project.adapter, "SELECT arr FROM postgres_db_array.public.array_table WHERE id = 1;", - fetch="one" + fetch="one", ) assert res[0] == [1, 2, 3], "The array should be fetched as a list [1, 2, 3]." @@ -321,10 +337,11 @@ def test_postgres_plugin_error_handling(self, project): # Try to attach without a DSN with pytest.raises(Exception) as exc_info: run_sql_with_adapter( - project.adapter, - "ATTACH '' AS postgres_db_error (TYPE POSTGRES);" + project.adapter, "ATTACH '' AS postgres_db_error (TYPE POSTGRES);" ) - assert "unable to connect to postgres" in str(exc_info.value).lower(), "Should raise exception for missing 'dsn'" + assert ( + "unable to connect to postgres" in str(exc_info.value).lower() + ), "Should raise exception for missing 'dsn'" # Provide a valid DSN for the next test db_host = os.getenv("DB_HOST") @@ -348,6 +365,8 @@ def test_postgres_plugin_error_handling(self, project): f""" ATTACH '{dsn}' AS postgres_db_error (TYPE POSTGRES); SET invalid_setting = 'true'; - """ + """, ) - assert "unrecognized configuration parameter" in str(exc_info.value).lower(), "Should raise syntax error for invalid setting" + assert ( + "unrecognized configuration parameter" in str(exc_info.value).lower() + ), "Should raise syntax error for invalid setting"