harlequin-postgres 0.1.2__py3-none-any.whl

harlequin_postgres/__init__.py

@@ -0,0 +1,3 @@
+from harlequin_postgres.adapter import HarlequinPostgresAdapter
+
+__all__ = ["HarlequinPostgresAdapter"]

harlequin_postgres/adapter.py

@@ -0,0 +1,352 @@
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from harlequin.adapter import HarlequinAdapter, HarlequinConnection, HarlequinCursor
+from harlequin.catalog import Catalog, CatalogItem
+from harlequin.exception import HarlequinConnectionError, HarlequinQueryError
+from psycopg2 import connect
+from psycopg2.extensions import connection, cursor
+from textual_fastdatatable.backend import AutoBackendType
+
+from harlequin_postgres.cli_options import POSTGRES_OPTIONS
+
+
+class HarlequinPostgresCursor(HarlequinCursor):
+    def __init__(self, conn: HarlequinPostgresConnection, cur: cursor) -> None:
+        self.conn = conn
+        self.cur = cur
+        self._limit: int | None = None
+
+    def columns(self) -> list[tuple[str, str]]:
+        assert self.cur.description is not None
+        return [
+            (col.name, self.conn._get_short_type_from_oid(col.type_code))
+            for col in self.cur.description
+        ]
+
+    def set_limit(self, limit: int) -> HarlequinPostgresCursor:
+        self._limit = limit
+        return self
+
+    def fetchall(self) -> AutoBackendType:
+        try:
+            if self._limit is None:
+                return self.cur.fetchall()
+            else:
+                return self.cur.fetchmany(self._limit)
+        except Exception as e:
+            raise HarlequinQueryError(
+                msg=str(e),
+                title="Harlequin encountered an error while executing your query.",
+            ) from e
+        finally:
+            self.cur.close()
+
+
+class HarlequinPostgresConnection(HarlequinConnection):
+    def __init__(
+        self,
+        conn_str: Sequence[str],
+        *_: Any,
+        init_message: str = "",
+        options: dict[str, Any],
+    ) -> None:
+        self.init_message = init_message
+        try:
+            if conn_str and conn_str[0]:
+                self.conn: connection = connect(dsn=conn_str[0], **options)
+            else:
+                self.conn = connect(**options)
+        except Exception as e:
+            raise HarlequinConnectionError(
+                msg=str(e), title="Harlequin could not connect to Postgres."
+            ) from e
+
+    def execute(self, query: str) -> HarlequinCursor | None:
+        try:
+            with self.conn:  # autocommit transaction
+                cur = self.conn.cursor()
+                cur.execute(query=query)
+        except Exception as e:
+            raise HarlequinQueryError(
+                msg=str(e),
+                title="Harlequin encountered an error while executing your query.",
+            ) from e
+        else:
+            if cur.description is not None:
+                return HarlequinPostgresCursor(self, cur)
+            else:
+                cur.close()
+                return None
+
+    def get_catalog(self) -> Catalog:
+        databases = self._get_databases()
+        db_items: list[CatalogItem] = []
+        for (db,) in databases:
+            schemas = self._get_schemas(db)
+            schema_items: list[CatalogItem] = []
+            for (schema,) in schemas:
+                relations = self._get_relations(db, schema)
+                rel_items: list[CatalogItem] = []
+                for rel, rel_type in relations:
+                    cols = self._get_columns(db, schema, rel)
+                    col_items = [
+                        CatalogItem(
+                            qualified_identifier=f'"{db}"."{schema}"."{rel}"."{col}"',
+                            query_name=f'"{col}"',
+                            label=col,
+                            type_label=self._get_short_type(col_type),
+                        )
+                        for col, col_type in cols
+                    ]
+                    rel_items.append(
+                        CatalogItem(
+                            qualified_identifier=f'"{db}"."{schema}"."{rel}"',
+                            query_name=f'"{db}"."{schema}"."{rel}"',
+                            label=rel,
+                            type_label="v" if rel_type == "VIEW" else "t",
+                            children=col_items,
+                        )
+                    )
+                schema_items.append(
+                    CatalogItem(
+                        qualified_identifier=f'"{db}"."{schema}"',
+                        query_name=f'"{db}"."{schema}"',
+                        label=schema,
+                        type_label="s",
+                        children=rel_items,
+                    )
+                )
+            db_items.append(
+                CatalogItem(
+                    qualified_identifier=f'"{db}"',
+                    query_name=f'"{db}"',
+                    label=db,
+                    type_label="db",
+                    children=schema_items,
+                )
+            )
+        return Catalog(items=db_items)
+
+    def _get_databases(self) -> list[tuple[str]]:
+        with self.conn.cursor() as cur:
+            cur.execute(
+                """
+                select datname
+                from pg_database
+                where
+                    datistemplate is false
+                    and datallowconn is true
+                ;"""
+            )
+            results = cur.fetchall()
+        return results
+
+    def _get_schemas(self, dbname: str) -> list[tuple[str]]:
+        with self.conn.cursor() as cur:
+            cur.execute(
+                f"""
+                select schema_name
+                from information_schema.schemata
+                where
+                    catalog_name = '{dbname}'
+                    and schema_name != 'information_schema'
+                    and schema_name not like 'pg_%'
+                ;"""
+            )
+            results = cur.fetchall()
+        return results
+
+    def _get_relations(self, dbname: str, schema: str) -> list[tuple[str, str]]:
+        with self.conn.cursor() as cur:
+            cur.execute(
+                f"""
+                select table_name, table_type
+                from information_schema.tables
+                where
+                    table_catalog = '{dbname}'
+                    and table_schema = '{schema}'
+                ;"""
+            )
+            results = cur.fetchall()
+        return results
+
+    def _get_columns(
+        self, dbname: str, schema: str, relation: str
+    ) -> list[tuple[str, str]]:
+        with self.conn.cursor() as cur:
+            cur.execute(
+                f"""
+                select column_name, data_type
+                from information_schema.columns
+                where
+                    table_catalog = '{dbname}'
+                    and table_schema = '{schema}'
+                    and table_name = '{relation}'
+                ;"""
+            )
+            results = cur.fetchall()
+        return results
+
+    @staticmethod
+    def _get_short_type(type_name: str) -> str:
+        MAPPING = {
+            "bigint": "##",
+            "bigserial": "##",
+            "bit": "010",
+            "boolean": "t/f",
+            "box": "□",
+            "bytea": "b",
+            "character": "s",
+            "cidr": "ip",
+            "circle": "○",
+            "date": "d",
+            "double": "#.#",
+            "inet": "ip",
+            "integer": "#",
+            "interval": "|-|",
+            "json": "{}",
+            "jsonb": "b{}",
+            "line": "—",
+            "lseg": "-",
+            "macaddr": "mac",
+            "macaddr8": "mac",
+            "money": "$$",
+            "numeric": "#.#",
+            "path": "╭",
+            "pg_lsn": "lsn",
+            "pg_snapshot": "snp",
+            "point": "•",
+            "polygon": "▽",
+            "real": "#.#",
+            "smallint": "#",
+            "smallserial": "#",
+            "serial": "#",
+            "text": "s",
+            "time": "t",
+            "timestamp": "ts",
+            "tsquery": "tsq",
+            "tsvector": "tsv",
+            "txid_snapshot": "snp",
+            "uuid": "uid",
+            "xml": "xml",
+            "array": "[]",
+        }
+        return MAPPING.get(type_name.split("(")[0].split(" ")[0], "?")
+
+    @staticmethod
+    def _get_short_type_from_oid(oid: int) -> str:
+        MAPPING = {
+            16: "t/f",
+            17: "b",
+            18: "s",
+            19: "s",
+            20: "##",
+            21: "#",
+            22: "[#]",
+            23: "#",
+            25: "s",
+            26: "oid",
+            114: "{}",
+            142: "xml",
+            600: "•",
+            601: "-",
+            602: "╭",
+            603: "□",
+            604: "▽",
+            628: "—",
+            651: "[ip]",
+            700: "#.#",
+            701: "#.#",
+            704: "|-|",
+            718: "○",
+            790: "$$",
+            829: "mac",
+            869: "ip",
+            650: "ip",
+            774: "mac",
+            1000: "[t/f]",
+            1001: "[b]",
+            1002: "[s]",
+            1003: "[s]",
+            1009: "[s]",
+            1013: "[oid]",
+            1014: "[s]",
+            1015: "[s]",
+            1016: "[#]",
+            1021: "[#.#]",
+            1022: "[#.#]",
+            1028: "[oid]",
+            1040: "[mac]",
+            1041: "[ip]",
+            1042: "s",
+            1043: "s",
+            1082: "d",
+            1083: "t",
+            1114: "ts",
+            1115: "[ts]",
+            1182: "[d]",
+            1183: "[t]",
+            1184: "ts",
+            1185: "[ts]",
+            1186: "|-|",
+            1187: "[|-|]",
+            1231: "[#.#]",
+            1266: "t",
+            1270: "[t]",
+            1560: "010",
+            1562: "010",
+            1700: "#.#",
+            2950: "uid",
+            3614: "tsv",
+            3615: "tsq",
+            3802: "b{}",
+        }
+        return MAPPING.get(oid, "?")
+
+
+class HarlequinPostgresAdapter(HarlequinAdapter):
+    ADAPTER_OPTIONS = POSTGRES_OPTIONS
+
+    def __init__(
+        self,
+        conn_str: Sequence[str],
+        host: str | None = None,
+        port: str | None = None,
+        dbname: str | None = None,
+        user: str | None = None,
+        password: str | None = None,
+        passfile: str | None = None,
+        require_auth: str | None = None,
+        channel_binding: str | None = None,
+        connect_timeout: int | None = None,
+        sslmode: str | None = None,
+        sslcert: str | None = None,
+        sslkey: str | None = None,
+        **_: Any,
+    ) -> None:
+        self.conn_str = conn_str
+        self.options = {
+            "host": host,
+            "port": port,
+            "dbname": dbname,
+            "user": user,
+            "password": password,
+            "passfile": passfile,
+            "require_auth": require_auth,
+            "channel_binding": channel_binding,
+            "connect_timeout": connect_timeout,
+            "sslmode": sslmode,
+            "sslcert": sslcert,
+            "sslkey": sslkey,
+        }
+
+    def connect(self) -> HarlequinPostgresConnection:
+        if len(self.conn_str) > 1:
+            raise HarlequinConnectionError(
+                "Cannot provide multiple connection strings to the Postgres adapter. "
+                f"{self.conn_str}"
+            )
+        conn = HarlequinPostgresConnection(self.conn_str, options=self.options)
+        return conn

