macrostrat.database 1.0.2__tar.gz → 3.0.0b1__tar.gz

macrostrat_database-3.0.0b1/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.1
+Name: macrostrat.database
+Version: 3.0.0b1
+Summary: A SQLAlchemy-based database toolkit.
+Author: Daven Quinn
+Author-email: dev@davenquinn.com
+Requires-Python: >=3.8,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: GeoAlchemy2 (>=0.14.0,<0.15.0)
+Requires-Dist: SQLAlchemy (>=2.0.18,<3.0.0)
+Requires-Dist: SQLAlchemy-Utils (>=0.41.1,<0.42.0)
+Requires-Dist: click (>=8.1.3,<9.0.0)
+Requires-Dist: macrostrat.utils (>=1.0.0,<2.0.0)
+Requires-Dist: psycopg2-binary (>=2.9.6,<3.0.0)
+Requires-Dist: sqlparse (>=0.4.4,<0.5.0)

{macrostrat.database-1.0.2 → macrostrat_database-3.0.0b1}/macrostrat/database/__init__.py

@@ -1,21 +1,23 @@
+import warnings
 from contextlib import contextmanager
-from pathlib import Path
 from typing import Optional
-from click import secho
 
 from sqlalchemy import create_engine, inspect, MetaData, text
-from sqlalchemy.orm import sessionmaker, scoped_session
-from sqlalchemy.schema import ForeignKey, Column
-from sqlalchemy.types import Integer
+from sqlalchemy.orm import sessionmaker, scoped_session, Session
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy.orm.exc import FlushError
-from macrostrat.utils import get_logger, relative_path
+from macrostrat.utils import get_logger
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.sql.expression import Insert
 
-from .utils import run_sql_file, run_query, get_or_create, run_sql_query_file
+
+from .utils import (
+    run_sql,
+    get_or_create,
+    reflect_table,
+    get_dataframe,
+)
 from .mapper import DatabaseMapper
-from .postgresql import on_conflict, prefix_inserts
+from .postgresql import prefix_inserts, on_conflict  # noqa
 
 
 metadata = MetaData()
@@ -25,9 +27,11 @@ log = get_logger(__name__)
 
 class Database(object):
     mapper: Optional[DatabaseMapper] = None
+    metadata: MetaData
+    session: Session
     __inspector__ = None
 
-    def __init__(self, db_conn, app=None, echo_sql=False, **kwargs):
+    def __init__(self, db_conn, echo_sql=False, **kwargs):
         """
         We can pass a connection string, a **Flask** application object
         with the appropriate configuration, or nothing, in which
@@ -38,11 +42,8 @@ class Database(object):
         compiles(Insert, "postgresql")(prefix_inserts)
 
         log.info(f"Setting up database connection '{db_conn}'")
-        self.engine = create_engine(
-            db_conn, executemany_mode="batch", echo=echo_sql, **kwargs
-        )
-        metadata.create_all(bind=self.engine)
-        self.meta = metadata
+        self.engine = create_engine(db_conn, echo=echo_sql, **kwargs)
+        self.metadata = kwargs.get("metadata", metadata)
 
         # Scoped session for database
         # https://docs.sqlalchemy.org/en/13/orm/contextual.html#unitofwork-contextual
@@ -51,10 +52,16 @@ class Database(object):
         self.session = scoped_session(self._session_factory)
         # Use the self.session_scope function to more explicitly manage sessions.
 
-    def automap(self):
+    def create_tables(self):
+        """
+        Create all tables described by the database's metadata instance.
+        """
+        metadata.create_all(bind=self.engine)
+
+    def automap(self, **kwargs):
         log.info("Automapping the database")
         self.mapper = DatabaseMapper(self)
-        self.mapper.automap_database()
+        self.mapper.reflect_database(**kwargs)
 
     @contextmanager
     def session_scope(self, commit=True):
@@ -85,26 +92,18 @@ class Database(object):
                 session.rollback()
                 log.debug(err)
 
-    def exec_sql_text(self, statement, *args, **kwargs):
-        """
-        Executes a sql command, in string on the database
-        Easy way to load data into a test database instance
-        """
-        connection = self.engine.connect()
-        connection.execute(text(statement), *args, **kwargs)
-
-    def exec_sql(self, fn, params=None):
+    def run_sql(self, fn, **kwargs):
         """Executes SQL files passed"""
-        # TODO: refactor this to exec_sql_file
-        secho(Path(fn).name, fg="cyan", bold=True)
-        run_sql_file(self.session, str(fn), params)
+        return iter(run_sql(self.session, fn, **kwargs))
 
-    def exec_sql_query(self, fn, params=None):
-        return run_sql_query_file(self.session, fn, params)
+    def exec_sql(self, sql, **kwargs):
+        """Executes SQL files passed"""
+        warnings.warn("exec_sql is deprecated. Use run_sql instead", DeprecationWarning)
+        return self.run_sql(sql, **kwargs)
 
-    def exec_query(self, *args):
+    def get_dataframe(self, *args):
         """Returns a Pandas DataFrame from a SQL query"""
-        return run_query(self.engine, *args)
+        return get_dataframe(self.engine, *args)
 
     @property
     def inspector(self):
@@ -134,12 +133,27 @@ class Database(object):
             model = getattr(self.model, model)
         return get_or_create(self.session, model, **kwargs)
 
+    def reflect_table(self, *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 using
+        `self.automap()`. Then, tables can be accessed using the
+        `self.table` 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
+        """
+        return reflect_table(self.engine, *args, **kwargs)
+
     @property
     def table(self):
         """
         Map of all tables in the database as SQLAlchemy table objects
         """
