macrostrat.database 3.1.0__tar.gz → 3.1.2__tar.gz

{macrostrat_database-3.1.0 → macrostrat_database-3.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: macrostrat.database
-Version: 3.1.0
+Version: 3.1.2
 Summary: A SQLAlchemy-based database toolkit.
 Author: Daven Quinn
 Author-email: dev@davenquinn.com

{macrostrat_database-3.1.0 → macrostrat_database-3.1.2}/macrostrat/database/postgresql.py

@@ -60,20 +60,6 @@ def prefix_inserts(insert, compiler, **kw):
     return compiler.visit_insert(insert, **kw)
 
 
-_psycopg2_setup_was_run = ContextVar("psycopg2-setup-was-run", default=False)
-
-
-def _setup_psycopg2_wait_callback():
-    """Set up the wait callback for PostgreSQL connections. This allows for query cancellation with Ctrl-C."""
-    # TODO: we might want to do this only once on engine creation
-    # https://github.com/psycopg/psycopg2/issues/333
-    val = _psycopg2_setup_was_run.get()
-    if val:
-        return
-    psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select)
-    _psycopg2_setup_was_run.set(True)
-
-
 def table_exists(db: Database, table_name: str, schema: str = "public") -> bool:
     """Check if a table exists in a PostgreSQL database."""
     sql = """SELECT EXISTS (

{macrostrat_database-3.1.0 → macrostrat_database-3.1.2}/macrostrat/database/utils.py

@@ -6,6 +6,8 @@ from typing import IO, Union
 from warnings import warn
 
 from click import echo, secho
+from psycopg2.extensions import set_wait_callback
+from psycopg2.extras import wait_select
 from psycopg2.sql import SQL, Composable, Composed
 from sqlalchemy import MetaData, create_engine, text
 from sqlalchemy.engine import Connection, Engine
@@ -23,8 +25,6 @@ from sqlparse import format, split
 
 from macrostrat.utils import cmd, get_logger
 
-from .postgresql import _setup_psycopg2_wait_callback
-
 log = get_logger(__name__)
 
 
@@ -150,7 +150,7 @@ def _get_queries(sql, interpret_as_file=None):
     if isinstance(sql, Path):
         sql = sql.read_text()
 
-    return split(sql)
+    return split(format(sql, strip_comments=True))
 
 
 def _is_prebind_param(param):
@@ -239,8 +239,6 @@ def _run_sql(connectable, sql, params=None, **kwargs):
             yield from _run_sql(conn, sql, params, **kwargs)
             return
 
-    _setup_psycopg2_wait_callback()
-
     stop_on_error = kwargs.pop("stop_on_error", False)
     raise_errors = kwargs.pop("raise_errors", False)
     has_server_binds = kwargs.pop("has_server_binds", None)
@@ -266,6 +264,9 @@ def _run_sql(connectable, sql, params=None, **kwargs):
 
     for query, params in zip(queries, params):
         trans = None
+        # This only does something for postgresql, but it's harmless to run it for other engines
+        set_wait_callback(wait_select)
+
         try:
             trans = connectable.begin()
         except InvalidRequestError:
@@ -320,6 +321,8 @@ def _run_sql(connectable, sql, params=None, **kwargs):
             log.error(err)
             if raise_errors:
                 raise err
+        finally:
+            set_wait_callback(None)
 
 
 def run_sql_file(connectable, filename, params=None, **kwargs):

{macrostrat_database-3.1.0 → macrostrat_database-3.1.2}/pyproject.toml

@@ -3,7 +3,7 @@ authors = ["Daven Quinn <dev@davenquinn.com>"]
 description = "A SQLAlchemy-based database toolkit."
 name = "macrostrat.database"
 packages = [{ include = "macrostrat" }]
-version = "3.1.0"
+version = "3.1.2"
 
 [tool.poetry.dependencies]
 GeoAlchemy2 = "^0.14.0"