assets.py

import time
from datetime import UTC, datetime
from io import BytesIO
from pathlib import Path

import pandas as pd
from dagster_pyspark import PySparkResource
from delta import DeltaTable
from models.file_upload import FileUpload
from pyspark import sql
from pyspark.sql import (
    SparkSession,
    functions as f,
)
from pyspark.sql.types import StringType, StructField, StructType
from sqlalchemy import select
from src.constants import DataTier
from src.data_quality_checks.utils import (
    aggregate_report_json,
    aggregate_report_spark_df,
    aggregate_report_statistics,
    dq_geolocation_extract_relevant_columns,
    dq_split_failed_rows,
    dq_split_passed_rows,
    row_level_checks,
)
from src.internal.common_assets.staging import StagingMode, StagingStep
from src.resources import ResourceKey
from src.schemas.file_upload import FileUploadConfig
from src.spark.config_expectations import config as config_expectations
from src.spark.transform_functions import (
    add_missing_columns,
    create_bronze_layer_columns_updated,
)
from src.utils.adls import (
    ADLSFileClient,
)
from src.utils.datahub.emit_dataset_metadata import (
    datahub_emit_metadata_with_exception_catcher,
)
from src.utils.db.primary import get_db_context
from src.utils.delta import check_table_exists, create_delta_table, create_schema
from src.utils.metadata import get_output_metadata, get_table_preview
from src.utils.op_config import FileConfig
from src.utils.pandas import pandas_loader
from src.utils.schema import (
    construct_full_table_name,
    construct_schema_name_for_tier,
    get_schema_columns,
    get_schema_columns_datahub,
    get_schema_table,
)
from src.utils.send_email_dq_report import send_email_dq_report_with_config
from src.utils.sentry import capture_op_exceptions

from dagster import (
    AssetOut,
    MetadataValue,
    OpExecutionContext,
    Output,
    asset,
    multi_asset,
)


@asset(io_manager_key=ResourceKey.ADLS_PASSTHROUGH_IO_MANAGER.value)
@capture_op_exceptions
def geolocation_raw(
    context: OpExecutionContext,
    adls_file_client: ADLSFileClient,
    config: FileConfig,
    spark: PySparkResource,
) -> Output[bytes]:
    raw = adls_file_client.download_raw(config.filepath)
    datahub_emit_metadata_with_exception_catcher(
        context=context,
        config=config,
        spark=spark,
    )
    return Output(raw, metadata=get_output_metadata(config))


@asset
def geolocation_metadata(
    context: OpExecutionContext,
    geolocation_raw: bytes,
    config: FileConfig,
    spark: PySparkResource,
):
    s: SparkSession = spark.spark_session

    context.log.info("Get upload details")
    file_size_bytes = config.file_size_bytes
    metadata = config.metadata
    data_source = metadata.get("data_source")
    if not (data_source is None or data_source == "giga_sync"):
        context.log.info("Data is not from Giga Sync, skipping metadata table update")
        return Output(None)

    file_path = config.filepath
    country_code = config.country_code
    schema_name = config.metastore_schema
    file_name = Path(file_path).name
    giga_sync_id = file_name.split("_")[0]
    giga_sync_uploaded_at = datetime.strptime(
        file_name.split(".")[0].split("_")[-1], "%Y%m%d-%H%M%S"
    )

    upload_details = {
        "giga_sync_id": giga_sync_id,
        "country_code": country_code,
        "giga_sync_uploaded_at": giga_sync_uploaded_at,
        "schema_name": schema_name,
        "raw_file_path": file_path,
        "file_size_bytes": file_size_bytes,
    }

    context.log.info("Create upload details dataframe")
    df = pd.DataFrame([upload_details])

    context.log.info("Create giga sync metadata dataframe")
    metadata_df = pd.DataFrame([metadata])

    context.log.info("Combine dataframes")
    metadata_df = pd.concat([df, metadata_df], axis="columns")
    metadata_df["created_at"] = pd.Timestamp.now()

    context.log.info("Create spark dataframe")
    metadata_df = s.createDataFrame(metadata_df)

    table_name = "school_geolocation_metadata"
    table_schema_name = "pipeline_tables"

    context.log.info("Get schema columns for metadata table")
    try:
        table_columns = get_schema_columns(s, "school_geolocation_metadata")
    except Exception:
        context.log.warning(
            "Schema table schemas.school_geolocation_metadata not found; "
            "using DataFrame schema for metadata table creation"
        )
        table_columns = list(metadata_df.schema.fields)

    context.log.info("Create the schema and table if they do not exist")
    metadata_df = add_missing_columns(metadata_df, table_columns)
    metadata_df = metadata_df.select(*StructType(table_columns).fieldNames())

    create_schema(s, table_schema_name)
    create_delta_table(
        s,
        table_schema_name,
        table_name,
        table_columns,
        context,
        if_not_exists=True,
    )

    context.log.info("Upsert the metadata from giga sync into the table")
    s.catalog.refreshTable(construct_full_table_name(table_schema_name, table_name))
    current_metadata_table = DeltaTable.forName(
        s, construct_full_table_name(table_schema_name, table_name)
    )

    (
        current_metadata_table.alias("metadata_current")
        .merge(
            metadata_df.alias("metadata_updates"),
            "metadata_current.giga_sync_id = metadata_updates.giga_sync_id",
        )
        .whenMatchedUpdateAll()
        .whenNotMatchedInsertAll()
        .execute()
    )
    context.log.info("Upsert operation completed")

    return Output(None)