-        if self.mapper._tables is None:
+        if self.mapper is None or self.mapper._tables is None:
             self.automap()
         return self.mapper._tables
 
@@ -150,7 +164,7 @@ class Database(object):
 
         https://docs.sqlalchemy.org/en/latest/orm/extensions/automap.html
         """
-        if self.mapper._models is None:
+        if self.mapper is None or self.mapper._models is None:
             self.automap()
         return self.mapper._models
 

{macrostrat.database-1.0.2 → macrostrat_database-3.0.0b1}/macrostrat/database/mapper/__init__.py

@@ -1,12 +1,13 @@
-from sqlalchemy.schema import Table
-from sqlalchemy import MetaData
+from distutils.log import warn
+from macrostrat.database.utils import reflect_table
 from sqlalchemy.ext.automap import generate_relationship
 from macrostrat.utils.logs import get_logger
-from .cache import DatabaseModelCache
+from warnings import warn
 
 # Drag in geographic types for database reflection
 from geoalchemy2 import Geometry, Geography
 
+from .cache import DatabaseModelCache
 from .utils import (
     ModelCollection,
     TableCollection,
@@ -18,12 +19,15 @@ from .utils import (
 
 log = get_logger(__name__)
 
+
 class AutomapError(Exception):
     pass
 
+
 model_builder = DatabaseModelCache()
 BaseModel = model_builder.automap_base()
 
+
 class DatabaseMapper:
     automap_base = BaseModel
     automap_error = None
@@ -38,10 +42,16 @@ class DatabaseMapper:
 
         # This stuff should be placed outside of core (one likely extension point).
         self.reflection_kwargs = dict(
-            name_for_scalar_relationship=kwargs.get("name_for_scalar_relationship", name_for_scalar_relationship),
-            name_for_collection_relationship=kwargs.get("name_for_collection_relationship", name_for_collection_relationship),
+            name_for_scalar_relationship=kwargs.get(
+                "name_for_scalar_relationship", name_for_scalar_relationship
+            ),
+            name_for_collection_relationship=kwargs.get(
+                "name_for_collection_relationship", name_for_collection_relationship
+            ),
             classname_for_table=kwargs.get("classname_for_table", _classname_for_table),
-            generate_relationship=kwargs.get("generate_relationship", generate_relationship),
+            generate_relationship=kwargs.get(
+                "generate_relationship", generate_relationship
+            ),
         )
 
         self._models = ModelCollection(self.automap_base.classes)
@@ -66,45 +76,26 @@ class DatabaseMapper:
     def reflect_schema(self, schema, use_cache=True):
         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(
-                self.db.engine, schema=schema, **self.reflection_kwargs
-            )
+            self.automap_base.prepare(schema=schema, **self.reflection_kwargs)
             return
         log.info(f"Reflecting schema {schema}")
         if schema == "public":
-            self.automap_base.prepare(
-                self.db.engine, reflect=True, schema=None, **self.reflection_kwargs
-            )
-        else:
-            # Reflect tables in schemas we care about
-            # Note: this will not reflect views because they don't have primary keys.
-            self.automap_base.metadata.reflect(bind=self.db.engine, schema=schema, **self.reflection_kwargs)
+            schema = None
+        # Reflect tables in schemas we care about
+        # Note: this will not reflect views because they don't have primary keys.
+        self.automap_base.prepare(
+            autoload_with=self.db.engine, schema=schema, **self.reflection_kwargs
+        )
         self._models = ModelCollection(self.automap_base.classes)
         self._tables = TableCollection(self._models)
 
     def reflect_table(self, 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)
-        tables = Table(
-            tablename,
-            meta,
-            *column_args,
-            autoload=True,
-            autoload_with=self.db.engine,
-            **kwargs,
+        # Warn that this method is deprecated
+        warn(
+            "DatabaseMapper.reflect_table is deprecated. Use Database.reflect_table instead",
+            DeprecationWarning,
         )
-        return tables
+        return reflect_table(self.db.engine, tablename, *column_args, **kwargs)
 
     def reflect_view(self, tablename, *column_args, **kwargs):
         pass

{macrostrat.database-1.0.2 → macrostrat_database-3.0.0b1}/macrostrat/database/mapper/cache.py

@@ -7,16 +7,17 @@ from .base import ModelHelperMixins
 
 log = get_logger(__name__)
 
+
 class AutomapError(Exception):
     pass
 
+
 class DatabaseModelCache(object):
     cache_file = None
 
     def __init__(self, cache_file=None):
         self.cache_file = cache_file
 
-
     @property
     def _metadata_cache_filename(self):
         return self.cache_file
@@ -36,9 +37,7 @@ class DatabaseModelCache(object):
             log.info(f"Cached database models to {self.cache_file}")
         except IOError:
             # couldn't write the file for some reason
-            log.info(
-                f"Could not cache database models to {self.cache_file}"
-            )
+            log.info(f"Could not cache database models to {self.cache_file}")
 
     def _load_database_map(self):
         # We have hard-coded the cache file for now
@@ -54,6 +53,8 @@ class DatabaseModelCache(object):
             log.info(
                 f"Could not find database model cache ({self._metadata_cache_filename})"
             )
+        except Exception as exc:
+            log.error(f"Error loading database model cache: {exc}")
         return cached_metadata
 
     def automap_base(self):
@@ -66,4 +67,3 @@ class DatabaseModelCache(object):
             base.loaded_from_cache = True
         base.builder = self
         return base
-

{macrostrat.database-1.0.2 → macrostrat_database-3.0.0b1}/macrostrat/database/mapper/utils.py

@@ -6,6 +6,7 @@ def primary_key(instance):
     prop_list = [mapper.get_property_by_column(column) for column in mapper.primary_key]
     return {prop.key: getattr(instance, prop.key) for prop in prop_list}
 
+
 def classname_for_table(table):
     if table.schema is not None:
         return f"{table.schema}_{table.name}"
@@ -83,6 +84,14 @@ class ModelCollection(BaseCollection):
     def keys(self):
         return [k for k in self.__models.keys()]
 
+    # Support for dict-like access
+    def __getitem__(self, key):
+        return self.__models[key]
+
+    # Support 'in' operator
+    def __contains__(self, key):
+        return key in self.__models
+
 
 class TableCollection(BaseCollection):
     """

