macrostrat.database 3.5.0__tar.gz → 3.5.2__tar.gz

{macrostrat_database-3.5.0 → macrostrat_database-3.5.2}/PKG-INFO

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

{macrostrat_database-3.5.0 → macrostrat_database-3.5.2}/macrostrat/database/__init__.py

@@ -5,13 +5,14 @@ from typing import Optional, Union
 
 from psycopg2.errors import InvalidSavepointSpecification
 from psycopg2.sql import Identifier
-from sqlalchemy import URL, MetaData, create_engine, inspect, Engine
+from sqlalchemy import URL, Engine, MetaData, create_engine, inspect
 from sqlalchemy.exc import IntegrityError, InternalError
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 from sqlalchemy.sql.expression import Insert
 
 from macrostrat.utils import get_logger
+
 from .mapper import DatabaseMapper
 from .postgresql import on_conflict, prefix_inserts  # noqa
 from .utils import (  # noqa
@@ -128,7 +129,7 @@ class Database(object):
         Returns: Iterator of results from the query.
         """
         params = self._setup_params(params, kwargs)
-        return iter(run_sql(self.session, fn, params, **kwargs))
+        return run_sql(self.session, fn, params, **kwargs)
 
     def run_query(self, sql, params=None, **kwargs):
         """Run a single query on the database object, returning the result.
@@ -184,6 +185,29 @@ class Database(object):
             self.__inspector__ = inspect(self.engine)
         return self.__inspector__
 
+    def refresh_schema(self, *, automap=None):
+        """
+        Refresh the current database connection
+
+        - closes the session and flushes
+        - removes the inspector
+
+        If automap is True, will automap the database after refreshing.
+        If automap is False, will not automap the database after refreshing.
+        If automap is None, it will re-map the database if it was previously mapped.
+        """
+        # Close the session
+        self.session.flush()
+        self.session.close()
+        # Remove the inspector
+        self.__inspector__ = None
+
+        if automap is None:
+            automap = self.mapper is not None
+
+        if automap:
+            self.automap()
+
     def entity_names(self, **kwargs):
         """
         Returns an iterator of names of *schema objects*

{macrostrat_database-3.5.0 → macrostrat_database-3.5.2}/macrostrat/database/transfer/dump_database.py

@@ -7,6 +7,7 @@ import aiofiles
 from sqlalchemy.engine import Engine
 
 from macrostrat.utils import get_logger
+
 from .stream_utils import print_stdout, print_stream_progress
 from .utils import _create_command
 

{macrostrat_database-3.5.0 → macrostrat_database-3.5.2}/macrostrat/database/transfer/move_tables.py

@@ -3,6 +3,7 @@ import asyncio
 from sqlalchemy.engine import Engine
 
 from macrostrat.utils import get_logger
+
 from .dump_database import pg_dump
 from .restore_database import pg_restore
 from .stream_utils import print_stdout, print_stream_progress

{macrostrat_database-3.5.0 → macrostrat_database-3.5.2}/macrostrat/database/transfer/restore_database.py

@@ -7,6 +7,7 @@ from rich.console import Console
 from sqlalchemy.engine import Engine
 
 from macrostrat.utils import get_logger
+
 from .stream_utils import print_stdout, print_stream_progress
 from .utils import _create_command, _create_database_if_not_exists
 

{macrostrat_database-3.5.0 → macrostrat_database-3.5.2}/macrostrat/database/transfer/stream_utils.py

@@ -5,6 +5,7 @@ import zlib
 from aiofiles.threadpool import AsyncBufferedIOBase
 
 from macrostrat.utils import get_logger
+
 from .utils import console
 
 log = get_logger(__name__)

{macrostrat_database-3.5.0 → macrostrat_database-3.5.2}/macrostrat/database/transfer/utils.py

@@ -5,7 +5,7 @@ from sqlalchemy.engine import Engine
 from sqlalchemy.engine.url import URL
 from sqlalchemy_utils import create_database, database_exists, drop_database
 
-from macrostrat.utils import get_logger, ApplicationError
+from macrostrat.utils import ApplicationError, get_logger
 
 console = Console()
 

{macrostrat_database-3.5.0 → macrostrat_database-3.5.2}/macrostrat/database/utils.py

@@ -1,15 +1,17 @@
 from contextlib import contextmanager
+from enum import Enum
 from pathlib import Path
 from re import search
+from sys import stderr
 from time import sleep
 from typing import IO, Union
 from warnings import warn
 
+import psycopg2.errors
 from click import echo, secho
 from psycopg2.extensions import set_wait_callback
 from psycopg2.extras import wait_select
 from psycopg2.sql import SQL, Composable, Composed
-import psycopg2.errors
 from rich.console import Console
 from sqlalchemy import MetaData, create_engine, text
 from sqlalchemy.engine import Connection, Engine
@@ -17,8 +19,8 @@ from sqlalchemy.exc import (
     IntegrityError,
     InternalError,
     InvalidRequestError,
+    OperationalError,
     ProgrammingError,
-    OperationalError
 )
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.schema import Table
@@ -46,23 +48,11 @@ def infer_is_sql_text(_string: str) -> bool:
     if isinstance(_string, bytes):
         _string = _string.decode("utf-8")
 
-    keywords = [
-        "SELECT",
-        "INSERT",
-        "UPDATE",
-        "CREATE",
-        "DROP",
-        "DELETE",
-        "ALTER",
-        "SET",
-        "GRANT",
-        "WITH",
-    ]
     lines = _string.split("\n")
     if len(lines) > 1:
         return True
     _string = _string.lower()
-    for i in keywords:
+    for i in _sql_keywords:
         if _string.strip().startswith(i.lower() + " "):
             return True
     return False
@@ -97,26 +87,40 @@ def get_dataframe(connectable, filename_or_query, **kwargs):
 
 
 def pretty_print(sql, **kwargs):
+    """Print and optionally summarize an SQL query"""
+    summarize = kwargs.pop("summarize", True)
+    if summarize:
+        sql = summarize_statement(sql)
+    secho(sql, **kwargs)
+
+
+_sql_keywords = [
+    "SELECT",
+    "INSERT",
+    "UPDATE",
+    "CREATE",
+    "DROP",
+    "DELETE",
+    "ALTER",
+    "SET",
+    "GRANT",
+    "WITH",
+    "NOTIFY",
+    "COPY",
+]
+
+
+def summarize_statement(sql):
     for line in sql.split("\n"):
-        for i in [
-            "SELECT",
-            "INSERT",
-            "UPDATE",
-            "CREATE",
-            "DROP",
-            "DELETE",
-            "ALTER",
-            "SET",
-            "GRANT",
-            "WITH",
-            "NOTIFY",
-            "COPY",
-        ]:
+        for i in _sql_keywords:
             if not line.startswith(i):
                 continue
-            start = line.split("(")[0].strip().rstrip(";").replace(" AS", "")
-            secho(start, **kwargs)
-            return
+            return line.split("(")[0].strip().rstrip(";").replace(" AS", "")
+
+
+class DevNull(object):
+    def write(self, *_):
+        pass
 
 
 def get_sql_text(sql, interpret_as_file=None, echo_file_name=True):
@@ -232,6 +236,16 @@ def infer_has_server_binds(sql):
     return "%s" in sql or search(r"%\(\w+\)s", sql)
 
 
+_default_statement_filter = lambda sql_text, params: True
+
+
+class PrintMode(Enum):
+    NONE = "none"
+    ERRORS = "errors"
+    SUMMARY = "summary"
+    ALL = "all"
+
+
 def _run_sql(connectable, sql, params=None, **kwargs):
     """
     Internal function for running a query on a SQLAlchemy connectable,
@@ -247,6 +261,12 @@ def _run_sql(connectable, sql, params=None, **kwargs):
     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 = kwargs.pop("output_mode", PrintMode.SUMMARY)
+    output_file = kwargs.pop("output_file", stderr)
+
+    if output_mode == PrintMode.NONE:
+        output_file = DevNull()
 
     if stop_on_error:
         raise_errors = True
@@ -279,6 +299,7 @@ def _run_sql(connectable, sql, params=None, **kwargs):
             query = _render_query(query, connectable)
 
         sql_text = str(query)
+
         if isinstance(query, str):
             sql_text = format(query, strip_comments=True).strip()
             if sql_text == "":
@@ -288,6 +309,21 @@ def _run_sql(connectable, sql, params=None, **kwargs):
             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 != PrintMode.ALL:
+            sql_text = summarize_statement(sql_text)
+
+        if not should_run:
+            secho(
+                sql_text,
+                dim=True,
+                strikethrough=True,
+                file=output_file,
+            )
+            continue
+
         # This only does something for postgresql, but it's harmless to run it for other engines
         set_wait_callback(wait_select)
 
@@ -309,7 +345,7 @@ def _run_sql(connectable, sql, params=None, **kwargs):
                 trans.commit()
             elif hasattr(connectable, "commit"):
                 connectable.commit()
-            pretty_print(sql_text, dim=True)
+            secho(sql_text, dim=True, file=output_file)
         except Exception as err:
             if trans is not None:
                 trans.rollback()
@@ -318,14 +354,16 @@ def _run_sql(connectable, sql, params=None, **kwargs):
             if raise_errors or _should_raise_query_error(err):
                 raise err
 
-            _print_error(sql_text, err)
+            _print_error(sql_text, err, file=output_file)
         finally:
             set_wait_callback(None)
 
 
 def _should_raise_query_error(err):
     """Determine if an error should be raised for a query or not."""
-    if not isinstance(err, (ProgrammingError, IntegrityError, InternalError, OperationalError)):
+    if not isinstance(
+        err, (ProgrammingError, IntegrityError, InternalError, OperationalError)
+    ):
         return True
 
     orig_err = getattr(err, "orig", None)
@@ -336,23 +374,27 @@ def _should_raise_query_error(err):
     # We might want to change this behavior in the future, or support more graceful handling of errors from other
     # database backends.
     # Ideally we could handle operational errors more gracefully
-    if isinstance(orig_err, psycopg2.errors.QueryCanceled) or getattr(orig_err, "pgcode", None) == "57014":
+    if (
+        isinstance(orig_err, psycopg2.errors.QueryCanceled)
+        or getattr(orig_err, "pgcode", None) == "57014"
+    ):
         return True
 
     return False
 
 
-def _print_error(sql_text, err):
+def _print_error(sql_text, err, **kwargs):
     if orig := getattr(err, "orig", None):
         _err = str(orig)
     else:
         _err = str(err)
     _err = _err.strip()
-    dim = "already exists" in _err
-    pretty_print(sql_text, fg=None if dim else "red", dim=True)
+    # Decide whether error should be dimmed
+    dim = kwargs.pop("dim", "already exists" in _err)
+    secho(sql_text, fg=None if dim else "red", dim=True, **kwargs)
     if dim:
         _err = "  " + _err
-    secho(_err, fg="red", dim=dim)
+    secho(_err, fg="red", dim=dim, **kwargs)
     log.error(err)
 
 
@@ -444,6 +486,9 @@ def run_sql(*args, **kwargs):
         returning a list after completion.
     ensure_single_query : bool
         If True, raise an error if multiple queries are passed when only one is expected.
+    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.
     """
     res = _run_sql(*args, **kwargs)
     if kwargs.pop("yield_results", False):
@@ -451,7 +496,9 @@ def run_sql(*args, **kwargs):
     return list(res)
 
 
-def execute(connectable, sql, params=None, stop_on_error=False):
+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
@@ -460,17 +507,12 @@ def execute(connectable, sql, params=None, stop_on_error=False):
         res = connectable.execute(text(sql), params=params)
         if hasattr(connectable, "commit"):
             connectable.commit()
-        pretty_print(sql, dim=True)
+        pretty_print(sql, dim=True, file=output_file, mode=output_mode)
         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)
+        _print_error(sql, dim=True, file=output_file, mode=output_mode)
         if stop_on_error:
             return
     finally:

{macrostrat_database-3.5.0 → macrostrat_database-3.5.2}/pyproject.toml

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