sql2iceberg 0.1.0__py3-none-any.whl

sql2iceberg/publish/aws.py

@@ -0,0 +1,108 @@
+"""AWS: make Iceberg tables visible in the Glue Data Catalog (so Athena, EMR, Redshift Spectrum can query them).
+
+If the tables were written through a Glue catalog (``catalog_type="glue"``) they are already registered and
+``plan()`` only adds a verification step. If they live in a Hadoop catalog on S3, they are registered in Glue
+with Iceberg's ``register_table`` procedure, run in a Spark session that has a Glue catalog configured.
+Needs ``pip install sql2iceberg[aws]`` (boto3) for verification. Credentials come from the usual AWS
+profile / environment, never from this library's config.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .base import Action, IcebergTable, Publisher, PublishResult, ident, lit
+
+
+@dataclass
+class GlueTarget:
+    database: str                      # Glue database (Iceberg namespace)
+    region: str = ""
+    already_in_glue: bool = False      # tables were written through a Glue catalog
+    spark_glue_catalog: str = "glue"   # name of the Spark catalog configured with GlueCatalog
+    athena_workgroup: str = "primary"
+    athena_output: str = ""            # s3://bucket/prefix/ for Athena results (needed to run the count check)
+
+    def __post_init__(self) -> None:
+        ident(self.database, "glue database")
+        ident(self.spark_glue_catalog, "spark catalog name")
+
+
+class GluePublisher(Publisher):
+    target = "aws"
+
+    def __init__(self, target: GlueTarget, spark=None, session=None) -> None:
+        super().__init__()
+        self.t, self._spark, self._session = target, spark, session
+
+    def plan(self, tables: list[IcebergTable]) -> PublishResult:
+        t = self.t
+        res = PublishResult()
+        add = res.actions.append
+        if not t.already_in_glue:
+            add(Action("aws", "sql", "Create Glue database via Spark",
+                       f"CREATE NAMESPACE IF NOT EXISTS {t.spark_glue_catalog}.{t.database};"))
+            for tb in tables:
+                if not tb.metadata_location.startswith("s3"):
+                    raise ValueError(f"{tb.name}: metadata_location must be an s3:// URI to register in Glue")
+                add(Action("aws", "sql", f"Register {tb.name} in Glue",
+                           f"CALL {t.spark_glue_catalog}.system.register_table(table => "
+                           f"{lit(t.database + '.' + ident(tb.name, 'table name'))}, "
+                           f"metadata_file => {lit(tb.metadata_location)});"))
+        add(Action("aws", "api", "Verify tables in Glue",
+                   "glue.get_table for: " + ", ".join(f"{t.database}.{tb.name}" for tb in tables)
+                   + " (expects Parameters.table_type == ICEBERG)"))
+        if t.athena_output:
+            add(Action("aws", "sql", "Count rows with Athena (billed per data scanned)",
+                       "\n".join(f'SELECT \'{tb.name}\' AS table_name, COUNT(*) FROM "{t.database}"."{tb.name}";'
+                                 for tb in tables), costs_money=True))
+        return res
+
+    def _boto(self, service: str):
+        if self._session is None:
+            try:
+                import boto3
+            except ImportError as exc:
+                raise ImportError("pip install 'sql2iceberg[aws]'") from exc
+            self._session = boto3.Session(region_name=self.t.region or None)
+        return self._session.client(service)
+
+    def _run(self, action: Action) -> str | None:
+        if action.kind == "sql" and action.costs_money:
+            return self._athena(action.text)
+        elif action.kind == "sql":
+            if self._spark is None:
+                raise RuntimeError("pass spark=<SparkSession with a Glue catalog> to register tables")
+            for stmt in [s for s in action.text.split(";\n") if s.strip()]:
+                self._spark.sql(stmt.rstrip(";"))
+        else:
+            self._verify_glue(action.text)
+
+    def _verify_glue(self, text: str) -> None:
+        glue = self._boto("glue")
+        names = text.split("for: ")[1].split(" (")[0].split(", ")
+        for full in names:
+            db, name = full.split(".")
+            tbl = glue.get_table(DatabaseName=db, Name=name)["Table"]
+            if tbl.get("Parameters", {}).get("table_type", "").upper() != "ICEBERG":
+                raise RuntimeError(f"{full} exists in Glue but is not an Iceberg table")
+
+    def _athena(self, sql_text: str) -> str:
+        import time
+
+        athena = self._boto("athena")
+        lines: list[str] = []
+        for stmt in [s for s in sql_text.split(";\n") if s.strip()]:
+            qid = athena.start_query_execution(
+                QueryString=stmt.rstrip(";"), WorkGroup=self.t.athena_workgroup,
+                ResultConfiguration={"OutputLocation": self.t.athena_output})["QueryExecutionId"]
+            while True:
+                state = athena.get_query_execution(QueryExecutionId=qid)["QueryExecution"]["Status"]["State"]
+                if state in ("SUCCEEDED", "FAILED", "CANCELLED"):
+                    break
+                time.sleep(1)
+            if state != "SUCCEEDED":
+                raise RuntimeError(f"Athena query {state}: {stmt[:80]}")
+            rows = athena.get_query_results(QueryExecutionId=qid)["ResultSet"]["Rows"][1:]
+            lines.extend(", ".join(c.get("VarCharValue", "") for c in r["Data"]) for r in rows)
+        return "\n".join(lines)

sql2iceberg/publish/base.py

@@ -0,0 +1,120 @@
+"""Publishers expose existing Iceberg tables to another platform (Snowflake, AWS Glue/Athena, BigQuery).
+
+Every publisher first builds a plan (a list of Actions: SQL or an API call described as text) and only
+changes anything in the target account when ``apply()`` is called. ``plan()`` never connects anywhere.
+"""
+
+from __future__ import annotations
+
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+
+from ..tracing import span
+
+_IDENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*$")
+
+
+def ident(value: str, what: str = "identifier") -> str:
+    """Validate a SQL object name. Names are interpolated into DDL, so reject anything unusual."""
+    if not _IDENT.match(value or ""):
+        raise ValueError(f"unsafe or empty {what}: {value!r} (letters, digits, _ and $ only)")
+    return value
+
+
+def lit(value: str) -> str:
+    """Quote a SQL string literal."""
+    return "'" + str(value).replace("\\", "\\\\").replace("'", "''") + "'"
+
+
+@dataclass
+class IcebergTable:
+    name: str                       # table name (no catalog/namespace)
+    namespace: str = "db"
+    metadata_location: str = ""     # full URI of the current metadata.json, e.g. s3://bucket/db/t/metadata/v3.metadata.json
+    format_version: int = 2
+
+
+@dataclass
+class Action:
+    target: str                     # snowflake | aws | gcp
+    kind: str                       # sql | api
+    description: str
+    text: str                       # SQL text, or a description of the API call
+    costs_money: bool = False       # warehouse / query usage possible
+    done: bool = False
+    error: str = ""
+    output: str = ""                # rows returned by a check, as text
+
+
+@dataclass
+class PublishResult:
+    actions: list[Action] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+
+    @property
+    def ok(self) -> bool:
+        return all(a.done and not a.error for a in self.actions)
+
+    def to_markdown(self) -> str:
+        out = [f"# Publish plan ({len(self.actions)} actions)\n"]
+        out += [f"- WARNING: {w}" for w in self.warnings]
+        for i, a in enumerate(self.actions, 1):
+            state = "applied" if a.done else ("FAILED: " + a.error if a.error else "planned")
+            out.append(f"\n## {i}. {a.description} [{a.target}, {state}]")
+            out.append(f"```{'sql' if a.kind == 'sql' else 'text'}\n{a.text}\n```")
+            if a.output:
+                out.append(f"Result:\n```\n{a.output}\n```")
+        return "\n".join(out) + "\n"
+
+
+class Publisher(ABC):
+    target = ""
+
+    def __init__(self) -> None:
+        self.warnings: list[str] = []
+
+    @abstractmethod
+    def plan(self, tables: list[IcebergTable]) -> PublishResult: ...
+
+    @abstractmethod
+    def _run(self, action: Action) -> str | None:
+        """Execute one action against the real account. Raise on failure; return any rows as text."""
+
+    def apply(self, tables: list[IcebergTable], *, confirm: bool = False) -> PublishResult:
+        """Run the plan. ``confirm=True`` is required so an accidental call cannot change an account."""
+        if not confirm:
+            raise PermissionError("apply() changes a real account: call it with confirm=True after reading plan()")
+        result = self.plan(tables)
+        for action in result.actions:
+            try:
+                with span("publish.action", target=action.target, step=action.description):
+                    action.output = self._run(action) or ""
+                action.done = True
+            except Exception as exc:  # noqa: BLE001  report and stop: later steps depend on earlier ones
+                action.error = f"{type(exc).__name__}: {exc}"[:1500]
+                break
+        return result
+
+
+def publisher_from_config(cfg: dict) -> tuple[Publisher, list[IcebergTable]]:
+    """Build a publisher and its tables from a plain dict (the JSON file the CLI reads).
+
+    {"target": "snowflake" | "aws" | "gcp", "<target>": {...target fields...},
+     "tables": [{"name": ..., "namespace": ..., "metadata_location": ..., "format_version": 2}]}
+    """
+    from .aws import GluePublisher, GlueTarget
+    from .gcp import BigQueryPublisher, BigQueryTarget
+    from .snowflake import SnowflakePublisher, SnowflakeTarget
+
+    kind = cfg.get("target")
+    tables = [IcebergTable(**t) for t in cfg.get("tables", [])]
+    if not tables:
+        raise ValueError("config has no tables")
+    if kind == "snowflake":
+        return SnowflakePublisher(SnowflakeTarget(**cfg["snowflake"])), tables
+    if kind == "aws":
+        return GluePublisher(GlueTarget(**cfg["aws"])), tables
+    if kind == "gcp":
+        return BigQueryPublisher(BigQueryTarget(**cfg["gcp"])), tables
+    raise ValueError("target must be snowflake, aws or gcp")

sql2iceberg/publish/gcp.py

@@ -0,0 +1,72 @@
+"""GCP: BigLake external Iceberg tables in BigQuery.
+
+BigQuery reads Iceberg data from Cloud Storage, so the warehouse must be on ``gs://``. Each table is created
+from its current metadata.json (a snapshot: re-run after writes). Needs a BigQuery connection resource with
+read access to the bucket, and ``pip install sql2iceberg[gcp]`` to apply. Authentication is Application Default
+Credentials (``gcloud auth application-default login`` or a service account), never stored here.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .base import Action, IcebergTable, Publisher, PublishResult, ident, lit
+
+
+@dataclass
+class BigQueryTarget:
+    project: str
+    dataset: str
+    connection: str                  # "<project>.<location>.<connection_id>"
+    location: str = "US"
+
+    def __post_init__(self) -> None:
+        for v, what in ((self.project.replace("-", "_"), "project"), (self.dataset, "dataset")):
+            ident(v, what)
+        if len(self.connection.split(".")) != 3:
+            raise ValueError("connection must look like project.location.connection_id")
+
+
+class BigQueryPublisher(Publisher):
+    target = "gcp"
+
+    def __init__(self, target: BigQueryTarget, client=None) -> None:
+        super().__init__()
+        self.t, self._client = target, client
+
+    def plan(self, tables: list[IcebergTable]) -> PublishResult:
+        t = self.t
+        res = PublishResult()
+        res.warnings.append("BigLake Iceberg tables are snapshots of one metadata file: re-run after each write.")
+        if any(tb.format_version >= 3 for tb in tables):
+            res.warnings.append("Some tables are Iceberg format-version 3; check BigQuery's supported versions "
+                                "and write with --iceberg-version 2 if unsure.")
+        res.actions.append(Action("gcp", "sql", "Create dataset",
+                                  f"CREATE SCHEMA IF NOT EXISTS `{t.project}.{t.dataset}` OPTIONS (location = {lit(t.location)});"))
+        for tb in tables:
+            if not tb.metadata_location.startswith("gs://"):
+                raise ValueError(f"{tb.name}: BigQuery needs a gs:// metadata_location, got {tb.metadata_location!r}")
+            res.actions.append(Action(
+                "gcp", "sql", f"Expose {tb.name}",
+                f"CREATE OR REPLACE EXTERNAL TABLE `{t.project}.{t.dataset}.{ident(tb.name, 'table name')}`\n"
+                f"  WITH CONNECTION `{t.connection}`\n"
+                f"  OPTIONS (format = 'ICEBERG', uris = [{lit(tb.metadata_location)}]);"))
+        res.actions.append(Action(
+            "gcp", "sql", "Count rows (billed per bytes processed)",
+            "\n".join(f"SELECT '{tb.name}' AS table_name, COUNT(*) AS row_count "
+                      f"FROM `{t.project}.{t.dataset}.{tb.name}`;" for tb in tables), costs_money=True))
+        return res
+
+    def _run(self, action: Action) -> str | None:
+        if self._client is None:
+            try:
+                from google.cloud import bigquery
+            except ImportError as exc:
+                raise ImportError("pip install 'sql2iceberg[gcp]'") from exc
+            self._client = bigquery.Client(project=self.t.project, location=self.t.location)
+        lines = []
+        for stmt in [s for s in action.text.split(";\n") if s.strip()]:
+            rows = self._client.query(stmt.rstrip(";")).result()
+            if stmt.lstrip().upper().startswith("SELECT"):
+                lines.extend(", ".join(str(v) for v in tuple(r)) for r in rows)
+        return "\n".join(lines) or None