{macrostrat.database-1.0.2 → macrostrat_database-3.0.0b1}/macrostrat/database/postgresql.py

@@ -5,9 +5,12 @@ from sqlalchemy.exc import CompileError
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.sql.expression import Insert
 from sqlalchemy.dialects import postgresql
+import psycopg2
+
 
 _import_mode = ContextVar("import-mode", default="do-nothing")
 
+
 # https://stackoverflow.com/questions/33307250/postgresql-on-conflict-in-sqlalchemy/62305344#62305344
 @contextmanager
 def on_conflict(action="restrict"):
@@ -47,3 +50,17 @@ def prefix_inserts(insert, compiler, **kw):
             index_elements=insert.table.primary_key
         )
     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)

macrostrat_database-3.0.0b1/macrostrat/database/utils.py

@@ -0,0 +1,450 @@
+from click import echo, secho
+from sqlalchemy.exc import ProgrammingError, IntegrityError, InternalError
+from sqlparse import split, format
+from sqlalchemy.sql.elements import TextClause, ClauseElement
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.engine import Engine, Connection
+from sqlalchemy.schema import Table
+from sqlalchemy import MetaData, create_engine, text
+from contextlib import contextmanager
+from sqlalchemy_utils import create_database, database_exists, drop_database
+from sqlalchemy.exc import InvalidRequestError
+from macrostrat.utils import cmd, get_logger
+from time import sleep
+from typing import Union, IO
+from pathlib import Path
+from warnings import warn
+from psycopg2.sql import SQL, Composable, Composed
+from re import search
+from macrostrat.utils import get_logger
+from .postgresql import _setup_psycopg2_wait_callback
+
+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,
+    false if it should be interpreted as a file path.
+    """
+    # If it's a byte string, decode it
+    if isinstance(_string, bytes):
+        _string = _string.decode("utf-8")
+
+    keywords = [
+        "SELECT",
+        "INSERT",
+        "UPDATE",
+        "CREATE",
+        "DROP",
+        "DELETE",
+        "ALTER",
+        "SET",
+    ]
+    lines = _string.split("\n")
+    if len(lines) > 1:
+        return True
+    _string = _string.lower()
+    for i in keywords:
+        if _string.strip().startswith(i.lower()):
+            return True
+    return False
+
+
+def canonicalize_query(file_or_text: Union[str, Path, IO]) -> Union[str, Path]:
+    if isinstance(file_or_text, Path):
+        return file_or_text
+    # If it's a file-like object, read it
+    if hasattr(file_or_text, "read"):
+        return file_or_text.read()
+    # Otherwise, assume it's a string
+    if infer_is_sql_text(file_or_text):
+        return file_or_text
+    pth = Path(file_or_text)
+    if pth.exists() and pth.is_file():
+        return pth
+    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):
+    for line in sql.split("\n"):
+        for i in ["SELECT", "INSERT", "UPDATE", "CREATE", "DROP", "DELETE", "ALTER"]:
+            if not line.startswith(i):
+                continue
+            start = line.split("(")[0].strip().rstrip(";").replace(" AS", "")
+            secho(start, **kwargs)
+            return
+
+
+def get_sql_text(sql, interpret_as_file=None, echo_file_name=True):
+    if interpret_as_file:
+        sql = Path(sql).read_text()
+    elif interpret_as_file is None:
+        sql = canonicalize_query(sql)
+
+    if isinstance(sql, Path):
+        if echo_file_name:
+            secho(sql.name, fg="cyan", bold=True)
+        sql = sql.read_text()
+
+    return sql
+
+
+def _get_queries(sql, interpret_as_file=None):
+    if isinstance(sql, (list, tuple)):
+        queries = []
+        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):
+        return [sql]
+
+    if sql in [None, ""]:
+        return
+    if interpret_as_file:
+        sql = Path(sql).read_text()
+    elif interpret_as_file is None:
+        sql = canonicalize_query(sql)
+
+    if isinstance(sql, Path):
+        sql = sql.read_text()
+
+    return split(sql)
+
+
+def _is_prebind_param(param):
+    return isinstance(param, Composable)
+
+
+def _split_params(params):
+    if params is None:
+        return None, None
+    new_params = []
+    new_bind_params = []
+    if isinstance(params, (list, tuple)):
+        for i in params:
+            if _is_prebind_param(i):
+                new_bind_params.append(i)
+            else:
+                new_params.append(i)
+    elif isinstance(params, dict):
+        new_params = {}
+        new_bind_params = {}
+        for k, v in params.items():
+            if _is_prebind_param(v):
+                new_bind_params[k] = v
+            else:
+                new_params[k] = v
+    if len(new_bind_params) == 0:
+        new_bind_params = None
+    return new_params, new_bind_params
+
+
+def _get_cursor(connectable):
+    if isinstance(connectable, Engine):
+        conn = connectable.connect()
+
+    # Find a connection or cursor object for the connectable
+    conn = connectable
+    if hasattr(conn, "raw_connection"):
+        conn = conn.raw_connection()
+    while hasattr(conn, "driver_connection") or hasattr(conn, "connection"):
+        if hasattr(conn, "driver_connection"):
+            conn = conn.driver_connection
+        else:
+            conn = conn.connection
+        if callable(conn):
+            conn = conn()
+    if hasattr(conn, "cursor"):
+        conn = conn.cursor()
+
+    return conn
+
+
+def _get_connection(connectable) -> Connection:
+    if isinstance(connectable, Engine):
+        return connectable.connect()
+    if isinstance(connectable, Connection):
+        return connectable
+    if not hasattr(connectable, "connection"):
+        return connectable
+    conn = connectable.connection
+    if callable(conn):
+        return conn()
+    return conn
+
+
+def _render_query(query: Union[SQL, Composed], connectable: Union[Engine, Connection]):
+    """Render a query to a SQL string."""
+    if not isinstance(query, (Composed, SQL)):
+        return query
+    # Find a connection or cursor object for the connectable
+    conn = _get_cursor(connectable)
+    return query.as_string(conn)
+
+
+def infer_has_server_binds(sql):
+    return "%s" in sql or search(r"%\(\w+\)s", sql)
+
+
+def _run_sql(connectable, sql, **kwargs):
+    """
+    Internal function for running a query on a SQLAlchemy connectable,
+    which always returns an iterator. The wrapper function adds the option
+    to return a list of results.
+    """
+    if isinstance(connectable, Engine):
+        with connectable.connect() as conn:
+            yield from _run_sql(conn, sql, **kwargs)
+            return
+
+    _setup_psycopg2_wait_callback()
+
+    params = kwargs.pop("params", None)
+    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)
+
+    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)
+
+    queries = _get_queries(sql, interpret_as_file=interpret_as_file)
+
+    if queries is None:
+        return
+
+    # 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):
+        trans = None
+        try:
+            trans = connectable.begin()
+        except InvalidRequestError:
+            trans = None
+        try:
+            params, pre_bind_params = _split_params(params)
+
+            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)
+
+            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()
+                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)
+
+            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()
+            pretty_print(sql_text, dim=True)
+        except (ProgrammingError, IntegrityError, InternalError) as err:
+            _err = str(err.orig).strip()
+            dim = "already exists" in _err
+            if trans is not None:
+                trans.rollback()
+            elif hasattr(connectable, "rollback"):
+                connectable.rollback()
+            pretty_print(sql_text, fg=None if dim else "red", dim=True)
+            if dim:
+                _err = "  " + _err
+            secho(_err, fg="red", dim=dim)
+            log.error(err)
+            if raise_errors:
+                raise err
+
+
+def run_sql_file(connectable, filename, **kwargs):
+    return run_sql(connectable, filename, interpret_as_file=True, **kwargs)
+
+
+def run_sql(*args, **kwargs):
+    """
+    Run a query on a SQLAlchemy connectable.
+
+    Parameters
+    ----------
+    connectable : Union[Engine, Connection]
+        A SQLAlchemy engine or connection object.
+    sql : Union[str, Path, IO, SQL, Composed]
+        A SQL query, or a file containing a SQL query.
+    params : Union[dict, list, tuple]
+        Parameters to bind to the query. If a list or tuple, the parameters
+        will be bound to the query in order. If a dict, the parameters will
+        be bound to the query by name.
+    stop_on_error : bool
+        If True, stop running queries if an error is encountered.
+    raise_errors : bool
+        If True, raise errors encountered while running queries.
+    has_server_binds : bool
+        Interpret the query to have server-side bind parameters (requiring execution
+        with the backend driver). By default, this is inferred from the query string,
+        but inference is not always reliable.
+    interpret_as_file : bool
+        If True, force interpreting the query as a file path.
+    yield_results : bool
+        If True, yield the results of the query as they are executed, rather than
+        returning a list after completion.
+    """
+    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):
+    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)
+        return res
+    except (ProgrammingError, IntegrityError) as err:
+        err = str(err.orig).strip()
+        dim = "already exists" in err
+        if hasattr(connectable, "rollback"):
+            connectable.rollback()
+        pretty_print(sql, fg=None if dim else "red", dim=True)
+        if dim:
+            err = "  " + err
+        secho(err, fg="red", dim=dim)
+        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."""
+    if ensure_empty:
+        drop_database(conn_string)
+    if not database_exists(conn_string):
+        create_database(conn_string)
+    try:
+        yield create_engine(conn_string)
+    finally:
+        if drop:
+            drop_database(conn_string)
+
+
+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-3.0.0b1/pyproject.toml