@asset(io_manager_key=ResourceKey.ADLS_SPARK_IO_MANAGER.value)
@capture_op_exceptions
def geolocation_bronze(
    context: OpExecutionContext,
    geolocation_raw: bytes,
    config: FileConfig,
    spark: PySparkResource,
) -> Output[sql.DataFrame]:
    s: SparkSession = spark.spark_session
    country_code = config.country_code
    mode = config.metadata["mode"]

    with get_db_context() as db:
        file_upload = db.scalar(
            select(FileUpload).where(FileUpload.id == config.filename_components.id),
        )
        if file_upload is None:
            raise FileNotFoundError(
                f"Database entry for FileUpload with id `{config.filename_components.id}` was not found",
            )

        file_upload = FileUploadConfig.from_orm(file_upload)

    column_to_schema_mapping = file_upload.column_to_schema_mapping

    string_col_mapping = {
        column_name: str
        for column_name, schema_name in column_to_schema_mapping.items()
        if schema_name in ("school_id_govt", "latitude", "longitude")
    }

    t0 = time.time()
    with BytesIO(geolocation_raw) as buffer:
        buffer.seek(0)
        pdf = pandas_loader(
            buffer,
            config.filepath,
            dtype_mapping=string_col_mapping,
            context=context,
        ).map(str)
    context.log.info(
        f"pandas_loader completed in {time.time() - t0:.2f}s — {len(pdf)} rows"
    )

    pdf.rename(lambda name: name.strip(), axis="columns", inplace=True)
    column_mapping_filtered = {
        k.strip(): v
        for k, v in column_to_schema_mapping.items()
        if (k is not None) and (v is not None)
    }
    pdf = pdf[column_to_schema_mapping.keys()]
    pdf.rename(column_mapping_filtered, axis="columns", inplace=True)

    t1 = time.time()
    df = s.createDataFrame(pdf)
    uploaded_columns = df.columns
    context.log.info(f"createDataFrame completed in {time.time() - t1:.2f}s")

    df = df.withColumn("school_id_govt", f.col("school_id_govt").cast(StringType()))

    df = create_bronze_layer_columns_updated(
        df, mode, uploaded_columns, country_code, s
    )

    t2 = time.time()
    datahub_emit_metadata_with_exception_catcher(
        context=context,
        config=config,
        spark=spark,
        schema_reference=df,
    )
    context.log.info(f"datahub_emit completed in {time.time() - t2:.2f}s")

    for column in config_expectations.TITLE_CASE_COLUMNS:
        if column in df.columns:
            df = df.withColumn(column, f.initcap(f.col(column)))

    df.cache()
    row_count = df.count()

    return Output(
        df,
        metadata={
            **get_output_metadata(config),
            "row_count": row_count,
            "column_mapping": column_mapping_filtered,
            "preview": get_table_preview(df),
        },
    )


