dataconvert updates

soleprint/station/tools/dataconvert/output.py

@@ -1,11 +1,14 @@
 """
-Output: one .sql file per table or sheet, named after it.
+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 render_table, sanitize_identifier
+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:
@@ -20,28 +23,127 @@ def table_filename(raw_name: str, source_name: str, sheet_count: int, bare_prefi
     return f"{clean}.sql"
 
 
-def write_tables(dfs: dict, out_dir: Path, source_name: str, max_rows=None, report=None, bare_prefixes=()):
-    """Write individual .sql files per table/sheet into the output directory."""
+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)
-        sql, total, full_bytes, exact = render_table(df, table, max_rows)
-        out_file = out_dir / filename
-        # Created on the first file, so a run that finds nothing to convert
-        # leaves no empty directory behind.
-        out_dir.mkdir(parents=True, exist_ok=True)
+        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)
 
-        # Several sources can feed the same table; they accumulate in one file.
-        mode = "a" if out_file.exists() else "w"
-        with open(out_file, mode, encoding="utf-8") as f:
-            f.write(sql)
+        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, filename, df, total, full_bytes, exact, shown)
+            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)"
-        say(f"Generated: {out_file}{suffix}")
+        for path in table_file.touched:
+            say(f"Generated: {path}{suffix}")