"""
iceberg_ops.autogen
===================
Autogenerate migration files from a schema diff.
This module powers ``iceberg-ops revision --autogenerate``: it compares an
existing Iceberg table against a target schema (``pa.Schema``,
:class:`~iceberg_ops.core.schema.SchemaContract`, or
:class:`pyiceberg.schema.Schema`) and produces a list of
:class:`MigrationOp`, which the renderer then turns into runnable
migration source code.
Public API
----------
:class:`MigrationOp`
Dataclass describing one op (kind + args).
:func:`schema_diff_to_ops(*, current, target, rename_hints=None, format_version=None, spec=None, sort_order=None)`
Compute the ops that turn ``current`` into ``target`` for column
schema (add / rename / delete / type_change).
:func:`spec_diff_to_ops(*, current, target, schema)`
Compute the ops that turn ``current`` partition spec into
``target``.
:func:`sort_order_diff_to_ops(*, current, target, schema)`
Emit one ``set_sort_order`` op when the orders differ.
:func:`properties_diff_to_ops(*, current, target)`
Emit ``set_property`` / ``remove_property`` ops.
:func:`render_migration_body(ops, *, target_table)`
Translate a ``list[MigrationOp]`` into ``(upgrade_body,
downgrade_body)`` Python source strings ready to drop into a
:class:`~iceberg_ops.migration.BaseMigration` skeleton.
:func:`render_migration_file(*, version, description, ops, target_table)`
Top-level helper that produces the full migration file source.
:func:`write_autogenerated_revision(*, version, description, ops, target_table, out_path, overwrite=False)`
Render and write the file to disk (UTF-8, LF line endings). Refuses
to replace an existing file unless ``overwrite=True``.
Scope
-----
- Every name, value and description rendered into a generated file is
written as a Python literal (``repr``), so a generated file always
compiles and cannot run code smuggled in through a column name or a
``-m`` message.
- Schema column ops handle primitive Iceberg types. Type changes are
emitted only for Iceberg type promotions (``int -> long``,
``float -> double``, decimal precision widening at the same scale).
- Column names containing a dot are refused: the ``MigrationContext``
helpers read a dot as a nested struct path.
- Steps that cannot run as generated (incompatible type changes, nested
types, dropping a column the partition spec, sort order or identifier
fields reference) or cannot be reversed (column drops, type widening,
partition-field removal, sort-order replacement, property removal with
an unknown old value) make the affected ``upgrade`` / ``downgrade``
raise :class:`NotImplementedError` before any step runs. The generated
calls are kept as comments so the author can finish the method by
hand.
- Partition specs and sort orders are diffed for programmatic use; the
CLI front-end (``revision --autogenerate``) is schema-only.
- Nullability and doc changes on existing columns are not emitted yet.
"""
from __future__ import annotations
import ast
import logging
import re
from collections.abc import Callable, Iterable
from dataclasses import dataclass, field
from pathlib import Path
from string import Template
from typing import Any, Literal
import pyarrow as pa
from pyiceberg import types as ice_types
from pyiceberg.partitioning import PartitionSpec
from pyiceberg.schema import Schema as IcebergSchema
from pyiceberg.table.sorting import NullOrder, SortDirection, SortOrder
from pyiceberg.transforms import IdentityTransform, VoidTransform
from iceberg_ops.core.schema import SchemaContract, diff_schemas
from iceberg_ops.core.types import arrow_type_to_iceberg, is_allowed_promotion
from iceberg_ops.scaffolding import _docstring_text
logger = logging.getLogger(__name__)
# All op kinds emitted by the diff functions and consumed by the renderer.
MigrationOpKind = Literal[
"add_column",
"rename_column",
"delete_column",
"type_change",
"add_partition_field",
"remove_partition_field",
"set_sort_order",
"set_property",
"remove_property",
]
[docs]
@dataclass
class MigrationOp:
"""A single migration operation that the autogen renderer will emit.
Attributes
----------
kind : str
One of the supported op kinds. See :data:`MigrationOpKind`.
args : dict
Op-specific arguments. The schema is documented per-kind in
the diff functions below; the renderer reads these args to
generate the matching Python source.
"""
kind: MigrationOpKind
args: dict[str, Any] = field(default_factory=dict)
# ===========================================================================
# Schema column diffs
# ===========================================================================
def _iceberg_type_to_source(t: ice_types.IcebergType) -> str | None:
"""Render an Iceberg type as its Python source form (qualified by
the ``ice_types.`` alias the renderer always imports).
Examples
--------
>>> _iceberg_type_to_source(ice_types.LongType())
'ice_types.LongType()'
>>> _iceberg_type_to_source(ice_types.DecimalType(10, 2))
'ice_types.DecimalType(precision=10, scale=2)'
Returns ``None`` for nested types
(:class:`~pyiceberg.types.ListType` / :class:`~pyiceberg.types.MapType` /
:class:`~pyiceberg.types.StructType`) and for any type whose
``repr`` is not a plain constructor call with literal arguments; the
renderer blocks such steps instead of emitting them.
"""
if not isinstance(t, ice_types.PrimitiveType):
return None
src = repr(t)
try:
node = ast.parse(src, mode="eval").body
except SyntaxError:
return None
if not (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and getattr(ice_types, node.func.id, None) is type(t)
and all(isinstance(arg, ast.Constant) for arg in node.args)
and all(isinstance(kw.value, ast.Constant) for kw in node.keywords)
):
return None
return f"ice_types.{src}"
def _align_target_to_current(
target_fields: list[tuple[str, ice_types.IcebergType, str | None, bool]],
current: IcebergSchema,
rename_hints: dict[str, str],
) -> IcebergSchema:
"""Build an :class:`IcebergSchema` from raw target fields, reusing
``current`` field-IDs by name (with rename-hint indirection).
Each entry of ``target_fields`` is ``(name, type, doc, required)``.
The aligned schema has the same column order as ``target_fields``;
field-IDs are taken from ``current`` for matching names, fresh IDs
are issued for genuinely new columns.
Rename mechanics: a hint ``{"old_name": "new_name"}`` makes the
target's ``new_name`` field carry ``current["old_name"].field_id``,
so :func:`~iceberg_ops.core.schema.diff_schemas` recognises the
pair as a rename rather than drop+add.
"""
current_by_name = {f.name: f for f in current.fields}
rev_hints = {new: old for old, new in rename_hints.items()}
next_id = max((f.field_id for f in current.fields), default=0) + 1
nested: list[ice_types.NestedField] = []
for name, ice_type, doc, required in target_fields:
if name in rev_hints and rev_hints[name] in current_by_name:
fid = current_by_name[rev_hints[name]].field_id
elif name in current_by_name:
fid = current_by_name[name].field_id
else:
fid = next_id
next_id += 1
nested.append(
ice_types.NestedField(
field_id=fid,
name=name,
field_type=ice_type,
required=required,
doc=doc,
)
)
return IcebergSchema(*nested)
def _arrow_field_to_iceberg_type(
f: pa.Field,
current: IcebergSchema,
rename_hints: dict[str, str],
) -> ice_types.IcebergType:
"""Map one target :class:`pa.Field` to its Iceberg type.
Primitive fields go through
:func:`iceberg_ops.core.types.arrow_type_to_iceberg`. A nested field
(list / struct / map) whose column (or rename-hint source) already
exists as a nested column reuses the current type: it is treated as
unchanged and a warning says it was not compared. A nested field
for a new column, or one replacing a primitive column, raises.
"""
if not pa.types.is_nested(f.type):
return arrow_type_to_iceberg(f.type)
source = {new: old for old, new in rename_hints.items()}.get(f.name, f.name)
current_by_name = {cf.name: cf for cf in current.fields}
existing = current_by_name.get(source)
if existing is not None and not isinstance(existing.field_type, ice_types.PrimitiveType):
logger.warning(
"Nested column %r in the pa.Schema target was not compared; "
"it is treated as unchanged.",
f.name,
)
return existing.field_type
if existing is not None:
raise ValueError(
f"Column {f.name!r} changes to a nested type in the pa.Schema target; "
"autogenerate cannot render this, use a hand-written migration."
)
raise ValueError(
f"Nested column {f.name!r} in a pa.Schema target cannot be added by "
"autogenerate; use a hand-written migration or an IcebergSchema target."
)
def _normalise_target(
target: pa.Schema | IcebergSchema | SchemaContract,
current: IcebergSchema,
rename_hints: dict[str, str],
) -> IcebergSchema:
"""Convert any supported ``target`` representation into an
ID-aligned :class:`IcebergSchema` ready for :func:`diff_schemas`.
- :class:`IcebergSchema`: passed through (caller is responsible for IDs).
- :class:`pa.Schema`: each field's type is mapped via
:func:`_arrow_field_to_iceberg_type`; ``required`` is the inverse
of Arrow's ``nullable`` flag; ``doc`` is taken from Arrow's field
metadata under the ``b"doc"`` key when present.
- :class:`SchemaContract`: ``ColumnSpec`` already carries
everything; we just translate to the tuple form.
"""
if isinstance(target, IcebergSchema):
return target
fields: list[tuple[str, ice_types.IcebergType, str | None, bool]] = []
if isinstance(target, pa.Schema):
for f in target:
doc = None
if f.metadata and b"doc" in f.metadata:
doc = f.metadata[b"doc"].decode("utf-8")
fields.append(
(
f.name,
_arrow_field_to_iceberg_type(f, current, rename_hints),
doc,
not f.nullable,
)
)
elif isinstance(target, SchemaContract):
for col in target.columns:
fields.append((col.name, col.type, col.doc, col.required))
else:
raise TypeError(
f"Unsupported target type {type(target).__name__}: "
"expected pa.Schema, IcebergSchema, or SchemaContract."
)
return _align_target_to_current(fields, current, rename_hints)
def _target_names(target: pa.Schema | IcebergSchema | SchemaContract) -> list[str]:
"""Top-level column names of any supported target, in order."""
if isinstance(target, IcebergSchema):
return [f.name for f in target.fields]
if isinstance(target, pa.Schema):
return list(target.names)
if isinstance(target, SchemaContract):
return [c.name for c in target.columns]
raise TypeError(
f"Unsupported target type {type(target).__name__}: "
"expected pa.Schema, IcebergSchema, or SchemaContract."
)
def _validate_rename_hints(
hints: dict[str, str],
current_names: Iterable[str],
target_names: Iterable[str],
) -> None:
"""Reject rename hints that would silently lose or corrupt columns.
Raises
------
ValueError
When a hint's old name is not a current column, its new name is
not a target column, the old name is still in the target, the
new name already exists in the current table (swaps and chained
renames need two migrations), or two hints share a new name.
"""
current_set, target_set = set(current_names), set(target_names)
claimed: dict[str, str] = {}
for old, new in hints.items():
label = f"rename hint {old!r} -> {new!r}"
if old not in current_set:
raise ValueError(f"{label}: {old!r} is not a column of the current table.")
if new not in target_set:
raise ValueError(f"{label}: {new!r} is not a column of the target schema.")
if old in target_set:
raise ValueError(
f"{label}: {old!r} is still in the target schema, so it is not being renamed."
)
if new in current_set:
raise ValueError(
f"{label}: {new!r} already exists in the current table; swaps and "
"chained renames have to be done in two migrations."
)
if new in claimed:
raise ValueError(
f"{label}: {new!r} is also the target of the hint for {claimed[new]!r}."
)
claimed[new] = old
def _drops_unlisted(target: pa.Schema | IcebergSchema | SchemaContract) -> bool:
"""Whether columns missing from *target* should be dropped."""
return not isinstance(target, SchemaContract) or target.strict
def _unlisted_columns_kept(
current: IcebergSchema,
target: pa.Schema | IcebergSchema | SchemaContract,
rename_hints: dict[str, str] | None = None,
) -> list[str]:
"""Current columns a non-strict :class:`SchemaContract` target keeps.
Returns the current column names that are neither listed in the
contract nor the source of a rename hint; ``[]`` for any other
target (those drop unlisted columns).
"""
if _drops_unlisted(target):
return []
listed = set(_target_names(target)) | set(rename_hints or {})
return [f.name for f in current.fields if f.name not in listed]
def _uncompared_nested_columns(
current: IcebergSchema,
target: pa.Schema | IcebergSchema | SchemaContract,
rename_hints: dict[str, str] | None = None,
) -> list[str]:
"""Nested columns of a :class:`pa.Schema` target that are not compared.
A nested (struct / list / map) target field whose column, or
rename-hint source, already exists as a nested column is treated
as unchanged (see :func:`_arrow_field_to_iceberg_type`). Returns the
target names of those columns; ``[]`` for any other target type.
"""
if not isinstance(target, pa.Schema):
return []
sources = {new: old for old, new in (rename_hints or {}).items()}
current_by_name = {f.name: f for f in current.fields}
skipped: list[str] = []
for f in target:
if not pa.types.is_nested(f.type):
continue
existing = current_by_name.get(sources.get(f.name, f.name))
if existing is not None and not isinstance(existing.field_type, ice_types.PrimitiveType):
skipped.append(f.name)
return skipped
def _ignored_attribute_changes(
current: IcebergSchema,
target: pa.Schema | IcebergSchema | SchemaContract,
rename_hints: dict[str, str] | None = None,
) -> list[str]:
"""Describe nullability / doc differences autogenerate does not emit.
Compares ``required`` and ``doc`` for every column present in both
schemas (matched by field id, so renames count), and flags new
columns the target marks required (they are added as optional).
A target doc of ``None`` means "unspecified" and is not reported.
Returns
-------
list[str]
One entry per affected column, e.g.
``"id (optional -> required)"``.
"""
hints = dict(rename_hints or {})
target_iceberg = _normalise_target(target, current, hints)
current_by_id = {f.field_id: f for f in current.fields}
notes: list[str] = []
for tf in target_iceberg.fields:
cf = current_by_id.get(tf.field_id)
if cf is None:
if tf.required:
notes.append(f"{tf.name} (new column added as optional, target is required)")
continue
changes: list[str] = []
if cf.required != tf.required:
old_req = "required" if cf.required else "optional"
new_req = "required" if tf.required else "optional"
changes.append(f"{old_req} -> {new_req}")
if tf.doc is not None and tf.doc != cf.doc:
changes.append("doc differs")
if changes:
notes.append(f"{tf.name} ({', '.join(changes)})")
return notes
def _required_format_version(t: ice_types.IcebergType) -> int:
"""Lowest table format version that can hold *t* (nested fields included)."""
from pyiceberg.schema import index_by_id
versions = [int(t.minimum_format_version())]
if not t.is_primitive:
versions.extend(
int(f.field_type.minimum_format_version()) for f in index_by_id(t).values()
)
return max(versions)
def _layout_references(
field: ice_types.NestedField,
current: IcebergSchema,
spec: PartitionSpec | None,
sort_order: SortOrder | None,
) -> list[str]:
"""Describe what in the table layout references *field* or its children.
Checks the identifier fields of ``current`` and, when given, the
non-void fields of ``spec`` and the fields of ``sort_order`` -
the references PyIceberg refuses to leave dangling when a column is
deleted.
"""
from pyiceberg.schema import index_by_id
ids = {field.field_id, *index_by_id(field.field_type)}
refs: list[str] = []
if spec is not None:
refs.extend(
f"partition field {pf.name!r}"
for pf in spec.fields
if pf.source_id in ids and not isinstance(pf.transform, VoidTransform)
)
if sort_order is not None and any(sf.source_id in ids for sf in sort_order.fields):
refs.append("the sort order")
if ids & set(current.identifier_field_ids):
refs.append("the identifier fields")
return refs
def _check_op_names(ops: list[MigrationOp]) -> None:
"""Refuse ops on column names that contain a dot.
Raises
------
ValueError
If any op's ``name`` / ``old_name`` / ``new_name`` contains
``"."``.
"""
dotted = sorted(
{
str(op.args[key])
for op in ops
for key in ("name", "old_name", "new_name")
if "." in str(op.args.get(key) or "")
}
)
if dotted:
raise ValueError(
f"Column name(s) {', '.join(repr(n) for n in dotted)} contain '.', which "
"MigrationContext helpers read as a nested struct path; write this step "
"by hand."
)
def _check_format_version(ops: list[MigrationOp], format_version: int) -> None:
"""Refuse added or changed column types the table's format version cannot hold.
Raises
------
ValueError
If an ``add_column`` / ``type_change`` type needs a newer table
format version than *format_version*.
"""
for op in ops:
key = {"add_column": "iceberg_type", "type_change": "new_iceberg_type"}.get(op.kind)
if key is None:
continue
ice_type = op.args[key]
needed = _required_format_version(ice_type)
if needed > format_version:
raise ValueError(
f"column {op.args['name']!r}: {ice_type} needs table format version "
f"{needed}; the table is version {format_version}."
)
[docs]
def schema_diff_to_ops(
*,
current: IcebergSchema,
target: pa.Schema | IcebergSchema | SchemaContract,
rename_hints: dict[str, str] | None = None,
format_version: int | None = None,
spec: PartitionSpec | None = None,
sort_order: SortOrder | None = None,
) -> list[MigrationOp]:
"""
Compute the list of :class:`MigrationOp` that turns ``current`` into
``target`` for column schema only (add / rename / delete /
type_change).
Uses :func:`~iceberg_ops.core.schema.diff_schemas` under the hood.
Renames are disambiguated by ``rename_hints`` (a mapping
``{old_name: new_name}``); when no hint covers a pair, a column
rename is emitted as drop + add.
Parameters
----------
current : IcebergSchema
Schema of the existing table - must carry stable field-IDs from
a real Iceberg table (e.g. ``table.schema()``).
target : pa.Schema | IcebergSchema | SchemaContract
Desired schema. :class:`IcebergSchema` is passed through;
:class:`pa.Schema` and :class:`SchemaContract` are normalised by
reusing ``current`` field-IDs for fields with the same name (or
via rename hints). Columns of ``current`` missing from the
target are dropped, except for a non-strict
:class:`SchemaContract` (``strict=False``, the default), which
only reconciles the columns it lists and keeps the rest.
Nested columns of a :class:`pa.Schema` target are not compared
when they already exist, and cannot be added.
rename_hints : dict[str, str], optional
Map of ``old_column_name`` -> ``new_column_name``. Each entry
causes a single ``rename_column`` op instead of drop + add.
Not accepted with an :class:`IcebergSchema` target, whose field
IDs already identify renames.
format_version : int, optional
The table's format version (``table.format_version``). When
given, an added or changed column whose type needs a newer
version (``timestamp_ns``, ``timestamptz_ns``, ``unknown``,
``geometry`` and ``geography`` need version 3) raises
``ValueError``. ``None`` (default) skips the check.
spec : PartitionSpec, optional
The partition spec in effect when the deletes run (normally
``table.spec()``). Deleting a column one of its non-void fields
is sourced from is rendered as an ``upgrade`` that raises.
sort_order : SortOrder, optional
The sort order in effect when the deletes run (normally
``table.sort_order()``); deleting a column it sorts by is
rendered as an ``upgrade`` that raises. The identifier fields
of ``current`` are always checked the same way.
Returns
-------
list[MigrationOp]
Ordered list of ops. Order: removed -> renamed -> type-changed
-> added; this keeps downstream renders deterministic.
Raises
------
ValueError
If a rename hint is inconsistent with the two schemas (unknown
old or new name, old name kept, new name already present, two
hints sharing a new name), if hints are combined with an
:class:`IcebergSchema` target, if the target repeats a column
name, if a target column type cannot be mapped, if an op names
a column containing ``"."`` (the ``MigrationContext`` helpers
read a dot as a nested struct path), or if *format_version* is
given and an added or changed type needs a newer one.
TypeError
If ``target`` is not one of the supported types.
Notes
-----
Op kinds emitted:
- ``add_column`` - args ``{"name", "iceberg_type", "doc"}``
- ``rename_column`` - args ``{"old_name", "new_name"}``
- ``delete_column`` - args ``{"name"}``, plus ``"referenced_by"``
(list of str) when the identifier fields, *spec* or *sort_order*
reference the column; the renderer then blocks the ``upgrade``.
- ``type_change`` - args ``{"name", "old_iceberg_type", "new_iceberg_type"}``
The ``iceberg_type`` values are :class:`pyiceberg.types.IcebergType`
instances; the renderer converts them to source code. A
``type_change`` that is not an Iceberg type promotion is still
emitted so the diff stays visible; the renderer turns it into an
``upgrade`` that raises.
Nullability and doc changes on existing columns are not emitted
yet (PyIceberg supports them via ``UpdateSchema.make_column_optional``
and ``update_column(doc=...)``). New columns are added as optional,
with the target doc when one is set.
"""
hints = dict(rename_hints or {})
target_names = _target_names(target)
duplicates = sorted({n for n in target_names if target_names.count(n) > 1})
if duplicates:
raise ValueError(f"Target schema repeats column name(s): {', '.join(duplicates)}.")
if hints:
if isinstance(target, IcebergSchema):
raise ValueError(
"rename_hints cannot be used with an IcebergSchema target: "
"its field IDs already identify renames."
)
_validate_rename_hints(hints, (f.name for f in current.fields), target_names)
target_iceberg = _normalise_target(target, current, hints)
diff = diff_schemas(current, target_iceberg)
target_by_id = {f.field_id: f for f in target_iceberg.fields}
current_by_id = {f.field_id: f for f in current.fields}
ops: list[MigrationOp] = []
# Removed columns first - emit before adds so generated migrations
# don't briefly carry both old and new versions of a soon-to-rename
# column when no rename hint covers it.
if _drops_unlisted(target):
for removed in diff.removed:
assert removed.field_id is not None
args: dict[str, Any] = {"name": removed.name}
refs = _layout_references(current_by_id[removed.field_id], current, spec, sort_order)
if refs:
args["referenced_by"] = refs
ops.append(MigrationOp(kind="delete_column", args=args))
elif diff.removed:
logger.info(
"Keeping %d column(s) not listed in the non-strict contract: %s",
len(diff.removed),
", ".join(r.name for r in diff.removed),
)
# Renames (changed entries with change_type == "renamed").
for changed in diff.changed:
if changed.change_type == "renamed":
ops.append(
MigrationOp(
kind="rename_column",
args={
"old_name": changed.old_name,
"new_name": changed.name,
},
)
)
# Type changes - keep the actual IcebergType instances so the
# renderer can check the promotion rule and nested types.
for changed in diff.changed:
if changed.change_type == "type_changed":
# FieldChange.field_id is typed Optional[int]; the
# diff_schemas path only emits "type_changed" entries for
# fields present in both schemas, so the id is guaranteed
# populated here.
assert changed.field_id is not None
old_type = current_by_id[changed.field_id].field_type
new_type = target_by_id[changed.field_id].field_type
ops.append(
MigrationOp(
kind="type_change",
args={
"name": changed.name,
"old_iceberg_type": old_type,
"new_iceberg_type": new_type,
},
)
)
# Adds last - for type_change semantics described above.
for added in diff.added:
assert added.field_id is not None
target_field = target_by_id[added.field_id]
ops.append(
MigrationOp(
kind="add_column",
args={
"name": added.name,
"iceberg_type": target_field.field_type,
"doc": target_field.doc,
},
)
)
_check_op_names(ops)
if format_version is not None:
_check_format_version(ops, format_version)
return ops
# ===========================================================================
# Partition spec / sort order / properties diffs
# ===========================================================================
def _transform_to_source(transform: Any) -> str:
"""Render a partition transform as Python source.
Every PyIceberg transform has a constructor-style ``repr`` (for
example ``BucketTransform(num_buckets=16)``); prefixing the
``transforms.`` module alias, which the file renderer imports, makes
it a valid expression in the generated migration.
"""
return f"transforms.{transform!r}"
[docs]
def spec_diff_to_ops(
*,
current: PartitionSpec,
target: PartitionSpec,
schema: IcebergSchema,
) -> list[MigrationOp]:
"""
Compute add / remove ops to evolve ``current`` partition spec into
``target``.
Compares fields by their ``name`` (the spec-level partition name,
not the source column). Each ``target`` field absent from
``current`` becomes an ``add_partition_field`` op; each ``current``
field absent from ``target`` becomes a ``remove_partition_field``
op.
Parameters
----------
current : PartitionSpec
``table.spec()`` from the existing table.
target : PartitionSpec
Desired spec.
schema : IcebergSchema
Used to resolve each partition field's source column name from
its ``source_id`` so the generated migration can pass the
correct argument to :meth:`MigrationContext.update_partition_spec`.
A nested source is rendered as its full dotted path
(``"addr.city"``).
Returns
-------
list[MigrationOp]
Ops in removal-first then add order.
Raises
------
ValueError
If a partition field keeps its name but changes its source
column or transform; autogenerate cannot evolve a field in
place, so give the new field a different name in the target.
Also if an added field's source id is not in *schema*.
Op kinds emitted
----------------
- ``remove_partition_field`` - args ``{"partition_name"}``
- ``add_partition_field`` - args
``{"source_column", "transform_source", "partition_name"}``, where
``transform_source`` is Python source such as
``"transforms.DayTransform()"``.
"""
current_by_name = {f.name: f for f in current.fields}
target_by_name = {f.name: f for f in target.fields}
for name, tf in target_by_name.items():
cf = current_by_name.get(name)
if cf is None or (cf.source_id, cf.transform) == (tf.source_id, tf.transform):
continue
if cf.transform != tf.transform:
change = f"{cf.transform} -> {tf.transform}"
else:
old_src = schema.find_column_name(cf.source_id)
new_src = schema.find_column_name(tf.source_id)
change = f"source column {old_src!r} -> {new_src!r}"
raise ValueError(
f"partition field {name!r} changes {change}; autogenerate cannot evolve "
"a field in place; rename it in the target spec."
)
ops: list[MigrationOp] = []
for name in current_by_name:
if name not in target_by_name:
ops.append(
MigrationOp(
kind="remove_partition_field",
args={"partition_name": name},
)
)
for name, pf in target_by_name.items():
if name not in current_by_name:
source_col = schema.find_column_name(pf.source_id)
if source_col is None:
raise ValueError(
f"partition field {name!r}: source id {pf.source_id} is not in the schema."
)
ops.append(
MigrationOp(
kind="add_partition_field",
args={
"source_column": source_col,
"transform_source": _transform_to_source(pf.transform),
"partition_name": name,
},
)
)
return ops
[docs]
def sort_order_diff_to_ops(
*,
current: SortOrder,
target: SortOrder,
schema: IcebergSchema,
) -> list[MigrationOp]:
"""
Emit a single ``set_sort_order`` op when the orders differ.
The renderer translates this to
``ctx.replace_sort_order(target_table, sort_fields=[(col, dir),
...])`` - replacing the entire sort order rather than adding /
removing fields, which matches PyIceberg's
``replace_sort_order`` semantics.
Two orders are considered identical iff they have the same number
of fields, in the same order, with the same ``source_id``,
``transform``, ``direction`` and ``null_order``.
Parameters
----------
current : SortOrder
``table.sort_order()`` from the existing table.
target : SortOrder
Desired sort order.
schema : IcebergSchema
Used to resolve each sort field's source column name; a nested
source is rendered as its full dotted path (``"addr.city"``).
Returns
-------
list[MigrationOp]
``[]`` when the orders are identical, otherwise one
``set_sort_order`` op with args ``{"sort_fields": [(column,
"asc" | "desc"), ...]}``.
Raises
------
ValueError
If the orders differ and a target field uses a non-identity
transform or a non-default null order (``asc`` must be
``NULLS_FIRST`` and ``desc`` ``NULLS_LAST``), which
:meth:`MigrationContext.replace_sort_order` cannot express, or
a target field's source id is not in *schema*.
"""
def _key(f: Any) -> tuple[Any, ...]:
return (f.source_id, f.transform, f.direction, f.null_order)
if len(current.fields) == len(target.fields) and all(
_key(a) == _key(b) for a, b in zip(current.fields, target.fields, strict=True)
):
return []
sort_fields: list[tuple[str, str]] = []
for f in target.fields:
col = schema.find_column_name(f.source_id)
if col is None:
raise ValueError(f"sort field source id {f.source_id} is not in the schema.")
ascending = f.direction == SortDirection.ASC
if not isinstance(f.transform, IdentityTransform):
raise ValueError(
f"sort field {col!r} uses {f.transform}; autogenerate only renders "
"identity sort fields. Write this sort order by hand."
)
default_nulls = NullOrder.NULLS_FIRST if ascending else NullOrder.NULLS_LAST
if f.null_order != default_nulls:
raise ValueError(
f"sort field {col!r} uses {f.null_order} with {f.direction}; autogenerate "
"only renders asc NULLS_FIRST and desc NULLS_LAST. Write this sort order by hand."
)
sort_fields.append((col, "asc" if ascending else "desc"))
return [MigrationOp(kind="set_sort_order", args={"sort_fields": sort_fields})]
[docs]
def properties_diff_to_ops(
*,
current: dict[str, str],
target: dict[str, str],
) -> list[MigrationOp]:
"""
Emit ``set_property`` and ``remove_property`` ops to reconcile the
table-property dicts.
Keys present in ``target`` but absent from ``current`` (or with a
different value) become ``set_property`` ops. Keys present in
``current`` but absent from ``target`` become ``remove_property``
ops; the old value is preserved in the args so the generated
downgrade body can restore it.
Parameters
----------
current : dict[str, str]
``table.properties`` of the existing table.
target : dict[str, str]
Desired properties.
Returns
-------
list[MigrationOp]
``remove_property`` ops (args ``{"key", "old_value"}``) first,
then ``set_property`` ops (args ``{"key", "value",
"old_value"}``), each group sorted by key.
"""
ops: list[MigrationOp] = []
# Removals first so the generated upgrade reads "remove old, set new".
for key in sorted(current):
if key not in target:
ops.append(
MigrationOp(
kind="remove_property",
args={"key": key, "old_value": current[key]},
)
)
for key in sorted(target):
new_value = target[key]
if current.get(key) != new_value:
ops.append(
MigrationOp(
kind="set_property",
args={
"key": key,
"value": new_value,
"old_value": current.get(key),
},
)
)
return ops
# ===========================================================================
# Migration-body renderer
# ===========================================================================
_INDENT = " " * 8 # method body inside ``class Migration`` (two indent levels)
@dataclass(frozen=True)
class _Blocked:
"""One side of a rendered op that must not run as generated.
``reason`` explains why; ``lines`` are the generated calls, kept as
comments in the output so the author can finish the method by hand.
"""
reason: str
lines: tuple[str, ...] = ()
# A rendered side of one op: runnable source lines, or a block.
_Side = list[str] | _Blocked
_ALLOWED_PROMOTIONS = (
"allowed: int->long, float->double, decimal precision widening at the same scale"
)
def _render_add_column(op: MigrationOp) -> tuple[_Side, _Side]:
name = op.args["name"]
ice_type = op.args["iceberg_type"]
doc = op.args.get("doc")
doc_src = f", doc={doc!r}" if doc is not None else ""
type_src = _iceberg_type_to_source(ice_type)
downgrade: _Side = [f"ctx.delete_column(target_table, {name!r})"]
if type_src is None:
return (
_Blocked(
f"column {name!r}: type {ice_type} is not auto-renderable; "
"write this step by hand",
(f"ctx.enforce_column(target_table, {name!r}, ...{doc_src})",),
),
downgrade,
)
return [f"ctx.enforce_column(target_table, {name!r}, {type_src}{doc_src})"], downgrade
def _render_rename_column(op: MigrationOp) -> tuple[_Side, _Side]:
old, new = op.args["old_name"], op.args["new_name"]
return (
[f"ctx.enforce_rename(target_table, {old!r}, {new!r})"],
[f"ctx.enforce_rename(target_table, {new!r}, {old!r})"],
)
def _render_delete_column(op: MigrationOp) -> tuple[_Side, _Side]:
name = op.args["name"]
line = f"ctx.delete_column(target_table, {name!r})"
downgrade = _Blocked(
f"delete_column {name!r} cannot be reversed: re-adding creates a new, "
"empty column (old data only via time travel)"
)
refs = op.args.get("referenced_by")
if refs:
return (
_Blocked(
f"column {name!r} is referenced by {', '.join(refs)}; remove those "
"references first",
(line,),
),
downgrade,
)
return [line], downgrade
def _render_type_change(op: MigrationOp) -> tuple[_Side, _Side]:
name = op.args["name"]
old_type = op.args["old_iceberg_type"]
new_type = op.args["new_iceberg_type"]
new_src = _iceberg_type_to_source(new_type)
downgrade = _Blocked(
f"type change on {name!r} ({old_type} -> {new_type}) cannot be narrowed back; "
"Iceberg only widens types"
)
line = f"ctx.enforce_column(target_table, {name!r}, {new_src or '...'})"
if new_src is None or _iceberg_type_to_source(old_type) is None:
return (
_Blocked(
f"column {name!r}: nested type change {old_type} -> {new_type} is not "
"auto-renderable; write this step by hand",
(line,),
),
downgrade,
)
if not is_allowed_promotion(old_type, new_type):
return (
_Blocked(
f"column {name!r}: {old_type} -> {new_type} is not an Iceberg type "
f"promotion ({_ALLOWED_PROMOTIONS})",
(line,),
),
downgrade,
)
return [line], downgrade
def _render_add_partition_field(op: MigrationOp) -> tuple[_Side, _Side]:
src_col = op.args["source_column"]
transform = op.args["transform_source"]
part_name = op.args["partition_name"]
return (
[
"ctx.update_partition_spec(",
" target_table,",
f" add_fields=[({src_col!r}, {transform}, {part_name!r})],",
")",
],
[
"ctx.update_partition_spec(",
" target_table,",
f" remove_fields=[{part_name!r}],",
")",
],
)
def _render_remove_partition_field(op: MigrationOp) -> tuple[_Side, _Side]:
part_name = op.args["partition_name"]
return (
[
"ctx.update_partition_spec(",
" target_table,",
f" remove_fields=[{part_name!r}],",
")",
],
_Blocked(
f"remove_partition_field {part_name!r} cannot be reversed: the original "
"source column and transform are not recorded"
),
)
def _render_set_sort_order(op: MigrationOp) -> tuple[_Side, _Side]:
fields_repr = ", ".join(repr((col, direction)) for col, direction in op.args["sort_fields"])
return (
[
"ctx.replace_sort_order(",
" target_table,",
f" sort_fields=[{fields_repr}],",
")",
],
_Blocked("set_sort_order cannot be reversed: the previous sort order is not recorded"),
)
def _render_set_property(op: MigrationOp) -> tuple[_Side, _Side]:
key = op.args["key"]
value = op.args["value"]
old_value = op.args.get("old_value")
upgrade = [f"ctx.set_table_properties(target_table, {{{key!r}: {value!r}}})"]
if old_value is None:
downgrade = [f"ctx.unset_table_properties(target_table, [{key!r}])"]
else:
downgrade = [f"ctx.set_table_properties(target_table, {{{key!r}: {old_value!r}}})"]
return upgrade, downgrade
def _render_remove_property(op: MigrationOp) -> tuple[_Side, _Side]:
key = op.args["key"]
old_value = op.args.get("old_value")
upgrade: _Side = [f"ctx.unset_table_properties(target_table, [{key!r}])"]
if old_value is None:
return upgrade, _Blocked(
f"remove_property {key!r} cannot be reversed: the old value was not recorded"
)
return upgrade, [f"ctx.set_table_properties(target_table, {{{key!r}: {old_value!r}}})"]
_RENDERERS: dict[str, Callable[[MigrationOp], tuple[_Side, _Side]]] = {
"add_column": _render_add_column,
"rename_column": _render_rename_column,
"delete_column": _render_delete_column,
"type_change": _render_type_change,
"add_partition_field": _render_add_partition_field,
"remove_partition_field": _render_remove_partition_field,
"set_sort_order": _render_set_sort_order,
"set_property": _render_set_property,
"remove_property": _render_remove_property,
}
def _render_pairs(ops: list[MigrationOp]) -> list[tuple[_Side, _Side]]:
"""Run each op's renderer, in op order."""
pairs: list[tuple[_Side, _Side]] = []
for op in ops:
renderer = _RENDERERS.get(op.kind)
if renderer is None:
raise ValueError(f"Unknown MigrationOp kind: {op.kind!r}")
pairs.append(renderer(op))
return pairs
def _blocking_reasons(ops: list[MigrationOp]) -> tuple[list[str], list[str]]:
"""Reasons the rendered ``upgrade`` / ``downgrade`` will raise.
Returns
-------
(upgrade_reasons, downgrade_reasons) : tuple of list[str]
Empty lists mean that method runs as generated.
"""
pairs = _render_pairs(ops)
up = [side.reason for side, _ in pairs if isinstance(side, _Blocked)]
down = [side.reason for _, side in reversed(pairs) if isinstance(side, _Blocked)]
return up, down
def _comment(text: str) -> str:
"""Render *text* as a single comment line (no line breaks can escape)."""
return "# " + "".join(ch if ch.isprintable() else " " for ch in text)
def _assemble(method: str, fragments: list[_Side], target_repr: str) -> list[str]:
"""Build one method body from rendered fragments.
Without blocks the fragments are emitted in order. With any block
the body raises :class:`NotImplementedError` straight after
``target_table`` is set - no ctx call precedes it - and every
fragment follows as comments, in order.
"""
lines = [f"target_table = {target_repr}", ""]
reasons = [frag.reason for frag in fragments if isinstance(frag, _Blocked)]
if not reasons:
for frag in fragments:
assert not isinstance(frag, _Blocked)
lines.extend(frag)
return lines
message = (
f"Autogenerated {method}() must be completed by hand before it can run; "
"no step has been applied. " + " ".join(f"[{r}]" for r in reasons)
)
lines.append(f"raise NotImplementedError({message!r})")
lines.append("")
lines.append("# Generated steps, in order, kept for manual completion:")
for frag in fragments:
if isinstance(frag, _Blocked):
if frag.lines:
lines.extend(_comment(line) for line in frag.lines)
else:
lines.append(_comment(f"(irreversible) {frag.reason}"))
else:
lines.extend(_comment(line) for line in frag)
return lines
[docs]
def render_migration_body(
ops: list[MigrationOp],
*,
target_table: str,
) -> tuple[str, str]:
"""
Translate a list of :class:`MigrationOp` into ``upgrade`` and
``downgrade`` method body source strings.
Parameters
----------
ops : list[MigrationOp]
Output of :func:`schema_diff_to_ops` and friends.
target_table : str
Table identifier used in the generated calls. Prefer a bare
name (``"orders"``): the migration context qualifies it with
the namespace the migration runs against, so one file serves
every namespace. A dotted name (``"sales.orders"``) pins the
namespace - the migration always modifies that table.
Returns
-------
(upgrade_body, downgrade_body) : tuple of str
Each is a ``\\n``-joined block of indented Python source ready
to drop into a :class:`~iceberg_ops.migration.BaseMigration`
method.
Raises
------
ValueError
If an op has an unknown ``kind``.
Notes
-----
The downgrade order is the reverse of upgrade - matches the
classic schema-migration convention. When any step of a method
cannot run as generated (an incompatible or nested type change, or
dropping a column the table layout references, on upgrade; a column
drop, type widening, partition-field removal,
sort-order replacement or property removal without a recorded old
value on downgrade), that method raises
:class:`NotImplementedError` before any step runs and keeps every
generated call as a comment.
"""
upgrade_lines: list[str]
downgrade_lines: list[str]
if not ops:
upgrade_lines = ["pass # autogenerated: no schema differences detected"]
downgrade_lines = ["pass # autogenerated: nothing to revert"]
return _indent_block(upgrade_lines), _indent_block(downgrade_lines)
target_repr = repr(target_table)
pairs = _render_pairs(ops)
upgrade_lines = _assemble("upgrade", [up for up, _ in pairs], target_repr)
# Downgrade is the reverse order of ops.
downgrade_lines = _assemble("downgrade", [down for _, down in reversed(pairs)], target_repr)
return _indent_block(upgrade_lines), _indent_block(downgrade_lines)
def _indent_block(lines: list[str]) -> str:
"""Prefix each line with the migration-method indent."""
return "\n".join(_INDENT + line if line else "" for line in lines)
# ===========================================================================
# Full migration-file renderer
# ===========================================================================
_AUTOGEN_TEMPLATE = '''\
"""
Migration $version - $doc_description.
.. warning::
This migration was autogenerated by iceberg-ops from a schema diff.
**Review every line before running** - autogenerate is a starting
point, not a finished implementation. In particular:
* Each step commits on its own, so a step that fails leaves the
earlier steps applied.
* Steps known not to run (incompatible or nested type changes,
dropping a column the partition spec, sort order or identifier
fields reference) and steps that cannot be reversed (column drops,
type widening, partition-field removal, sort-order replacement,
property removal with an unknown old value) make the method raise
``NotImplementedError`` before any step runs; the generated calls
are kept as comments for manual completion.
* upgrade() is idempotent and adopts columns that already exist;
downgrade() still removes every column this migration lists, even
in a namespace whose table already had it. Review downgrade()
before rolling back a namespace whose table differed from the one
that was diffed.
* Nullability and doc changes on existing columns are not emitted.
* Renames are emitted only for explicit rename hints
(``--rename-hint old:new``); otherwise a rename appears as a drop
plus an add.
"""
from __future__ import annotations
$imports
from iceberg_ops.migration import BaseMigration, MigrationContext
class Migration(BaseMigration):
version = $version_literal
description = $description_literal
def upgrade(self, ctx: MigrationContext) -> None:
$upgrade_body
def downgrade(self, ctx: MigrationContext) -> None:
$downgrade_body
'''
[docs]
def render_migration_file(
*,
version: str,
description: str,
ops: list[MigrationOp],
target_table: str,
) -> str:
"""
Render a complete migration source file from autogen ops.
Combines :func:`render_migration_body` with the autogen-specific
docstring (which carries the "Review every line before running"
warning) and the imports the ops need.
Parameters
----------
version : str
Four-digit version, e.g. ``"0007"``.
description : str
Free-text description (e.g. the ``-m`` message). Stored verbatim
as the ``description`` class attribute; a whitespace-collapsed
copy goes into the module docstring.
ops : list[MigrationOp]
Ops to render.
target_table : str
Table identifier for the generated calls; see
:func:`render_migration_body` (prefer a bare name).
Returns
-------
str
Full Python source ready to write to disk.
Raises
------
ValueError
If ``version`` is not four digits or an op has an unknown kind.
"""
if not re.fullmatch(r"\d{4}", version):
raise ValueError(f"version must be four digits, e.g. '0007' (got {version!r}).")
upgrade_body, downgrade_body = render_migration_body(ops, target_table=target_table)
imports: list[str] = []
if any(op.kind in {"add_column", "type_change"} for op in ops):
imports.append("from pyiceberg import types as ice_types")
if any(op.kind == "add_partition_field" for op in ops):
imports.append("from pyiceberg import transforms")
imports_block = "\n".join(imports) + "\n" if imports else ""
return Template(_AUTOGEN_TEMPLATE).substitute(
version=version,
doc_description=_docstring_text(description),
version_literal=repr(version),
description_literal=repr(description),
imports=imports_block,
upgrade_body=upgrade_body,
downgrade_body=downgrade_body,
)
[docs]
def write_autogenerated_revision(
*,
version: str,
description: str,
ops: list[MigrationOp],
target_table: str,
out_path: str | Path,
overwrite: bool = False,
) -> str:
"""
Render and write an autogenerated migration to *out_path*.
The file is written as UTF-8 with LF line endings on every platform.
Parameters
----------
version : str
Four-digit version, e.g. ``"0007"``.
description : str
Free-text description of the migration.
ops : list[MigrationOp]
Ops to render.
target_table : str
Table identifier for the generated calls; prefer a bare name so
the migration follows the namespace it runs against.
out_path : str or pathlib.Path
Destination file.
overwrite : bool, optional
Replace an existing file at *out_path*. Default ``False``.
Returns
-------
str
The absolute path of the written file (string for symmetry
with :func:`iceberg_ops.scaffolding.create_revision`).
Raises
------
FileExistsError
If *out_path* exists and ``overwrite`` is false.
ValueError
If ``version`` is not four digits or an op has an unknown kind.
"""
body = render_migration_file(
version=version,
description=description,
ops=ops,
target_table=target_table,
)
path = Path(out_path)
if overwrite:
path.write_text(body, encoding="utf-8", newline="\n")
return str(path.resolve())
# Exclusive create: a file created concurrently (another revision
# run picking the same version) is never overwritten.
try:
with open(path, "x", encoding="utf-8", newline="\n") as fh:
fh.write(body)
except FileExistsError:
raise FileExistsError(
f"{path} already exists; pass overwrite=True to replace it (hand edits are lost)."
) from None
return str(path.resolve())
__all__ = [
"MigrationOp",
"MigrationOpKind",
"properties_diff_to_ops",
"render_migration_body",
"render_migration_file",
"schema_diff_to_ops",
"sort_order_diff_to_ops",
"spec_diff_to_ops",
"write_autogenerated_revision",
]