@asset(io_manager_key=ResourceKey.ADLS_SPARK_IO_MANAGER.value)
@capture_op_exceptions
def geolocation_data_quality_results(
    context: OpExecutionContext,
    config: FileConfig,
    geolocation_bronze: sql.DataFrame,
    spark: PySparkResource,
) -> Output[sql.DataFrame]:
    s: SparkSession = spark.spark_session
    country_code = config.country_code
    schema_name = config.metastore_schema
    id = config.filename_components.id
    dataset_type = "geolocation"

    current_timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")

    columns = get_schema_columns(s, schema_name)
    schema = StructType(columns)

    if check_table_exists(s, schema_name, country_code, DataTier.SILVER):
        silver_tier_schema_name = construct_schema_name_for_tier(
            "school_geolocation", DataTier.SILVER
        )
        silver_table_name = construct_full_table_name(
            silver_tier_schema_name, country_code
        )
        s.catalog.refreshTable(silver_table_name)
        silver = DeltaTable.forName(s, silver_table_name).alias("silver").toDF()
    else:
        silver = s.createDataFrame(s.sparkContext.emptyRDD(), schema=schema)

    casted_silver = silver.withColumn(
        "school_id_govt", f.col("school_id_govt").cast(StringType())
    )
    casted_bronze = geolocation_bronze.withColumn(
        "school_id_govt", f.col("school_id_govt").cast(StringType())
    )

    renamed_bronze = casted_bronze.withColumnRenamed("signature", "dq_signature")

    dq_results = row_level_checks(
        df=renamed_bronze,
        silver=casted_silver,
        dataset_type=dataset_type,
        _country_code_iso3=country_code,
        mode=config.metadata["mode"],
        context=context,
    )

    dq_results = dq_results.withColumnRenamed("dq_signature", "signature")

    # Collapse all individual dq_ check columns into a single map<string, int> column.
    # This reduces ~120+ columns to one, cutting the DataFrame width by ~60 %.
    # dq_has_critical_error and failure_reason remain as top-level columns because
    # they are used for row-level filtering throughout the pipeline.
    # In Trino the map is queryable as: dq_results['is_null_optional-latitude']
    dq_flag_cols = [
        c
        for c in dq_results.columns
        if c.startswith("dq_") and c != "dq_has_critical_error"
    ]
    map_args = []
    for col_name in dq_flag_cols:
        map_args.extend([f.lit(col_name[len("dq_") :]), f.col(col_name).cast("int")])
    dq_results = dq_results.withColumn("dq_results", f.create_map(*map_args)).drop(
        *dq_flag_cols
    )

    dq_results_schema_name = f"{schema_name}_dq_results"
    # Replace hyphens with underscores so the identifier is valid in Spark SQL
    safe_id = id.replace("-", "_")
    table_name = f"{safe_id}_{country_code}_{current_timestamp}"

    schema_columns = [
        StructField(field.name, field.dataType, nullable=True)
        for field in dq_results.schema.fields
    ]

    dq_results_table_name = construct_full_table_name(
        dq_results_schema_name,
        table_name,
    )

    create_schema(s, dq_results_schema_name)
    create_delta_table(
        s,
        dq_results_schema_name,
        table_name,
        schema_columns,
        context,
        if_not_exists=True,
    )
    dq_results.cache()
    dq_results.write.format("delta").mode("append").saveAsTable(dq_results_table_name)

    datahub_emit_metadata_with_exception_catcher(
        context=context,
        config=config,
        spark=spark,
    )

    return Output(
        dq_results.coalesce(1),
        metadata={
            **get_output_metadata(config),
            "row_count": dq_results.count(),
            "preview": get_table_preview(dq_results),
        },
    )