@@ -0,0 +1,23 @@
+[tool.poetry]
+authors = ["Daven Quinn <dev@davenquinn.com>"]
+description = "A SQLAlchemy-based database toolkit."
+name = "macrostrat.database"
+packages = [{ include = "macrostrat" }]
+version = "3.0.0-beta1"
+
+[tool.poetry.dependencies]
+GeoAlchemy2 = "^0.14.0"
+SQLAlchemy = "^2.0.18"
+SQLAlchemy-Utils = "^0.41.1"
+click = "^8.1.3"
+"macrostrat.utils" = "^1.0.0"
+psycopg2-binary = "^2.9.6"
+python = "^3.8"
+sqlparse = "^0.4.4"
+
+[tool.poetry.dev-dependencies]
+"macrostrat.utils" = { path = "../utils", develop = true }
+
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = ["poetry-core>=1.0.0"]

macrostrat.database-1.0.2/PKG-INFO

@@ -1,20 +0,0 @@
-Metadata-Version: 2.1
-Name: macrostrat.database
-Version: 1.0.2
-Summary: A small library based on SQLAlchemy to assist with common database tasks.
-Author: Daven Quinn
-Author-email: dev@davenquinn.com
-Requires-Python: >=3.8,<4.0
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Requires-Dist: GeoAlchemy2 (>=0.9.4,<0.10.0)
-Requires-Dist: SQLAlchemy (>=1.4.26,<2.0.0)
-Requires-Dist: SQLAlchemy-Utils (>=0.37.0,<0.38.0)
-Requires-Dist: click (>=8.1.3,<9.0.0)
-Requires-Dist: macrostrat.utils (>=1.0.0,<2.0.0)
-Requires-Dist: migra (>=3.0.1621480950,<4.0.0)
-Requires-Dist: psycopg2-binary (>=2.9.1,<3.0.0)
-Requires-Dist: schemainspect (>=3.0.1616029793,<4.0.0)
-Requires-Dist: sqlparse (>=0.4.0,<0.5.0)