harlequin_postgres/cli_options.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+from harlequin.options import (
+    FlagOption,  # noqa
+    ListOption,  # noqa
+    PathOption,  # noqa
+    SelectOption,  # noqa
+    TextOption,
+)
+
+host = TextOption(
+    name="host",
+    description=(
+        "Specifies the host name of the machine on which the server is running. "
+        "If the value begins with a slash, it is used as the directory for the "
+        "Unix-domain socket."
+    ),
+    short_decls=["-h"],
+    default="localhost",
+)
+
+
+port = TextOption(
+    name="port",
+    description=(
+        "Port number to connect to at the server host, or socket file name extension "
+        "for Unix-domain connections."
+    ),
+    short_decls=["-p"],
+    default="5432",
+)
+
+
+dbname = TextOption(
+    name="dbname",
+    description=(
+        "Port number to connect to at the server host, or socket file name extension "
+        "for Unix-domain connections."
+    ),
+    short_decls=["-d"],
+    default="postgres",
+)
+
+
+user = TextOption(
+    name="user",
+    description=("PostgreSQL user name to connect as."),
+    short_decls=["-u", "--username", "-U"],
+)
+
+
+password = TextOption(
+    name="password",
+    description=("Password to be used if the server demands password authentication."),
+)
+
+
+passfile = PathOption(
+    name="passfile",
+    description=(
+        "Specifies the name of the file used to store passwords. Defaults to "
+        "~/.pgpass, or %APPDATA%\postgresql\pgpass.conf on Windows. (No error is "
+        "reported if this file does not exist.)"
+    ),
+    resolve_path=True,
+    exists=False,
+    file_okay=True,
+    dir_okay=False,
+)
+
+require_auth = SelectOption(
+    name="require_auth",
+    description=(
+        "Specifies the authentication method that the client requires from the server. "
+        "If the server does not use the required method to authenticate the client, or "
+        "if the authentication handshake is not fully completed by the server, the "
+        "connection will fail."
+    ),
+    choices=["password", "md5", "gss", "sspi", "scram-sha-256", "none"],
+)
+
+channel_binding = SelectOption(
+    name="channel_binding",
+    description=(
+        "This option controls the client's use of channel binding. A setting of "
+        "require means that the connection must employ channel binding, prefer "
+        "means that the client will choose channel binding if available, and "
+        "disable prevents the use of channel binding. The default is prefer if "
+        "PostgreSQL is compiled with SSL support; otherwise the default is disable."
+    ),
+    choices=["require", "prefer", "disable"],
+)
+
+
+def _int_validator(s: str | None) -> tuple[bool, str]:
+    if s is None:
+        return True, ""
+    try:
+        _ = int(s)
+    except ValueError:
+        return False, f"Cannot convert {s} to an int!"
+    else:
+        return True, ""
+
+
+connect_timeout = TextOption(
+    name="connect_timeout",
+    description=(
+        "Maximum time to wait while connecting, in seconds (write as an integer, "
+        "e.g., 10)."
+    ),
+    validator=_int_validator,
+)
+
+sslmode = SelectOption(
+    name="sslmode",
+    description=(
+        "Determines whether or with what priority a secure SSL TCP/IP connection will "
+        "be negotiated with the server."
+    ),
+    choices=["disable", "allow", "prefer", "require", "verify-ca", "verify-full"],
+    default="prefer",
+)
+
+sslcert = PathOption(
+    name="sslcert",
+    description=(
+        "Specifies the file name of the client SSL certificate. "
+        "Ignored if an SSL connection is not made."
+    ),
+    default="~/.postgresql/postgresql.crt",
+)
+
+sslkey = TextOption(
+    name="sslkey",
+    description=(
+        "Specifies the location for the secret key used for the client certificate. "
+        "It can either specify a file name that will be used instead of the default "
+        "~/.postgresql/postgresql.key, or it can specify a key obtained from an "
+        "external engine. An external engine specification should consist of a "
+        "colon-separated engine name and an engine-specific key identifier. This "
+        "parameter is ignored if an SSL connection is not made."
+    ),
+)
+
+
+POSTGRES_OPTIONS = [
+    host,
+    port,
+    dbname,
+    user,
+    password,
+    passfile,
+    require_auth,
+    channel_binding,
+    connect_timeout,
+    sslmode,
+    sslcert,
+    sslkey,
+]

