from contextlib import contextmanager from dataclasses import dataclass from typing import Optional, Union from dbt.adapters.contracts.connection import AdapterResponse, Credentials from dbt.adapters.events.logging import AdapterLogger from dbt.adapters.events.types import TypeCodeNotFound from dbt.adapters.postgres.record import PostgresRecordReplayHandle from dbt.adapters.sql import SQLConnectionManager from dbt_common.exceptions import DbtDatabaseError, DbtRuntimeError from dbt_common.events.functions import warn_or_error from dbt_common.helper_types import Port from dbt_common.record import get_record_mode_from_env, RecorderMode from mashumaro.jsonschema.annotations import Maximum, Minimum import psycopg2 from typing_extensions import Annotated logger = AdapterLogger("Postgres") @dataclass class PostgresCredentials(Credentials): host: str user: str # Annotated is used by mashumaro for jsonschema generation port: Annotated[Port, Minimum(0), Maximum(65535)] password: str # on postgres the password is mandatory 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" retries: int = 1 _ALIASES = {"dbname": "database", "pass": "password"} @property def type(self): return "postgres" @property def unique_field(self): return self.host def _connection_keys(self): return ( "host", "port", "user", "database", "schema", "connect_timeout", "role", "search_path", "keepalives_idle", "sslmode", "sslcert", "sslkey", "sslrootcert", "application_name", "retries", ) class PostgresConnectionManager(SQLConnectionManager): TYPE = "postgres" @contextmanager def exception_handler(self, sql): try: yield except psycopg2.DatabaseError as e: logger.debug("Postgres error: {}".format(str(e))) try: self.rollback_if_open() except psycopg2.Error: logger.debug("Failed to release connection!") pass raise DbtDatabaseError(str(e).strip()) from e except Exception as e: logger.debug("Error running SQL: {}", sql) logger.debug("Rolling back transaction.") self.rollback_if_open() if isinstance(e, DbtRuntimeError): # during a sql query, an internal to dbt exception was raised. # this sounds a lot like a signal handler and probably has # useful information, so raise it without modification. raise raise DbtRuntimeError(e) from e @classmethod def open(cls, connection): if connection.state == "open": logger.debug("Connection is already open, skipping open.") return connection credentials = cls.get_credentials(connection.credentials) kwargs = {} # we don't want to pass 0 along to connect() as postgres will try to # call an invalid setsockopt() call (contrary to the docs). if credentials.keepalives_idle: kwargs["keepalives_idle"] = credentials.keepalives_idle # psycopg2 doesn't support search_path officially, # see https://github.com/psycopg/psycopg2/issues/465 search_path = credentials.search_path if search_path is not None and search_path != "": # see https://postgresql.org/docs/9.5/libpq-connect.html kwargs["options"] = "-c search_path={}".format(search_path.replace(" ", "\\ ")) if credentials.sslmode: kwargs["sslmode"] = credentials.sslmode if credentials.sslcert is not None: kwargs["sslcert"] = credentials.sslcert if credentials.sslkey is not None: kwargs["sslkey"] = credentials.sslkey if credentials.sslrootcert is not None: kwargs["sslrootcert"] = credentials.sslrootcert if credentials.application_name: kwargs["application_name"] = credentials.application_name def connect(): handle = None # In replay mode, we won't connect to a real database at all, while # in record and diff modes we do, but insert an intermediate handle # object which monitors native connection activity. rec_mode = get_record_mode_from_env() if rec_mode != RecorderMode.REPLAY: 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 rec_mode is not None: # If using the record/replay mechanism, regardless of mode, we # use a wrapper. handle = PostgresRecordReplayHandle(handle, connection) if credentials.role: handle.cursor().execute("set role {}".format(credentials.role)) 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, ] def exponential_backoff(attempt: int): return attempt * attempt return cls.retry_connection( connection, connect=connect, logger=logger, retry_limit=credentials.retries, retry_timeout=exponential_backoff, retryable_exceptions=retryable_exceptions, ) def cancel(self, connection): connection_name = connection.name try: pid = connection.handle.get_backend_pid() except psycopg2.InterfaceError as exc: # if the connection is already closed, not much to cancel! if "already closed" in str(exc): logger.debug(f"Connection {connection_name} was already closed") return # probably bad, re-raise it raise sql = "select pg_terminate_backend({})".format(pid) logger.debug("Cancelling query '{}' ({})".format(connection_name, pid)) _, cursor = self.add_query(sql) res = cursor.fetchone() logger.debug("Cancel query '{}': {}".format(connection_name, res)) def add_begin_query(self): pass @classmethod def get_credentials(cls, credentials): return credentials @classmethod def get_response(cls, cursor) -> AdapterResponse: message = str(cursor.statusmessage) rows = cursor.rowcount status_message_parts = message.split() if message is not None else [] status_messsage_strings = [part for part in status_message_parts if not part.isdigit()] code = " ".join(status_messsage_strings) return AdapterResponse(_message=message, code=code, rows_affected=rows) @classmethod def data_type_code_to_name(cls, type_code: Union[int, str]) -> str: if type_code in psycopg2.extensions.string_types: return psycopg2.extensions.string_types[type_code].name else: warn_or_error(TypeCodeNotFound(type_code=type_code)) return f"unknown type_code {type_code}"