macrostrat.database 3.5.4__tar.gz → 4.0.0__tar.gz

{macrostrat_database-3.5.4 → macrostrat_database-4.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: macrostrat.database
-Version: 3.5.4
+Version: 4.0.0
 Summary: A SQLAlchemy-based database toolkit.
 Author: Daven Quinn
 Author-email: dev@davenquinn.com
@@ -9,12 +9,14 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: GeoAlchemy2 (>=0.15.2,<0.16.0)
 Requires-Dist: SQLAlchemy (>=2.0.18,<3.0.0)
 Requires-Dist: SQLAlchemy-Utils (>=0.41.1,<0.42.0)
 Requires-Dist: aiofiles (>=23.2.1,<24.0.0)
 Requires-Dist: click (>=8.1.3,<9.0.0)
 Requires-Dist: macrostrat.utils (>=1.3.0,<2.0.0)
-Requires-Dist: psycopg2-binary (>=2.9.6,<3.0.0)
+Requires-Dist: psycopg (>=3.2.1,<4.0.0)
+Requires-Dist: psycopg2 (>=2.9.11,<3.0.0)
 Requires-Dist: rich (>=13.7.1,<14.0.0)
 Requires-Dist: sqlparse (>=0.5.1,<0.6.0)

{macrostrat_database-3.5.4 → macrostrat_database-4.0.0}/macrostrat/database/__init__.py

@@ -3,16 +3,15 @@ from contextlib import contextmanager
 from pathlib import Path
 from typing import Optional, Union
 
-from psycopg2.errors import InvalidSavepointSpecification
-from psycopg2.sql import Identifier
-from sqlalchemy import URL, Engine, MetaData, create_engine, inspect
-from sqlalchemy.exc import IntegrityError, InternalError
+from psycopg.errors import InvalidSavepointSpecification
+from psycopg.sql import Identifier
+from sqlalchemy import URL, Engine, MetaData, inspect
+from sqlalchemy.exc import IntegrityError, OperationalError
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 from sqlalchemy.sql.expression import Insert
 
 from macrostrat.utils import get_logger
-
 from .mapper import DatabaseMapper
 from .postgresql import on_conflict, prefix_inserts  # noqa
 from .utils import (  # noqa
@@ -25,6 +24,7 @@ from .utils import (  # noqa
     run_fixtures,
     run_query,
     run_sql,
+    create_engine,
 )
 
 metadata = MetaData()
@@ -60,12 +60,8 @@ class Database(object):
 
         self.instance_params = kwargs.pop("instance_params", {})
 
-        if isinstance(db_conn, Engine):
-            log.info(f"Set up database connection with engine {db_conn.url}")
-            self.engine = db_conn
-        else:
-            log.info(f"Setting up database connection with URL '{db_conn}'")
-            self.engine = create_engine(db_conn, echo=echo_sql, **kwargs)
+        self.engine = create_engine(db_conn, echo=echo_sql, **kwargs)
+
         self.metadata = kwargs.get("metadata", metadata)
 
         # Scoped session for database
@@ -334,6 +330,7 @@ class Database(object):
 
         if connection is None:
             connection = self.session.connection()
+
         params = {"name": Identifier(name)}
         run_query(connection, "SAVEPOINT {name}", params)
         should_rollback = rollback == "always"
@@ -356,7 +353,7 @@ def _clear_savepoint(connection, name, rollback=True):
             run_query(connection, "ROLLBACK TO SAVEPOINT {name}", params)
         else:
             run_query(connection, "RELEASE SAVEPOINT {name}", params)
-    except InternalError as err:
+    except OperationalError as err:
         if isinstance(err.orig, InvalidSavepointSpecification):
             log.warning(
                 f"Savepoint {name} does not exist; we may have already rolled back."

{macrostrat_database-3.5.4 → macrostrat_database-4.0.0}/macrostrat/database/postgresql.py

@@ -4,11 +4,11 @@ from contextlib import contextmanager
 from contextvars import ContextVar
 from typing import TYPE_CHECKING
 
-import psycopg2
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.exc import CompileError
 from sqlalchemy.ext.compiler import compiles
-from sqlalchemy.sql.expression import Insert, text
+from sqlalchemy.sql.dml import Insert
+from sqlalchemy.sql.expression import text
 
 if TYPE_CHECKING:
     from ..database import Database
@@ -26,9 +26,10 @@ def on_conflict(action="restrict"):
         _insert_mode.reset(token)
 
 
-# @compiles(Insert, "postgresql")
+@compiles(Insert, "postgresql")
 def prefix_inserts(insert, compiler, **kw):
     """Conditionally adapt insert statements to use on-conflict resolution (a PostgreSQL feature)"""
+
     if insert._post_values_clause is not None:
         return compiler.visit_insert(insert, **kw)
 
@@ -63,7 +64,7 @@ def prefix_inserts(insert, compiler, **kw):
 def table_exists(db: Database, table_name: str, schema: str = "public") -> bool:
     """Check if a table exists in a PostgreSQL database."""
     sql = """SELECT EXISTS (
-        SELECT FROM information_schema.tables 
+        SELECT FROM information_schema.tables
         WHERE table_schema = :schema
           AND table_name = :table_name
     );"""

{macrostrat_database-3.5.4 → macrostrat_database-4.0.0}/macrostrat/database/utils.py

@@ -4,18 +4,17 @@ from enum import Enum
 from pathlib import Path
 from re import search
 from sys import stderr
-from time import sleep
 from typing import IO, Union
 from warnings import warn
 
-import psycopg2.errors
 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 psycopg.errors import QueryCanceled
+from psycopg.sql import SQL, Composable, Composed
 from rich.console import Console
-from sqlalchemy import MetaData, create_engine, text
+from sqlalchemy import MetaData, text
+from sqlalchemy import create_engine as base_create_engine
 from sqlalchemy.engine import Connection, Engine
+from sqlalchemy.engine.url import make_url
 from sqlalchemy.exc import (
     IntegrityError,
     InternalError,
@@ -29,6 +28,7 @@ from sqlalchemy.sql.elements import ClauseElement, TextClause
 from sqlalchemy_utils import create_database as _create_database
 from sqlalchemy_utils import database_exists, drop_database
 from sqlparse import format, split
+from time import sleep
 
 from macrostrat.utils import cmd, get_logger
 
@@ -196,6 +196,8 @@ def _get_cursor(connectable):
     while hasattr(conn, "driver_connection") or hasattr(conn, "connection"):
         if hasattr(conn, "driver_connection"):
             conn = conn.driver_connection
+        elif conn.connection == conn:
+            break
         else:
             conn = conn.connection
         if callable(conn):
@@ -296,7 +298,7 @@ def _run_sql(connectable, sql, params=None, **kwargs):
         if pre_bind_params is not None:
             if not isinstance(query, SQL):
                 query = SQL(query)
-            # Pre-bind the parameters using PsycoPG2
+            # Pre-bind the parameters using psycopg
             query = query.format(**pre_bind_params)
 
         if isinstance(query, (SQL, Composed)):
@@ -328,9 +330,6 @@ def _run_sql(connectable, sql, params=None, **kwargs):
             )
             continue
 
-        # 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:
@@ -360,7 +359,7 @@ def _run_sql(connectable, sql, params=None, **kwargs):
 
             _print_error(sql_text, err, file=output_file)
         finally:
-            set_wait_callback(None)
+            pass
 
 
 def _should_raise_query_error(err):
@@ -379,7 +378,7 @@ def _should_raise_query_error(err):
     # database backends.
     # Ideally we could handle operational errors more gracefully
     if (
-        isinstance(orig_err, psycopg2.errors.QueryCanceled)
+        isinstance(orig_err, QueryCanceled)
         or getattr(orig_err, "pgcode", None) == "57014"
     ):
         return True
@@ -602,6 +601,26 @@ def create_database(url, **kwargs):
     _create_database(url, **kwargs)
 
 
+def create_engine(db_conn, **kwargs):
+    if isinstance(db_conn, Engine):
+        log.info(f"Set up database connection with engine {db_conn.url}")
+        if db_conn.driver == "psycopg2":
+            log.warning(
+                "The psycopg2 driver is deprecated. Please use psycopg3 instead."
+            )
+        return db_conn
+    else:
+        log.info(f"Setting up database connection with URL '{db_conn}'")
+        url = db_conn
+        if isinstance(url, str):
+            url = make_url(url)
+        # Set the driver to psycopg if not already set
+        if url.drivername != "postgresql+psycopg":
+            url = url.set(drivername="postgresql+psycopg")
+
+        return base_create_engine(url, **kwargs)
+
+
 def connection_args(engine):
     """Get PostgreSQL connection arguments for an engine"""
     _psql_flags = {"-U": "username", "-h": "host", "-p": "port", "-P": "password"}
@@ -617,15 +636,24 @@ def connection_args(engine):
     return flags, engine.url.database
 
 
-def db_isready(engine_or_url):
-    args, _ = connection_args(engine_or_url)
-    c = cmd("pg_isready", args, capture_output=True)
-    return c.returncode == 0
+def db_isready(engine_or_url, use_shell_command=False):
+    if use_shell_command:
+        args, _ = connection_args(engine_or_url)
+        c = cmd("pg_isready", args, capture_output=True)
+        return c.returncode == 0
+    # Use a more typical sqlalchemy connection approach
+    engine = create_engine(engine_or_url)
+    try:
+        with engine.connect() as conn:
+            conn.execute(text("SELECT 1"))
+        return True
+    except OperationalError:
+        return False
 
 
-def wait_for_database(engine_or_url, quiet=False):
+def wait_for_database(engine_or_url, *, quiet=False, use_shell_command=False):
     msg = "Waiting for database..."
-    while not db_isready(engine_or_url):
+    while not db_isready(engine_or_url, use_shell_command=use_shell_command):
         if not quiet:
             echo(msg, err=True)
         log.info(msg)

{macrostrat_database-3.5.4 → macrostrat_database-4.0.0}/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.5.4"
+version = "4.0.0"
 
 [tool.poetry.dependencies]
 GeoAlchemy2 = "^0.15.2"
@@ -11,13 +11,14 @@ SQLAlchemy = "^2.0.18"
 SQLAlchemy-Utils = "^0.41.1"
 click = "^8.1.3"
 "macrostrat.utils" = "^1.3.0"
-psycopg2-binary = "^2.9.6"
 python = "^3.10"
 sqlparse = "^0.5.1"
 aiofiles = "^23.2.1"
 rich = "^13.7.1"
+psycopg = "^3.2.1"
+psycopg2 = "^2.9.11"
 
-[tool.poetry.dev-dependencies]
+[tool.poetry.group.dev.dependencies]
 "macrostrat.utils" = { path = "../utils", develop = true }
 
 [build-system]