harlequin_postgres-0.1.2.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Ted Conbeer
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

harlequin_postgres-0.1.2.dist-info/METADATA

@@ -0,0 +1,88 @@
+Metadata-Version: 2.1
+Name: harlequin-postgres
+Version: 0.1.2
+Summary: A Harlequin adapter for Postgres.
+License: MIT
+Author: Ted Conbeer
+Author-email: tconbeer@users.noreply.github.com
+Requires-Python: >=3.8.1,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: harlequin (>=1.4,<2.0)
+Requires-Dist: psycopg2-binary (>=2.9.9,<3.0.0)
+Description-Content-Type: text/markdown
+
+# harlequin-postgres
+
+This repo provides the Harlequin adapter for Postgres.
+
+## Installation
+
+`harlequin-postgres` depends on `harlequin`, so installing this package will also install Harlequin.
+
+### Using pip
+
+To install this adapter into an activated virtual environment:
+```bash
+pip install harlequin-postgres
+```
+
+### Using poetry
+
+```bash
+poetry add harlequin-postgres
+```
+
+### Using pipx
+
+If you do not already have Harlequin installed:
+
+```bash
+pip install harlequin-postgres
+```
+
+If you would like to add the Postgres adapter to an existing Harlequin installation:
+
+```bash
+pipx inject harlequin harlequin-postgres
+```
+
+### As an Extra
+Alternatively, you can install Harlequin with the `postgres` extra:
+
+```bash
+pip install harlequin[postgres]
+```
+
+```bash
+poetry add harlequin[postgres]
+```
+
+```bash
+pipx install harlequin[postgres]
+```
+
+## Usage and Configuration
+
+You can open Harlequin with the Postgres adapter by selecting it with the `-a` option and passing a [Posgres DSN](https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING):
+
+```bash
+harlequin -a postgres "postgres://my-user:my-pass@localhost:5432/my-database"
+```
+
+You can also pass all or parts of the connection string as separate options. The following is equivalent to the above DSN:
+
+```bash
+harlequin -a postgres -h localhost -p 5432 -U my-user --password my-pass -d my-database
+```
+
+Many more options are available; to see the full list, run:
+
+```bash
+harlequin --help
+```
+
+For more information, see the [Harlequin Docs](https://harlequin.sh/docs/postgres/index).

harlequin_postgres-0.1.2.dist-info/RECORD

@@ -0,0 +1,9 @@
+harlequin_postgres/__init__.py,sha256=ZUrCa_NVFJRbvxKXuzBhW2VS1RlR4hTJFL210L0r-x8,104
+harlequin_postgres/adapter.py,sha256=WbAnZWkQmDJ6YOD_TFsFzzjdlW0F6SRDtFzTEG4Ejqk,11202
+harlequin_postgres/cli_options.py,sha256=xHgSrN_BFmODCdDLUEYvRQ_X8salEqzhTGhQ3CE2-n0,4344
+harlequin_postgres/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+harlequin_postgres-0.1.2.dist-info/LICENSE,sha256=Wdf1mXEX48JTe3Unchs3u8M7sUOB6v0fGsMI5WP7bB4,1068
+harlequin_postgres-0.1.2.dist-info/METADATA,sha256=krvCYcA0Uwj-JJFX_Xa0Da9llKZRRbqgTLdNXljBMkc,2161
+harlequin_postgres-0.1.2.dist-info/WHEEL,sha256=d2fvjOD7sXsVzChCqf0Ty0JbHKBaLYwDbGQDwQTnJ50,88
+harlequin_postgres-0.1.2.dist-info/entry_points.txt,sha256=wGOeVVp9meAYkVPIny5_dApH0h6Wg3ktne7TlRNnxyI,74
+harlequin_postgres-0.1.2.dist-info/RECORD,,

harlequin_postgres-0.1.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 1.7.0
+Root-Is-Purelib: true
+Tag: py3-none-any

harlequin_postgres-0.1.2.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[harlequin.adapter]
+postgres=harlequin_postgres:HarlequinPostgresAdapter
+