macrostrat.database 3.4.1__tar.gz → 3.5.1__tar.gz

{macrostrat_database-3.4.1 → macrostrat_database-3.5.1}/PKG-INFO

@@ -1,13 +1,11 @@
 Metadata-Version: 2.1
 Name: macrostrat.database
-Version: 3.4.1
+Version: 3.5.1
 Summary: A SQLAlchemy-based database toolkit.
 Author: Daven Quinn
 Author-email: dev@davenquinn.com
-Requires-Python: >=3.8,<4.0
+Requires-Python: >=3.10,<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
 Classifier: Programming Language :: Python :: 3.12
@@ -16,7 +14,7 @@ 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.0.0,<2.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.4.1 → macrostrat_database-3.5.1}/macrostrat/database/__init__.py

@@ -128,7 +128,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.

{macrostrat_database-3.4.1 → macrostrat_database-3.5.1}/macrostrat/database/transfer/dump_database.py

@@ -1,4 +1,5 @@
 import asyncio
+import sys
 from pathlib import Path
 from typing import Optional
 
@@ -6,8 +7,8 @@ import aiofiles
 from sqlalchemy.engine import Engine
 
 from macrostrat.utils import get_logger
-
-from .utils import _create_command, print_stdout, print_stream_progress
+from .stream_utils import print_stdout, print_stream_progress
+from .utils import _create_command
 
 log = get_logger(__name__)
 
@@ -45,11 +46,22 @@ async def pg_dump(
     )
 
 
-async def pg_dump_to_file(dumpfile: Path, *args, **kwargs):
-    proc = await pg_dump(*args, **kwargs)
+async def pg_dump_to_file(engine: Engine, dumpfile: Path | None, **kwargs):
+    proc = await pg_dump(engine, **kwargs)
+    if dumpfile is None or dumpfile == sys.stdout:
+        # If we have no dumpfile, just print to stdout
+        await _monitor_stdout(proc)
+        return
     # Open dump file as an async stream
     async with aiofiles.open(dumpfile, mode="wb") as dest:
         await asyncio.gather(
             asyncio.create_task(print_stream_progress(proc.stdout, dest)),
             asyncio.create_task(print_stdout(proc.stderr)),
         )
+
+
+async def _monitor_stdout(proc):
+    await asyncio.gather(
+        asyncio.create_task(print_stdout(proc.stdout)),
+        asyncio.create_task(print_stream_progress(proc.stderr, None)),
+    )

{macrostrat_database-3.4.1 → macrostrat_database-3.5.1}/macrostrat/database/transfer/move_tables.py

@@ -3,10 +3,9 @@ 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 .utils import print_stdout, print_stream_progress
+from .stream_utils import print_stdout, print_stream_progress
 
 log = get_logger(__name__)
 

{macrostrat_database-3.4.1 → macrostrat_database-3.5.1}/macrostrat/database/transfer/restore_database.py

@@ -7,13 +7,8 @@ from rich.console import Console
 from sqlalchemy.engine import Engine
 
 from macrostrat.utils import get_logger
-
-from .utils import (
-    _create_command,
-    _create_database_if_not_exists,
-    print_stdout,
-    print_stream_progress,
-)
+from .stream_utils import print_stdout, print_stream_progress
+from .utils import _create_command, _create_database_if_not_exists
 
 console = Console()
 
@@ -56,11 +51,13 @@ async def pg_restore(
     )
 
 
-async def pg_restore_from_file(dumpfile: Path, *args, **kwargs):
-    proc = await pg_restore(*args, **kwargs)
+async def pg_restore_from_file(dumpfile: Path, engine: Engine, **kwargs):
+    proc = await pg_restore(engine, **kwargs)
     # Open dump file as an async stream
     async with aiofiles.open(dumpfile, mode="rb") as source:
         await asyncio.gather(
-            asyncio.create_task(print_stream_progress(source, proc.stdin)),
+            asyncio.create_task(
+                print_stream_progress(source, proc.stdin, prefix="Restored")
+            ),
             asyncio.create_task(print_stdout(proc.stderr)),
         )

macrostrat_database-3.5.1/macrostrat/database/transfer/stream_utils.py

