Streaming Logistics Cost Monitoring with Spark Structured Streaming

# Databricks notebook source
# MAGIC %md
# MAGIC ### Setup: Create Unity Catalog Schema and Mock Lookup Tables
# MAGIC
# MAGIC This notebook initializes the Unity Catalog schema and creates a mock `silver.fuel_prices` table for our streaming pipeline.
# MAGIC It also defines the expected schema for incoming Kafka messages.

# COMMAND ----------

import logging
from pyspark.sql import SparkSession
from pyspark.sql.types import (
    StructType, StructField, StringType, TimestampType, DoubleType, DateType, IntegerType
)
from delta.tables import DeltaTable

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# COMMAND ----------

# Configuration
CATALOG_NAME = "supply_chain_lakehouse"
SCHEMA_NAME = "bronze" # For raw events
SCHEMA_NAME_SILVER = "silver" # For enriched data and lookup tables
SCHEMA_NAME_GOLD = "gold" # For tariff analysis lookup

# COMMAND ----------

# MAGIC %md
# MAGIC #### 1. Create Unity Catalog Schemas
# MAGIC
# MAGIC We'll ensure the necessary schemas exist in Unity Catalog.

# COMMAND ----------

try:
    spark.sql(f"CREATE CATALOG IF NOT EXISTS {CATALOG_NAME}")
    spark.sql(f"USE CATALOG {CATALOG_NAME}")
    spark.sql(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA_NAME}")
    spark.sql(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA_NAME_SILVER}")
    spark.sql(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA_NAME_GOLD}") # Ensure gold schema from DLT exists
    logger.info(f"Unity Catalog schemas '{CATALOG_NAME}.{SCHEMA_NAME}', '{CATALOG_NAME}.{SCHEMA_NAME_SILVER}', '{CATALOG_NAME}.{SCHEMA_NAME_GOLD}' ensured.")
except Exception as e:
    logger.error(f"Error creating Unity Catalog schemas: {e}")
    raise

# COMMAND ----------

# MAGIC %md
# MAGIC #### 2. Define Expected Kafka Message Schema
# MAGIC
# MAGIC This is the schema we expect for the `logistics_events` Kafka topic.

# COMMAND ----------

logistics_event_schema = StructType([
    StructField("event_id", StringType(), False),
    StructField("shipment_id", StringType(), False),
    StructField("event_type", StringType(), False), # e.g., 'SHIPMENT_UPDATE', 'FUEL_STOP'
    StructField("event_timestamp", TimestampType(), False),
    StructField("origin_location", StringType(), True),
    StructField("destination_location", StringType(), True),
    StructField("product_hs_code", StringType(), True), # For tariff correlation
    StructField("distance_km", DoubleType(), True),
    StructField("fuel_consumed_liters", DoubleType(), True),
    StructField("driver_id", StringType(), True),
    StructField("current_status", StringType(), True)
])

logger.info("Logistics event schema defined.")

# COMMAND ----------

# MAGIC %md
# MAGIC #### 3. Create or Update `silver.fuel_prices` Lookup Table
# MAGIC
# MAGIC We'll create a mock fuel prices table. In a real-world scenario, this would be populated by another data pipeline (e.g., daily batch job or another stream).

# COMMAND ----------

fuel_prices_table_path = f"{CATALOG_NAME}.{SCHEMA_NAME_SILVER}.fuel_prices"

# Data for the mock fuel prices table
fuel_price_data = [
    ("Diesel", "North America", 1.25, "USD", "2025-12-19", "2025-12-19 08:00:00"),
    ("Diesel", "Europe", 1.50, "EUR", "2025-12-19", "2025-12-19 08:00:00"),
    ("Diesel", "Asia", 1.10, "USD", "2025-12-19", "2025-12-19 08:00:00"),
    ("Gasoline", "North America", 1.15, "USD", "2025-12-19", "2025-12-19 08:00:00"),
    ("Gasoline", "Europe", 1.40, "EUR", "2025-12-19", "2025-12-19 08:00:00"),
]

