sql2iceberg 0.1.0__py3-none-any.whl

sql2iceberg/pyspark_gen.py

@@ -0,0 +1,251 @@
+"""Generate a runnable PySpark script: ingest source tables over JDBC into Iceberg, then run the curation SQL.
+
+The script reads credentials from environment variables when it runs. Nothing secret is written into it.
+Iceberg table names are ``<catalog>.<namespace>.<table>``; source schema/database qualifiers are dropped.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+
+import sqlglot
+from sqlglot import exp
+from sqlglot.errors import ParseError
+
+from .connectors import REGISTRY
+from .convert import UNPARSED, UNSUPPORTED, convert_statement
+from .lineage import map_statements
+from .tracing import span
+
+DEFAULT_ICEBERG_PACKAGE = "org.apache.iceberg:iceberg-spark-runtime-4.1_2.13:1.12.0"
+CATALOG_TYPES = ("hadoop", "glue", "rest", "hive", "jdbc")
+
+
+@dataclass
+class MigrationConfig:
+    platform: str = "snowflake"
+    catalog: str = "lake"
+    namespace: str = "db"
+    catalog_type: str = "hadoop"            # hadoop | glue | rest | hive | jdbc (use jdbc, not hadoop, on S3)
+    region: str = "us-east-2"               # S3 region when the warehouse is s3://
+    gcs_project: str = ""                   # GCP project when the warehouse is gs://
+    gcs_token_from_env: bool = False        # gs:// only: read GCS_OAUTH_TOKEN (+ _EXPIRES_AT_MS) at run time, else ADC
+    warehouse: str = "./warehouse"          # path or s3://bucket/prefix
+    catalog_uri: str = ""                   # rest / hive / jdbc (e.g. jdbc:sqlite:catalog.db)
+    format_version: int = 2
+    iceberg_package: str = DEFAULT_ICEBERG_PACKAGE
+    extra_packages: list[str] = field(default_factory=list)
+    jdbc_driver: str = ""                   # default: the connector's driver class
+    jdbc_package: str | None = None         # default: the connector's Maven coordinate; "" for none
+    jars: list[str] = field(default_factory=list)   # local jars (offline machines); use "" iceberg_package then
+    jdbc_url: str = ""                      # default: the connector's own URL, resolved at generation time
+    source_schema: str = ""                 # prefix for unqualified source tables, e.g. "PUBLIC"
+    # per source table: {"partition_column": "id", "lower": 1, "upper": 1000000, "num_partitions": 8}
+    table_options: dict[str, dict] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if self.catalog_type not in CATALOG_TYPES:
+            raise ValueError(f"catalog_type must be one of {CATALOG_TYPES}")
+        if self.platform not in REGISTRY:
+            raise ValueError(f"unknown platform {self.platform!r}")
+
+
+def target_name(name: str) -> str:
+    """Last part of a table name, lower-cased: the Iceberg table it maps to."""
+    return name.split(".")[-1].strip('"`[]').lower()
+
+
+def _catalog_conf(cfg: MigrationConfig) -> dict[str, str]:
+    p = f"spark.sql.catalog.{cfg.catalog}"
+    conf = {p: "org.apache.iceberg.spark.SparkCatalog", f"{p}.warehouse": cfg.warehouse}
+    if cfg.warehouse.startswith("s3"):
+        conf[f"{p}.io-impl"] = "org.apache.iceberg.aws.s3.S3FileIO"
+        conf[f"{p}.client.region"] = cfg.region
+    if cfg.warehouse.startswith("gs"):
+        conf[f"{p}.io-impl"] = "org.apache.iceberg.gcp.gcs.GCSFileIO"
+        if cfg.gcs_project:
+            conf[f"{p}.gcs.project-id"] = cfg.gcs_project
+        if cfg.gcs_token_from_env:
+            conf[f"{p}.gcs.oauth2.token"] = "$ENV:GCS_OAUTH_TOKEN"
+            conf[f"{p}.gcs.oauth2.token-expires-at"] = "$ENV:GCS_OAUTH_TOKEN_EXPIRES_AT_MS"
+    if cfg.catalog_type == "jdbc":
+        conf[f"{p}.catalog-impl"] = "org.apache.iceberg.jdbc.JdbcCatalog"
+        conf[f"{p}.uri"] = cfg.catalog_uri
+    elif cfg.catalog_type == "glue":
+        conf[f"{p}.catalog-impl"] = "org.apache.iceberg.aws.glue.GlueCatalog"
+        conf[f"{p}.io-impl"] = "org.apache.iceberg.aws.s3.S3FileIO"
+    else:
+        conf[f"{p}.type"] = cfg.catalog_type
+        if cfg.catalog_uri:
+            conf[f"{p}.uri"] = cfg.catalog_uri
+    return conf
+
+
+def _conf_line(key: str, value: str) -> str:
+    if value.startswith("$ENV:"):  # tokens are read when the script runs, never written into it
+        return f"    .config({key!r}, os.environ[{value[5:]!r}])"
+    return f"    .config({key!r}, {value!r})"
+
+
+def rewrite_for_iceberg(spark_sql: str, cfg: MigrationConfig) -> str | None:
+    """Point every table at the Iceberg catalog; make CTAS create Iceberg tables.
+
+    Returns None for a plain CREATE TABLE (no query): ingestion or the generated DDL creates those.
+    """
+    tree = sqlglot.parse_one(spark_sql, read="spark")
+    if isinstance(tree, exp.Create) and not tree.expression and (tree.args.get("kind") or "").upper() == "TABLE":
+        return None
+    ctes = {c.alias for c in tree.find_all(exp.CTE)}
+
+    def fix(node):
+        if isinstance(node, exp.Table) and isinstance(node.this, exp.Identifier) and node.name not in ctes:
+            node = node.copy()
+            node.set("this", exp.to_identifier(target_name(node.name)))
+            node.set("db", exp.to_identifier(cfg.namespace))
+            node.set("catalog", exp.to_identifier(cfg.catalog))
+        return node
+
+    text = tree.transform(fix).sql(dialect="spark", comments=False)
+    props = f"USING iceberg TBLPROPERTIES ('format-version'='{cfg.format_version}')"
+    return re.sub(r"^(CREATE (?:OR REPLACE )?TABLE (?:IF NOT EXISTS )?\S+)\s+AS\b",
+                  lambda m: f"{m.group(1)} {props} AS", text, count=1, flags=re.IGNORECASE)
+
+
+def _source_names(statements, cfg: MigrationConfig) -> dict[str, str]:
+    """target table -> table name to read over JDBC, for tables the code reads but never creates."""
+    out: dict[str, str] = {}
+    for name in map_statements(statements)["source_tables"]:
+        src = name if "." in name or not cfg.source_schema else f"{cfg.source_schema}.{name}"
+        out[target_name(name)] = src
+    return out
+
+
+_HEADER = '''"""Generated by sql2iceberg. Review every statement before running it on real data.
+
+Set before running: {envs}
+Run with: spark-submit {submit}pipeline_pyspark.py
+"""
+import os
+
+from pyspark.sql import SparkSession
+
+CATALOG = {catalog!r}
+NAMESPACE = {namespace!r}
+FORMAT_VERSION = {format_version!r}
+JDBC_URL = {jdbc_url!r}
+JDBC_DRIVER = {jdbc_driver!r}
+# Iceberg table -> (source table to read, partitioned-read options)
+SOURCES = {sources}
+
+spark = (
+    SparkSession.builder.appName("sql2iceberg-migration")
+{jars_conf}{pkgs_conf}    .config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
+{conf}
+    .getOrCreate()
+)
+spark.sql(f"CREATE NAMESPACE IF NOT EXISTS {{CATALOG}}.{{NAMESPACE}}")
+
+
+def ingest(target, source, opts):
+    """Copy one source table into Iceberg (create or replace)."""
+    reader = (
+        spark.read.format("jdbc")
+        .option("url", JDBC_URL)
+        .option("driver", JDBC_DRIVER)
+        .option("dbtable", source)
+        .option("user", os.environ[{user_env!r}])
+    )
+    key_file = os.environ.get({key_env!r}) if {key_env!r} else None
+    if key_file:  # key-pair login
+        reader = reader.option({key_option!r}, key_file)
+    else:
+        reader = reader.option("password", os.environ[{pw_env!r}])
+    if opts.get("partition_column"):
+        reader = (
+            reader.option("partitionColumn", opts["partition_column"])
+            .option("lowerBound", opts["lower"])
+            .option("upperBound", opts["upper"])
+            .option("numPartitions", opts.get("num_partitions", 4))
+        )
+    (
+        reader.load()
+        .writeTo(f"{{CATALOG}}.{{NAMESPACE}}.{{target}}")
+        .using("iceberg")
+        .tableProperty("format-version", str(FORMAT_VERSION))
+        .createOrReplace()
+    )
+
+
+for _target, (_source, _opts) in SOURCES.items():
+    ingest(_target, _source, _opts)
+
+# ---- curation ----
+'''
+
+
+def pyspark_pipeline(statements, cfg: MigrationConfig, dialect: str | None = None) -> tuple[str, list[dict]]:
+    with span("pyspark_gen", platform=cfg.platform, catalog_type=cfg.catalog_type, statements=len(statements)):
+        return _pyspark_pipeline(statements, cfg, dialect)
+
+
+def _pyspark_pipeline(statements, cfg: MigrationConfig, dialect: str | None = None) -> tuple[str, list[dict]]:
+    """Return (script text, per-statement report). Unsupported or dynamic statements are commented out."""
+    conn = REGISTRY[cfg.platform]
+    dialect = dialect or conn.sqlglot_dialect
+    sources = _source_names(statements, cfg) if statements else {}
+    table_src = {t: (s, cfg.table_options.get(s) or cfg.table_options.get(t) or {}) for t, s in sources.items()}
+    jdbc_pkg = conn.jdbc_package if cfg.jdbc_package is None else cfg.jdbc_package
+    packages = ",".join(p for p in (cfg.iceberg_package, jdbc_pkg, *cfg.extra_packages) if p)
+    jdbc_url = cfg.jdbc_url
+    if not jdbc_url:
+        try:
+            jdbc_url = conn().jdbc_url()
+        except KeyError as exc:
+            raise ValueError(f"set MigrationConfig.jdbc_url or the environment variable {exc.args[0]}") from None
+    envs = ", ".join(e for e in (conn.jdbc_user_env, conn.jdbc_password_env) if e)
+    if conn.jdbc_key_env:
+        envs += f" (or {conn.jdbc_key_env} instead of the password)"
+    pkg_arg = f"--packages {packages} " if packages else ""
+    joined_jars = ",".join(cfg.jars)
+    jars_conf = f"    .config('spark.jars', {joined_jars!r})\n" if cfg.jars else ""
+    pkgs_conf = f"    .config('spark.jars.packages', {packages!r})\n" if packages else ""
+    lines = [_HEADER.format(
+        envs=envs, submit=pkg_arg, catalog=cfg.catalog, namespace=cfg.namespace,
+        format_version=cfg.format_version, jdbc_url=jdbc_url, jdbc_driver=cfg.jdbc_driver or conn.jdbc_driver,
+        jars_conf=jars_conf,
+        sources="{\n" + "".join(f"    {t!r}: ({s!r}, {o!r}),\n" for t, (s, o) in table_src.items()) + "}",
+        pkgs_conf=pkgs_conf, user_env=conn.jdbc_user_env, pw_env=conn.jdbc_password_env,
+        key_env=conn.jdbc_key_env, key_option=conn.jdbc_key_option,
+        conf="\n".join(_conf_line(k, v) for k, v in _catalog_conf(cfg).items()),
+    )]
+    report: list[dict] = []
+    for s in statements:
+        entry = {"id": s.id, "status": "skipped", "note": ""}
+        report.append(entry)
+        if s.dynamic or s.parse_error:
+            entry["note"] = "built from variables" if s.dynamic else "could not parse"
+            lines.append(f"# {s.id}: skipped ({entry['note']})")
+            continue
+        r = convert_statement(s.sql, dialect=dialect)
+        entry["status"] = r.status
+        lines.append(f"# {s.id} [{r.status}]")
+        lines.extend(f"#   {f.code}: {f.message}" for f in r.findings)
+        if r.status in (UNSUPPORTED, UNPARSED) or r.spark_sql is None:
+            lines.append("# NOT GENERATED: rewrite by hand")
+            lines.append("# " + (r.spark_sql or s.sql).replace("\n", "\n# "))
+            continue
+        try:
+            sql = rewrite_for_iceberg(r.spark_sql, cfg)
+        except ParseError as exc:
+            entry["status"], entry["note"] = UNPARSED, str(exc)[:120]
+            lines.append(f"# NOT GENERATED: converted SQL did not parse as Spark ({entry['note']})")
+            continue
+        if sql is None:
+            entry["note"] = "plain CREATE TABLE: table is created by ingestion"
+            lines.append("# plain CREATE TABLE skipped: ingestion or generated DDL creates it")
+        else:
+            lines.append(f'spark.sql("""{sql}""")')
+        lines.append("")
+    lines.append("spark.stop()")
+    return "\n".join(lines) + "\n", report

sql2iceberg/pytest_plugin.py

@@ -0,0 +1,74 @@
+"""pytest plugin: fixtures and markers for validating converted SQL.
+
+Installed through the `pytest11` entry point. Spark fixtures need `pip install 'sql2iceberg[pytest]'`.
+Set SQL2ICE_ICEBERG=1 (and optionally ICEBERG_PKG) to add the Iceberg runtime and a Hadoop catalog.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+from .convert import convert_statement
+
+DEFAULT_ICEBERG_PKG = "org.apache.iceberg:iceberg-spark-runtime-3.5_2.12:1.9.2"
+
+
+def pytest_addoption(parser):
+    group = parser.getgroup("sql2iceberg")
+    group.addoption("--live-snowflake", action="store_true", default=False,
+                    help="run tests marked live_snowflake against a real Snowflake account (spends credits)")
+
+
+def pytest_collection_modifyitems(config, items):
+    if config.getoption("--live-snowflake"):
+        return
+    skip = pytest.mark.skip(reason="needs --live-snowflake (spends Snowflake credits)")
+    for item in items:
+        if "live_snowflake" in item.keywords:
+            item.add_marker(skip)
+
+
+@pytest.fixture(scope="session")
+def sf2i_has_iceberg() -> bool:
+    return os.environ.get("SQL2ICE_ICEBERG") == "1"
+
+
+@pytest.fixture(scope="session")
+def sf2i_spark(tmp_path_factory, sf2i_has_iceberg):
+    pyspark_sql = pytest.importorskip("pyspark.sql")
+    warehouse = str(tmp_path_factory.mktemp("warehouse"))
+    builder = (
+        pyspark_sql.SparkSession.builder.master("local[2]")
+        .appName("sql2iceberg-tests")
+        .config("spark.ui.enabled", "false")
+        .config("spark.sql.shuffle.partitions", "2")
+        .config("spark.sql.warehouse.dir", warehouse + "/spark")
+    )
+    if sf2i_has_iceberg:
+        pkg = os.environ.get("ICEBERG_PKG", DEFAULT_ICEBERG_PKG)
+        jar = os.environ.get("ICEBERG_JAR")  # local jar, for machines that cannot reach Maven Central
+        builder = (
+            builder.config("spark.jars", jar) if jar else builder.config("spark.jars.packages", pkg)
+        )
+        builder = (
+            builder.config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
+            .config("spark.sql.catalog.lake", "org.apache.iceberg.spark.SparkCatalog")
+            .config("spark.sql.catalog.lake.type", "hadoop")
+            .config("spark.sql.catalog.lake.warehouse", warehouse)
+            .config("spark.sql.defaultCatalog", "lake")
+        )
+    session = builder.getOrCreate()
+    session.sparkContext.setLogLevel("ERROR")
+    if sf2i_has_iceberg:
+        session.sql("CREATE NAMESPACE IF NOT EXISTS lake.db")
+        session.sql("USE lake.db")
+    yield session
+    session.stop()
+
+
+@pytest.fixture()
+def sf2i_convert():
+    """Return a function that converts one Snowflake statement and returns a Result."""
+    return convert_statement

