Source code for sqltrack.engines.postgres

"""
This module provides functionality to enable using PostgreSQL as backend.
"""

from pathlib import Path

from ..engine import Engine
from ..json import jsonb

import psycopg as pg


DATA_DIR = (Path(__file__).parent).absolute()


# monkey patch psycopg cursor to add executescript method like sqlite
pg.Cursor.executescript = pg.Cursor.execute


class PostgresEngine(Engine):
    def __init__(self, config):
        # make a copy of the config so we can change it
        config = dict(config)
        # dbpath, engine are not a valid postgres connection option
        config.pop("dbpath", None)
        config.pop("engine")
        # pop schema and put as search path
        if "schema" in config:
            schema = config.pop("schema")
            self.schema = schema
            config["options"] = f"--search_path={schema}"
        # the finished DSN
        self._dsn = " ".join(f"{k}={pg.sql.quote(v)}" for k, v in config.items())
        self._oid_mapping = {}

    def data_dir(self):
        return DATA_DIR

    def connect(self, cursor_factory=pg.Cursor):
        return pg.connect(self._dsn, cursor_factory=cursor_factory)

    def map_type(self, client, typ):
        if not self._oid_mapping:
            with client.cursor() as cur:
                # get oid mapping
                cur.execute("SELECT oid, typname FROM pg_type;")
                for oid, typename in cur.fetchall():
                    self._oid_mapping[oid] = typename
        dumper = pg.adapters.get_dumper(typ, pg.adapt.PyFormat.AUTO)
        if not dumper or not dumper.oid:
            raise ValueError(f"cannot detect postgres type for {typ!r}")
        return self._oid_mapping[dumper.oid]