macrostrat.database 3.5.4__tar.gz → 4.0.1__tar.gz

macrostrat_database-4.0.1/PKG-INFO

@@ -0,0 +1,23 @@
+Metadata-Version: 2.3
+Name: macrostrat.database
+Version: 4.0.1
+Summary: A SQLAlchemy-based database toolkit.
+Author: Daven Quinn
+Author-email: Daven Quinn <dev@davenquinn.com>
+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
+Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: geoalchemy2>=0.15.2,<0.16
+Requires-Dist: sqlalchemy>=2.0.18,<3
+Requires-Dist: sqlalchemy-utils>=0.41.1,<0.42
+Requires-Dist: click>=8.1.3,<9
+Requires-Dist: macrostrat-utils>=1.3.3,<2
+Requires-Dist: sqlparse>=0.5.1,<0.6
+Requires-Dist: aiofiles>=23.2.1,<24
+Requires-Dist: rich>=13.7.1,<14
+Requires-Dist: psycopg>=3.2.1,<4
+Requires-Dist: psycopg2>=2.9.11,<3
+Requires-Python: >=3.10, <4

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

@@ -3,10 +3,10 @@ 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
@@ -17,6 +17,7 @@ from .mapper import DatabaseMapper
 from .postgresql import on_conflict, prefix_inserts  # noqa
 from .utils import (  # noqa
     create_database,
+    create_engine,
     database_exists,
     drop_database,
     get_dataframe,
@@ -60,12 +61,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 +331,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 +354,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-4.0.1/macrostrat/database/compat.py

@@ -0,0 +1,36 @@
+from warnings import warn
+
+import psycopg.sql as psql3
+import psycopg2.sql as psql2
+
+
+def update_legacy_identifier(identifier):
+    """
+    For backwards compatibility with current code, we need to map psycopg2 identifiers to their equivalents in psycopg3,
+    while printing a warning that the mapping is deprecated.
+    :param identifier:
+    :return: psycopg3 equivalent of identifier
+    """
+    new_identifier = _map_psycopg2_identifier_to_psycopg3_identifier_internal(
+        identifier
+    )
+    if new_identifier is not identifier:
+        warn(
+            "psycopg2 identifiers are deprecated. Please use psycopg3 identifiers instead.",
+            DeprecationWarning,
+        )
+    return new_identifier
+
+
+def _map_psycopg2_identifier_to_psycopg3_identifier_internal(identifier):
+    if isinstance(identifier, psql2.Identifier):
+        return psql3.Identifier(*identifier._wrapped)
+    if isinstance(identifier, psql2.SQL):
+        return psql3.SQL(identifier._wrapped)
+    if isinstance(identifier, psql2.Literal):
+        return psql3.Literal(identifier._wrapped)
+    if isinstance(identifier, psql2.Placeholder):
+        return psql3.Placeholder(identifier._obj)
+    if isinstance(identifier, psql2.Composed):
+        return psql3.Composed(identifier._obj)
+    return identifier

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

@@ -2,18 +2,26 @@ from __future__ import annotations
 
 from contextlib import contextmanager
 from contextvars import ContextVar
-from typing import TYPE_CHECKING
+from enum import Enum
+from typing import Any, Sequence, 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
 
-_insert_mode = ContextVar("insert-mode", default="do-nothing")
+
+class OnConflictAction(str, Enum):
+    DO_NOTHING = "do-nothing"
+    DO_UPDATE = "do-update"
+    RESTRICT = "restrict"
+
+
+_insert_mode = ContextVar("insert-mode", default="restrict")
 
 
 # https://stackoverflow.com/questions/33307250/postgresql-on-conflict-in-sqlalchemy/62305344#62305344
@@ -26,9 +34,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)
 
@@ -60,10 +69,36 @@ def prefix_inserts(insert, compiler, **kw):
     return compiler.visit_insert(insert, **kw)
 
 
+def upsert(
+    table,
+    values: dict[str, Any],
+    index_elements: Sequence[str],
+    *,
+    on_conflict: OnConflictAction = "do-update",
+):
+    stmt = postgresql.insert(table).values(values)
+    if on_conflict == "restrict":
+        return stmt
+
+    if on_conflict == "do-nothing":
+        return stmt.on_conflict_do_nothing(index_elements=list(index_elements))
+
+    update_values = {
+        column.name: getattr(stmt.excluded, column.name)
+        for column in table.columns
+        if column.name in values and column.name not in index_elements
+    }
+
+    return stmt.on_conflict_do_update(
+        index_elements=list(index_elements),
+        set_=update_values,
+    )
+
+
 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/utils.py → macrostrat_database-4.0.1/macrostrat/database/query.py