@@ -0,0 +1,118 @@
+import asyncio
+import sys
+import zlib
+
+from aiofiles.threadpool import AsyncBufferedIOBase
+
+from macrostrat.utils import get_logger
+from .utils import console
+
+log = get_logger(__name__)
+
+
+async def print_stream_progress(
+    input: asyncio.StreamReader | asyncio.subprocess.Process,
+    out_stream: asyncio.StreamWriter | None | AsyncBufferedIOBase = None,
+    *,
+    verbose: bool = False,
+    chunk_size: int = 1024,
+    prefix: str = None,
+):
+    """This should be unified with print_stream_progress, but there seem to be
+    slight API differences between aiofiles and asyncio.StreamWriter APIs.?"""
+    in_stream = input
+    if isinstance(in_stream, asyncio.subprocess.Process):
+        in_stream = input.stdout
+
+    megabytes_written = 0
+    i = 0
+
+    # Iterate over the stream by chunks
+    try:
+        while True:
+            chunk = await in_stream.read(chunk_size)
+            if not chunk:
+                log.info("End of stream")
+                break
+            if verbose:
+                log.info(chunk)
+            megabytes_written += len(chunk) / 1_000_000
+            if isinstance(out_stream, AsyncBufferedIOBase):
+                await out_stream.write(chunk)
+                await out_stream.flush()
+            elif out_stream is not None:
+                out_stream.write(chunk)
+                await out_stream.drain()
+            i += 1
+            if i == 100:
+                i = 0
+                _print_progress(megabytes_written, end="\r", prefix=prefix)
+    except asyncio.CancelledError:
+        pass
+    finally:
+        _print_progress(megabytes_written, prefix=prefix)
+
+        if isinstance(out_stream, AsyncBufferedIOBase):
+            out_stream.close()
+        elif out_stream is not None:
+            out_stream.close()
+            await out_stream.wait_closed()
+
+
+def _print_progress(megabytes: float, **kwargs):
+    prefix = kwargs.pop("prefix", None)
+    if prefix is None:
+        prefix = "Dumped"
+    progress = f"{prefix} {megabytes:.1f} MB"
+    kwargs["file"] = sys.stderr
+    print(progress, **kwargs)
+
+
+async def print_stdout(stream: asyncio.StreamReader):
+    async for line in stream:
+        log.info(line)
+        console.print(line.decode("utf-8"), style="dim")
+
+
+class DecodingStreamReader(asyncio.StreamReader):
+    """A StreamReader that decompresses gzip files (if compressed)"""
+
+    # https://ejosh.co/de/2022/08/stream-a-massive-gzipped-json-file-in-python/
+
+    def __init__(self, stream, encoding="utf-8", errors="strict"):
+        super().__init__()
+        self.stream = stream
+        self._is_gzipped = None
+        self.d = zlib.decompressobj(zlib.MAX_WBITS | 16)
+
+    def decompress(self, input: bytes) -> bytes:
+        decompressed = self.d.decompress(input)
+        data = b""
+        while self.d.unused_data != b"":
+            buf = self.d.unused_data
+            self.d = zlib.decompressobj(zlib.MAX_WBITS | 16)
+            data = self.d.decompress(buf)
+        return decompressed + data
+
+    def transform_data(self, data):
+        if self._is_gzipped is None:
+            self._is_gzipped = data[:2] == b"\x1f\x8b"
+            log.info("is_gzipped: %s", self._is_gzipped)
+        if self._is_gzipped:
+            # Decompress the data
+            data = self.decompress(data)
+        return data
+
+    async def read(self, n=-1):
+        data = await self.stream.read(n)
+        return self.transform_data(data)
+
+    async def readline(self):
+        res = b""
+        while res == b"":
+            # Read next line
+            line = await self.stream.readline()
+            if not line:
+                break
+            res += self.transform_data(line)
+        return res

macrostrat_database-3.5.1/macrostrat/database/transfer/utils.py

@@ -0,0 +1,97 @@
+from urllib.parse import quote
+
+from rich.console import Console
+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
+
+console = Console()
+
+log = get_logger(__name__)
+
+
+def _docker_local_run_args(postgres_container: str = "postgres:15"):
+    return [
+        "docker",
+        "run",
+        "-i",
+        "--network",
+        "host",
+        "--attach",
+        "stdin",
+        "--attach",
+        "stdout",
+        "--attach",
+        "stderr",
+        "--log-driver",
+        "none",
+        "--rm",
+        postgres_container,
+    ]
+
+
+def _create_database_if_not_exists(
+    _url: URL, *, create=False, allow_exists=True, overwrite=False
+):
+    database = _url.database
+    if overwrite:
+        create = True
+    db_exists = database_exists(_url)
+    if db_exists:
+        msg = f"Database [bold underline]{database}[/] already exists"
+        if overwrite:
+            console.print(f"{msg}, overwriting")
+            drop_database(_url)
+            db_exists = False
+        elif not allow_exists:
+            raise ApplicationError(msg, details="Use `--overwrite` to overwrite")
+        else:
+            console.print(msg)
+
+    if create and not db_exists:
+        console.print(f"Creating database [bold cyan]{database}[/]")
+        create_database(_url)
+
+    if not db_exists and not create:
+        raise ApplicationError(
+            f"Database [bold cyan]{database}[/] does not exist. ",
+            "Use `--create` to create it.",
+        )
+
+
+def _create_command(
+    engine: Engine,
+    *command,
+    args=None | list[str],
+    prefix=None | list[str],
+    container="postgres:16",
+):
+    if args is None:
+        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,
+        ]
+
+    return _cmd
+
+
+def raw_database_url(url: URL):
+    """Replace the password asterisks with the actual password, in order to pass to other commands."""
+    _url = str(url)
+    if "***" not in _url or url.password is None:
+        return _url
+    return _url.replace("***", quote(url.password, safe=""))