sql2iceberg/publish/metadata.py

@@ -0,0 +1,63 @@
+"""Find the current metadata.json of Iceberg tables."""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+from .base import IcebergTable
+
+_VERSION = re.compile(r"^(?:v)?(\d+)(?:-[0-9a-f-]+)?\.(?:gz\.)?metadata\.json$")
+
+
+def latest_local_metadata(warehouse: str | Path, namespace: str, table: str) -> Path:
+    """Newest metadata file of a Hadoop-catalog table on a local or mounted path."""
+    folder = Path(warehouse) / namespace / table / "metadata"
+    best: tuple[int, Path] | None = None
+    for f in folder.glob("*.metadata.json*"):
+        m = _VERSION.match(f.name)
+        if m and (best is None or int(m.group(1)) > best[0]):
+            best = (int(m.group(1)), f)
+    if best is None:
+        raise FileNotFoundError(f"no metadata.json under {folder}")
+    return best[1]
+
+
+def metadata_from_spark(spark, catalog: str, namespace: str, table: str) -> str:
+    """Current metadata location from Iceberg's metadata_log_entries (works for any catalog type)."""
+    rows = spark.sql(
+        f"SELECT file FROM {catalog}.{namespace}.{table}.metadata_log_entries ORDER BY timestamp DESC LIMIT 1"
+    ).collect()
+    if not rows:
+        raise LookupError(f"{catalog}.{namespace}.{table} has no metadata log")
+    return rows[0][0]
+
+
+def tables_from_spark(spark, catalog: str, namespace: str) -> list[IcebergTable]:
+    """Every table in a namespace with its current metadata location and format version."""
+    out = []
+    for row in spark.sql(f"SHOW TABLES IN {catalog}.{namespace}").collect():
+        name = row["tableName"]
+        loc = metadata_from_spark(spark, catalog, namespace, name)
+        props = {r[0]: r[1] for r in spark.sql(f"SHOW TBLPROPERTIES {catalog}.{namespace}.{name}").collect()}
+        out.append(IcebergTable(name, namespace, loc, int(props.get("format-version", 2))))
+    return out
+
+
+def tables_from_jdbc_catalog(sqlite_file: str | Path, namespace: str, catalog: str = "lake",
+                             format_version: int = 2) -> list[IcebergTable]:
+    """Tables and current metadata locations from an Iceberg JdbcCatalog stored in a SQLite file.
+
+    Use this when the catalog is the SQLite file written by ``catalog_type="jdbc"``. ``format_version`` is not
+    stored in the catalog table, so pass the version the tables were written with.
+    """
+    import sqlite3
+
+    con = sqlite3.connect(str(sqlite_file))
+    try:
+        rows = con.execute(
+            "SELECT table_name, metadata_location FROM iceberg_tables "
+            "WHERE catalog_name = ? AND table_namespace = ? ORDER BY table_name", (catalog, namespace)).fetchall()
+    finally:
+        con.close()
+    return [IcebergTable(n, namespace, loc, format_version) for n, loc in rows]