@@ -1,20 +1,18 @@
 import os
-from contextlib import contextmanager
+from dataclasses import dataclass
 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 typing import Callable, Any, 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
+import psycopg2.sql as psql2
+from click import secho
+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 text
 from sqlalchemy.engine import Connection, Engine
 from sqlalchemy.exc import (
     IntegrityError,
@@ -23,23 +21,17 @@ from sqlalchemy.exc import (
     OperationalError,
     ProgrammingError,
 )
-from sqlalchemy.orm import sessionmaker
-from sqlalchemy.schema import Table
-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 sqlalchemy.sql.elements import TextClause
 from sqlparse import format, split
 
-from macrostrat.utils import cmd, get_logger
+from macrostrat.database.compat import (
+    update_legacy_identifier,
+)
+from macrostrat.utils import get_logger
 
 log = get_logger(__name__)
 
 
-def db_session(engine):
-    factory = sessionmaker(bind=engine)
-    return factory()
-
-
 def infer_is_sql_text(_string: str) -> bool:
     """
     Return True if the string is a valid SQL query,
@@ -74,19 +66,6 @@ def canonicalize_query(file_or_text: Union[str, Path, IO]) -> Union[str, Path]:
     return file_or_text
 
 
-def get_dataframe(connectable, filename_or_query, **kwargs):
-    """
-    Run a query on a SQL database (represented by
-    a SQLAlchemy database object) and turn it into a
-    `Pandas` dataframe.
-    """
-    from pandas import read_sql
-
-    sql = get_sql_text(filename_or_query)
-
-    return read_sql(sql, connectable, **kwargs)
-
-
 def pretty_print(sql, **kwargs):
     """Print and optionally summarize an SQL query"""
     summarize = kwargs.pop("summarize", True)
@@ -117,6 +96,7 @@ def summarize_statement(sql):
             if not line.startswith(i):
                 continue
             return line.split("(")[0].strip().rstrip(";").replace(" AS", "")
+    return sql.strip().split("\n")[0].strip().rstrip(";")
 
 
 def get_sql_text(sql, interpret_as_file=None, echo_file_name=True):
@@ -139,9 +119,7 @@ def _get_queries(sql, interpret_as_file=None):
         for i in sql:
             queries.extend(_get_queries(i, interpret_as_file=interpret_as_file))
         return queries
-    if isinstance(sql, TextClause):
-        return [sql]
-    if isinstance(sql, SQL):
+    if isinstance(sql, (SQL, psql2.SQL, TextClause)):
         return [sql]
 
     if sql in [None, ""]:
@@ -158,7 +136,7 @@ def _get_queries(sql, interpret_as_file=None):
 
 
 def _is_prebind_param(param):
-    return isinstance(param, Composable)
+    return isinstance(param, (Composable, psql2.Composable))
 
 
 def _split_params(params):
@@ -169,7 +147,7 @@ def _split_params(params):
     if isinstance(params, (list, tuple)):
         for i in params:
             if _is_prebind_param(i):
-                new_bind_params.append(i)
+                new_bind_params.append(update_legacy_identifier(i))
             else:
                 new_params.append(i)
     elif isinstance(params, dict):
@@ -177,7 +155,7 @@ def _split_params(params):
         new_bind_params = {}
         for k, v in params.items():
             if _is_prebind_param(v):
-                new_bind_params[k] = v
+                new_bind_params[k] = update_legacy_identifier(v)
             else:
                 new_params[k] = v
     if len(new_bind_params) == 0:
@@ -196,6 +174,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):
@@ -228,8 +208,11 @@ def _render_query(query: Union[SQL, Composed], connectable: Union[Engine, Connec
     return query.as_string(conn)
 
 
-def infer_has_server_binds(sql):
-    return "%s" in sql or search(r"%\(\w+\)s", sql)
+def infer_has_server_binds(sql) -> bool:
+    if "%s" in sql:
+        return True
+    res = search(r"%\(\w+\)s", sql)
+    return res is not None
 
 
 _default_statement_filter = lambda sql_text, params: True
@@ -254,7 +237,40 @@ def _normalize_output_args(kwargs):
     return output_mode, output_file
 
 
-def _run_sql(connectable, sql, params=None, **kwargs):
+@dataclass
+class StatementResult:
+    query: Any
+    params: Any = None
+    skip: bool = False
+    label: str | None = None
+
+    @classmethod
+    def skipped(cls, query=None, params=None) -> "StatementResult":
+        return cls(query=query, params=params, skip=True)
+
+
+@dataclass
+class StatementContext:
+    index: int
+    query: Any
+    params: Any
+    sql_text: str
+
+
+TransformFn = Callable[[StatementContext], list[StatementResult] | None]
+Connectable = Union[Engine, Connection]
+
+
+def _statement_filter_to_transform(statement_filter) -> TransformFn:
+    def transform(ctx: StatementContext) -> list[StatementResult] | None:
+        if not statement_filter(ctx.sql_text, ctx.params):
+            return [StatementResult.skipped(query=ctx.query, params=ctx.params)]
+        return None
+
+    return transform
+
+
+def _run_sql(connectable, sql, params=None, *, print_skipped=True, **kwargs):
     """
     Internal function for running a query on a SQLAlchemy connectable,
     which always returns an iterator. The wrapper function adds the option
@@ -267,17 +283,30 @@ def _run_sql(connectable, sql, params=None, **kwargs):
 
     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)
     ensure_single_query = kwargs.pop("ensure_single_query", False)