fuel_prices_df = spark.createDataFrame(
    fuel_price_data,
    schema=["fuel_type", "region", "price_per_unit", "currency", "effective_date", "update_time"]
)

try:
    # Use MERGE INTO for idempotent updates
    # This ensures that if the table exists, it's updated, otherwise created.
    # For a simple mock, we can just overwrite or create.
    # For production, a MERGE is better if data changes frequently.
    if spark.catalog.tableExists(fuel_prices_table_path):
        logger.info(f"Table {fuel_prices_table_path} exists. Overwriting with mock data.")
        fuel_prices_df.write \
            .format("delta") \
            .mode("overwrite") \
            .option("mergeSchema", "true") \
            .saveAsTable(fuel_prices_table_path)
    else:
        logger.info(f"Table {fuel_prices_table_path} does not exist. Creating with mock data.")
        fuel_prices_df.write \
            .format("delta") \
            .mode("append") \
            .saveAsTable(fuel_prices_table_path)

    logger.info(f"Mock fuel prices table '{fuel_prices_table_path}' created/updated successfully.")
    spark.sql(f"SELECT * FROM {fuel_prices_table_path}").display()

except Exception as e:
    logger.error(f"Error creating/updating fuel prices table: {e}")
    raise

# COMMAND ----------

# MAGIC %md
# MAGIC #### 4. Verify Tariff Analysis Table (from DLT)
# MAGIC
# MAGIC We'll verify the `gold.tariff_analysis` table from previous DLT chapters exists and has data.
# MAGIC This table is crucial for joining.

# COMMAND ----------

tariff_analysis_table_path = f"{CATALOG_NAME}.{SCHEMA_NAME_GOLD}.tariff_analysis"

try:
    if spark.catalog.tableExists(tariff_analysis_table_path):
        logger.info(f"Tariff analysis table '{tariff_analysis_table_path}' found.")
        # Display a sample to confirm data presence
        spark.sql(f"SELECT * FROM {tariff_analysis_table_path} LIMIT 5").display()
    else:
        logger.warning(f"Tariff analysis table '{tariff_analysis_table_path}' not found. Please ensure DLT pipelines from previous chapters are run.")
        # Create a dummy table for local testing if not found, for demonstration purposes.
        # In production, this would be a hard dependency.
        dummy_tariff_data = [
            ("870323", "USA", "Mexico", 0.05, "2025-01-01"),
            ("870323", "Germany", "France", 0.02, "2025-01-01"),
            ("847130", "China", "USA", 0.10, "2025-01-01"),
        ]
        dummy_tariff_df = spark.createDataFrame(
            dummy_tariff_data,
            schema=["hs_code", "origin_country", "destination_country", "tariff_rate", "effective_date"]
        )
        dummy_tariff_df.write \
            .format("delta") \
            .mode("overwrite") \
            .saveAsTable(tariff_analysis_table_path)
        logger.info(f"Created dummy tariff analysis table '{tariff_analysis_table_path}' for demonstration.")
        spark.sql(f"SELECT * FROM {tariff_analysis_table_path} LIMIT 5").display()
except Exception as e:
    logger.error(f"Error verifying/creating tariff analysis table: {e}")
    raise

# COMMAND ----------

# MAGIC %md
# MAGIC #### 5. Kafka Configuration
# MAGIC
# MAGIC We'll store Kafka connection details in a configuration file or Databricks secrets. For this tutorial, we'll use a Python dictionary, but strongly recommend Databricks Secrets for production.

# COMMAND ----------

# File: conf/kafka_config.py
# This file would typically be managed by Databricks Secrets or a configuration management system.
# For tutorial purposes, we define it here.

