pgpack-dumper 0.3.0.0__cp314-cp314-macosx_11_0_arm64.whl

pgpack_dumper/__init__.py

@@ -0,0 +1,46 @@
+"""Library for read and write PGPack format between PostgreSQL and file."""
+
+from pgcopylib import (
+    PGCopyReader,
+    PGCopyWriter,
+)
+from pgpack import (
+    CompressionMethod,
+    PGPackReader,
+    PGPackWriter,
+)
+
+from .common import (
+    PGConnector,
+    CopyBuffer,
+    CopyBufferError,
+    CopyBufferObjectError,
+    CopyBufferTableNotDefined,
+    PGPackDumperError,
+    PGPackDumperReadError,
+    PGPackDumperWriteError,
+    PGPackDumperWriteBetweenError,
+)
+from .dumper import PGPackDumper
+from .version import __version__
+
+
+__all__ = (
+    "__version__",
+    "CompressionMethod",
+    "CopyBuffer",
+    "CopyBufferError",
+    "CopyBufferObjectError",
+    "CopyBufferTableNotDefined",
+    "PGConnector",
+    "PGCopyReader",
+    "PGCopyWriter",
+    "PGPackDumper",
+    "PGPackDumperError",
+    "PGPackDumperReadError",
+    "PGPackDumperWriteError",
+    "PGPackDumperWriteBetweenError",
+    "PGPackReader",
+    "PGPackWriter",
+)
+__author__ = "0xMihalich"

pgpack_dumper/common/__init__.py

@@ -0,0 +1,58 @@
+"""Common functions and classes."""
+
+from .columns import make_columns
+from .connector import PGConnector
+from .copy import CopyBuffer
+from .diagram import (
+    DBMetadata,
+    format_table,
+    transfer_diagram,
+)
+from .errors import (
+    CopyBufferError,
+    CopyBufferObjectError,
+    CopyBufferTableNotDefined,
+    PGPackDumperError,
+    PGPackDumperReadError,
+    PGPackDumperWriteError,
+    PGPackDumperWriteBetweenError,
+)
+from .logger import DumperLogger
+from .metadata import read_metadata
+from .query import (
+    chunk_query,
+    query_path,
+    query_template,
+    random_name,
+    search_object,
+)
+from .reader import CopyReader
+from .stream import StreamReader
+from .structs import PGObject
+
+
+__all__ = (
+    "CopyBuffer",
+    "CopyBufferError",
+    "CopyBufferObjectError",
+    "CopyBufferTableNotDefined",
+    "CopyReader",
+    "DBMetadata",
+    "DumperLogger",
+    "PGConnector",
+    "PGObject",
+    "PGPackDumperError",
+    "PGPackDumperReadError",
+    "PGPackDumperWriteBetweenError",
+    "PGPackDumperWriteError",
+    "StreamReader",
+    "chunk_query",
+    "format_table",
+    "make_columns",
+    "query_path",
+    "query_template",
+    "random_name",
+    "read_metadata",
+    "search_object",
+    "transfer_diagram",
+)

pgpack_dumper/common/columns.py

@@ -0,0 +1,30 @@
+from collections import OrderedDict
+
+from pgcopylib import PGOid
+from pgpack.common import PGParam
+
+
+def make_columns(
+    list_columns: list[str],
+    pgtypes: list[PGOid],
+    pgparam: list[PGParam],
+) -> OrderedDict[str, str]:
+    """Make DBMetadata.columns dictionary."""
+
+    columns = OrderedDict()
+
+    for col_name, pgtype, param in zip(
+        list_columns,
+        pgtypes,
+        pgparam,
+    ):
+        col_type = pgtype.name
+
+        if pgtype is PGOid.bpchar:
+            col_type = f"{col_type}({param.length})"
+        elif pgtype is PGOid.numeric:
+            col_type = f"{col_type}({param.length}, {param.scale})"
+
+        columns[col_name] = col_type
+
+    return columns

pgpack_dumper/common/connector.py

