macrostrat.database 4.2.0__tar.gz → 4.3.0__tar.gz

{macrostrat_database-4.2.0 → macrostrat_database-4.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: macrostrat.database
-Version: 4.2.0
+Version: 4.3.0
 Summary: A SQLAlchemy-based database toolkit.
 Author: Daven Quinn
 Author-email: Daven Quinn <dev@davenquinn.com>

macrostrat_database-4.3.0/macrostrat/database/__init__.py

@@ -0,0 +1,14 @@
+from .core import Database
+from .mapper import DatabaseMapper
+from .postgresql import on_conflict, prefix_inserts  # noqa
+from .query import run_fixtures, run_query, run_sql, execute  # noqa
+from .utils import (  # noqa
+    create_database,
+    create_engine,
+    database_exists,
+    drop_database,
+    get_dataframe,
+    get_or_create,
+    reflect_table,
+    get_database_url,
+)

macrostrat_database-4.2.0/macrostrat/database/__init__.py → macrostrat_database-4.3.0/macrostrat/database/core.py

@@ -13,16 +13,14 @@ 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 .query import run_fixtures, run_query, run_sql, execute  # noqa
-from .utils import (  # noqa
-    create_database,
+from .postgresql import prefix_inserts
+from .query import run_fixtures, run_query, run_sql
+from .utils import (
     create_engine,
-    database_exists,
-    drop_database,
     get_dataframe,
     get_or_create,
     reflect_table,
+    DatabaseInput,
 )
 
 metadata = MetaData()
@@ -30,6 +28,28 @@ metadata = MetaData()
 log = get_logger(__name__)
 
 
+def _parse_table_name(name, schema=None):
+    """Parse a table name into (schema, table_name).
+
+    Accepts "table", "schema.table", or ("schema", "table").
+    An explicit schema kwarg takes precedence. Defaults to "public".
+    """
+    if isinstance(name, tuple):
+        parsed_schema, table_name = name[0], name[1]
+    elif "." in str(name):
+        parsed_schema, table_name = str(name).split(".", 1)
+    else:
+        parsed_schema, table_name = None, str(name)
+    return (schema or parsed_schema or "public"), table_name
+
+
+def _model_key(schema, table_name):
+    """Return the ModelCollection key used by automap for a given table."""
+    if schema == "public":
+        return table_name
+    return f"{schema}_{table_name}"
+
+
 class Database(object):
     mapper: Optional[DatabaseMapper] = None
     metadata: MetaData
@@ -38,7 +58,7 @@ class Database(object):
 
     __inspector__ = None
 
-    def __init__(self, db_conn: Union[str, URL, Engine], *, echo_sql=False, **kwargs):
+    def __init__(self, db_conn: DatabaseInput, *, echo_sql=False, **kwargs):
         """
         Wrapper for interacting with a database using SQLAlchemy.
         Optimized for use with PostgreSQL, but usable with SQLite
@@ -58,7 +78,10 @@ class Database(object):
 
         self.instance_params = kwargs.pop("instance_params", {})
 
-        self.engine = create_engine(db_conn, echo=echo_sql, **kwargs)
+        if echo_sql:
+            kwargs["echo"] = True
+
+        self.engine = create_engine(db_conn, **kwargs)
 
         self.metadata = kwargs.get("metadata", metadata)
 
@@ -68,6 +91,7 @@ class Database(object):
         self._session_factory = sessionmaker(bind=self.engine)
         self.session = scoped_session(self._session_factory)
         # Use the self.session_scope function to more explicitly manage sessions.
+        self._table_cache: dict = {}
 
     def create_tables(self):
         """
@@ -80,6 +104,10 @@ class Database(object):
         self.mapper = DatabaseMapper(self)
         self.mapper.reflect_database(**kwargs)
 
+    def get_server_version(self):
+        with self.engine.connect():
+            return self.engine.dialect.server_version_info
+
     @contextmanager
     def session_scope(self, commit=True):
         """Provide a transactional scope around a series of operations."""
@@ -343,6 +371,95 @@ class Database(object):
             self.session.close()
             self.session = _prev_session
 
+    def get_table(self, name, *, schema=None):
+        """Return a reflected SQLAlchemy Table object, with per-instance caching.
+
+        After the first call the result is cached; subsequent calls for the
+        same table are instant.  If automap has already been run the mapper's
+        existing Table is reused, avoiding a second round-trip.
+
+        Args:
+            name: Table name as ``"table"``, ``"schema.table"``, or
+                  ``("schema", "table")``.
+            schema: Explicit schema override (default ``"public"``).
+        """
+        schema_, table_name = _parse_table_name(name, schema)
+        cache_key = (schema_, table_name)
+        if cache_key in self._table_cache:
+            return self._table_cache[cache_key]
+
+        # Reuse the already-reflected Table from automap when available
+        if self.mapper is not None:
+            model_key = _model_key(schema_, table_name)
+            if model_key in self.mapper._models:
+                tbl = self.mapper._models[model_key].__table__
+                self._table_cache[cache_key] = tbl
+                return tbl
+
+        # Per-table reflection; "public" → None matches how automap stores it
+        reflect_schema = None if schema_ == "public" else schema_
+        tbl = reflect_table(self.engine, table_name, schema=reflect_schema)
+        self._table_cache[cache_key] = tbl
+        return tbl
+
+    def get_model(self, name, *, schema=None, automap=True):
+        """Return the ORM model class for a table.
+
+        If the target schema has not yet been reflected and ``automap=True``
+        (the default), it is reflected lazily before the lookup.  Set
+        ``automap=False`` to raise ``LookupError`` instead, which is useful
+        when you want strict control over when reflection happens.
+
+        Args:
+            name: Table name as ``"table"``, ``"schema.table"``, or
+                  ``("schema", "table")``.
+            schema: Explicit schema override (default ``"public"``).
+            automap: Lazily reflect the schema if not yet mapped.
+
+        Raises:
+            LookupError: When the model is not found.
+        """
+        schema_, table_name = _parse_table_name(name, schema)
+        model_key = _model_key(schema_, table_name)
+
+        if self.mapper is not None and model_key in self.mapper._models:
+            return self.mapper._models[model_key]
+
+        if not automap:
+            raise LookupError(
+                f"No ORM model found for {schema_}.{table_name}. "
+                "Call db.automap() first, or use get_table() for a Table object."
+            )
+
+        # Lazy automap: reflect only the needed schema
+        if self.mapper is None:
+            self.automap(schemas=[schema_])
+        elif schema_ not in self.mapper._reflected_schemas:
+            self.mapper.reflect_schema(schema_)
+
+        if model_key in self.mapper._models:
+            return self.mapper._models[model_key]
+
+        raise LookupError(
+            f"No ORM model found for {schema_}.{table_name} after reflecting "
+            f"schema '{schema_}'. Verify the table exists, or use get_table()."
+        )
+
+    def __getitem__(self, name):
+        """Subscript shorthand for get_table()."""
+        return self.get_table(name)
+
+    # Destroy engine on cleanup
+    def cleanup(self):
+        try:
+            self.session.close()
+        except OperationalError:
+            pass
+        self.engine.dispose()
+
+    def __del__(self):
+        self.cleanup()
+
 
 def _clear_savepoint(connection, name, rollback=True):
     params = {"name": Identifier(name)}

{macrostrat_database-4.2.0 → macrostrat_database-4.3.0}/macrostrat/database/mapper/__init__.py

@@ -33,6 +33,7 @@ class DatabaseMapper:
     automap_error = None
     _models = None
     _tables = None
+    _reflected_schemas: set
 
     def __init__(self, db, **kwargs):
         # https://docs.sqlalchemy.org/en/13/orm/extensions/automap.html#sqlalchemy.ext.automap.AutomapBase.prepare
@@ -56,6 +57,7 @@ class DatabaseMapper:
 
         self._models = ModelCollection(self.automap_base.classes)
         self._tables = TableCollection(self._models)
+        self._reflected_schemas = set()
 
     def reflect_database(self, schemas=["public"], use_cache=True):
         # This stuff should be placed outside of core (one likely extension point).
@@ -74,6 +76,7 @@ class DatabaseMapper:
         self.automap_base.builder._cache_database_map(self.automap_base.metadata)
 
     def reflect_schema(self, schema, use_cache=True):
+        self._reflected_schemas.add(schema or "public")
         if use_cache and self.automap_base.loaded_from_cache:
             log.info("Database models for %s have been loaded from cache", schema)
             self.automap_base.prepare(schema=schema, **self.reflection_kwargs)

{macrostrat_database-4.2.0 → macrostrat_database-4.3.0}/macrostrat/database/postgresql.py

@@ -12,7 +12,7 @@ from sqlalchemy.sql.dml import Insert
 from sqlalchemy.sql.expression import text
 
 if TYPE_CHECKING:
-    from ..database import Database
+    from .core import Database
 
 
 class OnConflictAction(str, Enum):

{macrostrat_database-4.2.0 → macrostrat_database-4.3.0}/macrostrat/database/query.py

@@ -10,10 +10,6 @@ from warnings import warn
 
 import psycopg2.sql as psql2
 from click import secho
-from macrostrat.database.compat import (
-    update_legacy_identifier,
-)
-from macrostrat.utils import get_logger
 from psycopg.errors import QueryCanceled
 from psycopg.sql import SQL, Composable, Composed
 from rich.console import Console
@@ -29,6 +25,11 @@ from sqlalchemy.exc import (
 from sqlalchemy.sql.elements import TextClause
 from sqlparse import format, split
 
+from macrostrat.database.compat import (
+    update_legacy_identifier,
+)
+from macrostrat.utils import get_logger
+
 log = get_logger(__name__)
 
 
@@ -270,7 +271,15 @@ def _statement_filter_to_transform(statement_filter) -> TransformFn:
     return transform
 
 
-def _run_sql(connectable, sql, params=None, *, print_skipped=True, **kwargs):
+def _run_sql(
+    connectable,
+    sql,
+    params=None,
+    *,
+    print_skipped=True,
+    use_transaction=True,
+    **kwargs,
+):
     """
     Internal function for running a query on a SQLAlchemy connectable,
     which always returns an iterator. The wrapper function adds the option
@@ -341,6 +350,7 @@ def _run_sql(connectable, sql, params=None, *, print_skipped=True, **kwargs):
                 output_mode=output_mode,
                 print_skipped=print_skipped,
                 has_server_binds=has_server_binds,
+                use_transaction=use_transaction,
             )
 
 
@@ -389,6 +399,7 @@ def _execute_one(
     output_mode: OutputMode = OutputMode.SUMMARY,
     has_server_binds: bool | None = None,
     print_skipped: bool = True,
+    use_transaction: bool = True,
 ):
     params = result.params
 
@@ -398,20 +409,24 @@ def _execute_one(
 
     if result.label is not None:
         display_text = result.label
+    elif output_mode == OutputMode.NONE:
+        display_text = None
     elif output_mode != OutputMode.ALL:
         display_text = summarize_statement(str(query))
     else:
         display_text = str(query)
 
     if result.skip:
-        if print_skipped:
+        if print_skipped and display_text is not None:
             secho(display_text, dim=True, strikethrough=True, file=output_file)
         return
 
-    try:
-        trans = connectable.begin()
-    except InvalidRequestError:
-        trans = None
+    trans = None
+    if use_transaction:
+        try:
+            trans = connectable.begin()
+        except InvalidRequestError:
+            pass
 
     try:
         log.debug("Executing SQL: \n %s", query)
@@ -430,7 +445,8 @@ def _execute_one(
         elif hasattr(connectable, "commit"):
             connectable.commit()
 
-        secho(display_text, dim=True, file=output_file)
+        if display_text is not None:
+            secho(display_text, dim=True, file=output_file)
 
     except Exception as err:
         if trans is not None:
@@ -439,7 +455,8 @@ def _execute_one(
             connectable.rollback()
         if raise_errors or _should_raise_query_error(err):
             raise err
-        _print_error(display_text, err, file=output_file)
+        if display_text is not None:
+            _print_error(display_text, err, file=output_file)
 
 
 def _should_raise_query_error(err):
@@ -590,6 +607,8 @@ def run_sql(*args, **kwargs):
         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.
+    use_transaction: bool
+        Whether to run the query in a transaction block
     """
     res = _run_sql(*args, **kwargs)
     if kwargs.pop("yield_results", False):

{macrostrat_database-4.2.0 → macrostrat_database-4.3.0}/macrostrat/database/transfer/utils.py

@@ -72,20 +72,14 @@ def _create_command(
         args = []
 
     command_prefix = prefix or _docker_local_run_args(container)
-    _cmd = [*command_prefix, *command, str(engine.url), *args]
 
-    log.info(" ".join(_cmd))
-
-    # Replace asterisks with the real password (if any). This is kind of backwards
-    # but it works.
-    if "***" in str(engine.url) and engine.url.password is not None:
-        _cmd = [
-            *command_prefix,
-            *command,
-            raw_database_url(engine.url),
-            *args,
-        ]
+    # Strip the SQLAlchemy dialect suffix (e.g. +psycopg) — pg_dump/pg_restore
+    # only understand plain postgresql:// URLs.  Also expand any *** password
+    # masking so the subprocess receives the real credentials.
+    pg_url = raw_database_url(engine.url.set(drivername="postgresql"))
 
+    _cmd = [*command_prefix, *command, pg_url, *args]
+    log.info(" ".join(_cmd))
     return _cmd
 
 

macrostrat_database-4.3.0/macrostrat/database/utils.py

@@ -0,0 +1,348 @@
+from contextlib import contextmanager
+from time import sleep
+from typing import Union
+from uuid import uuid4
+from warnings import warn
+
+from click import echo
+from psycopg.errors import AdminShutdown
+from psycopg.sql import Identifier
+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, URL
+from sqlalchemy.exc import (
+    OperationalError,
+)
+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,
+    database_exists,
+    drop_database as _drop_database,
+)
+
+from macrostrat.utils import cmd, get_logger
+from .query import get_sql_text, execute  # noqa
+
+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
+
+DatabaseInput = Union["Database", Engine, str, URL]
+
+
+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 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(*args, **kwargs):
+    warn(
+        "temp_database is deprecated, use temporary_database instead",
+        DeprecationWarning,
+    )
+    with temporary_database(*args, **kwargs) as engine:
+        yield engine
+
+
+@contextmanager
+def temporary_database(
+    _input: DatabaseInput,
+    *,
+    drop=True,
+    ensure_empty=False,
+    exists_ok=True,
+    template=None,
+    force_drop=False,
+):
+    """Create a temporary database and tear it down after tests."""
+    url = get_database_url(_input)
+    create_database(url, exists_ok=exists_ok, replace=ensure_empty, template=template)
+    engine = create_engine(url)
+    try:
+        yield engine
+        engine.dispose()
+    finally:
+        if drop:
+            drop_database(engine, force=force_drop)
+
+
+def drop_database(_input: DatabaseInput, force=None, allow_missing=False):
+    """Drop a database.
+
+    Parameters
+    ----------
+    engine : Database, Engine or str
+        A SQLAlchemy engine or database URL.
+    force: bool
+        If true, use the `force` parameter
+    """
+    url = get_database_url(_input)
+    if not database_exists(url):
+        if not allow_missing:
+            raise ValueError(f"Database {url} does not exist")
+    elif "postgres" in url.drivername and force is not False:
+        # Check if we can force-drop and do so if we can
+        _force_drop_postgresql_database(url)
+    else:
+        # Drop the database without force
+        _drop_database(url)
+
+
+def _force_drop_postgresql_database(url):
+    # Check if we can force-drop and do so if we can
+    database_name = url.database
+    user_url = url._replace(database=None)
+    user_engine = create_engine(
+        user_url, execution_options={"isolation_level": "AUTOCOMMIT"}
+    )
+    # Get postgresql version from engine
+    major_version = 0
+    with allow_shutdown(user_engine) as conn:
+        conn.autocommit = True
+        pg_version = user_engine.dialect.server_version_info
+        major_version = pg_version[0]
+        can_use_modern_force = major_version >= 13
+        sql = "DROP DATABASE {database_name}"
+        params = dict(database_name=Identifier(database_name))
+        if can_use_modern_force:
+            sql += " WITH (FORCE)"
+        else:
+            close_all_connections(user_engine, database=database_name)
+        run_sql(conn, sql, params=params, raise_errors=True, use_transaction=False)
+    user_engine.dispose()
+
+
+@contextmanager
+def allow_shutdown(engine):
+    with engine.connect() as conn:
+        conn.autocommit = True
+        try:
+            yield conn
+        except AdminShutdown:
+            pass
+        except OperationalError as exc:
+            if isinstance(exc.orig, AdminShutdown):
+                pass
+            else:
+                raise exc
+
+
+def close_all_connections(engine: Engine, database: str = None):
+    """Close all connections to the database."""
+    if database is None:
+        database = engine.url.database
+    sql = "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = :database"
+    params = dict(database=database)
+    with allow_shutdown(engine) as conn:
+        run_sql(conn, sql, params=params, raise_errors=True, use_transaction=False)
+
+
+def get_database_url(_input: DatabaseInput) -> URL:
+    from .core import Database
+
+    if isinstance(_input, Database):
+        return _input.engine.url
+    elif isinstance(_input, Engine):
+        return _input.url
+    elif isinstance(_input, str) or isinstance(_input, URL):
+        return make_url(_input)
+    else:
+        raise ValueError(f"Invalid input type: {_input}")
+
+
+@contextmanager
+def template_database(
+    _input: DatabaseInput,
+    *,
+    name: str = None,
+    force_drop=True,
+    close_source_connections=False,
+):
+    """Create a temporary template database using an existing database as a template."""
+
+    url = get_database_url(_input)
+    if close_source_connections:
+        engine = create_engine(_input)
+        close_all_connections(engine)
+        engine.dispose()
+
+    db_name = url.database
+    template_db_name = name
+    if name is None:
+        uid = str(uuid4())[:8]
+        template_db_name = db_name + "_template_" + uid
+    # Close connection to the database so we can create a new one based on the template
+    new_db_url = url.set(database=template_db_name)
+    with temporary_database(
+        new_db_url, drop=True, exists_ok=False, template=db_name, force_drop=force_drop
+    ) as engine:
+        yield engine
+
+
+def create_database(_input: DatabaseInput, **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`.
+    """
+    url = get_database_url(_input)
+    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(_input: DatabaseInput, **kwargs):
+    from .core import Database
+
+    recreate = False
+    # If we specify engine options, we should recreate the engine
+    if len(kwargs) > 0:
+        recreate = True
+    db_conn = _input
+    if isinstance(_input, Database):
+        db_conn = _input.engine
+    elif isinstance(_input, str):
+        db_conn = make_url(_input)
+    elif isinstance(_input, URL):
+        db_conn = _input
+
+    if isinstance(db_conn, Engine):
+        if recreate:
+            # Reuse the existing 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:
+            db_conn = db_conn.url
+
+    if not isinstance(db_conn, URL):
+        raise ValueError(f"Invalid input type: {_input}")
+    url = db_conn
+
+    log.info(f"Setting up database connection with URL '{url}'")
+    # Set the driver to psycopg if not already set
+    if "postgres" in url.drivername:
+        url = url.set(drivername="postgresql+psycopg")
+
+    return base_create_engine(url, **kwargs)
+
+
+def connection_args(_input: DatabaseInput, with_password=False):
+    """Get PostgreSQL connection arguments for an engine"""
+    _psql_flags = {"-U": "username", "-h": "host", "-p": "port", "-P": "password"}
+    url = get_database_url(_input)
+
+    flags = ""
+    for flag, _attr in _psql_flags.items():
+        val = getattr(url, _attr)
+        if flag == "-P" and not with_password:
+            continue
+        if val is not None:
+            flags += f" {flag} {val}"
+    return flags, url.database
+
+
+def db_isready(_input: DatabaseInput, use_shell_command=False):
+    if use_shell_command:
+        args, _ = connection_args(_input, 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(_input)
+    try:
+        with engine.connect() as conn:
+            conn.execute(text("SELECT 1"))
+        return True
+    except OperationalError:
+        return False
+
+
+def wait_for_database(_input: DatabaseInput, *, quiet=False, use_shell_command=False):
+    msg = "Waiting for database..."
+    while not db_isready(_input, 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.2.0 → macrostrat_database-4.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "macrostrat.database"
-version = "4.2.0"
+version = "4.3.0"
 description = "A SQLAlchemy-based database toolkit."
 authors = [{ name = "Daven Quinn", email = "dev@davenquinn.com" }]
 requires-python = ">=3.10,<4"

macrostrat_database-4.2.0/macrostrat/database/utils.py

@@ -1,221 +0,0 @@
-import warnings
-from contextlib import contextmanager
-from time import sleep
-from uuid import uuid4
-
-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 (
-    OperationalError,
-)
-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 macrostrat.utils import cmd, get_logger
-from .query import get_sql_text, execute  # noqa
-
-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 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 temporary_database(
-    conn_string, *, drop=True, ensure_empty=False, exists_ok=True, template=None
-):
-    """Create a temporary database and tear it down after tests."""
-    create_database(
-        conn_string, exists_ok=exists_ok, replace=ensure_empty, template=template
-    )
-    try:
-        engine = create_engine(conn_string)
-        yield engine
-        engine.dispose()
-    finally:
-        if drop:
-            drop_database(conn_string)
-
-
-@contextmanager
-def temp_database(*args, **kwargs):
-    warnings.warn(
-        "temp_database is deprecated, use temporary_database instead",
-        DeprecationWarning,
-    )
-    with temporary_database(*args, **kwargs) as engine:
-        yield engine
-
-
-@contextmanager
-def template_database(engine: Engine, *, name: str = None):
-    """Create a temporary template database using an existing database as a template."""
-    db_name = engine.url.database
-    template_db_name = name
-    if name is None:
-        uid = str(uuid4())[:8]
-        template_db_name = db_name + "_template_" + uid
-    # Close connection to the database so we can create a new one based on the template
-    new_db_url = engine.url.set(database=template_db_name)
-    engine.dispose()
-    with temporary_database(
-        new_db_url, drop=True, exists_ok=False, template=db_name
-    ) as engine:
-        yield engine
-
-
-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 "postgres" in url.drivername:
-            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)