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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: harlequin-postgres
-Version: 0.2.2
+Version: 0.3.0
 Summary: A Harlequin adapter for Postgres.
 License: MIT
 Author: Ted Conbeer
@@ -11,7 +11,7 @@ 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: harlequin (>=1.20,<2.0)
 Requires-Dist: psycopg2-binary (>=2.9.9,<3.0.0)
 Description-Content-Type: text/markdown
 
@@ -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.3.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.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "harlequin-postgres"
-version = "0.2.2"
+version = "0.3.0"
 description = "A Harlequin adapter for Postgres."
 authors = ["Ted Conbeer <tconbeer@users.noreply.github.com>"]
 license = "MIT"
@@ -14,13 +14,13 @@ postgres = "harlequin_postgres:HarlequinPostgresAdapter"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
-harlequin = "^1.7"
+harlequin = "^1.20"
 psycopg2-binary = "^2.9.9"
 
 [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"
@@ -32,9 +32,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.3.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,10 +8,11 @@ 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.extensions import TRANSACTION_STATUS_IDLE, connection, cursor
 from psycopg2.pool import ThreadedConnectionPool
 from textual_fastdatatable.backend import AutoBackendType
 
@@ -62,23 +64,43 @@ class HarlequinPostgresConnection(HarlequinConnection):
         try:
             if conn_str and conn_str[0]:
                 self.pool: ThreadedConnectionPool = ThreadedConnectionPool(
-                    1, 5, dsn=conn_str[0], **options
+                    2, 5, dsn=conn_str[0], **options
                 )
             else:
-                self.pool = ThreadedConnectionPool(1, 5, **options)
+                self.pool = ThreadedConnectionPool(2, 5, **options)
         except Exception as e:
             raise HarlequinConnectionError(
                 msg=str(e), title="Harlequin could not connect to Postgres."
             ) from e
 
+        self._main_conn: connection = self.pool.getconn(key="main")
+        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 == TRANSACTION_STATUS_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 Exception as e:
             cur.close()
+            self.rollback()
             raise HarlequinQueryError(
                 msg=str(e),
                 title="Harlequin encountered an error while executing your query.",
@@ -89,8 +111,12 @@ class HarlequinPostgresConnection(HarlequinConnection):
             else:
                 cur.close()
                 return None
-        finally:
-            self.pool.putconn(conn, key="main")
+
+    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()
@@ -146,7 +172,30 @@ class HarlequinPostgresConnection(HarlequinConnection):
         completions = _get_completions(conn)
         self.pool.putconn(conn)
         return completions
-        ...
+
+    def close(self) -> None:
+        self.pool.putconn(self._main_conn, key="main")
+        self.pool.closeall()
+
+    @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()

{harlequin_postgres-0.2.2 → harlequin_postgres-0.3.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",
 )