{macrostrat_database-3.4.1 → macrostrat_database-3.5.1}/macrostrat/database/utils.py

@@ -5,11 +5,11 @@ 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
@@ -18,7 +18,7 @@ from sqlalchemy.exc import (
     InternalError,
     InvalidRequestError,
     ProgrammingError,
-    OperationalError
+    OperationalError,
 )
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.schema import Table
@@ -232,6 +232,9 @@ def infer_has_server_binds(sql):
     return "%s" in sql or search(r"%\(\w+\)s", sql)
 
 
+_default_statement_filter = lambda sql_text, params: True
+
+
 def _run_sql(connectable, sql, params=None, **kwargs):
     """
     Internal function for running a query on a SQLAlchemy connectable,
@@ -247,6 +250,7 @@ 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)
 
     if stop_on_error:
         raise_errors = True
@@ -288,6 +292,11 @@ 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)
+        if not should_run:
+            pretty_print(sql_text, dim=True, strikethrough=True)
+            continue
+
         # This only does something for postgresql, but it's harmless to run it for other engines
         set_wait_callback(wait_select)
 
@@ -325,7 +334,9 @@ def _run_sql(connectable, sql, params=None, **kwargs):
 
 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,7 +347,10 @@ 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
@@ -444,6 +458,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):

{macrostrat_database-3.4.1 → macrostrat_database-3.5.1}/pyproject.toml

@@ -3,16 +3,16 @@ authors = ["Daven Quinn <dev@davenquinn.com>"]
 description = "A SQLAlchemy-based database toolkit."
 name = "macrostrat.database"
 packages = [{ include = "macrostrat" }]
-version = "3.4.1"
+version = "3.5.1"
 
 [tool.poetry.dependencies]
 GeoAlchemy2 = "^0.15.2"
 SQLAlchemy = "^2.0.18"
 SQLAlchemy-Utils = "^0.41.1"
 click = "^8.1.3"
-"macrostrat.utils" = "^1.0.0"
+"macrostrat.utils" = "^1.3.0"
 psycopg2-binary = "^2.9.6"
-python = "^3.8"
+python = "^3.10"
 sqlparse = "^0.5.1"
 aiofiles = "^23.2.1"
 rich = "^13.7.1"

macrostrat_database-3.4.1/macrostrat/database/transfer/utils.py

@@ -1,107 +0,0 @@
-import asyncio
-from urllib.parse import quote
-
-from aiofiles.threadpool.binary import AsyncBufferedIOBase
-from rich.console import Console
-from sqlalchemy.engine import Engine
-from sqlalchemy.engine.url import URL
-from sqlalchemy_utils import create_database, database_exists
-
-from macrostrat.utils import get_logger
-
-console = Console()
-
-log = get_logger(__name__)
-
-
-def _docker_local_run_args(postgres_container: str = "postgres:15"):
-    return [
-        "docker",
-        "run",
-        "-i",
-        "--rm",
-        "--network",
-        "host",
-        postgres_container,
-    ]
-
-
-def _create_database_if_not_exists(_url: URL, create=False):
-    database = _url.database
-    db_exists = database_exists(_url)
-    if db_exists:
-        console.print(f"Database [bold cyan]{database}[/] already exists")
-
-    if create and not db_exists:
-        console.print(f"Creating database [bold cyan]{database}[/]")
-        create_database(_url)
-
-    if not db_exists and not create:
-        raise ValueError(
-            f"Database [bold cyan]{database}[/] does not exist. "
-            "Use `--create` to create it."
-        )
-
-
-def _create_command(
-    engine: Engine,
-    *command,
-    args=[],
-    prefix=None | list[str],
-    container="postgres:16",
-):
-    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,
-        ]
-
-    return _cmd
-
-
-async def print_stream_progress(
-    in_stream: asyncio.StreamReader,
-    out_stream: asyncio.StreamWriter | AsyncBufferedIOBase,
-):
-    """This should be unified with print_stream_progress, but there seem to be
-    slight API differences between aiofiles and asyncio.StreamWriter APIs.?"""
-    megabytes_written = 0
-    i = 0
-    async for line in in_stream:
-        megabytes_written += len(line) / 1_000_000
-        if isinstance(out_stream, AsyncBufferedIOBase):
-            await out_stream.write(line)
-            await out_stream.flush()
-        else:
-            out_stream.write(line)
-            await out_stream.drain()
-        i += 1
-        if i == 1000:
-            i = 0
-            _print_progress(megabytes_written, end="\r")
-
-    out_stream.close()
-    _print_progress(megabytes_written)
-
-
-def _print_progress(megabytes: float, **kwargs):
-    progress = f"Dumped {megabytes:.1f} MB"
-    print(progress, **kwargs)
-
-
-async def print_stdout(stream: asyncio.StreamReader):
-    async for line in stream:
-        console.print(line.decode("utf-8"), style="dim")
-
-
-def raw_database_url(url: URL):
-    return str(url).replace("***", quote(url.password, safe=""))