-    statement_filter = kwargs.pop("statement_filter", _default_statement_filter)
     output_mode, output_file = _normalize_output_args(kwargs)
+    has_server_binds = kwargs.pop("has_server_binds", None)
+
+    statement_filter = kwargs.pop("statement_filter", None)
+    transform_statement: TransformFn | None = kwargs.pop("transform_statement", None)
 
     if stop_on_error:
         raise_errors = True
         warn(DeprecationWarning("stop_on_error is deprecated, use raise_errors"))
 
-    interpret_as_file = kwargs.pop("interpret_as_file", None)
+    if statement_filter is not None:
+        warn(
+            DeprecationWarning(
+                "statement_filter is deprecated, use transform_statement"
+            )
+        )
+        if transform_statement is not None:
+            raise ValueError(
+                "Cannot specify both statement_filter and transform_statement"
+            )
+        transform_statement = _statement_filter_to_transform(statement_filter)
 
+    interpret_as_file = kwargs.pop("interpret_as_file", None)
     queries = _get_queries(sql, interpret_as_file=interpret_as_file)
 
     if queries is None:
@@ -286,81 +315,116 @@ def _run_sql(connectable, sql, params=None, **kwargs):
     if ensure_single_query and len(queries) > 1:
         raise ValueError("Multiple queries passed when only one was expected")
 
-    # check if parameters is a list of the same length as the number of queries
     if not isinstance(params, list) or not len(params) == len(queries):
         params = [params] * len(queries)
 
-    for query, _params in zip(queries, params):
-        params, pre_bind_params = _split_params(_params)
+    for index, (query, _params) in enumerate(zip(queries, params)):
+        _query, sql_text = _render_query_text(connectable, query, _params)
+        if sql_text == "":
+            continue
 
-        if pre_bind_params is not None:
-            if not isinstance(query, SQL):
-                query = SQL(query)
-            # Pre-bind the parameters using PsycoPG2
-            query = query.format(**pre_bind_params)
+        ctx = StatementContext(
+            index=index, query=query, params=_params, sql_text=sql_text
+        )
 
-        if isinstance(query, (SQL, Composed)):
-            query = _render_query(query, connectable)
+        results = transform_statement(ctx) if transform_statement is not None else None
 
-        sql_text = str(query)
+        if results is None:
+            results = [StatementResult(query=query, params=_params)]
 
