soleprint/soleprint/station/tools/dataconvert/output.py

"""
Output: one .sql file per table or sheet, named after it, split into numbered
parts when a size limit is set and the table outgrows it.
"""

import glob
import re
from pathlib import Path

from progress import say
from sqlgen import DEFAULT_BATCH_ROWS, Shape, render_sample, sanitize_identifier, utf8_len


def table_filename(raw_name: str, source_name: str, sheet_count: int, bare_prefixes=()) -> str:
    """
    Sheets of a multi-sheet workbook are prefixed with the workbook, so two
    workbooks cannot collide, unless the config names the sheet as already
    unique (bare_sheet_prefixes).
    """
    clean = sanitize_identifier(raw_name)
    if sheet_count > 1 and not clean.startswith(tuple(bare_prefixes)):
        return f"{sanitize_identifier(source_name)}_{clean}.sql"
    return f"{clean}.sql"


class TableFile:
    """
    Where one table's statements go: <table>.sql, or once it passes max_bytes,
    <table>.001.sql, <table>.002.sql, ... in load order.

    Several sources can feed the same table and they accumulate: writing resumes
    at the last part that exists, as the unsplit file always did.
    """

    def __init__(self, out_dir: Path, filename: str, max_bytes=None):
        self.out_dir = out_dir
        self.stem = filename[:-len(".sql")]
        self.max_bytes = max_bytes
        self.handle = None
        self.touched = []
        parts = self.parts()
        self.index = int(parts[-1].name[len(self.stem) + 1:-len(".sql")]) if parts else 0
        self.size = self.path.stat().st_size if self.path.exists() else 0

    @property
    def path(self) -> Path:
        name = f"{self.stem}.sql" if self.index == 0 else f"{self.stem}.{self.index:03d}.sql"
        return self.out_dir / name

    def parts(self):
        pattern = re.compile(re.escape(self.stem) + r"\.\d{3}\.sql$")
        if not self.out_dir.is_dir():
            return []
        return sorted(p for p in self.out_dir.glob(glob.escape(self.stem) + ".*.sql") if pattern.match(p.name))

    def write(self, text: str):
        if self.handle is None:
            # Created on the first write, so a run that finds nothing to
            # convert leaves no empty directory behind.
            self.out_dir.mkdir(parents=True, exist_ok=True)
            self.handle = open(self.path, "a", encoding="utf-8")
            if self.path not in self.touched:
                self.touched.append(self.path)
        self.handle.write(text)
        self.size += utf8_len(text)

    def full(self, adding: int) -> bool:
        return self.max_bytes is not None and self.size > 0 and self.size + adding > self.max_bytes

    def next_part(self):
        """Move on to the next numbered part; the unnumbered file becomes part 001."""
        self.close()
        if self.index == 0:
            base, first = self.path, self.out_dir / f"{self.stem}.001.sql"
            base.rename(first)
            self.touched = [first if p == base else p for p in self.touched]
            self.index = 1
        self.index += 1
        self.size = 0

    def close(self):
        if self.handle is not None:
            self.handle.close()
            self.handle = None


def write_full(shape: Shape, df, table_file: TableFile) -> int:
    """Every row, streamed; a new part, as its own transaction, whenever the limit is reached."""
    opening, closing = shape.open(), shape.close()
    open_bytes, close_bytes = utf8_len(opening), utf8_len(closing)
    if table_file.full(open_bytes + close_bytes):
        table_file.next_part()
    table_file.write(opening)
    written = open_bytes
    units_in_part = 0
    for unit in shape.units(df):
        size = utf8_len(unit)
        # Only between units, and never leaving a part empty: a single unit
        # bigger than the limit gets a part of its own instead of looping.
        if units_in_part and table_file.full(size + close_bytes):
            table_file.write(closing)
            table_file.next_part()
            table_file.write(opening)
            written += close_bytes + open_bytes
            units_in_part = 0
        table_file.write(unit)
        written += size
        units_in_part += 1
    table_file.write(closing)
    return written + close_bytes


def write_tables(dfs: dict, out_dir: Path, source_name: str, max_rows=None, report=None, bare_prefixes=(),
                 fmt="insert", batch_rows=DEFAULT_BATCH_ROWS, max_bytes=None, folder=""):
    """
    Write individual .sql files per table/sheet into the output directory, or
    into its `folder` subdirectory when the source folders are kept.
    """
    target = out_dir / folder if folder else out_dir
    for raw_name, df in dfs.items():
        if df.empty:
            continue

        table = sanitize_identifier(raw_name)
        filename = table_filename(raw_name, source_name, len(dfs), bare_prefixes)
        shape = Shape(fmt, table, df.columns, batch_rows)
        total = len(df)
        exact = max_rows is None or max_rows >= total
        # A sample is never split: it is small, and it is not for loading.
        table_file = TableFile(target, filename, max_bytes if exact else None)

        try:
            if exact:
                full_bytes = write_full(shape, df, table_file)
            else:
                text, full_bytes = render_sample(shape, df, max_rows)
                table_file.write(text)
        finally:
            table_file.close()

        shown = total if exact else max_rows
        parts = len(table_file.parts()) or 1
        if report is not None:
            report.add(source_name, table, f"{folder}/{filename}" if folder else filename,
                       df, total, full_bytes, exact, shown, parts)

        suffix = "" if exact else f" ({shown} of {total} rows)"
        for path in table_file.touched:
            say(f"Generated: {path}{suffix}")