macrostrat.database 4.1.0__tar.gz → 4.2.0__tar.gz

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

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

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

@@ -12,9 +12,9 @@ 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 .query import run_fixtures, run_query, run_sql, execute  # noqa
 from .utils import (  # noqa
     create_database,
     create_engine,
@@ -23,9 +23,6 @@ from .utils import (  # noqa
     get_dataframe,
     get_or_create,
     reflect_table,
-    run_fixtures,
-    run_query,
-    run_sql,
 )
 
 metadata = MetaData()

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

@@ -10,6 +10,10 @@ 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
@@ -25,11 +29,6 @@ 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__)
 
 
@@ -376,7 +375,7 @@ def escape_postgresql_cast_parameters(sql_text):
     for res in re.findall(regex, sql_text):
         param, cast_type = res
         sql_text = sql_text.replace(
-            f":{param}::{cast_type}", f":{param}\:\:{cast_type}"
+            f":{param}::{cast_type}", ":" + param + r"\:\:" + cast_type
         )
     return sql_text
 
@@ -596,3 +595,27 @@ def 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()

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

@@ -1,5 +1,7 @@
+import warnings
 from contextlib import contextmanager
 from time import sleep
+from uuid import uuid4
 
 from click import echo
 from sqlalchemy import MetaData
@@ -8,19 +10,16 @@ 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
+from .query import get_sql_text, execute  # noqa
 
 log = get_logger(__name__)
 
@@ -46,30 +45,6 @@ def db_session(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
@@ -97,9 +72,13 @@ def get_db_model(db, model_name: str):
 
 
 @contextmanager
-def temp_database(conn_string, drop=True, ensure_empty=False):
+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=True, replace=ensure_empty)
+    create_database(
+        conn_string, exists_ok=exists_ok, replace=ensure_empty, template=template
+    )
     try:
         engine = create_engine(conn_string)
         yield engine
@@ -109,6 +88,33 @@ def temp_database(conn_string, drop=True, ensure_empty=False):
             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.
 

{macrostrat_database-4.1.0 → macrostrat_database-4.2.0}/pyproject.toml

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