@@ -0,0 +1,11 @@
+from typing import NamedTuple
+
+
+class PGConnector(NamedTuple):
+    """Connector for PostgreSQL."""
+
+    host: str
+    dbname: str
+    user: str
+    password: str
+    port: int

pgpack_dumper/common/copy.py

@@ -0,0 +1,157 @@
+from logging import Logger
+from typing import (
+    Generator,
+    Iterator,
+)
+
+from psycopg import (
+    Copy,
+    Cursor,
+)
+
+from .errors import (
+    CopyBufferObjectError,
+    CopyBufferTableNotDefined,
+)
+from .query import (
+    query_template,
+    search_object,
+)
+from .structs import PGObject
+from .metadata import read_metadata
+
+
+class CopyBuffer:
+
+    def __init__(
+        self,
+        cursor: Cursor,
+        logger: Logger,
+        query: str | None = None,
+        table_name: str | None = None,
+    ) -> None:
+        """Class initialization."""
+
+        self.cursor = cursor
+        self.logger = logger
+        self.query = query
+        self.table_name = table_name
+
+    @property
+    def metadata(self) -> bytes:
+        """Get metadata as bytes."""
+
+        host = self.cursor.connection.info.host
+        self.logger.info(f"Start read metadata from host {host}.")
+        metadata = read_metadata(
+            self.cursor,
+            self.query,
+            self.table_name,
+        )
+        self.logger.info(f"Read metadata from host {host} done.")
+        return metadata
+
+    def copy_to(self) -> Iterator[Copy]:
+        """Get copy object from PostgreSQL."""
+
+        if not self.query and not self.table_name:
+            error_msg = "Query or table not defined."
+            self.logger.error(f"CopyBufferTableNotDefined: {error_msg}")
+            raise CopyBufferTableNotDefined(error_msg)
+
+        host = self.cursor.connection.info.host
+
+        if not self.query:
+            self.logger.info(f"Start read from {host}.{self.table_name}.")
+            self.cursor.execute(query_template("relkind").format(
+                table_name=self.table_name,
+            ))
+            relkind = self.cursor.fetchone()[0]
+            pg_object = PGObject[relkind]
+            if not pg_object.is_readable:
+                error_msg = f"Read from {pg_object} not support."
+                self.logger.error(f"CopyBufferObjectError: {error_msg}")
+                raise CopyBufferObjectError(error_msg)
+            self.logger.info(f"Use method read from {pg_object}.")
+            if not pg_object.is_readobject:
+                self.table_name = f"(select * from {self.table_name})"
+        elif self.query:
+            self.logger.info(f"Start read query from {host}.")
+            self.logger.info("Use method read from select.")
+            self.table_name = f"({self.query})"
+
+        return self.cursor.copy(
+            query_template("copy_to").format(table_name=self.table_name)
+        )
+
+    def copy_from(
+        self,
+        copyobj: Iterator[bytes],
+    ) -> None:
+        """Write PGCopy dump into PostgreSQL."""
+
+        if not self.table_name:
+            error_msg = "Table not defined."
+            self.logger.error(f"CopyBufferTableNotDefined: {error_msg}")
+            raise CopyBufferTableNotDefined(error_msg)
+
+        host = self.cursor.connection.info.host
+        size = 0
+        self.logger.info(f"Start write into {host}.{self.table_name}.")
+
+        with self.cursor.copy(
+            query_template("copy_from").format(table_name=self.table_name)
+        ) as cp:
+            for bytes_data in copyobj:
+                size += len(bytes_data)
+                cp.write(bytes_data)
+
+        self.logger.info(f"Successfully sending {size} bytes.")
+        self.logger.info(f"Write into {host}.{self.table_name} done.")
+
+    def copy_between(
+        self,
+        copy_buffer: "CopyBuffer",
+    ) -> None:
+        """Write from PostgreSQL into PostgreSQL."""
+
+        with copy_buffer.copy_to() as copy_to:
+            destination_host = self.cursor.connection.info.host
+            source_host = copy_buffer.cursor.connection.info.host
+            source_object = search_object(
+                copy_buffer.table_name,
+                copy_buffer.query,
+            )
+            size = 0
+            self.logger.info(
+                f"Copy {source_object} from {source_host} into "
+                f"{destination_host}.{self.table_name} started."
+            )
+
+            with self.cursor.copy(
+                query_template("copy_from").format(table_name=self.table_name)
+            ) as copy_from:
+                for data in copy_to:
+                    size += len(data)
+                    copy_from.write(data)
+
+            self.logger.info(f"Successfully sending {size} bytes.")
+            self.logger.info(
+                f"Copy {source_object} from {source_host}"
+                f"into {destination_host}.{self.table_name} done."
+            )
+
+    def copy_reader(self) -> Generator[bytes, None, None]:
+        """Read bytes from copy object."""
+
+        host = self.cursor.connection.info.host
+        source = search_object(
+            self.table_name,
+            self.query,
+        )
+
+        with self.copy_to() as copy_object:
+            for data in copy_object:
+                yield bytes(data)
+
+        self.logger.info(f"Read {source} from {host} done.")