macrostrat.database-1.0.2/macrostrat/database/utils.py

@@ -1,164 +0,0 @@
-from click import echo, secho
-from sqlalchemy.exc import ProgrammingError, IntegrityError
-from sqlparse import split, format
-from sqlalchemy.sql import ClauseElement
-from sqlalchemy import create_engine, text
-from sqlalchemy.orm import sessionmaker
-from contextlib import contextmanager
-from sqlalchemy_utils import create_database, database_exists, drop_database
-from macrostrat.utils import cmd, get_logger
-from time import sleep
-
-log = get_logger(__name__)
-
-
-def db_session(engine):
-    factory = sessionmaker(bind=engine)
-    return factory()
-
-
-def run_query(db, 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
-
-    if "SELECT" in str(filename_or_query):
-        # We are working with a query string instead of
-        # an SQL file.
-        sql = filename_or_query
-    else:
-        with open(filename_or_query) as f:
-            sql = f.read()
-
-    return read_sql(sql, db, **kwargs)
-
-
-def pretty_print(sql, **kwargs):
-    for line in sql.split("\n"):
-        for i in ["SELECT", "INSERT", "UPDATE", "CREATE", "DROP", "DELETE", "ALTER"]:
-            if not line.startswith(i):
-                continue
-            start = line.split("(")[0].strip().rstrip(";").replace(" AS", "")
-            secho(start, **kwargs)
-            return
-
-
-def run_sql(session, sql, params=None, stop_on_error=False):
-    queries = split(sql)
-    for q in queries:
-        sql = format(q, strip_comments=True).strip()
-        if sql == "":
-            continue
-        try:
-            session.execute(text(sql), params=params)
-            if hasattr(session, "commit"):
-                session.commit()
-            pretty_print(sql, dim=True)
-        except (ProgrammingError, IntegrityError) as err:
-            err = str(err.orig).strip()
-            dim = "already exists" in err
-            if hasattr(session, "rollback"):
-                session.rollback()
-            pretty_print(sql, fg=None if dim else "red", dim=True)
-            if dim:
-                err = "  " + err
-            secho(err, fg="red", dim=dim)
-            if stop_on_error:
-                return
-
-
-def _exec_raw_sql(engine, sql):
-    """Execute SQL unsafely on an sqlalchemy Engine"""
-    try:
-        engine.execute(text(sql))
-        pretty_print(sql, dim=True)
-    except (ProgrammingError, IntegrityError) as err:
-        err = str(err.orig).strip()
-        dim = "already exists" in err
-        pretty_print(sql, fg=None if dim else "red", dim=True)
-        if dim:
-            err = "  " + err
-        secho(err, fg="red", dim=dim)
-
-
-def run_sql_file(session, sql_file, params=None):
-    sql = open(sql_file).read()
-    run_sql(session, sql, params=params)
-
-
-def run_sql_query_file(session, sql_file, params=None):
-    sql = open(sql_file).read()
-    return session.execute(sql, params)
-
-
-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."""
-    if ensure_empty:
-        drop_database(conn_string)
-    if not database_exists(conn_string):
-        create_database(conn_string)
-    try:
-        yield create_engine(conn_string)
-    finally:
-        if drop:
-            drop_database(conn_string)
-
-
-def connection_args(engine):
-    """Get PostgreSQL connection arguments for a 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)

