spark-sql-migrations 0.0.6__tar.gz → 0.0.8__tar.gz

{spark_sql_migrations-0.0.6 → spark_sql_migrations-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spark_sql_migrations
-Version: 0.0.6
+Version: 0.0.8
 Summary: Chained, idempotent SQL migrations for Spark and Databricks.
 Home-page: https://github.com/mcjug2015/spark_sql_migrations
 Author: Victor Semenov

{spark_sql_migrations-0.0.6 → spark_sql_migrations-0.0.8}/setup.py

@@ -299,5 +299,5 @@ Apache License 2.0 — see [LICENSE](LICENSE).
     },
     'python_requires': '>=3.10',
     'url': 'https://github.com/mcjug2015/spark_sql_migrations',
-    'version': '0.0.6',
+    'version': '0.0.8',
 })

{spark_sql_migrations-0.0.6 → spark_sql_migrations-0.0.8}/spark_sql_migrations/spark_sql/spark_sql.py

@@ -27,8 +27,11 @@ import os
 import re
 import uuid
 from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
 
 from jinja2 import Environment, FileSystemLoader, PackageLoader, select_autoescape
+from pyspark.sql.session import SparkSession
+from pyspark.sql.types import Row
 
 from spark_sql_migrations.spark_utils import get_spark, is_dbr
 
@@ -51,17 +54,17 @@ _COMMENT_OR_BLANK_RE = re.compile(r"^\s*(--.*)?$")
 @dataclass(frozen=True)
 class Migration:
     revision_id: str
-    prev_revision_id: str | None
+    prev_revision_id: Optional[str]
     template_name: str
 
 
-def get_default_template_path():
+def get_default_template_path() -> str:
     """the template this package ships; callers may pass their own instead."""
     return os.path.join(os.path.dirname(__file__), "..", TEMPLATES_PACKAGE_PATH, DEFAULT_TEMPLATE_NAME)
 
 
-def _parse_migration(migrations_dir, template_name):
-    headers = {}
+def _parse_migration(migrations_dir: str, template_name: str) -> Migration:
+    headers: Dict[str, Any | str] = {}
     migration_file_path = os.path.join(migrations_dir, template_name)
     if os.path.getsize(migration_file_path) == 0:
         raise ValueError(f"{template_name} is empty;")
@@ -84,8 +87,8 @@ def _parse_migration(migrations_dir, template_name):
     )
 
 
-def get_ordered_migration_objs(migration_objs):
-    by_revision = {}
+def get_ordered_migration_objs(migration_objs: List[Migration]) -> List[Migration]:
+    by_revision: Dict[str, Migration] = {}
     for migration in migration_objs:
         clash = by_revision.get(migration.revision_id)
         if clash:
@@ -102,7 +105,7 @@ def get_ordered_migration_objs(migration_objs):
             f"found {[m.template_name for m in roots]};"
         )
 
-    next_by_revision = {}
+    next_by_revision: Dict[str, Migration] = {}
     for migration in migration_objs:
         if not migration.prev_revision_id:
             continue
@@ -119,7 +122,7 @@ def get_ordered_migration_objs(migration_objs):
         next_by_revision[migration.prev_revision_id] = migration
 
     ordered = []
-    current = roots[0]
+    current: Optional[Migration] = roots[0]
     while current:
         ordered.append(current)
         current = next_by_revision.get(current.revision_id)
@@ -134,8 +137,14 @@ def get_ordered_migration_objs(migration_objs):
     return ordered
 
 
-def get_migrations_list(migrations_dir):
-    """every migration in one directory, walked from its root to its head."""
+def get_migrations_list(migrations_dir: str, max_revision_id: Optional[str] = None) -> List[Migration]:
+    """every migration in one directory, walked from its root to its head.
+
+    ``revision_id`` cuts the walk short: the chain is returned up to and including
+    the migration carrying that revision. A revision no migration in the directory
+    carries is not an error -- the whole chain comes back, same as for no revision
+    at all.
+    """
     if not os.path.isdir(migrations_dir):
         raise ValueError(f"no migrations dir at {migrations_dir};")
 
@@ -147,10 +156,18 @@ def get_migrations_list(migrations_dir):
     if not migrations:
         return []
 
