native-dumper 0.3.5.0__cp310-cp310-macosx_11_0_arm64.whl

native_dumper/__init__.py

@@ -0,0 +1,33 @@
+"""Library for read and write Native format between Clickhouse and file."""
+
+from light_compressor import CompressionMethod
+
+from .common import (
+    CHConnector,
+    ClickhouseServerError,
+    DumperLogger,
+    HTTPCursor,
+    NativeDumperError,
+    NativeDumperReadError,
+    NativeDumperValueError,
+    NativeDumperWriteError,
+)
+
+from .dumper import NativeDumper
+from .version import __version__
+
+
+__all__ = (
+    "__version__",
+    "CHConnector",
+    "ClickhouseServerError",
+    "CompressionMethod",
+    "DumperLogger",
+    "HTTPCursor",
+    "NativeDumper",
+    "NativeDumperError",
+    "NativeDumperReadError",
+    "NativeDumperValueError",
+    "NativeDumperWriteError",
+)
+__author__ = "0xMihalich"

native_dumper/common/__init__.py

@@ -0,0 +1,58 @@
+"""Common utilities."""
+
+from .columns import make_columns
+from .connector import CHConnector
+from .cursor import HTTPCursor
+from .defines import (
+    CHUNK_SIZE,
+    DBMS_DEFAULT_TIMEOUT_SEC,
+    DEFAULT_DATABASE,
+    DEFAULT_PASSWORD,
+    DEFAULT_PORT,
+    DEFAULT_USER,
+)
+from .diagram import (
+    DBMetadata,
+    format_table,
+    transfer_diagram,
+)
+from .errors import (
+    ClickhouseServerError,
+    NativeDumperError,
+    NativeDumperReadError,
+    NativeDumperValueError,
+    NativeDumperWriteError,
+)
+from .logger import DumperLogger
+from .multiquery import chunk_query
+from .pyo3http import (
+    HttpResponse,
+    HttpSession,
+)
+from .writer import file_writer
+
+
+__all__ = (
+    "CHUNK_SIZE",
+    "DBMS_DEFAULT_TIMEOUT_SEC",
+    "DEFAULT_DATABASE",
+    "DEFAULT_PASSWORD",
+    "DEFAULT_PORT",
+    "DEFAULT_USER",
+    "CHConnector",
+    "ClickhouseServerError",
+    "DBMetadata",
+    "DumperLogger",
+    "HTTPCursor",
+    "HttpResponse",
+    "HttpSession",
+    "NativeDumperError",
+    "NativeDumperReadError",
+    "NativeDumperValueError",
+    "NativeDumperWriteError",
+    "chunk_query",
+    "file_writer",
+    "format_table",
+    "make_columns",
+    "transfer_diagram",
+)

native_dumper/common/columns.py

@@ -0,0 +1,30 @@
+from collections import OrderedDict
+
+from nativelib import Column
+
+
+def make_columns(
+    column_list: list[Column],
+) -> OrderedDict[str, str]:
+    """Make DBMetadata.columns dictionary."""
+
+    columns = OrderedDict()
+
+    for column in column_list:
+        col_type = column.dtype.name
+        info = column.info
+
+        if col_type == "FixedString":
+            col_type = f"{col_type}({info.length})"
+        elif col_type == "Decimal":
+            col_type = f"{col_type}({info.precision}, {info.scale})"
+        elif col_type == "DateTime64":
+            col_type = f"{col_type}({info.precision}, {info.tzinfo})"
+        elif col_type in ("Enum8", "Enum16"):
+            col_type = f"{col_type}({info.enumcase})"
+        elif col_type == "Time64":
+            col_type = f"{col_type}({info.precision})"
+
+        columns[column.column] = col_type
+
+    return columns

native_dumper/common/connector.py

@@ -0,0 +1,18 @@
+from typing import NamedTuple
+
+from .defines import (
+    DEFAULT_DATABASE,
+    DEFAULT_USER,
+    DEFAULT_PASSWORD,
+    DEFAULT_PORT,
+)
+
+
+class CHConnector(NamedTuple):
+    """Connector for Clickhouse."""
+
+    host: str
+    dbname: str = DEFAULT_DATABASE
+    user: str = DEFAULT_USER
+    password: str = DEFAULT_PASSWORD
+    port: int = DEFAULT_PORT