@multi_asset(
    outs={
        "geolocation_dq_schools_passed_human_readable": AssetOut(
            io_manager_key=ResourceKey.ADLS_SPARK_SINGLE_FILE_IO_MANAGER.value,
        ),
        "geolocation_dq_schools_failed_human_readable": AssetOut(
            io_manager_key=ResourceKey.ADLS_SPARK_SINGLE_FILE_IO_MANAGER.value,
        ),
    },
)
@capture_op_exceptions
def geolocation_data_quality_results_human_readable(
    context: OpExecutionContext,
    geolocation_data_quality_results: sql.DataFrame,
    config: FileConfig,
):
    context.log.info("Get the file upload object from the database")
    with get_db_context() as db:
        file_upload = db.scalar(
            select(FileUpload).where(FileUpload.id == config.filename_components.id),
        )
        if file_upload is None:
            raise FileNotFoundError(
                f"Database entry for FileUpload with id `{config.filename_components.id}` was not found",
            )

    context.log.info("Obtain the list of uploaded columns")
    file_upload = FileUploadConfig.from_orm(file_upload)
    column_mapping = file_upload.column_to_schema_mapping
    uploaded_columns = list(column_mapping.values())
    context.log.info(f"The list of uploaded columns is: {uploaded_columns}")
    mode = config.metadata["mode"]

    context.log.info("Create a new dataframe with only the relevant columns")
    df, human_readable_mappings = dq_geolocation_extract_relevant_columns(
        geolocation_data_quality_results, uploaded_columns, mode
    )
    for map_key, human_name in human_readable_mappings.items():
        df = df.withColumn(
            human_name,
            f.when(f.element_at(f.col("dq_results"), map_key) == 1, "No").otherwise(
                f.when(f.element_at(f.col("dq_results"), map_key) == 0, "Yes")
            ),
        )
    df = df.drop("dq_results")

    # Cache once — both filters read from the same plan
    df.cache()

    df_passed = df.filter(df.dq_has_critical_error == 0).drop(
        "dq_has_critical_error", "failure_reason"
    )
    df_failed = df.filter(df.dq_has_critical_error == 1).drop("dq_has_critical_error")

    output_metadata = get_output_metadata(config)

    yield Output(
        df_passed,
        output_name="geolocation_dq_schools_passed_human_readable",
        metadata={
            **output_metadata,
            "row_count": df_passed.count(),
            "preview": get_table_preview(df_passed),
        },
    )
    yield Output(
        df_failed,
        output_name="geolocation_dq_schools_failed_human_readable",
        metadata={
            **output_metadata,
            "row_count": df_failed.count(),
            "preview": get_table_preview(df_failed),
        },
    )


@asset(io_manager_key=ResourceKey.ADLS_JSON_IO_MANAGER.value)
@capture_op_exceptions
async def geolocation_data_quality_results_summary(
    context: OpExecutionContext,
    geolocation_bronze: sql.DataFrame,
    geolocation_data_quality_results: sql.DataFrame,
    spark: PySparkResource,
    config: FileConfig,
) -> Output[dict]:
    with get_db_context() as db:
        file_upload = db.scalar(
            select(FileUpload).where(FileUpload.id == config.filename_components.id),
        )
        if file_upload is None:
            raise FileNotFoundError(
                f"Database entry for FileUpload with id `{config.filename_components.id}` was not found",
            )
    file_upload = FileUploadConfig.from_orm(file_upload)
    column_mapping = file_upload.column_to_schema_mapping
    uploaded_columns = list(column_mapping.values())
    mode = config.metadata["mode"]
    context.log.info(f"The list of uploaded columns is: {uploaded_columns}")

    dq_results, _ = dq_geolocation_extract_relevant_columns(
        geolocation_data_quality_results, uploaded_columns, mode=mode
    )

    dq_summary_statistics = aggregate_report_json(
        df_aggregated=aggregate_report_spark_df(
            spark.spark_session,
            dq_results,
        ),
        df_bronze=geolocation_bronze,
        df_data_quality_checks=dq_results,
    )

    datahub_emit_metadata_with_exception_catcher(
        context=context,
        config=config,
        spark=spark,
    )

    await send_email_dq_report_with_config(
        dq_results=dq_summary_statistics,
        config=config,
        context=context,
    )
    return Output(dq_summary_statistics, metadata=get_output_metadata(config))