sql2iceberg/report.py

@@ -0,0 +1,49 @@
+"""Build the conversion report (JSON and readable Markdown)."""
+
+from __future__ import annotations
+
+import json
+from collections import Counter
+
+from .convert import Result
+from .detect import Statement
+
+
+def build(statements: list[Statement], results: dict[str, Result]) -> dict:
+    items = []
+    for s in statements:
+        item = {"id": s.id, "file": s.file, "line": s.line}
+        if s.dynamic:
+            item.update(status="dynamic", note="built from variables; needs parameters")
+        elif s.parse_error:
+            item.update(status="unparsed", note=s.parse_error)
+        else:
+            item.update(results[s.id].to_dict())
+        items.append(item)
+    counts = Counter(i["status"] for i in items)
+    return {"summary": dict(counts), "total": len(items), "statements": items}
+
+
+def to_json(report: dict) -> str:
+    return json.dumps(report, indent=2)
+
+
+def to_markdown(report: dict) -> str:
+    lines = ["# sql2iceberg conversion report", "", f"Statements found: {report['total']}", ""]
+    for status, n in sorted(report["summary"].items()):
+        lines.append(f"- {status}: {n}")
+    lines.append("")
+    for item in report["statements"]:
+        lines.append(f"## {item['id']} ({item['status']})")
+        if item.get("source_sql"):
+            lines += ["", "```sql", item["source_sql"], "```"]
+        if item.get("spark_sql"):
+            lines += ["", "Spark SQL:", "```sql", item["spark_sql"], "```"]
+        for f in item.get("findings", []):
+            lines.append(f"- **{f['code']}** ({f['severity']}): {f['message']}")
+        if item.get("note"):
+            lines.append(f"- {item['note']}")
+        lines.append("")
+    lines.append("Statuses mean: converted = no known problems found (still verify); "
+                 "review = semantics may differ from Snowflake; unsupported = known not to run in Spark.")
+    return "\n".join(lines)