kafka_bootstrap_servers = "<YOUR_KAFKA_BOOTSTRAP_SERVERS>" # e.g., "pkc-lxxxx.us-east-1.aws.confluent.cloud:9092"
kafka_topic_logistics_events = "logistics_events"
kafka_sasl_jaas_config = "org.apache.kafka.common.security.plain.PlainLoginModule required username='<YOUR_KAFKA_API_KEY>' password='<YOUR_KAFKA_API_SECRET>';"
kafka_security_protocol = "SASL_SSL"
kafka_sasl_mechanism = "PLAIN"

# IMPORTANT: In a production environment, NEVER hardcode credentials.
# Use Databricks Secrets (dbutils.secrets.get) for secure access.
# Example:
# kafka_bootstrap_servers = dbutils.secrets.get(scope="kafka_scope", key="bootstrap_servers")
# kafka_api_key = dbutils.secrets.get(scope="kafka_scope", key="api_key")
# kafka_api_secret = dbutils.secrets.get(scope="kafka_scope", key="api_secret")
# kafka_sasl_jaas_config = f"org.apache.kafka.common.security.plain.PlainLoginModule required username='{kafka_api_key}' password='{kafka_api_secret}';"

logger.info("Kafka configuration loaded (ensure secrets are handled securely in production).")

# COMMAND ----------

# MAGIC %md
# MAGIC #### Setup Complete
# MAGIC
# MAGIC The necessary Unity Catalog schemas are in place, mock lookup tables are populated, and Kafka configuration is defined. We are now ready to build our streaming jobs.

# COMMAND ----------

# End of 01_setup_schemas_and_lookups.py

# Databricks notebook source
# MAGIC %md
# MAGIC ### Stream Raw Logistics Events to Bronze Delta Table
# MAGIC
# MAGIC This notebook sets up a Spark Structured Streaming job to ingest raw logistics events from a Kafka topic
# MAGIC and land them into the `bronze.raw_logistics_events` Delta table.

# COMMAND ----------

import logging
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, current_timestamp, from_json, lit, expr
from pyspark.sql.types import (
    StructType, StructField, StringType, TimestampType, DoubleType, DateType, IntegerType
)
from delta.tables import DeltaTable
from pyspark.sql.streaming import StreamingQuery
from pyspark.sql.utils import StreamingQueryException

# Import Kafka configuration
# In a production setup, this would be a dedicated config module or Databricks Secrets
try:
    from supply_chain_project.08_structured_streaming_cost_monitoring.conf.kafka_config import (
        kafka_bootstrap_servers, kafka_topic_logistics_events,
        kafka_sasl_jaas_config, kafka_security_protocol, kafka_sasl_mechanism
    )
except ImportError:
    # Fallback for direct notebook execution if conf/kafka_config.py is not in path or run separately
    logging.warning("kafka_config.py not found as a module. Assuming variables are defined in current scope or hardcoded for testing.")
    kafka_bootstrap_servers = "<YOUR_KAFKA_BOOTSTRAP_SERVERS>"
    kafka_topic_logistics_events = "logistics_events"
    kafka_sasl_jaas_config = "org.apache.kafka.common.security.plain.PlainLoginModule required username='<YOUR_KAFKA_API_KEY>' password='<YOUR_KAFKA_API_SECRET>';"
    kafka_security_protocol = "SASL_SSL"
    kafka_sasl_mechanism = "PLAIN"


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# COMMAND ----------

# Configuration
CATALOG_NAME = "supply_chain_lakehouse"
SCHEMA_NAME = "bronze"
BRONZE_TABLE_NAME = "raw_logistics_events"
BRONZE_TABLE_PATH = f"{CATALOG_NAME}.{SCHEMA_NAME}.{BRONZE_TABLE_NAME}"
CHECKPOINT_LOCATION = f"/mnt/{CATALOG_NAME}/checkpoints/{SCHEMA_NAME}/{BRONZE_TABLE_NAME}" # Use ADLS/S3 mount for checkpoints