-        if isinstance(query, str):
-            sql_text = format(query, strip_comments=True).strip()
-            if sql_text == "":
-                continue
-            # Check for server-bound parameters in sql native style. If there are none, use
-            # the SQLAlchemy text() function, otherwise use the raw query string
-            if has_server_binds is None:
-                has_server_binds = infer_has_server_binds(sql_text)
-
-        should_run = statement_filter(sql_text, params)
-
-        # Shorten summary text for printing
-        if output_mode != OutputMode.ALL:
-            sql_text = summarize_statement(sql_text)
-
-        if not should_run:
-            secho(
-                sql_text,
-                dim=True,
-                strikethrough=True,
-                file=output_file,
+        for result in results:
+            yield from _execute_one(
+                connectable,
+                result,
+                output_file,
+                raise_errors=raise_errors,
+                output_mode=output_mode,
+                print_skipped=print_skipped,
+                has_server_binds=has_server_binds,
             )
-            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:
-            trans = None
-        try:
-            log.debug("Executing SQL: \n %s", query)
-            if has_server_binds:
-                conn = _get_connection(connectable)
-                res = conn.exec_driver_sql(query, params)
-            else:
-                if not isinstance(query, TextClause):
-                    query = text(query)
-                res = connectable.execute(query, params)
-            yield res
-            if trans is not None:
-                trans.commit()
-            elif hasattr(connectable, "commit"):
-                connectable.commit()
-            secho(sql_text, dim=True, file=output_file)
-        except Exception as err:
-            if trans is not None:
-                trans.rollback()
-            elif hasattr(connectable, "rollback"):
-                connectable.rollback()
-            if raise_errors or _should_raise_query_error(err):
-                raise err
-
-            _print_error(sql_text, err, file=output_file)
-        finally:
-            set_wait_callback(None)
+
+def _render_query_text(connectable, query, params):
+    params, pre_bind_params = _split_params(params)
+    if isinstance(query, (psql2.SQL, psql2.Composed)):
+        query = update_legacy_identifier(query)
+
+    if pre_bind_params is not None:
+        if not isinstance(query, SQL):
+            query = SQL(query)
+        # Pre-bind the parameters using psycopg
+        query = query.format(**pre_bind_params)
+
+    if isinstance(query, (SQL, Composed)):
+        query = _render_query(query, connectable)
+
+    sql_text = str(query)
+    if isinstance(query, str):
+        sql_text = format(query, strip_comments=True).strip()
+
+    return query, sql_text
+
+
+def _execute_one(
+    connectable,
+    result: StatementResult,
+    output_file: IO,
+    *,
+    raise_errors: bool = True,
+    output_mode: OutputMode = OutputMode.SUMMARY,
+    has_server_binds: bool | None = None,
+    print_skipped: bool = True,
+):
+    params = result.params
+
+    query, sql_text = _render_query_text(connectable, result.query, params)
+    if has_server_binds is None:
+        has_server_binds = infer_has_server_binds(sql_text)
+
+    if result.label is not None:
+        display_text = result.label
+    elif output_mode != OutputMode.ALL:
+        display_text = summarize_statement(str(query))
+    else:
+        display_text = str(query)
+
+    if result.skip:
+        if print_skipped:
+            secho(display_text, dim=True, strikethrough=True, file=output_file)
+        return
+
+    try:
+        trans = connectable.begin()
+    except InvalidRequestError:
+        trans = None
+
+    try:
+        log.debug("Executing SQL: \n %s", query)
+        if has_server_binds:
+            conn = _get_connection(connectable)
+            res = conn.exec_driver_sql(query, params)
+        else:
+            if not isinstance(query, TextClause):
+                query = text(query)
+            res = connectable.execute(query, params)
+
+        yield res
+
+        if trans is not None:
+            trans.commit()
+        elif hasattr(connectable, "commit"):
+            connectable.commit()
+
+        secho(display_text, dim=True, file=output_file)
+
+    except Exception as err:
+        if trans is not None:
+            trans.rollback()
+        elif hasattr(connectable, "rollback"):
+            connectable.rollback()
+        if raise_errors or _should_raise_query_error(err):
+            raise err
+        _print_error(display_text, err, file=output_file)
 
 
 def _should_raise_query_error(err):
@@ -379,7 +443,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
@@ -506,144 +570,13 @@ def run_sql(*args, **kwargs):
     statement_filter : Callable
         A function that takes a SQL statement and parameters and returns True if the statement
         should be run, and False if it should be skipped.
+    transform_statement: TransformFn | None
+        A function that takes a StatementContext and returns a list of StatementResult
+        objects, which can modify the query, parameters, and whether the statement
+        should be skipped or not. This allows for more complex logic than a simple
+        statement filter.
     """
     res = _run_sql(*args, **kwargs)
     if kwargs.pop("yield_results", False):
         return res
     return list(res)
-
-
-def execute(connectable, sql, params=None, stop_on_error=False, **kwargs):
-    output_file = kwargs.pop("output_file", None)
-    output_mode = kwargs.pop("output_mode", None)
-    sql = format(sql, strip_comments=True).strip()
-    if sql == "":
-        return
-    try:
-        connectable.begin()
-        res = connectable.execute(text(sql), params=params)
-        if hasattr(connectable, "commit"):
-            connectable.commit()
-        pretty_print(sql, dim=True, file=output_file, mode=output_mode)
-        return res
-    except (ProgrammingError, IntegrityError) as err:
-        if hasattr(connectable, "rollback"):
-            connectable.rollback()
-        _print_error(sql, dim=True, file=output_file, mode=output_mode)
-        if stop_on_error:
-            return
-    finally:
-        if hasattr(connectable, "close"):
-            connectable.close()
-
-
-def get_or_create(session, model, defaults=None, **kwargs):
-    """
-    Get an instance of a model, or create it if it doesn't
-    exist.
-
-    https://stackoverflow.com/questions/2546207
-    """
-    instance = session.query(model).filter_by(**kwargs).first()
-    if instance:
-        instance._created = False
-        return instance
-    else:
-        params = dict(
-            (k, v) for k, v in kwargs.items() if not isinstance(v, ClauseElement)
-        )
-        params.update(defaults or {})
-        instance = model(**params)
-        session.add(instance)
-        instance._created = True
-        return instance
-
-
-def get_db_model(db, model_name: str):
-    return getattr(db.model, model_name)
-
-
-@contextmanager
-def temp_database(conn_string, drop=True, ensure_empty=False):
-    """Create a temporary database and tear it down after tests."""
-    create_database(conn_string, exists_ok=True, replace=ensure_empty)
-    try:
-        yield create_engine(conn_string)
-    finally:
-        if drop:
-            drop_database(conn_string)
-
-
-def create_database(url, **kwargs):
-    """Create a database if it doesn't exist.
-
-    Parameters
-    ----------
-    url : str
-        A SQLAlchemy database URL.
-    exists_ok : bool
-        If True, don't raise an error if the database already exists.
-    replace : bool
-        If True, drop the database if it exists and create a new one.
-    kwargs : dict
-        Additional keyword arguments to pass to `sqlalchemy_utils.create_database`.
-    """
-    db_exists = database_exists(url)
-
-    should_replace = kwargs.pop("replace", False)
-    exists_ok = kwargs.pop("exists_ok", False)
-
-    if should_replace and db_exists:
-        drop_database(url)
-        db_exists = False
-
-    if exists_ok and db_exists:
-        return
-    _create_database(url, **kwargs)
-
-
-def connection_args(engine):
-    """Get PostgreSQL connection arguments for an engine"""
-    _psql_flags = {"-U": "username", "-h": "host", "-p": "port", "-P": "password"}
-
-    if isinstance(engine, str):
-        # We passed a connection url!
-        engine = create_engine(engine)
-    flags = ""
-    for flag, _attr in _psql_flags.items():
-        val = getattr(engine.url, _attr)
-        if val is not None:
-            flags += f" {flag} {val}"
-    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 wait_for_database(engine_or_url, quiet=False):
-    msg = "Waiting for database..."
-    while not db_isready(engine_or_url):
-        if not quiet:
-            echo(msg, err=True)
-        log.info(msg)
-        sleep(1)
-
-
-def reflect_table(engine, tablename, *column_args, **kwargs):
-    """
-    One-off reflection of a database table or view. Note: for most purposes,
-    it will be better to use the database tables automapped at runtime in the
-    `self.tables` object. However, this function can be useful for views (which
-    are not reflected automatically), or to customize type definitions for mapped
-    tables.
-
-    A set of `column_args` can be used to pass columns to override with the mapper, for
-    instance to set up foreign and primary key constraints.
-    https://docs.sqlalchemy.org/en/13/core/reflection.html#reflecting-views
-    """
-    schema = kwargs.pop("schema", "public")
-    meta = MetaData(schema=schema)
-    return Table(tablename, meta, *column_args, autoload_with=engine, **kwargs)

macrostrat_database-4.0.1/macrostrat/database/utils.py

@@ -0,0 +1,213 @@
+from contextlib import contextmanager
+from time import sleep
+
+from click import echo
+from sqlalchemy import MetaData
+from sqlalchemy import create_engine as base_create_engine
+from sqlalchemy import text
+from sqlalchemy.engine import Engine
+from sqlalchemy.engine.url import make_url
+from sqlalchemy.exc import (
+    IntegrityError,
+    OperationalError,
+    ProgrammingError,
+)
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.schema import Table
+from sqlalchemy.sql.elements import ClauseElement
+from sqlalchemy_utils import create_database as _create_database
+from sqlalchemy_utils import database_exists, drop_database
+from sqlparse import format
+
+from macrostrat.utils import cmd, get_logger
+from .query import get_sql_text
+
+log = get_logger(__name__)
+
+# Ensure that old import structure still works
+from .query import run_sql, run_query, run_sql_file, run_fixtures  # noqa: F401
+
+
+def get_dataframe(connectable, filename_or_query, **kwargs):
+    """
+    Run a query on a SQL database (represented by
+    a SQLAlchemy database object) and turn it into a
+    `Pandas` dataframe.
+    """
+    from pandas import read_sql
+
+    sql = get_sql_text(filename_or_query)
+
+    return read_sql(sql, connectable, **kwargs)
+
+
+def db_session(engine):
+    factory = sessionmaker(bind=engine)
+    return factory()
+
+
+def execute(connectable, sql, params=None, stop_on_error=False, **kwargs):
+    output_file = kwargs.pop("output_file", None)
+    output_mode = kwargs.pop("output_mode", None)
+    sql = format(sql, strip_comments=True).strip()
+    if sql == "":
+        return
+    try:
+        connectable.begin()
+        res = connectable.execute(text(sql), params=params)
+        if hasattr(connectable, "commit"):
+            connectable.commit()
+        pretty_print(sql, dim=True, file=output_file, mode=output_mode)
+        return res
+    except (ProgrammingError, IntegrityError) as err:
+        if hasattr(connectable, "rollback"):
+            connectable.rollback()
+        _print_error(sql, dim=True, file=output_file, mode=output_mode)
+        if stop_on_error:
+            return
+    finally:
+        if hasattr(connectable, "close"):
+            connectable.close()
+
+
+def get_or_create(session, model, defaults=None, **kwargs):
+    """
+    Get an instance of a model, or create it if it doesn't
+    exist.
+
+    https://stackoverflow.com/questions/2546207
+    """
+    instance = session.query(model).filter_by(**kwargs).first()
+    if instance:
+        instance._created = False
+        return instance
+    else:
+        params = dict(
+            (k, v) for k, v in kwargs.items() if not isinstance(v, ClauseElement)
+        )
+        params.update(defaults or {})
+        instance = model(**params)
+        session.add(instance)
+        instance._created = True
+        return instance
+
+
+def get_db_model(db, model_name: str):
+    return getattr(db.model, model_name)
+
+
+@contextmanager
+def temp_database(conn_string, drop=True, ensure_empty=False):
+    """Create a temporary database and tear it down after tests."""
+    create_database(conn_string, exists_ok=True, replace=ensure_empty)
+    try:
+        yield create_engine(conn_string)
+    finally:
+        if drop:
+            drop_database(conn_string)
+
+
+def create_database(url, **kwargs):
+    """Create a database if it doesn't exist.
+
+    Parameters
+    ----------
+    url : str
+        A SQLAlchemy database URL.
+    exists_ok : bool
+        If True, don't raise an error if the database already exists.
+    replace : bool
+        If True, drop the database if it exists and create a new one.
+    kwargs : dict
+        Additional keyword arguments to pass to `sqlalchemy_utils.create_database`.
+    """
+    db_exists = database_exists(url)
+
+    should_replace = kwargs.pop("replace", False)
+    exists_ok = kwargs.pop("exists_ok", False)
+
+    if should_replace and db_exists:
+        drop_database(url)
+        db_exists = False
+
+    if exists_ok and db_exists:
+        return
+    _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, with_password=False):
+    """Get PostgreSQL connection arguments for an engine"""
+    _psql_flags = {"-U": "username", "-h": "host", "-p": "port", "-P": "password"}
+
+    if isinstance(engine, str):
+        # We passed a connection url!
+        engine = create_engine(engine)
+    flags = ""
+    for flag, _attr in _psql_flags.items():
+        val = getattr(engine.url, _attr)
+        if flag == "-P" and not with_password:
+            continue
+        if val is not None:
+            flags += f" {flag} {val}"
+    return flags, engine.url.database
+
+
+def db_isready(engine_or_url, use_shell_command=False):
+    if use_shell_command:
+        args, _ = connection_args(engine_or_url, with_password=True)
+        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, use_shell_command=False):
+    msg = "Waiting for database..."
+    while not db_isready(engine_or_url, use_shell_command=use_shell_command):
+        if not quiet:
+            echo(msg, err=True)
+        log.info(msg)
+        sleep(1)
+
+
+def reflect_table(engine, tablename, *column_args, **kwargs):
+    """
+    One-off reflection of a database table or view. Note: for most purposes,
+    it will be better to use the database tables automapped at runtime in the
+    `self.tables` object. However, this function can be useful for views (which
+    are not reflected automatically), or to customize type definitions for mapped
+    tables.
+
+    A set of `column_args` can be used to pass columns to override with the mapper, for
+    instance to set up foreign and primary key constraints.
+    https://docs.sqlalchemy.org/en/13/core/reflection.html#reflecting-views
+    """
+    schema = kwargs.pop("schema", "public")
+    meta = MetaData(schema=schema)
+    return Table(tablename, meta, *column_args, autoload_with=engine, **kwargs)

macrostrat_database-4.0.1/pyproject.toml

@@ -0,0 +1,44 @@
+[project]
+name = "macrostrat.database"
+version = "4.0.1"
+description = "A SQLAlchemy-based database toolkit."
+authors = [{ name = "Daven Quinn", email = "dev@davenquinn.com" }]
+requires-python = ">=3.10,<4"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+]
+dependencies = [
+    "GeoAlchemy2>=0.15.2,<0.16",
+    "SQLAlchemy>=2.0.18,<3",
+    "SQLAlchemy-Utils>=0.41.1,<0.42",
+    "click>=8.1.3,<9",
+    "macrostrat.utils>=1.3.3,<2",
+    "sqlparse>=0.5.1,<0.6",
+    "aiofiles>=23.2.1,<24",
+    "rich>=13.7.1,<14",
+    "psycopg>=3.2.1,<4",
+    "psycopg2>=2.9.11,<3",
+]
+
+[dependency-groups]
+dev = ["macrostrat.utils"]
+
+[tool.uv]
+default-groups = "all"
+
+[tool.uv.sources]
+"macrostrat.utils" = { path = "../utils", editable = true }
+
+[tool.uv.build-backend]
+module-name = ["macrostrat"]
+module-root = ""
+namespace = true
+
+[build-system]
+requires = ["uv_build>=0.9.21,<0.12.0"]
+build-backend = "uv_build"

macrostrat_database-3.5.4/PKG-INFO

@@ -1,20 +0,0 @@
-Metadata-Version: 2.1
-Name: macrostrat.database
-Version: 3.5.4
-Summary: A SQLAlchemy-based database toolkit.
-Author: Daven Quinn
-Author-email: dev@davenquinn.com
-Requires-Python: >=3.10,<4.0
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-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: rich (>=13.7.1,<14.0.0)
-Requires-Dist: sqlparse (>=0.5.1,<0.6.0)

macrostrat_database-3.5.4/pyproject.toml

@@ -1,25 +0,0 @@
-[tool.poetry]
-authors = ["Daven Quinn <dev@davenquinn.com>"]
-description = "A SQLAlchemy-based database toolkit."
-name = "macrostrat.database"
-packages = [{ include = "macrostrat" }]
-version = "3.5.4"
-
-[tool.poetry.dependencies]
-GeoAlchemy2 = "^0.15.2"
-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"
-
-[tool.poetry.dev-dependencies]
-"macrostrat.utils" = { path = "../utils", develop = true }
-
-[build-system]
-build-backend = "poetry.core.masonry.api"
-requires = ["poetry-core>=1.0.0"]