native_dumper/common/cursor.py

@@ -0,0 +1,196 @@
+from typing import Iterable
+from uuid import uuid4
+
+from light_compressor import (
+    CompressionMethod,
+    define_reader,
+)
+from nativelib import (
+    Column,
+    NativeReader,
+)
+
+from ..version import __version__
+from .connector import CHConnector
+from .defines import CHUNK_SIZE
+from .errors import ClickhouseServerError
+from .logger import Logger
+from .pyo3http import (
+    HttpResponse,
+    HttpSession,
+)
+
+
+def string_error(data: bytes) -> str:
+    """Bytes to string decoder."""
+
+    return data.decode("utf-8", errors="replace").strip()
+
+
+class HTTPCursor:
+    """Class for send queryes to Clickhouse server
+    and read/write Native format."""
+
+    def __init__(
+        self,
+        connector: CHConnector,
+        compression_method: CompressionMethod,
+        logger: Logger,
+        timeout: int,
+    ) -> None:
+        """Class initialization."""
+
+        self.connector = connector
+        self.compression_method = compression_method
+        self.logger = logger
+        self.timeout = timeout
+        self.session = HttpSession(timeout=self.timeout)
+        self.is_connected = False
+        self.headers = {
+            "Accept": "*/*",
+            "Connection": "keep-alive",
+            "User-Agent": f"{self.__class__.__name__}/{__version__}",
+            "Accept-Encoding": self.compression_method.method,
+            "Content-Encoding": self.compression_method.method,
+            "X-ClickHouse-User": self.connector.user,
+            "X-ClickHouse-Key": self.connector.password,
+            "X-ClickHouse-Compression": self.compression_method.method,
+            "X-ClickHouse-Format": "Native",
+            "X-Content-Type-Options": "nosniff",
+        }
+        self.mode = {
+            443: "https",
+        }.get(int(self.connector.port), "http")
+        self.url = (
+            f"{self.mode}://{self.connector.host}:{self.connector.port}/"
+            "?enable_http_compression=1"
+        )
+        self.params = {
+            "database": connector.dbname,
+            "query": "",
+            "session_id": str(uuid4()),
+        }
+        self.check_length = {
+            CompressionMethod.NONE: 1024,
+        }
+        self.server_version = None
+
+    def send_hello(self) -> str:
+        """Get server version."""
+
+        reader = self.get_stream("SELECT version()")
+        server_version = tuple(reader.to_rows())[0][0]
+        self.is_connected = True
+        self.server_version = server_version
+        return self.server_version
+
+    def get_response(
+        self,
+        query: str,
+        data: Iterable[bytes] | None = None,
+    ) -> HttpResponse:
+        """Get response from clickhouse server."""
+
+        self.params["query"] = query
+
+        response = self.session.post(
+            url=self.url,
+            params=self.params,
+            headers=self.headers,
+            timeout=self.timeout,
+            data=data,
+        )
+        status = response.get_status()
+
+        if status != 200:
+
+            if not self.is_connected:
+                error = string_error(response.read())
+                response.close()
+            else:
+                bufferobj = define_reader(response, self.compression_method)
+                error = string_error(bufferobj.read(CHUNK_SIZE))
+                bufferobj.close()
+
+            self.logger.error(f"ClickhouseServerError: {error}")
+            raise ClickhouseServerError(error)
+
+        return response
+
+    def get_stream(
+        self,
+        query: str,
+    ) -> NativeReader:
+        """Get answer from server as unpacked stream file."""
+
+        stream = self.get_response(query)
+
+        try:
+            bufferobj = define_reader(stream, self.compression_method)
+            check_error = bufferobj.read(
+                self.check_length.get(self.compression_method, 4)
+            )[:4]
+            bufferobj.seek(0)
+        except EOFError:
+            error = (
+                "Code: 92. DB::Exception: (EMPTY_DATA_PASSED) "
+                f"(version {self.server_version} (official build))"
+            )
+            self.logger.error(f"ClickhouseServerError: {error}")
+            raise ClickhouseServerError(error)
+
+        if check_error == b"Code":
+            error = string_error(bufferobj.read(CHUNK_SIZE))
+            bufferobj.close()
+            self.logger.error(f"ClickhouseServerError: {error}")
+            raise ClickhouseServerError(error)
+
+        return NativeReader(bufferobj)
+
+    def upload_data(
+        self,
+        table: str,
+        data: Iterable[bytes],
+    ) -> None:
+        """Download data into table."""
+
+        self.get_response(
+            query=f"INSERT INTO {table} FORMAT Native",
+            data=data,
+        )
+
+    def metadata(
+        self,
+        table: str,
+    ) -> list[Column]:
+        """Get table metadata."""
+
+        reader = self.get_stream(f"DESCRIBE TABLE {table}")
+        return [
+            Column(*describe[:2])
+            for describe in reader.to_rows()
+        ]
+
+    def execute(
+        self,
+        query: str,
+    ) -> None:
+        """Simple exetute method without return."""
+
+        self.get_response(query)
+
+    def last_query(self) -> str:
+        """Show last query."""
+
+        return self.params["query"]
+
+    def refresh(self) -> None:
+        """Refresh Session ID."""
+
+        self.params["session_id"] = str(uuid4())
+
+    def close(self) -> None:
+        """Close HTTPCursor session."""
+
+        self.session.close()
+        self.is_connected = False