# Ensure mount point exists if using /mnt/
# dbutils.fs.mkdirs(f"/mnt/{CATALOG_NAME}/checkpoints")
# For Unity Catalog, external locations are preferred for managed checkpoints.
# For simplicity, we assume an external location or managed checkpointing.

# COMMAND ----------

# MAGIC %md
# MAGIC #### 1. Define Kafka Read Stream

# COMMAND ----------

try:
    # Read from Kafka using Spark Structured Streaming
    kafka_df = spark.readStream \
        .format("kafka") \
        .option("kafka.bootstrap.servers", kafka_bootstrap_servers) \
        .option("subscribe", kafka_topic_logistics_events) \
        .option("startingOffsets", "earliest") \
        .option("failOnDataLoss", "false") \
        .option("kafka.security.protocol", kafka_security_protocol) \
        .option("kafka.sasl.mechanism", kafka_sasl_mechanism) \
        .option("kafka.sasl.jaas.config", kafka_sasl_jaas_config) \
        .load()

    logger.info(f"Kafka read stream initialized for topic: {kafka_topic_logistics_events}")

except Exception as e:
    logger.error(f"Failed to initialize Kafka read stream: {e}")
    raise

# COMMAND ----------

# MAGIC %md
# MAGIC #### 2. Transform and Select Data for Bronze Table

# COMMAND ----------

# Add ingestion timestamp and cast value to string
bronze_df = kafka_df.selectExpr(
    "CAST(key AS BINARY) as key",
    "CAST(value AS STRING) as value",
    "topic",
    "partition",
    "offset",
    "timestamp",
    "current_timestamp() as ingestion_time"
)

logger.info("Data transformed for bronze layer.")

# COMMAND ----------

# MAGIC %md
# MAGIC #### 3. Write Stream to Delta Lake Bronze Table

# COMMAND ----------

# Define the write stream query
try:
    bronze_stream_query = bronze_df.writeStream \
        .format("delta") \
        .outputMode("append") \
        .option("checkpointLocation", CHECKPOINT_LOCATION) \
        .option("mergeSchema", "true") # Allows schema evolution
        .trigger(processingTime="30 seconds") # Process data every 30 seconds
        .toTable(BRONZE_TABLE_PATH)

    logger.info(f"Bronze stream query started for table: {BRONZE_TABLE_PATH}. Checkpoint: {CHECKPOINT_LOCATION}")
    logger.info("To stop the stream, run: bronze_stream_query.stop()")

    # For development, you might want to wait for termination for a short period
    # bronze_stream_query.awaitTermination(timeout=300) # Wait for 5 minutes

except StreamingQueryException as e:
    logger.error(f"Streaming query failed: {e}")
    # Handle specific streaming errors
    if "offset out of range" in str(e):
        logger.error("Kafka offset out of range. Consider resetting startingOffsets or checkpoint.")
    raise
except Exception as e:
    logger.error(f"An unexpected error occurred during stream writing: {e}")
    raise

# COMMAND ----------

# MAGIC %md
# MAGIC #### Monitoring and Stopping the Stream
# MAGIC
# MAGIC You can monitor the stream's progress in the Databricks UI under the "Streaming" tab of your cluster.
# MAGIC To manually stop the stream (e.g., for development or redeployment):

# COMMAND ----------

# To stop the stream, uncomment and run the following line
# if bronze_stream_query.isActive:
#     bronze_stream_query.stop()
#     logger.info("Bronze stream query stopped.")

# COMMAND ----------

# End of 02_stream_raw_logistics_events_bronze.py

# Databricks notebook source
# MAGIC %md
# MAGIC ### Stream Enriched Logistics Costs to Silver Delta Table
# MAGIC
# MAGIC This notebook sets up a Spark Structured Streaming job to read raw logistics events from the Bronze layer,
# MAGIC parse them, enrich with tariff and fuel price data, calculate total costs, and land them into the
# MAGIC `silver.enriched_logistics_costs` Delta table.

# COMMAND ----------

