|
| 1 | +import logging |
| 2 | +import os |
| 3 | +import sys |
| 4 | + |
| 5 | +import duckdb |
| 6 | +from pyspark.sql import SparkSession |
| 7 | + |
| 8 | +from databricks.labs.lakebridge.assessments.profiler_validator import EmptyTableValidationCheck, build_validation_report |
| 9 | + |
| 10 | +logger = logging.getLogger(__name__) |
| 11 | + |
| 12 | + |
| 13 | +def main(*argv) -> None: |
| 14 | + logger.debug(f"Arguments received: {argv}") |
| 15 | + assert len(sys.argv) == 4, f"Invalid number of arguments: {len(sys.argv)}" |
| 16 | + catalog_name = sys.argv[0] |
| 17 | + schema_name = sys.argv[1] |
| 18 | + extract_location = sys.argv[2] |
| 19 | + source_tech = sys.argv[3] |
| 20 | + logger.info(f"Validating {source_tech} profiler extract located at '{extract_location}'.") |
| 21 | + valid_extract = _validate_profiler_extract(extract_location) |
| 22 | + if valid_extract: |
| 23 | + _ingest_profiler_tables(catalog_name, schema_name, extract_location) |
| 24 | + else: |
| 25 | + raise ValueError("Corrupt or invalid profiler extract.") |
| 26 | + |
| 27 | + |
| 28 | +def _validate_profiler_extract(extract_location: str) -> bool: |
| 29 | + logger.info("Validating the profiler extract file.") |
| 30 | + validation_checks = [] |
| 31 | + try: |
| 32 | + with duckdb.connect(database=extract_location) as duck_conn: |
| 33 | + tables = duck_conn.execute("SHOW ALL TABLES").fetchall() |
| 34 | + for table in tables: |
| 35 | + fq_table_name = f"{table[0]}.{table[1]}.{table[2]}" |
| 36 | + empty_check = EmptyTableValidationCheck(fq_table_name) |
| 37 | + validation_checks.append(empty_check) |
| 38 | + report = build_validation_report(validation_checks, duck_conn) |
| 39 | + except duckdb.IOException as e: |
| 40 | + logger.exception(f"Could not access the profiler extract: '{extract_location}'.") |
| 41 | + raise e |
| 42 | + except Exception as e: |
| 43 | + logger.exception(f"Unable to validate the profiler extract: '{extract_location}'.") |
| 44 | + raise e |
| 45 | + |
| 46 | + if len(report) > 0: |
| 47 | + report_errors = list(filter(lambda x: x.outcome == "FAIL" and x.severity == "ERROR", report)) |
| 48 | + num_errors = len(report_errors) |
| 49 | + logger.info(f"There are {num_errors} validation errors in the profiler extract.") |
| 50 | + for error in report_errors: |
| 51 | + logging.info(error) |
| 52 | + else: |
| 53 | + raise ValueError("Profiler extract validation report is empty.") |
| 54 | + return num_errors == 0 |
| 55 | + |
| 56 | + |
| 57 | +def _ingest_profiler_tables(catalog_name: str, schema_name: str, extract_location: str) -> None: |
| 58 | + try: |
| 59 | + with duckdb.connect(database=extract_location) as duck_conn: |
| 60 | + tables_to_ingest = duck_conn.execute("SHOW ALL TABLES").fetchall() |
| 61 | + except duckdb.IOException as e: |
| 62 | + logger.error(f"Could not access the profiler extract: '{extract_location}': {e}") |
| 63 | + raise duckdb.IOException(f"Could not access the profiler extract: '{extract_location}'.") from e |
| 64 | + except Exception as e: |
| 65 | + logger.error(f"Unable to read tables from profiler extract: '{extract_location}': {e}") |
| 66 | + raise e |
| 67 | + |
| 68 | + if len(tables_to_ingest) == 0: |
| 69 | + raise ValueError("Profiler extract contains no tables.") |
| 70 | + |
| 71 | + successful_tables = [] |
| 72 | + unsuccessful_tables = [] |
| 73 | + for source_table in tables_to_ingest: |
| 74 | + try: |
| 75 | + fq_source_table_name = f"{source_table[0]}.{source_table[1]}.{source_table[2]}" |
| 76 | + fq_delta_table_name = f"{catalog_name}.{schema_name}.{source_table[2]}" |
| 77 | + logger.info(f"Ingesting profiler table: '{fq_source_table_name}'") |
| 78 | + _ingest_table(extract_location, fq_source_table_name, fq_delta_table_name) |
| 79 | + successful_tables.append(fq_source_table_name) |
| 80 | + except (ValueError, IndexError, TypeError) as e: |
| 81 | + logger.error(f"Failed to construct source and destination table names: {e}") |
| 82 | + unsuccessful_tables.append(source_table) |
| 83 | + except duckdb.Error as e: |
| 84 | + logger.error(f"Failed to ingest table from profiler database: {e}") |
| 85 | + unsuccessful_tables.append(source_table) |
| 86 | + logger.info(f"Ingested {len(successful_tables)} tables from profiler extract.") |
| 87 | + logger.info(",".join(successful_tables)) |
| 88 | + logger.info(f"Failed to ingest {len(unsuccessful_tables)} tables from profiler extract.") |
| 89 | + logger.info(",".join(unsuccessful_tables)) |
| 90 | + |
| 91 | + |
| 92 | +def _ingest_table(extract_location: str, source_table_name: str, target_table_name: str) -> None: |
| 93 | + """ |
| 94 | + Ingest a table from a DuckDB profiler extract into a managed Delta table in Unity Catalog. |
| 95 | + """ |
| 96 | + try: |
| 97 | + with duckdb.connect(database=extract_location, read_only=True) as duck_conn: |
| 98 | + query = f"SELECT * FROM {source_table_name}" |
| 99 | + pdf = duck_conn.execute(query).df() |
| 100 | + # Save table as a managed Delta table in Unity Catalog |
| 101 | + logger.info(f"Saving profiler table '{target_table_name}' to Unity Catalog.") |
| 102 | + spark = SparkSession.builder.getOrCreate() |
| 103 | + df = spark.createDataFrame(pdf) |
| 104 | + df.write.format("delta").mode("overwrite").saveAsTable(target_table_name) |
| 105 | + except duckdb.CatalogException as e: |
| 106 | + logger.error(f"Could not find source table '{source_table_name}' in profiler extract: {e}") |
| 107 | + raise duckdb.CatalogException(f"Could not find source table '{source_table_name}' in profiler extract.") from e |
| 108 | + except duckdb.IOException as e: |
| 109 | + logger.error(f"Could not access the profiler extract: '{extract_location}': {e}") |
| 110 | + raise duckdb.IOException(f"Could not access the profiler extract: '{extract_location}'.") from e |
| 111 | + except Exception as e: |
| 112 | + logger.error(f"Unable to ingest table '{source_table_name}' from profiler extract: {e}") |
| 113 | + raise e |
| 114 | + |
| 115 | + |
| 116 | +if __name__ == "__main__": |
| 117 | + # Ensure that the ingestion job is being run on a Databricks cluster |
| 118 | + if "DATABRICKS_RUNTIME_VERSION" not in os.environ: |
| 119 | + raise SystemExit("The Lakebridge profiler ingestion job is only intended to run in a Databricks Runtime.") |
| 120 | + main(*sys.argv) |
0 commit comments