native_dumper/common/defines.py

@@ -0,0 +1,6 @@
+CHUNK_SIZE = 16_384
+DEFAULT_DATABASE = ""
+DEFAULT_USER = "default"
+DEFAULT_PASSWORD = ""
+DEFAULT_PORT = 8123
+DBMS_DEFAULT_TIMEOUT_SEC = 300

native_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)
+    )

native_dumper/common/errors.py

@@ -0,0 +1,18 @@
+class ClickhouseServerError(ValueError):
+    """Clickhouse errors."""
+
+
+class NativeDumperError(Exception):
+    """NativeDumper base error."""
+
+
+class NativeDumperReadError(NativeDumperError):
+    """NativeDumper read error."""
+
+
+class NativeDumperWriteError(NativeDumperError):
+    """NativeDumper write error."""
+
+
+class NativeDumperValueError(ValueError):
+    """NativeDumper value error."""

native_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):
+    """NativeDumper logger."""
+
+    def __init__(
+            self,
+            level: int = DEBUG,
+            use_console: bool = True,
+    ) -> None:
+        """Class initialize."""
+
+        super().__init__("NativeDumper")
+
+        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()}/native_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

native_dumper/common/multiquery.py

@@ -0,0 +1,34 @@
+from re import split
+
+def chunk_query(query: str | None) -> tuple[list[str]]:
+    """Chunk multiquery to queryes."""
+
+    if not query:
+        return [], []
+
+    pattern = r";(?=(?:[^']*'[^']*')*[^']*$)"
+    parts = [
+        part.strip(";").strip()
+        for part in split(pattern, query)
+        if part.strip(";").strip()
+    ]
+
+    if not parts:
+        return [], []
+
+    first_part: list[str] = []
+    second_part: list[str] = []
+
+    for i, part in enumerate(parts):
+        first_part.append(part)
+
+        if (i + 1 < len(parts) and parts[i + 1].lower().startswith(
+                ('with ', 'select ')
+            )
+        ):
+            second_part = parts[i + 1:]
+            break
+    else:
+        second_part = []
+
+    return first_part, second_part

native_dumper/common/pyo3http/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "pyo3http"
+version = "0.1.0"
+edition = "2021"
+
+[lib]
+name = "pyo3http"
+crate-type = ["cdylib"]
+
+[dependencies]
+pyo3 = { version = "0.26.0", features = ["extension-module"] }
+reqwest = { version = "0.12.24", features = ["stream", "json"] }
+tokio = { version = "1.0", features = ["full"] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"