macrostrat.database-1.0.2/pyproject.toml

@@ -1,27 +0,0 @@
-[tool.poetry]
-authors = ["Daven Quinn <dev@davenquinn.com>"]
-description = "A small library based on SQLAlchemy to assist with common database tasks."
-name = "macrostrat.database"
-packages = [
-  {include = "macrostrat"},
-]
-version = "1.0.2"
-
-[tool.poetry.dependencies]
-GeoAlchemy2 = "^0.9.4"
-SQLAlchemy = "^1.4.26"
-SQLAlchemy-Utils = "^0.37.0"
-click = "^8.1.3"
-"macrostrat.utils" = "^1.0.0"
-migra = "^3.0.1621480950"
-psycopg2-binary = "^2.9.1"
-python = "^3.8"
-schemainspect = "^3.0.1616029793"
-sqlparse = "^0.4.0"
-
-[tool.poetry.dev-dependencies]
-"macrostrat.utils" = {path = "../utils", develop = true}
-
-[build-system]
-build-backend = "poetry.core.masonry.api"
-requires = ["poetry-core>=1.0.0"]

macrostrat.database-1.0.2/setup.py

@@ -1,38 +0,0 @@
-# -*- coding: utf-8 -*-
-from setuptools import setup
-
-packages = \
-['macrostrat', 'macrostrat.database', 'macrostrat.database.mapper']
-
-package_data = \
-{'': ['*']}
-
-install_requires = \
-['GeoAlchemy2>=0.9.4,<0.10.0',
- 'SQLAlchemy-Utils>=0.37.0,<0.38.0',
- 'SQLAlchemy>=1.4.26,<2.0.0',
- 'click>=8.1.3,<9.0.0',
- 'macrostrat.utils>=1.0.0,<2.0.0',
- 'migra>=3.0.1621480950,<4.0.0',
- 'psycopg2-binary>=2.9.1,<3.0.0',
- 'schemainspect>=3.0.1616029793,<4.0.0',
- 'sqlparse>=0.4.0,<0.5.0']
-
-setup_kwargs = {
-    'name': 'macrostrat.database',
-    'version': '1.0.2',
-    'description': 'A small library based on SQLAlchemy to assist with common database tasks.',
-    'long_description': None,
-    'author': 'Daven Quinn',
-    'author_email': 'dev@davenquinn.com',
-    'maintainer': None,
-    'maintainer_email': None,
-    'url': None,
-    'packages': packages,
-    'package_data': package_data,
-    'install_requires': install_requires,
-    'python_requires': '>=3.8,<4.0',
-}
-
-
-setup(**setup_kwargs)