@asset(io_manager_key=ResourceKey.ADLS_GENERIC_FILE_IO_MANAGER.value)
def geolocation_data_quality_report(
    context: OpExecutionContext,
    geolocation_data_quality_results: sql.DataFrame,
    geolocation_raw: bytes,
    config: FileConfig,
    spark: PySparkResource,
):
    with get_db_context() as db:
        file_upload = db.scalar(
            select(FileUpload).where(FileUpload.id == config.filename_components.id),
        )
        if file_upload is None:
            raise FileNotFoundError(
                f"Database entry for FileUpload with id `{config.filename_components.id}` was not found",
            )

        file_upload = FileUploadConfig.from_orm(file_upload)

    with BytesIO(geolocation_raw) as buffer:
        buffer.seek(0)
        original_df = pandas_loader(buffer, config.filepath, context=context).map(str)

    original_df_columns = original_df.columns
    uploaded_columns = file_upload.column_to_schema_mapping.values()

    uploaded_columns_not_used = list(set(original_df_columns) - set(uploaded_columns))

    schema = get_schema_table(spark.spark_session, config.metastore_schema)
    important_columns_df = schema.filter(f.col("is_important"))
    important_columns_list = [
        row[0] for row in important_columns_df.select("name").collect()
    ]

    important_columns_not_uploaded = list(
        set(important_columns_list) - set(uploaded_columns)
    )
    important_columns_not_uploaded = [
        col for col in important_columns_not_uploaded if not col.startswith("admin")
    ]

    upload_details = {
        "country_code": file_upload.country,
        "file_name": file_upload.original_filename,
        "uploaded_columns_not_used": uploaded_columns_not_used,
        "important_columns_not_uploaded": important_columns_not_uploaded,
    }

    dq_report = aggregate_report_statistics(
        geolocation_data_quality_results, upload_details
    )
    return Output(dq_report)


@asset(io_manager_key=ResourceKey.ADLS_SPARK_IO_MANAGER.value)
@capture_op_exceptions
def geolocation_dq_passed_rows(
    context: OpExecutionContext,
    geolocation_data_quality_results: sql.DataFrame,
    config: FileConfig,
    spark: PySparkResource,
) -> Output[sql.DataFrame]:
    df_passed = dq_split_passed_rows(
        geolocation_data_quality_results,
        config.dataset_type,
    )

    schema_reference = get_schema_columns_datahub(
        spark.spark_session,
        config.metastore_schema,
    )
    datahub_emit_metadata_with_exception_catcher(
        context=context,
        config=config,
        spark=spark,
        schema_reference=schema_reference,
    )

    df_passed.cache()
    row_count = df_passed.count()

    return Output(
        df_passed,
        metadata={
            **get_output_metadata(config),
            "row_count": row_count,
            "preview": get_table_preview(df_passed),
        },
    )


@asset(io_manager_key=ResourceKey.ADLS_SPARK_IO_MANAGER.value)
@capture_op_exceptions
def geolocation_dq_failed_rows(
    context: OpExecutionContext,
    geolocation_data_quality_results: sql.DataFrame,
    config: FileConfig,
    spark: PySparkResource,
) -> Output[sql.DataFrame]:
    df_failed = dq_split_failed_rows(
        geolocation_data_quality_results,
        config.dataset_type,
    )

    schema_reference = get_schema_columns_datahub(
        spark.spark_session,
        config.metastore_schema,
    )
    datahub_emit_metadata_with_exception_catcher(
        context=context,
        config=config,
        spark=spark,
        schema_reference=schema_reference,
        df_failed=df_failed,
    )

    df_failed.cache()
    row_count = df_failed.count()

    return Output(
        df_failed,
        metadata={
            **get_output_metadata(config),
            "row_count": row_count,
            "preview": get_table_preview(df_failed),
        },
    )


