macrostrat.database 3.4.1__tar.gz → 3.5.0__tar.gz

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

@@ -1,13 +1,11 @@
 Metadata-Version: 2.1
 Name: macrostrat.database
-Version: 3.4.1
+Version: 3.5.0
 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.0}/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.0}/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.0}/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.0/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.0/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.0}/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.0"
 
 [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=""))