import logging
from pyspark.sql import SparkSession
from pyspark.sql.functions import (
    col, current_timestamp, from_json, lit, expr, coalesce, when,
    date_trunc, to_date, year, month, dayofmonth,
    sha2, concat_ws, md5
)
from pyspark.sql.types import (
    StructType, StructField, StringType, TimestampType, DoubleType, DateType, IntegerType, LongType
)
from delta.tables import DeltaTable
from pyspark.sql.streaming import StreamingQuery
from pyspark.sql.utils import StreamingQueryException

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# COMMAND ----------

# Configuration
CATALOG_NAME = "supply_chain_lakehouse"
SCHEMA_NAME_BRONZE = "bronze"
SCHEMA_NAME_SILVER = "silver"
SCHEMA_NAME_GOLD = "gold"

BRONZE_TABLE_PATH = f"{CATALOG_NAME}.{SCHEMA_NAME_BRONZE}.raw_logistics_events"
SILVER_FUEL_PRICES_TABLE_PATH = f"{CATALOG_NAME}.{SCHEMA_NAME_SILVER}.fuel_prices"
GOLD_TARIFF_ANALYSIS_TABLE_PATH = f"{CATALOG_NAME}.{SCHEMA_NAME_GOLD}.tariff_analysis"
SILVER_ENRICHED_COSTS_TABLE_NAME = "enriched_logistics_costs"
SILVER_ENRICHED_COSTS_TABLE_PATH = f"{CATALOG_NAME}.{SCHEMA_NAME_SILVER}.{SILVER_ENRICHED_COSTS_TABLE_NAME}"
CHECKPOINT_LOCATION_SILVER = f"/mnt/{CATALOG_NAME}/checkpoints/{SCHEMA_NAME_SILVER}/{SILVER_ENRICHED_COSTS_TABLE_NAME}"

# Ensure mount point exists if using /mnt/
# dbutils.fs.mkdirs(f"/mnt/{CATALOG_NAME}/checkpoints")
# For Unity Catalog, external locations are preferred for managed checkpoints.

# COMMAND ----------

# MAGIC %md
# MAGIC #### 1. Define Expected Kafka Message Schema (for parsing Bronze `value` column)
# MAGIC
# MAGIC This schema must match the structure of your Kafka `logistics_events` messages.

# COMMAND ----------

# This schema is used to parse the 'value' (JSON string) column from the bronze table
logistics_event_schema = StructType([
    StructField("event_id", StringType(), False),
    StructField("shipment_id", StringType(), False),
    StructField("event_type", StringType(), False), # e.g., 'SHIPMENT_UPDATE', 'FUEL_STOP'
    StructField("event_timestamp", TimestampType(), False),
    StructField("origin_location", StringType(), True),
    StructField("destination_location", StringType(), True),
    StructField("product_hs_code", StringType(), True), # For tariff correlation
    StructField("distance_km", DoubleType(), True),
    StructField("fuel_consumed_liters", DoubleType(), True),
    StructField("driver_id", StringType(), True),
    StructField("current_status", StringType(), True)
])

logger.info("Logistics event schema for parsing defined.")

# COMMAND ----------

# MAGIC %md
# MAGIC #### 2. Read Stream from Bronze Delta Table

# COMMAND ----------

try:
    # Read from the bronze Delta table as a streaming source
    bronze_stream_df = spark.readStream \
        .format("delta") \
        .option("ignoreDeletes", "true") \
        .option("ignoreChanges", "true") \
        .table(BRONZE_TABLE_PATH)

    logger.info(f"Structured Stream initialized from Bronze table: {BRONZE_TABLE_PATH}")

except Exception as e:
    logger.error(f"Failed to initialize stream from Bronze Delta table: {e}")
    raise

# COMMAND ----------

# MAGIC %md
# MAGIC #### 3. Parse JSON, Flatten, and Apply Data Quality Checks

# COMMAND ----------