-    return get_ordered_migration_objs(migrations)
+    ordered = get_ordered_migration_objs(migrations)
+    if max_revision_id is None:
+        return ordered
+
+    for index, migration in enumerate(ordered):
+        if migration.revision_id == max_revision_id:
+            return ordered[: index + 1]
+
+    return ordered
 
 
-def get_table_version(spark, cat: str, schema: str):
+def get_table_version(spark: SparkSession, cat: str, schema: str) -> Optional[Row]:
     """the revision VERSION_TABLE records, or None if it holds no row yet."""
     version_table = f"{cat}.{schema}.{VERSION_TABLE}"
     if not spark.catalog.tableExists(version_table):
@@ -158,8 +175,10 @@ def get_table_version(spark, cat: str, schema: str):
     return spark.read.table(version_table).select("version_num").head()
 
 
-def get_unapplied_migrations_list(spark, migrations_dir, cat: str, schema: str):
-    ordered = get_migrations_list(migrations_dir)
+def get_unapplied_migrations_list(
+    spark: SparkSession, migrations_dir: str, cat: str, schema: str, max_revision_id: Optional[str] = None
+) -> List[Migration]:
+    ordered = get_migrations_list(migrations_dir, max_revision_id)
     current = get_table_version(spark, cat, schema)
     if not current:
         return ordered
@@ -180,7 +199,7 @@ def get_unapplied_migrations_list(spark, migrations_dir, cat: str, schema: str):
     raise ValueError(f"{VERSION_TABLE} is at {current_revision_id}, which matches no migration in {migrations_dir};")
 
 