sql2iceberg/tracing.py

@@ -0,0 +1,134 @@
+"""A small tracer: timed, nested spans for every step of a run, written as JSON lines.
+
+Enable it with ``--trace trace.jsonl`` on the CLI, ``SQL2ICEBERG_TRACE=trace.jsonl`` in the environment, or in code:
+
+    from sql2iceberg.tracing import Tracer, set_tracer
+    set_tracer(Tracer("trace.jsonl"))
+
+With no tracer configured every call is a no-op. Attribute values whose key looks secret (password, token,
+key, secret, credential) and anything that looks like a bearer/access token are replaced with ``***`` before
+they are stored, so a trace file is safe to attach to a GitHub issue after a quick read.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+import re
+import threading
+import time
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+
+_SECRET_KEY = re.compile(r"pass(word)?|token|secret|credential|private|(^|_)key($|_)", re.IGNORECASE)
+_SECRET_VALUE = re.compile(r"(ya29\.[\w.-]+|AKIA[0-9A-Z]{16}|-----BEGIN [A-Z ]*PRIVATE KEY-----|eyJ[\w-]{20,}\.[\w-]+\.[\w-]+)")
+
+
+def redact(key: str, value):
+    if _SECRET_KEY.search(key):
+        return "***"
+    if isinstance(value, str):
+        return _SECRET_VALUE.sub("***", value)[:500]
+    if isinstance(value, (int, float, bool)) or value is None:
+        return value
+    return _SECRET_VALUE.sub("***", str(value))[:500]
+
+
+@dataclass
+class Span:
+    id: str
+    parent: str | None
+    name: str
+    start: float
+    attrs: dict = field(default_factory=dict)
+    duration_ms: float = 0.0
+    status: str = "ok"
+    error: str = ""
+
+    def set(self, **attrs) -> None:
+        self.attrs.update({k: redact(k, v) for k, v in attrs.items()})
+
+    def to_dict(self) -> dict:
+        d = {"id": self.id, "parent": self.parent, "name": self.name, "start": round(self.start, 3),
+             "duration_ms": round(self.duration_ms, 1), "status": self.status, "attrs": self.attrs}
+        if self.error:
+            d["error"] = self.error
+        return d
+
+
+class Tracer:
+    def __init__(self, path: str | Path | None = None, run_id: str | None = None) -> None:
+        self.path = Path(path) if path else None
+        self.run_id = run_id or uuid.uuid4().hex[:12]
+        self.spans: list[Span] = []
+        self._local = threading.local()
+        self._lock = threading.Lock()
+        if self.path:
+            self.path.parent.mkdir(parents=True, exist_ok=True)
+
+    def _stack(self) -> list[Span]:
+        if not hasattr(self._local, "stack"):
+            self._local.stack = []
+        return self._local.stack
+
+    @contextlib.contextmanager
+    def span(self, name: str, **attrs):
+        stack = self._stack()
+        sp = Span(uuid.uuid4().hex[:8], stack[-1].id if stack else None, name, time.time())
+        sp.set(**attrs)
+        stack.append(sp)
+        t0 = time.perf_counter()
+        try:
+            yield sp
+        except BaseException as exc:
+            sp.status, sp.error = "error", f"{type(exc).__name__}: {exc}"[:500]
+            raise
+        finally:
+            sp.duration_ms = (time.perf_counter() - t0) * 1000
+            stack.pop()
+            with self._lock:
+                self.spans.append(sp)
+                if self.path:
+                    with self.path.open("a", encoding="utf-8") as f:
+                        f.write(json.dumps({"run": self.run_id, **sp.to_dict()}) + "\n")
+
+    def summary(self) -> str:
+        rows = ["| step | status | ms |", "|---|---|---|"]
+        for sp in sorted(self.spans, key=lambda s: s.start):
+            depth = 0
+            p = sp.parent
+            ids = {s.id: s for s in self.spans}
+            while p and p in ids:
+                depth += 1
+                p = ids[p].parent
+            rows.append(f"| {'  ' * depth}{sp.name} | {sp.status} | {sp.duration_ms:.0f} |")
+        return f"Trace {self.run_id}\n\n" + "\n".join(rows) + "\n"
+
+
+class _Noop:
+    @contextlib.contextmanager
+    def span(self, name: str, **attrs):
+        yield Span("", None, name, time.time())
+
+
+_current: Tracer | _Noop | None = None
+
+
+def set_tracer(tracer: Tracer | None) -> None:
+    global _current
+    _current = tracer
+
+
+def get_tracer() -> Tracer | _Noop:
+    global _current
+    if _current is None:
+        path = os.environ.get("SQL2ICEBERG_TRACE")
+        _current = Tracer(path) if path else _Noop()
+    return _current
+
+
+def span(name: str, **attrs):
+    """``with span("step", table="orders") as sp: ... sp.set(rows=3)``"""
+    return get_tracer().span(name, **attrs)