# Parse the JSON string 'value' column into structured columns
parsed_df = bronze_stream_df.withColumn(
    "parsed_data",
    from_json(col("value"), logistics_event_schema)
).select(
    col("parsed_data.*"),
    col("ingestion_time").alias("bronze_ingestion_time"),
    col("offset").alias("kafka_offset_id") # Keep offset for traceability
)

# Apply basic data quality checks
# For example, filter out records where 'event_id' or 'shipment_id' is null
# Or flag them for a Dead Letter Queue (DLQ) if using a more advanced pattern
clean_df = parsed_df.filter(
    col("event_id").isNotNull() & col("shipment_id").isNotNull()
).withColumn(
    "data_quality_status", lit("CLEAN")
).withColumn(
    "error_details", lit(None)
)

# For records failing schema parsing or basic checks, you might route them to a DLQ
# For simplicity, we're filtering out. In production, consider a more robust DLQ approach.
bad_schema_df = parsed_df.filter(
    col("event_id").isNull() | col("shipment_id").isNull()
).withColumn(
    "data_quality_status", lit("BAD_SCHEMA_OR_MISSING_KEY_FIELDS")
).withColumn(
    "error_details", lit("Missing event_id or shipment_id in raw JSON")
)

# Union clean and bad_schema records to write to a single table, but flag bad ones.
# For this tutorial, we'll proceed with `clean_df` directly for simplicity,
# assuming a separate DLQ stream for `bad_schema_df` in a full production system.
# If we wanted to write all, we'd union them and then write.
# For now, we'll only process clean_df to the silver table.

logger.info("JSON parsing and initial data quality checks applied.")

# COMMAND ----------

# MAGIC %md
# MAGIC #### 4. Load Lookup Tables (Tariff Analysis and Fuel Prices)
# MAGIC
# MAGIC We'll load these as batch DataFrames which Spark will handle efficiently (e.g., broadcasting for small tables) for joining with the stream.

# COMMAND ----------

try:
    # Load the tariff analysis table
    tariff_analysis_df = spark.read.table(GOLD_TARIFF_ANALYSIS_TABLE_PATH) \
        .select(
            col("hs_code").alias("tariff_hs_code"),
            col("origin_country").alias("tariff_origin_country"),
            col("destination_country").alias("tariff_destination_country"),
            col("tariff_rate"),
            col("effective_date").alias("tariff_effective_date")
        )

    logger.info(f"Tariff analysis lookup table '{GOLD_TARIFF_ANALYSIS_TABLE_PATH}' loaded.")

    # Load the fuel prices table
    fuel_prices_df = spark.read.table(SILVER_FUEL_PRICES_TABLE_PATH) \
        .select(
            col("fuel_type").alias("fuel_lookup_type"),
            col("region").alias("fuel_lookup_region"),
            col("price_per_unit").alias("fuel_price_per_unit"),
            col("currency").alias("fuel_currency"),
            col("effective_date").alias("fuel_effective_date")
        )

    logger.info(f"Fuel prices lookup table '{SILVER_FUEL_PRICES_TABLE_PATH}' loaded.")

except Exception as e:
    logger.error(f"Failed to load lookup tables: {e}")
    raise

# COMMAND ----------

# MAGIC %md
# MAGIC #### 5. Enrich Data with Lookups and Calculate Costs

# COMMAND ----------

# Determine region from origin_location for fuel price lookup (simplified)
# In reality, you'd need a more sophisticated geo-mapping service
enriched_df = clean_df.withColumn(
    "event_region",
    when(col("origin_location").contains("America"), "North America")
    .when(col("origin_location").contains("Europe"), "Europe")
    .when(col("origin_location").contains("Asia"), "Asia")
    .otherwise("Other")
)

