harlequin-postgres 0.2.2__tar.gz → 0.4.0__tar.gz

{harlequin_postgres-0.2.2 → harlequin_postgres-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: harlequin-postgres
-Version: 0.2.2
+Version: 0.4.0
 Summary: A Harlequin adapter for Postgres.
 License: MIT
 Author: Ted Conbeer
@@ -11,8 +11,8 @@ 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.7,<2.0)
-Requires-Dist: psycopg2-binary (>=2.9.9,<3.0.0)
+Requires-Dist: harlequin (>=1.20,<2.0)
+Requires-Dist: psycopg[binary,pool] (>=3.2,<4.0)
 Description-Content-Type: text/markdown
 
 # harlequin-postgres
@@ -85,4 +85,10 @@ Many more options are available; to see the full list, run:
 harlequin --help
 ```
 
+## Manual Transactions
+
+To use Manual transaction mode, click on the label in the Run Query Bar to toggle the transaction mode from Auto to Manual.
+
+## Further Documentation
+
 For more information, see the [Harlequin Docs](https://harlequin.sh/docs/postgres/index).

{harlequin_postgres-0.2.2 → harlequin_postgres-0.4.0}/README.md

@@ -68,4 +68,10 @@ Many more options are available; to see the full list, run:
 harlequin --help
 ```
 