sql2iceberg/publish/snowflake.py

@@ -0,0 +1,143 @@
+"""Snowflake: query Iceberg tables you manage elsewhere, as externally managed Iceberg tables.
+
+Two ways to tell Snowflake where the table is:
+  * ``glue``          catalog integration to AWS Glue; Snowflake follows the Glue table (CATALOG_TABLE_NAME)
+  * ``object_store``  catalog integration on the storage itself; each table points at a metadata.json file
+
+SQL follows the Snowflake documentation for CREATE EXTERNAL VOLUME, CREATE CATALOG INTEGRATION and
+CREATE ICEBERG TABLE (externally managed). The table DDL matches the shape of a table already created on a live
+account by an earlier session (object_store mode); Glue mode and the volume/integration DDL are unverified
+until the live test in tests/test_live.py has been run.
+Creating these objects needs a role with the matching privileges, and queries use warehouse credits.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from urllib.parse import urlparse
+
+from .base import Action, IcebergTable, Publisher, PublishResult, ident, lit
+
+_PROVIDERS = {"s3": "S3", "gs": "GCS", "gcs": "GCS", "azure": "AZURE", "abfss": "AZURE"}
+
+
+@dataclass
+class SnowflakeTarget:
+    database: str
+    schema: str
+    base_url: str                      # storage root that contains the warehouse, e.g. s3://bucket/warehouse/
+    mode: str = "object_store"         # object_store | glue
+    external_volume: str = "SF2ICE_VOL"
+    catalog_integration: str = "SF2ICE_CATALOG"
+    aws_role_arn: str = ""             # role Snowflake assumes to read the bucket (S3) and Glue (glue mode)
+    glue_catalog_id: str = ""          # AWS account id (glue mode)
+    glue_region: str = ""
+    glue_namespace: str = "db"
+    azure_tenant_id: str = ""
+    auto_refresh: bool = False         # glue mode only
+    create_volume: bool = True         # False: reuse an existing external volume with this name
+    create_integration: bool = True    # False: reuse an existing catalog integration with this name
+    create_schema: bool = True
+
+    def __post_init__(self) -> None:
+        for field_name in ("database", "schema", "external_volume", "catalog_integration"):
+            ident(getattr(self, field_name), field_name)
+        if self.mode not in ("object_store", "glue"):
+            raise ValueError("mode must be object_store or glue")
+        scheme = urlparse(self.base_url).scheme
+        if scheme not in _PROVIDERS:
+            raise ValueError(f"base_url must start with one of {sorted(_PROVIDERS)}://")
+        if scheme == "s3" and self.create_volume and not self.aws_role_arn:
+            raise ValueError("aws_role_arn is required for an S3 external volume")
+        if self.mode == "glue" and self.create_integration and not (
+                self.glue_catalog_id and self.glue_region and self.aws_role_arn):
+            raise ValueError("glue mode needs glue_catalog_id, glue_region and aws_role_arn")
+
+
+class SnowflakePublisher(Publisher):
+    target = "snowflake"
+
+    def __init__(self, target: SnowflakeTarget, connector=None) -> None:
+        super().__init__()
+        self.t = target
+        self._connector = connector
+
+    def _volume_sql(self) -> str:
+        t = self.t
+        scheme = urlparse(t.base_url).scheme
+        provider = _PROVIDERS[scheme]
+        base = self._base()
+        fields = ["NAME = 'loc1'", f"STORAGE_PROVIDER = '{provider}'", f"STORAGE_BASE_URL = {lit(base)}"]
+        if provider == "S3":
+            fields.append(f"STORAGE_AWS_ROLE_ARN = {lit(t.aws_role_arn)}")
+        if provider == "AZURE":
+            fields.append(f"AZURE_TENANT_ID = {lit(t.azure_tenant_id)}")
+        return (f"CREATE EXTERNAL VOLUME IF NOT EXISTS {t.external_volume}\n  STORAGE_LOCATIONS = (( "
+                + ", ".join(fields) + " ));")
+
+    def _integration_sql(self) -> str:
+        t = self.t
+        if t.mode == "glue":
+            return (f"CREATE CATALOG INTEGRATION IF NOT EXISTS {t.catalog_integration}\n"
+                    f"  CATALOG_SOURCE = GLUE\n  CATALOG_NAMESPACE = {lit(t.glue_namespace)}\n  TABLE_FORMAT = ICEBERG\n"
+                    f"  GLUE_AWS_ROLE_ARN = {lit(t.aws_role_arn)}\n  GLUE_CATALOG_ID = {lit(t.glue_catalog_id)}\n"
+                    f"  GLUE_REGION = {lit(t.glue_region)}\n  ENABLED = TRUE;")
+        return (f"CREATE CATALOG INTEGRATION IF NOT EXISTS {t.catalog_integration}\n"
+                f"  CATALOG_SOURCE = OBJECT_STORE\n  TABLE_FORMAT = ICEBERG\n  ENABLED = TRUE;")
+
+    def _base(self) -> str:
+        base = self.t.base_url if self.t.base_url.endswith("/") else self.t.base_url + "/"
+        return base.replace("gs://", "gcs://", 1)
+
+    def _relative_path(self, uri: str) -> str:
+        """Path of a metadata file relative to the external volume's base URL."""
+        uri, base = uri.replace("gs://", "gcs://", 1), self._base()
+        if not uri.startswith(base):
+            raise ValueError(f"{uri} is not under the external volume base {base}")
+        return uri[len(base):]
+
+    def _table_sql(self, tb: IcebergTable) -> str:
+        t = self.t
+        head = f"CREATE ICEBERG TABLE IF NOT EXISTS {t.database}.{t.schema}.{ident(tb.name, 'table name')}\n"
+        head += f"  EXTERNAL_VOLUME = {lit(t.external_volume)}\n  CATALOG = {lit(t.catalog_integration)}\n"
+        if t.mode == "glue":
+            return head + f"  CATALOG_TABLE_NAME = {lit(tb.name)}" + (
+                "\n  AUTO_REFRESH = TRUE" if t.auto_refresh else "") + ";"
+        if not tb.metadata_location:
+            raise ValueError(f"{tb.name}: metadata_location is required in object_store mode")
+        return head + f"  METADATA_FILE_PATH = {lit(self._relative_path(tb.metadata_location))};"
+
+    def plan(self, tables: list[IcebergTable]) -> PublishResult:
+        t = self.t
+        res = PublishResult()
+        if any(tb.format_version >= 3 for tb in tables):
+            res.warnings.append("Some tables are Iceberg format-version 3. Check the current Snowflake docs for "
+                                "v3 support in externally managed tables; use --iceberg-version 2 if unsure.")
+        if t.mode == "object_store":
+            res.warnings.append("object_store mode points at a fixed metadata file. Re-run after each write "
+                                "(ALTER ICEBERG TABLE ... REFRESH 'new/metadata.json') or use glue mode.")
+        add = res.actions.append
+        if t.create_volume:
+            add(Action("snowflake", "sql", "Create external volume", self._volume_sql()))
+        if t.create_integration:
+            add(Action("snowflake", "sql", "Create catalog integration", self._integration_sql()))
+        if t.create_schema:
+            add(Action("snowflake", "sql", "Create schema", f"CREATE SCHEMA IF NOT EXISTS {t.database}.{t.schema};"))
+        for tb in tables:
+            add(Action("snowflake", "sql", f"Expose {tb.name}", self._table_sql(tb)))
+        add(Action("snowflake", "sql", "Count rows (uses a warehouse)",
+                   "\n".join(f"SELECT '{tb.name}' AS table_name, COUNT(*) AS row_count "
+                             f"FROM {t.database}.{t.schema}.{tb.name};" for tb in tables), costs_money=True))
+        return res
+
+    def _run(self, action: Action) -> str | None:
+        if self._connector is None:
+            from ..connectors import get_connector
+
+            self._connector = get_connector("snowflake")
+        lines = []
+        for stmt in [x for x in action.text.split(";\n") if x.strip()]:
+            rows = self._connector.query(stmt.rstrip(";"))
+            if stmt.lstrip().upper().startswith("SELECT"):
+                lines.extend(", ".join(str(v) for v in r) for r in rows)
+        return "\n".join(lines) or None