# Join with Fuel Prices
# This is a stream-static join. Spark handles this by broadcasting the static DF.
enriched_df = enriched_df.join(
    fuel_prices_df,
    (enriched_df["event_region"] == fuel_prices_df["fuel_lookup_region"]) &
    (to_date(enriched_df["event_timestamp"]) >= fuel_prices_df["fuel_effective_date"]) & # Use effective date for prices
    (lit("Diesel") == fuel_prices_df["fuel_lookup_type"]), # Assuming 'Diesel' for simplicity, could be dynamic
    "leftouter"
)

# Join with Tariff Analysis
enriched_df = enriched_df.join(
    tariff_analysis_df,
    (enriched_df["product_hs_code"] == tariff_analysis_df["tariff_hs_code"]) &
    (to_date(enriched_df["event_timestamp"]) >= tariff_analysis_df["tariff_effective_date"]), # Use effective date for tariffs
    "leftouter"
)

# Calculate costs
final_silver_df = enriched_df.withColumn(
    "calculated_fuel_cost_usd",
    coalesce(col("fuel_consumed_liters") * col("fuel_price_per_unit"), lit(0.0))
).withColumn(
    "tariff_impact_usd",
    when(col("tariff_rate").isNotNull(), lit(100.0) * col("tariff_rate")) # Example: 100 USD base value for calculation
    .otherwise(lit(0.0))
).withColumn(
    "other_fixed_cost_usd", lit(50.0) # Example: a fixed handling cost per event
).withColumn(
    "total_logistics_cost_usd",
    col("calculated_fuel_cost_usd") + col("tariff_impact_usd") + col("other_fixed_cost_usd")
).withColumn(
    "processing_timestamp", current_timestamp()
).select(
    col("event_id"),
    col("shipment_id"),
    col("event_type"),
    col("event_timestamp"),
    col("origin_location"),
    col("destination_location"),
    col("product_hs_code"),
    col("distance_km"),
    col("fuel_consumed_liters"),
    col("driver_id"),
    col("current_status"),
    col("tariff_impact_usd"),
    col("fuel_price_per_unit"),
    col("calculated_fuel_cost_usd"),
    col("total_logistics_cost_usd"),
    col("processing_timestamp"),
    col("data_quality_status"),
    col("error_details")
)

logger.info("Data enriched and costs calculated.")

# COMMAND ----------

# MAGIC %md
# MAGIC #### 6. Write Stream to Delta Lake Silver Table

# COMMAND ----------

# Define the write stream query for the silver table
try:
    silver_stream_query = final_silver_df.writeStream \
        .format("delta") \
        .outputMode("append") \
        .option("checkpointLocation", CHECKPOINT_LOCATION_SILVER) \
        .option("mergeSchema", "true") \
        .trigger(processingTime="30 seconds") # Process data every 30 seconds
        .toTable(SILVER_ENRICHED_COSTS_TABLE_PATH)

    logger.info(f"Silver stream query started for table: {SILVER_ENRICHED_COSTS_TABLE_PATH}. Checkpoint: {CHECKPOINT_LOCATION_SILVER}")
    logger.info("To stop the stream, run: silver_stream_query.stop()")

    # For development, you might want to wait for termination for a short period
    # silver_stream_query.awaitTermination(timeout=300) # Wait for 5 minutes

except StreamingQueryException as e:
    logger.error(f"Streaming query failed: {e}")
    # Handle specific streaming errors
    raise
except Exception as e:
    logger.error(f"An unexpected error occurred during stream writing to silver: {e}")
    raise

# COMMAND ----------

# MAGIC %md
# MAGIC #### Monitoring and Stopping the Stream
# MAGIC
# MAGIC You can monitor the stream's progress in the Databricks UI under the "Streaming" tab of your cluster.
# MAGIC To manually stop the stream (e.g., for development or redeployment):

# COMMAND ----------

# To stop the stream, uncomment and run the following line
# if silver_stream_query.isActive:
#     silver_stream_query.stop()
#     logger.info("Silver stream query stopped.")

# COMMAND ----------

# End of 03_stream_enriched_logistics_costs_silver.py