pgpack_dumper/common/diagram.py

@@ -0,0 +1,78 @@
+from collections import OrderedDict
+from typing import NamedTuple
+
+
+class DBMetadata(NamedTuple):
+    """Database object."""
+
+    name: str
+    version: str
+    columns: OrderedDict
+
+
+def truncate_text(text: str, max_length: int) -> str:
+    """Truncate text and add ellipsis if too long."""
+
+    if len(text) > max_length:
+        return text[: max_length - 1] + "…"
+    return text
+
+
+def format_table(
+    metadata: DBMetadata,
+    direction: str,
+    table_width: int = 51,
+) -> list[str]:
+    """Format single table as list of lines."""
+
+    lines = []
+
+    title = f"{direction} [{metadata.name} {metadata.version}]"
+    lines.append(f"┌{''.ljust(table_width, '─')}┐")
+    lines.append(
+        f"│ {truncate_text(title, table_width - 1).ljust(table_width - 1)}│"
+    )
+    lines.append(f"╞{'═' * 25}╤{'═' * 25}╡")
+    lines.append(f"│ {'Column Name'.ljust(23)} │ {'Data Type'.ljust(23)} │")
+    lines.append(f"╞{'═' * 25}╪{'═' * 25}╡")
+
+    for i, (col_name, col_type) in enumerate(metadata.columns.items()):
+        truncated_name = truncate_text(col_name, 23)
+        truncated_type = truncate_text(str(col_type), 23)
+        lines.append(
+            f"│ {truncated_name.ljust(23)} │ {truncated_type.ljust(23)} │"
+        )
+        if i < len(metadata.columns) - 1:
+            lines.append(f"├{'─' * 25}┼{'─' * 25}┤")
+
+    lines.append(f"└{'─' * 25}┴{'─' * 25}┘")
+    return lines
+
+
+def transfer_diagram(source: DBMetadata, destination: DBMetadata) -> str:
+    """Make transfer diagram with two tables and arrow."""
+
+    src_lines = format_table(source, "Source")
+    dest_lines = format_table(destination, "Destination")
+    max_lines = max(len(src_lines), len(dest_lines), 9)
+
+    src_lines.extend([" " * 53] * (max_lines - len(src_lines)))
+    dest_lines.extend([" " * 53] * (max_lines - len(dest_lines)))
+
+    middle_line = max_lines // 2
+    arrow_config = [
+        (middle_line - 3, " │╲   "),
+        (middle_line - 2, " │ ╲  "),
+        (middle_line - 1, "┌┘  ╲ "),
+        (middle_line, "│    ╲"),
+        (middle_line + 1, "│    ╱"),
+        (middle_line + 2, "└┐  ╱ "),
+        (middle_line + 3, " │ ╱  "),
+        (middle_line + 4, " │╱   "),
+    ]
+    arrow_map = {line: arrow for line, arrow in arrow_config}
+
+    return "Transfer data diagram:\n" + "\n".join(
+        f"{src_lines[row]} {arrow_map.get(row, '      ')} {dest_lines[row]}"
+        for row in range(max_lines)
+    )