@asset
@capture_op_exceptions
def geolocation_error_table(
    context: OpExecutionContext,
    geolocation_dq_failed_rows: sql.DataFrame,
    config: FileConfig,
    spark: PySparkResource,
) -> Output[None]:
    s: SparkSession = spark.spark_session

    if geolocation_dq_failed_rows.isEmpty():
        context.log.info("No failed rows to write to aggregated error table.")
        return Output(None)

    file_id = config.filename_components.id
    file_name = Path(config.filepath).name
    country_code = config.country_code
    dataset_type = config.dataset_type

    df = geolocation_dq_failed_rows

    df = df.withColumn("giga_sync_file_id", f.lit(file_id))
    df = df.withColumn("giga_sync_file_name", f.lit(file_name))
    df = df.withColumn("dataset_type", f.lit(dataset_type))
    df = df.withColumn("country_code", f.lit(country_code))
    df = df.withColumn("created_at", f.current_timestamp())

    schema_name = "school_geolocation_error_table"
    table_name = country_code.lower()
    full_table_name = construct_full_table_name(schema_name, table_name)

    create_schema(s, schema_name)

    try:
        if s.catalog.tableExists(full_table_name):
            context.log.info(f"Deleting existing errors for file_id: {file_id}")
            delta_table = DeltaTable.forName(s, full_table_name)
            delta_table.delete(f.col("giga_sync_file_id") == f.lit(file_id))
        else:
            context.log.info(
                f"Table {full_table_name} does not exist. It will be created on write."
            )
    except Exception as exc:
        context.log.warning(f"Failed to delete existing rows: {exc}")

    context.log.info(f"Appending failed rows to {full_table_name}")
    row_count = df.count()

    (
        df.write.format("delta")
        .mode("append")
        .option("mergeSchema", "true")
        .saveAsTable(full_table_name)
    )

    return Output(
        None,
        metadata={
            **get_output_metadata(config),
            "row_count": row_count,
            "preview": get_table_preview(df),
        },
    )


@asset
@capture_op_exceptions
def geolocation_staging(
    context: OpExecutionContext,
    geolocation_dq_passed_rows: sql.DataFrame,
    adls_file_client: ADLSFileClient,
    spark: PySparkResource,
    config: FileConfig,
) -> Output[None]:
    if geolocation_dq_passed_rows.isEmpty():
        context.log.warning("Skipping staging as there are no rows passing DQ checks")
        return Output(None)

    schema_reference = get_schema_columns_datahub(
        spark.spark_session,
        config.metastore_schema,
    )
    datahub_emit_metadata_with_exception_catcher(
        context=context,
        config=config,
        spark=spark,
        schema_reference=schema_reference,
    )
    staging_step = StagingStep(
        context,
        config,
        adls_file_client,
        spark.spark_session,
        StagingMode.UPDATE,
    )
    pending = staging_step(geolocation_dq_passed_rows)

    if pending is None:
        return Output(
            None,
            metadata={
                **get_output_metadata(config),
                "insert_count": MetadataValue.int(0),
                "update_count": MetadataValue.int(0),
                "unchanged_count": MetadataValue.int(0),
                "delete_count": MetadataValue.int(0),
            },
        )

    counts = pending.groupBy("change_type").count().collect()
    count_map = {row["change_type"]: row["count"] for row in counts}
    return Output(
        None,
        metadata={
            **get_output_metadata(config),
            "insert_count": MetadataValue.int(count_map.get("INSERT", 0)),
            "update_count": MetadataValue.int(count_map.get("UPDATE", 0)),
            "unchanged_count": MetadataValue.int(count_map.get("UNCHANGED", 0)),
            "delete_count": MetadataValue.int(count_map.get("DELETE", 0)),
        },
    )


@asset
@capture_op_exceptions
def geolocation_delete_staging(
    context: OpExecutionContext,
    adls_file_client: ADLSFileClient,
    spark: PySparkResource,
    config: FileConfig,
) -> Output[None]:
    delete_row_ids = adls_file_client.download_json(config.filepath)
    if isinstance(delete_row_ids, list):
        # dedupe change IDs
        delete_row_ids = list(set(delete_row_ids))

    staging_step = StagingStep(
        context,
        config,
        adls_file_client,
        spark.spark_session,
        StagingMode.DELETE,
    )
    staging = staging_step(delete_row_ids)

    if staging is not None:
        datahub_emit_metadata_with_exception_catcher(
            context=context,
            config=config,
            spark=spark,
        )
        return Output(
            None,
            metadata={
                **get_output_metadata(config),
                "preview": get_table_preview(staging),
                "delete_row_ids": MetadataValue.json(delete_row_ids),
            },
        )

    return Output(
        None,
        metadata={
            **get_output_metadata(config),
            "delete_row_ids": MetadataValue.json(delete_row_ids),
        },
    )