-def gate_on_is_dbr(migration_content):
+def gate_on_is_dbr(migration_content: str) -> str:
     """wrap a migration's body in the jinja conditional, below its headers.
 
     the revision headers stay outside it: they are read off the file as comments,
@@ -199,7 +218,7 @@ def gate_on_is_dbr(migration_content):
 
 def create_new_migration(
     message: str, template_path: str, output_path: str, truncate_slug_length: int = 40, add_is_dbr: bool = False
-):
+) -> None:
     rev_id = uuid.uuid4().hex[-12:]
 
     date_prefix = datetime.datetime.today().strftime("%y%m%d")
@@ -228,32 +247,38 @@ def create_new_migration(
         migration_file.write(default_migration_content)
 
 
-def apply_template(output_dir, template, cat: str, schema: str):
+def apply_template(output_dir: str, template, cat: str, schema: str) -> str:
     result_sql = template.render(cat=cat, schema=schema, is_dbr=is_dbr())
     with open(os.path.join(output_dir, template.name.replace(".sql", "_primed.sql")), "w") as file_handle:
         file_handle.write(result_sql)
     return result_sql
 
 
-def get_ascending_letters_within_minute():
-    micros_since_minute = datetime.datetime.now() - datetime.datetime.now().replace(second=0, microsecond=0)
-    result = str(micros_since_minute.microseconds).translate(str.maketrans("0123456789", "ABCDEFGHIJ"))
-    return result
+def get_ascending_letters_within_minute() -> str:
+    """a letter stamp that sorts in the order it was taken, within one minute.
+
+    the full microseconds since the minute began, not timedelta.microseconds -- that
+    holds only the sub-second part and so starts over every second. zero-padded to a
+    fixed width, because these are compared as text once they are folder names.
+    """
+    now = datetime.datetime.now()
+    micros_since_minute = (now - now.replace(second=0, microsecond=0)) // datetime.timedelta(microseconds=1)
+    return f"{micros_since_minute:08d}".translate(str.maketrans("0123456789", "ABCDEFGHIJ"))
 
 
-def get_output_folder(output_parent_path):
+def get_output_folder(output_parent_path: str) -> str:
     folder_name = f"{datetime.datetime.today().strftime('%Y%m%d_%H%M')}_{get_ascending_letters_within_minute()}"
     return os.path.join(output_parent_path, folder_name)
 
 
-def has_no_statements(rendered_sql):
+def has_no_statements(rendered_sql: str) -> bool:
     for line in rendered_sql.splitlines():
         if not _COMMENT_OR_BLANK_RE.match(line):
             return False
     return True
 
 
-def render_and_apply(spark, output_folder, template, cat: str, schema: str):
+def render_and_apply(spark: SparkSession, output_folder: str, template, cat: str, schema: str) -> None:
     """render one migration and run it, unless this engine was gated out of it."""
     result_sql = apply_template(output_folder, template, cat=cat, schema=schema)
     if has_no_statements(result_sql):
@@ -262,7 +287,7 @@ def render_and_apply(spark, output_folder, template, cat: str, schema: str):
     spark.sql(result_sql)
 
 
-def migrate_initial(spark, output_folder, cat: str, schema: str):
+def migrate_initial(spark: SparkSession, output_folder: str, cat: str, schema: str) -> None:
     """apply the package's own bootstrap chain, in filename order rather than by revision."""
     env = Environment(
         loader=PackageLoader(package_name=PACKAGE_NAME, package_path=INITIAL_MIGRATIONS_PACKAGE_PATH),
@@ -276,14 +301,21 @@ def migrate_initial(spark, output_folder, cat: str, schema: str):
     logger.info(f"rendered {len(all_templates)} initial migrations;")
 
 
-def record_revision(spark, cat: str, schema: str, revision_id: str):
+def record_revision(spark: SparkSession, cat: str, schema: str, revision_id: str) -> None:
     version_table = f"{cat}.{schema}.{VERSION_TABLE}"
     revision_literal = revision_id.replace("'", "''")
     spark.sql(f"delete from {version_table}")
     spark.sql(f"insert into {version_table} (version_num) values ('{revision_literal}')")
 
 
-def migrate_w_rev(spark, output_folder, migrations_dir, cat: str, schema: str):
+def upgrade(
+    spark: SparkSession,
+    output_folder: str,
+    migrations_dir: str,
+    cat: str,
+    schema: str,
+    max_revision_id: Optional[str] = None,
+) -> None:
     """
     TODO XXX permit applying up to a revision below the head revision
     TODO XXX add backwards migrations, permit migrating backwards
@@ -291,7 +323,9 @@ def migrate_w_rev(spark, output_folder, migrations_dir, cat: str, schema: str):
     if not os.path.isdir(migrations_dir):
         logger.info(f"no chain at {migrations_dir}; skipping it;")
         return
-    unapplied = get_unapplied_migrations_list(spark, migrations_dir, cat=cat, schema=schema)
+    unapplied = get_unapplied_migrations_list(
+        spark, migrations_dir, cat=cat, schema=schema, max_revision_id=max_revision_id
+    )
 
     logger.info(
         f"{len(unapplied)} unapplied migrations; "
@@ -305,7 +339,14 @@ def migrate_w_rev(spark, output_folder, migrations_dir, cat: str, schema: str):
         logger.info(f"{migration.template_name} done; now at {migration.revision_id};")
 
 
-def run_migrations(spark, cat, schema, output_folder, migrations_dir):
+def run_upgrade(
+    spark: SparkSession,
+    cat: str,
+    schema: str,
+    output_folder: str,
+    migrations_dir: str,
+    max_revision_id: Optional[str] = None,
+) -> None:
     initial_output_folder = os.path.join(output_folder, INITIAL_MIGRATIONS_PACKAGE_PATH)
     client_output_folder = os.path.join(output_folder, CLIENT_OUTPUT_DIRNAME)
     os.makedirs(initial_output_folder, exist_ok=True)
@@ -315,42 +356,61 @@ def run_migrations(spark, cat, schema, output_folder, migrations_dir):
         logger.info(f"{version_table} already exists; skipping the initial migrations;")
     else:
         migrate_initial(spark, initial_output_folder, cat, schema)
-    migrate_w_rev(spark, client_output_folder, migrations_dir, cat, schema)
+    upgrade(spark, client_output_folder, migrations_dir, cat, schema, max_revision_id)
 
 
-def main(cat, schema, migrations_dir, output_parent_path=None):
+def main(
+    cat: str,
+    schema: str,
+    migrations_dir: str,
+    output_parent_path: Optional[str] = None,
+    max_revision_id: Optional[str] = None,
+) -> None:
     output_folder = get_output_folder(output_parent_path or os.path.join(os.getcwd(), "migrations_out"))
     os.makedirs(output_folder)
-    run_migrations(get_spark(), cat, schema, output_folder, migrations_dir)
+    run_upgrade(get_spark(), cat, schema, output_folder, migrations_dir, max_revision_id)
 
 
-def _cli_main(cat, schema, migrations_dir):  # pragma: no cover
-    main(cat, schema, migrations_dir)
+def _cli_main(
+    cat: str, schema: str, migrations_dir: str, max_revision_id: Optional[str] = None
+) -> None:  # pragma: no cover
+    main(cat, schema, migrations_dir, max_revision_id=max_revision_id)
 
 
-def _cli_create_new_migration(message, output_path, template_path, add_is_dbr):  # pragma: no cover
+def _cli_create_new_migration(
+    message: str, output_path: str, template_path: str, add_is_dbr: bool
+) -> None:  # pragma: no cover
     create_new_migration(message, template_path or get_default_template_path(), output_path, add_is_dbr=add_is_dbr)
 
 
-def build_parser():  # pragma: no cover
+def build_parser() -> argparse.ArgumentParser:  # pragma: no cover
     parser = argparse.ArgumentParser(description="spark_sql_migrations migration utilities")
     subparsers = parser.add_subparsers(dest="command", required=True)
 
-    p_run = subparsers.add_parser("run", help="apply all pending migrations")
-    p_run.add_argument("--cat", default="spark_catalog")
-    p_run.add_argument("--schema", default="default")
-    p_run.add_argument("--migrations-dir", required=True, help="the directory holding the migration chain")
+    p_run = subparsers.add_parser("run", help="upgrade to head/apply pending migrations")
+    p_run.add_argument("--cat", type=str, default="spark_catalog")
+    p_run.add_argument("--schema", type=str, default="default")
+    p_run.add_argument("--migrations-dir", type=str, required=True, help="the directory holding the migration chain")
+    p_run.add_argument(
+        "--max-revision-id",
+        type=str,
+        default=None,
+        help="apply the migration carrying this revision_id and stop; defaults to applying the chain to its head",
+    )
     p_run.set_defaults(func=_cli_main)
 
     p_create = subparsers.add_parser("create_new_migration", help="create a new migration from the template")
-    p_create.add_argument("--message", required=True)
+    p_create.add_argument("--message", type=str, required=True)
     p_create.add_argument(
         "--output-path",
+        type=str,
         required=True,
         help="the chain directory to write the new migration into. if sql files are in there, the last sql's rev_id"
         "will be set as the new migrations prev_rev_id.",
     )
-    p_create.add_argument("--template-path", default=None, help="defaults to the template shipped in this package")
+    p_create.add_argument(
+        "--template-path", type=str, default=None, help="defaults to the template shipped in this package"
+    )
     # the %% are argparse's: it interpolates help strings
     p_create.add_argument(
         "--add-is-dbr",
@@ -362,7 +422,7 @@ def build_parser():  # pragma: no cover
     return parser
 
 
-def cli(argv=None):  # pragma: no cover
+def cli(argv: Optional[list[str]] = None) -> None:  # pragma: no cover
     """the `spark-sql-migrations` console command.
 
     an application entry point, so unlike the rest of this package it configures

{spark_sql_migrations-0.0.6 → spark_sql_migrations-0.0.8}/spark_sql_migrations/spark_utils.py

@@ -11,7 +11,7 @@ METASTORE_DIR_ENV_VAR = "SPARK_METASTORE_DIR"
 DEFAULT_WAREHOUSE_DIRNAME = "spark-warehouse"
 
 
-def is_dbr():
+def is_dbr() -> bool:
     try:
         from pyspark.dbutils import DBUtils  # type: ignore # noqa: F401
 
@@ -21,7 +21,7 @@ def is_dbr():
         return False
 
 
-def get_warehouse_dir():
+def get_warehouse_dir() -> str:
     """where local (non-DBR) spark keeps its warehouse.
 
     resolved against the caller's cwd, never against __file__: installed from a
@@ -34,7 +34,7 @@ def get_warehouse_dir():
     )
 
 
-def get_spark(use_dbc=False):
+def get_spark(use_dbc: bool = False) -> SparkSession:
     if use_dbc or is_dbr():
         from databricks.connect.session import DatabricksSession  # type: ignore # pants: no-infer-dep
 

{spark_sql_migrations-0.0.6 → spark_sql_migrations-0.0.8}/spark_sql_migrations.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spark_sql_migrations
-Version: 0.0.6
+Version: 0.0.8
 Summary: Chained, idempotent SQL migrations for Spark and Databricks.
 Home-page: https://github.com/mcjug2015/spark_sql_migrations
 Author: Victor Semenov