pgpack_dumper/common/errors.py

@@ -0,0 +1,26 @@
+class CopyBufferError(Exception):
+    """CopyBuffer base error."""
+
+
+class CopyBufferObjectError(TypeError):
+    """Destination object not support."""
+
+
+class CopyBufferTableNotDefined(ValueError):
+    """Destination table not defined."""
+
+
+class PGPackDumperError(Exception):
+    """PGPackDumper base error."""
+
+
+class PGPackDumperReadError(PGPackDumperError):
+    """PGPackDumper read error."""
+
+
+class PGPackDumperWriteError(PGPackDumperError):
+    """PGPackDumper write error."""
+
+
+class PGPackDumperWriteBetweenError(PGPackDumperWriteError):
+    """PGPackDumper write between error."""

pgpack_dumper/common/logger.py

@@ -0,0 +1,70 @@
+from datetime import datetime
+from logging import (
+    DEBUG,
+    FileHandler,
+    Formatter,
+    Logger,
+    StreamHandler,
+)
+from os import makedirs
+from os.path import dirname
+from sys import stdout
+
+from ..version import __version__
+
+
+def root_dir() -> str:
+    """Get project directory."""
+
+    try:
+        import __main__
+
+        return dirname(__main__.__file__)
+    except AttributeError:
+        return ""
+
+
+class DumperLogger(Logger):
+    """PGPackDumper logger."""
+
+    def __init__(
+            self,
+            level: int = DEBUG,
+            use_console: bool = True,
+    ) -> None:
+        """Class initialize."""
+
+        super().__init__("PGPackDumper")
+
+        self.fmt = (
+            f"%(asctime)s | %(levelname)-8s | ver {__version__} "
+            "| %(funcName)s-%(filename)s-%(lineno)04d <%(message)s>"
+        )
+        self.setLevel(level)
+        self.log_path = f"{root_dir()}/pgpack_logs"
+        makedirs(self.log_path, exist_ok=True)
+
+        formatter = Formatter(
+            fmt=self.fmt,
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+
+        file_handler = FileHandler(
+            "{}/{:%Y-%m-%d}_{}.log".format(
+                self.log_path,
+                datetime.now(),
+                self.name,
+            ),
+            encoding="utf-8",
+        )
+        file_handler.setLevel(DEBUG)
+        file_handler.setFormatter(formatter)
+        self.addHandler(file_handler)
+
+        if use_console:
+            console_handler = StreamHandler(stdout)
+            console_handler.setLevel(level)
+            console_handler.setFormatter(formatter)
+            self.addHandler(console_handler)
+
+        self.propagate = False

pgpack_dumper/common/metadata.py

@@ -0,0 +1,38 @@
+from psycopg import Cursor
+
+from .query import (
+    query_template,
+    random_name,
+)
+
+
+def read_metadata(
+    cursor: Cursor,
+    query: str | None = None,
+    table_name: str | None = None,
+) -> bytes:
+    """Read metadata for query or table."""
+
+    if not query and not table_name:
+        raise ValueError()
+
+    if query:
+        session_name = random_name()
+        prepare_name = f"{session_name}_prepare"
+        table_name = f"{session_name}_temp"
+        cursor.execute(query_template("prepare").format(
+            prepare_name=prepare_name,
+            query=query,
+            table_name=table_name,
+        ))
+
+    cursor.execute(query_template("attributes").format(
+        table_name=table_name,
+    ))
+
+    metadata: bytes = cursor.fetchone()[0]
+
+    if query:
+        cursor.execute(f"drop table if exists {table_name};")
+
+    return metadata

pgpack_dumper/common/query.py

@@ -0,0 +1,62 @@
+from pathlib import Path
+from random import randbytes
+from re import match
+
+
+pattern = r"\(select \* from (.*)\)|(.*)"
+
+
+def search_object(table: str, query: str = "") -> str:
+    """Return current string for object."""
+
+    if query:
+        return "query"
+
+    return match(pattern, table).group(1) or table
+
+
+def random_name() -> str:
+    """Generate random name for prepare and temp table."""
+
+    return f"session_{randbytes(8).hex()}"  # noqa: S311
+
+
+def query_path() -> str:
+    """Path for queryes."""
+
+    return f"{Path(__file__).parent.absolute()}/queryes/{{}}.sql"
+
+
+def query_template(query_name: str) -> str:
+    """Get query template for his name."""
+
+    path = query_path().format(query_name)
+
+    with open(path, encoding="utf-8") as query:
+        return query.read()
+
+
+def chunk_query(query: str | None) -> tuple[list[str]]:
+    """Chunk multiquery to queryes."""
+
+    if not query:
+        return [], []
+
+    first_part: list[str] = [
+        part.strip()
+        for part in query.split(";")
+    ]
+    second_part: list[str] = []
+
+    for _ in first_part:
+        second_part.append(first_part.pop())
+        if any(
+            word == second_part[-1][:len(word)].lower()
+            for word in ("with", "select")
+        ):
+            second_part = list(reversed(second_part))
+            break
+
+    return first_part, second_part
+
+

pgpack_dumper/common/queryes/attributes.sql

@@ -0,0 +1,4 @@
+select json_agg(json_build_array(attnum, json_build_array(attname, atttypid::int4, 
+case when atttypid = 1042 then atttypmod - 4 when atttypid = 1700 then (atttypmod - 4) >> 16 else attlen end,
+case when atttypid = 1700 then (atttypmod - 4) & 65535 else 0 end, attndims)))::text::bytea as metadata
+from pg_attribute where attrelid = '{table_name}'::regclass and attnum > 0 and not attisdropped;

pgpack_dumper/common/queryes/copy_from.sql

@@ -0,0 +1 @@
+copy {table_name} from stdin with (format binary);

pgpack_dumper/common/queryes/copy_to.sql

@@ -0,0 +1 @@
+copy {table_name} to stdout with (format binary);

pgpack_dumper/common/queryes/dbname.sql

@@ -0,0 +1,2 @@
+select case when count(nspname) = 1 then 'greenplum' else 'postgres' end as dbname
+from pg_catalog.pg_namespace where nspname = 'gp_toolkit';

pgpack_dumper/common/queryes/gpversion.sql

@@ -0,0 +1 @@
+select substring(version() from 'Greenplum Database (.*?) build') as gp_version;

pgpack_dumper/common/queryes/prepare.sql

@@ -0,0 +1,4 @@
+prepare {prepare_name} as {query} limit 0;
+drop table if exists {table_name};
+create temporary table {table_name} as execute {prepare_name} (null);
+deallocate prepare {prepare_name};

pgpack_dumper/common/queryes/relkind.sql

@@ -0,0 +1 @@
+select relkind from pg_class where oid = '{table_name}'::regclass;

pgpack_dumper/common/reader.pxd

@@ -0,0 +1,11 @@
+cdef class CopyReader:
+
+    cdef object copyobj
+    cdef object iterator
+    cdef object bufferobj
+    cdef bint closed
+    cdef long long total_read
+
+    cpdef bytes read(self, long long size)
+    cpdef long long tell(self)
+    cpdef void close(self)

pgpack_dumper/common/reader.pyi

@@ -0,0 +1,35 @@
+from typing import Iterator, Iterable
+
+from psycopg import Copy
+
+
+class CopyReader:
+    """Read from iterable Copy object."""
+
+    def __init__(
+        self,
+        copyobj: Iterable[Copy],
+    ) -> None:
+        """Class initialization."""
+
+        self.copyobj: Iterable[Copy]
+        self.iterator: Iterator[bytearray]
+        self.bufferobj: bytearray
+        self.closed: bool
+        self.total_read: int
+        ...
+
+    def read(self, size: int) -> bytes:
+        """Read from copy."""
+
+        ...
+
+    def tell(self) -> int:
+        """Return the current stream position."""
+
+        ...
+
+    def close(self) -> None:
+        """Close CopyReader."""
+
+        ...