+## Manual Transactions
+
+To use Manual transaction mode, click on the label in the Run Query Bar to toggle the transaction mode from Auto to Manual.
+
+## Further Documentation
+
 For more information, see the [Harlequin Docs](https://harlequin.sh/docs/postgres/index).

{harlequin_postgres-0.2.2 → harlequin_postgres-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "harlequin-postgres"
-version = "0.2.2"
+version = "0.4.0"
 description = "A Harlequin adapter for Postgres."
 authors = ["Ted Conbeer <tconbeer@users.noreply.github.com>"]
 license = "MIT"
@@ -14,16 +14,15 @@ postgres = "harlequin_postgres:HarlequinPostgresAdapter"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
-harlequin = "^1.7"
-psycopg2-binary = "^2.9.9"
+harlequin = "^1.20"
+psycopg = { version = "^3.2", extras = ["binary", "pool"]}
 
 [tool.poetry.group.dev.dependencies]
-ruff = "^0.1.6"
+ruff = "^0.5"
 pytest = "^7.4.3"
-mypy = "^1.7.0"
+mypy = "^1.10.0"
 pre-commit = "^3.5.0"
 importlib_metadata = { version = ">=4.6.0", python = "<3.10.0" }
-types-psycopg2 = "^2.9.21.16"
 pyinstrument = "^4.6.1"
 
 [build-system]
@@ -32,9 +31,11 @@ build-backend = "poetry.core.masonry.api"
 
 
 [tool.ruff]
-select = ["A", "B", "E", "F", "I"]
 target-version = "py38"
 
+[tool.ruff.lint]
+select = ["A", "B", "E", "F", "I"]
+
 [tool.mypy]
 python_version = "3.8"
 files = [

{harlequin_postgres-0.2.2 → harlequin_postgres-0.4.0}/src/harlequin_postgres/adapter.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from itertools import cycle
 from typing import Any, Sequence
 
 from harlequin import (
@@ -7,11 +8,14 @@ from harlequin import (
     HarlequinCompletion,
     HarlequinConnection,
     HarlequinCursor,
+    HarlequinTransactionMode,
 )
 from harlequin.catalog import Catalog, CatalogItem
 from harlequin.exception import HarlequinConnectionError, HarlequinQueryError
-from psycopg2.extensions import connection, cursor
-from psycopg2.pool import ThreadedConnectionPool
+from psycopg import Connection, Cursor, conninfo
+from psycopg.errors import QueryCanceled
+from psycopg.pq import TransactionStatus
+from psycopg_pool import ConnectionPool
 from textual_fastdatatable.backend import AutoBackendType
 
 from harlequin_postgres.cli_options import POSTGRES_OPTIONS
@@ -19,16 +23,19 @@ from harlequin_postgres.completions import _get_completions
 
 
 class HarlequinPostgresCursor(HarlequinCursor):
-    def __init__(self, conn: HarlequinPostgresConnection, cur: cursor) -> None:
+    def __init__(self, conn: HarlequinPostgresConnection, cur: Cursor) -> None:
         self.conn = conn
         self.cur = cur
+        # we need to copy the description from the cursor in case the results are
+        # fetched and the cursor is closed before columns() is called.
+        assert cur.description is not None
+        self.description = cur.description.copy()
         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
+            for col in self.description
         ]
 
     def set_limit(self, limit: int) -> HarlequinPostgresCursor:
@@ -41,6 +48,8 @@ class HarlequinPostgresCursor(HarlequinCursor):
                 return self.cur.fetchall()
             else:
                 return self.cur.fetchmany(self._limit)
+        except QueryCanceled:
+            return []
         except Exception as e:
             raise HarlequinQueryError(
                 msg=str(e),
@@ -60,25 +69,73 @@ class HarlequinPostgresConnection(HarlequinConnection):
     ) -> None:
         self.init_message = init_message
         try:
-            if conn_str and conn_str[0]:
-                self.pool: ThreadedConnectionPool = ThreadedConnectionPool(
-                    1, 5, dsn=conn_str[0], **options
-                )
-            else:
-                self.pool = ThreadedConnectionPool(1, 5, **options)
+            self.conn_info = conninfo.conninfo_to_dict(
+                conninfo=conn_str[0] if conn_str else "", **options
+            )
+        except Exception as e:
+            raise HarlequinConnectionError(
+                msg=str(e),
+                title=(
+                    "Harlequin could not connect to Postgres. "
+                    "Invalid connection string."
+                ),
+            ) from e
+        try:
+            raw_timeout = self.conn_info.get("connect_timeout")
+            timeout = float(raw_timeout) if raw_timeout is not None else 30.0
+        except (TypeError, ValueError) as e:
+            raise HarlequinConnectionError(
+                msg=str(e),
+                title=(
+                    "Harlequin could not connect to Postgres. "
+                    "Invalid value for connection_timeout."
+                ),
+            ) from e
+        try:
+            self.pool: ConnectionPool = ConnectionPool(
+                conninfo=conn_str[0] if conn_str and conn_str[0] else "",
+                min_size=2,
+                max_size=5,
+                kwargs=options,
+                open=True,
+                timeout=timeout,
+            )
+            self._main_conn: Connection = self.pool.getconn()
         except Exception as e:
             raise HarlequinConnectionError(
                 msg=str(e), title="Harlequin could not connect to Postgres."
             ) from e
 
+        self._transaction_modes = cycle(
+            [
+                HarlequinTransactionMode(label="Auto"),
+                HarlequinTransactionMode(
+                    label="Manual",
+                    commit=self.commit,
+                    rollback=self.rollback,
+                ),
+            ]
+        )
+        self.toggle_transaction_mode()
+
     def execute(self, query: str) -> HarlequinCursor | None:
+        if (
+            self.transaction_mode.label != "Auto"
+            and self._main_conn.info.transaction_status == TransactionStatus.IDLE
+        ):
+            cur = self._main_conn.cursor()
+            cur.execute(query="begin;")
+            cur.close()
+
         try:
-            conn: connection = self.pool.getconn(key="main")
-            with conn:  # autocommit transaction
-                cur = conn.cursor()
-                cur.execute(query=query)
+            cur = self._main_conn.cursor()
+            cur.execute(query=query)
+        except QueryCanceled:
+            cur.close()
+            return None
         except Exception as e:
             cur.close()
+            self.rollback()
             raise HarlequinQueryError(
                 msg=str(e),
                 title="Harlequin encountered an error while executing your query.",
@@ -89,8 +146,15 @@ class HarlequinPostgresConnection(HarlequinConnection):
             else:
                 cur.close()
                 return None
-        finally:
-            self.pool.putconn(conn, key="main")
+            
+    def cancel(self) -> None:
+        self._main_conn.cancel_safe()
+
+    def commit(self) -> None:
+        self._main_conn.commit()
+
+    def rollback(self) -> None:
+        self._main_conn.rollback()
 
     def get_catalog(self) -> Catalog:
         databases = self._get_databases()
@@ -142,14 +206,37 @@ class HarlequinPostgresConnection(HarlequinConnection):
         return Catalog(items=db_items)
 
     def get_completions(self) -> list[HarlequinCompletion]:
-        conn: connection = self.pool.getconn()
+        conn: Connection = self.pool.getconn()
         completions = _get_completions(conn)
         self.pool.putconn(conn)
         return completions
-        ...
+
+    def close(self) -> None:
+        self.pool.putconn(self._main_conn)
+        self.pool.close()
+
+    @property
+    def transaction_mode(self) -> HarlequinTransactionMode:
+        return self._transaction_mode
+
+    def toggle_transaction_mode(self) -> HarlequinTransactionMode:
+        self._transaction_mode = next(self._transaction_modes)
+        self._sync_transaction_mode()
+        return self._transaction_mode
+
+    def _sync_transaction_mode(self) -> None:
+        """
+        Sync this class's transaction mode with the main connection
+        """
+        conn = self._main_conn
+        if self.transaction_mode.label == "Auto":
+            conn.autocommit = True
+            conn.commit()
+        else:
+            conn.autocommit = False
 
     def _get_databases(self) -> list[tuple[str]]:
-        conn: connection = self.pool.getconn()
+        conn: Connection = self.pool.getconn()
         with conn.cursor() as cur:
             cur.execute(
                 """
@@ -166,7 +253,7 @@ class HarlequinPostgresConnection(HarlequinConnection):
         return results
 
     def _get_schemas(self, dbname: str) -> list[tuple[str]]:
-        conn: connection = self.pool.getconn()
+        conn: Connection = self.pool.getconn()
         with conn.cursor() as cur:
             cur.execute(
                 f"""
@@ -184,7 +271,7 @@ class HarlequinPostgresConnection(HarlequinConnection):
         return results
 
     def _get_relations(self, dbname: str, schema: str) -> list[tuple[str, str]]:
-        conn: connection = self.pool.getconn()
+        conn: Connection = self.pool.getconn()
         with conn.cursor() as cur:
             cur.execute(
                 f"""
@@ -203,7 +290,7 @@ class HarlequinPostgresConnection(HarlequinConnection):
     def _get_columns(
         self, dbname: str, schema: str, relation: str
     ) -> list[tuple[str, str]]:
-        conn: connection = self.pool.getconn()
+        conn: Connection = self.pool.getconn()
         with conn.cursor() as cur:
             cur.execute(
                 f"""
@@ -339,6 +426,7 @@ class HarlequinPostgresConnection(HarlequinConnection):
 
 class HarlequinPostgresAdapter(HarlequinAdapter):
     ADAPTER_OPTIONS = POSTGRES_OPTIONS
+    IMPLEMENTS_CANCEL = True
 
     def __init__(
         self,
@@ -351,14 +439,14 @@ class HarlequinPostgresAdapter(HarlequinAdapter):
         passfile: str | None = None,
         require_auth: str | None = None,
         channel_binding: str | None = None,
-        connect_timeout: int | None = None,
+        connect_timeout: int | float | None = None,
         sslmode: str | None = None,
         sslcert: str | None = None,
         sslkey: str | None = None,
         **_: Any,
     ) -> None:
         self.conn_str = conn_str
-        self.options = {
+        self.options: dict[str, str | int | None] = {
             "host": host,
             "port": port,
             "dbname": dbname,
@@ -367,12 +455,30 @@ class HarlequinPostgresAdapter(HarlequinAdapter):
             "passfile": passfile,
             "require_auth": require_auth,
             "channel_binding": channel_binding,
-            "connect_timeout": connect_timeout,
+            "connect_timeout": connect_timeout,  # type: ignore[dict-item]
             "sslmode": sslmode,
             "sslcert": sslcert,
             "sslkey": sslkey,
         }
 
+    @property
+    def connection_id(self) -> str | None:
+        """
+        Use a simplified connection string, with only the host, port, and database
+        """
+        try:
+            conn_info = conninfo.conninfo_to_dict(
+                conninfo=self.conn_str[0] if self.conn_str else "",
+                **self.options,
+            )
+        except Exception:
+            return None
+
+        host = conn_info.get("host", "localhost")
+        port = conn_info.get("port", "5432")
+        dbname = conn_info.get("dbname", "postgres")
+        return f"{host}:{port}/{dbname}"
+
     def connect(self) -> HarlequinPostgresConnection:
         if len(self.conn_str) > 1:
             raise HarlequinConnectionError(

{harlequin_postgres-0.2.2 → harlequin_postgres-0.4.0}/src/harlequin_postgres/cli_options.py

@@ -33,10 +33,7 @@ port = TextOption(
 
 dbname = TextOption(
     name="dbname",
-    description=(
-        "Port number to connect to at the server host, or socket file name extension "
-        "for Unix-domain connections."
-    ),
+    description=("The database name to use when connecting with the Postgres server."),
     short_decls=["-d"],
     default="postgres",
 )

{harlequin_postgres-0.2.2 → harlequin_postgres-0.4.0}/src/harlequin_postgres/completions.py

@@ -4,10 +4,10 @@ import csv
 from pathlib import Path
 
 from harlequin import HarlequinCompletion
-from psycopg2.extensions import connection
+from psycopg import Connection
 
 
-def _get_completions(conn: connection) -> list[HarlequinCompletion]:
+def _get_completions(conn: Connection) -> list[HarlequinCompletion]:
     completions: list[HarlequinCompletion] = []
 
     # source: https://www.postgresql.org/docs/current/sql-keywords-appendix.html