{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "intro-md",
   "metadata": {},
   "source": [
    "# End-to-End Data Pipeline Notebook\n",
    "\n",
    "This notebook demonstrates a fully integrated data pipeline supporting both **batch** and **streaming** processing, including data ingestion, transformation, validation, storage, monitoring, governance, and ML tracking.\n",
    "\n",
    "It is designed for rapid deployment and customization using Docker Compose, Apache Airflow, Apache Spark, Kafka, MinIO, PostgreSQL, Prometheus, Grafana, and more.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "arch-md",
   "metadata": {},
   "source": [
    "## Architecture Overview\n",
    "\n",
    "Below is a text-based diagram of the pipeline architecture:\n",
    "\n",
    "```\n",
    "                              ┌─────────────────────────────┐\n",
    "                              │         Batch Source        │\n",
    "                              │   (MySQL, Files, etc.)      │\n",
    "                              └────────────┬────────────────┘\n",
    "                                           │\n",
    "                                           │ (Extract/Validate)\n",
    "                                           ▼\n",
    "                           ┌────────────────────────────────┐\n",
    "                           │      Airflow Batch DAG         │\n",
    "                           │ - Extracts data from MySQL     │\n",
    "                           │ - Validates with GreatExpect.  │\n",
    "                           │ - Uploads raw data to MinIO    │\n",
    "                           └────────────┬───────────────────┘\n",
    "                                        │ (spark-submit)\n",
    "                                        ▼\n",
    "                           ┌───────────────────────────────────┐\n",
    "                           │         Spark Batch Job           │\n",
    "                           │ - Transforms, cleans, enriches    │\n",
    "                           │ - Writes to PostgreSQL & MinIO    │\n",
    "                           └────────────┬──────────────────────┘\n",
    "                                        │\n",
    "                                        ▼\n",
    "                           ┌────────────────────────────────┐\n",
    "                           │       Processed Data Store     │\n",
    "                           │         (PostgreSQL)           │\n",
    "                           └────────────────────────────────┘\n",
    "                                        \n",
    "Streaming Side:\n",
    "                              ┌─────────────────────────────┐\n",
    "                              │       Streaming Source      │\n",
    "                              │         (Kafka)             │\n",
    "                              └────────────┬────────────────┘\n",
    "                                           │\n",
    "                                           ▼\n",
    "                           ┌───────────────────────────────────┐\n",
    "                           │    Spark Streaming Job            │\n",
    "                           │ - Consumes Kafka messages         │\n",
    "                           │ - Detects anomalies               │\n",
    "                           │ - Writes to PostgreSQL & MinIO    │\n",
    "                           └───────────────────────────────────┘\n",
    "\n",
    "Monitoring & Governance:\n",
    "                              ┌────────────────────────────────┐\n",
    "                              │   Monitoring & Governance      │\n",
    "                              │ - Prometheus & Grafana         │\n",
    "                              │ - Apache Atlas/OpenMetadata    │\n",
    "                              └────────────────────────────────┘\n",
    "\n",
    "ML & Serving:\n",
    "                              ┌────────────────────────────────┐\n",
    "                              │        AI/ML Serving           │\n",
    "                              │ - Feature Store (Feast)        │\n",
    "                              │ - MLflow Model Tracking        │\n",
    "                              │ - BI Dashboards                │\n",
    "                              └────────────────────────────────┘\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dir-md",
   "metadata": {},
   "source": [
    "## Directory Structure\n",
    "\n",
    "```\n",
    "end-to-end-pipeline/\n",
    "  ├── docker-compose.yaml            # Docker Compose for all services\n",
    "  ├── README.md                      # This documentation\n",
    "  ├── airflow/\n",
    "  │   ├── Dockerfile                 # Airflow image\n",
    "  │   ├── requirements.txt           # Python deps for Airflow\n",
    "  │   └── dags/\n",
    "  │       ├── batch_ingestion_dag.py # Batch ingestion DAG\n",
    "  │       └── streaming_monitoring_dag.py  # Streaming monitoring DAG\n",
    "  ├── spark/\n",
    "  │   ├── Dockerfile                 # Spark image\n",
    "  │   ├── spark_batch_job.py         # Spark batch ETL job\n",
    "  │   └── spark_streaming_job.py     # Spark streaming job\n",
    "  ├── kafka/\n",
    "  │   └── producer.py                # Kafka producer simulation\n",
    "  ├── great_expectations/\n",
    "  │   ├── great_expectations.yaml    # GE config\n",
    "  │   └── expectations/\n",
    "  │       └── raw_data_validation.py # GE suite\n",
    "  ├── governance/\n",
    "  │   └── atlas_stub.py              # Dataset lineage registration\n",
    "  ├── monitoring/\n",
    "  │   ├── monitoring.py              # Monitoring setup for Prometheus & Grafana\n",
    "  │   └── prometheus.yml             # Prometheus config\n",
    "  ├── ml/\n",
    "  │   ├── feature_store_stub.py      # Feature Store stub\n",
    "  │   └── mlflow_tracking.py         # MLflow tracking\n",
    "  └── scripts/\n",
    "      └── init_db.sql                # MySQL initialization script\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "id": "docker-compose-cell",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Docker Compose content\n",
    "docker_compose_yaml = '''\n",
    "version: \"3.8\"\n",
    "services:\n",
    "  zookeeper:\n",
    "    image: confluentinc/cp-zookeeper:7.3.2\n",
    "    container_name: zookeeper\n",
    "    environment:\n",
    "      ZOOKEEPER_CLIENT_PORT: 2181\n",
    "    ports:\n",
    "      - \"2181:2181\"\n",
    "\n",
    "  kafka:\n",
    "    image: confluentinc/cp-kafka:7.3.2\n",
    "    container_name: kafka\n",
    "    depends_on:\n",
    "      - zookeeper\n",
    "    ports:\n",
    "      - \"9092:9092\"\n",
    "    environment:\n",
    "      KAFKA_BROKER_ID: 1\n",
    "      KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181\n",
    "      KAFKA_LISTENERS: PLAINTEXT://0.0.0.0:9092\n",
    "      KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://kafka:9092\n",
    "\n",
    "  mysql:\n",
    "    image: mysql:8.0\n",
    "    container_name: mysql\n",
    "    environment:\n",
    "      MYSQL_ROOT_PASSWORD: root\n",
    "      MYSQL_DATABASE: source_db\n",
    "      MYSQL_USER: user\n",
    "      MYSQL_PASSWORD: pass\n",
    "    ports:\n",
    "      - \"3306:3306\"\n",
    "    volumes:\n",
    "      - ./scripts/init_db.sql:/docker-entrypoint-initdb.d/init_db.sql\n",
    "\n",
    "  postgres:\n",
    "    image: postgres:14\n",
    "    container_name: postgres\n",
    "    environment:\n",
    "      POSTGRES_DB: processed_db\n",
    "      POSTGRES_USER: user\n",
    "      POSTGRES_PASSWORD: pass\n",
    "    ports:\n",
    "      - \"5432:5432\"\n",
    "\n",
    "  minio:\n",
    "    image: minio/minio:latest\n",
    "    container_name: minio\n",
    "    environment:\n",
    "      MINIO_ROOT_USER: minio\n",
    "      MINIO_ROOT_PASSWORD: minio123\n",
    "    command: server /data --console-address \":9001\"\n",
    "    ports:\n",
    "      - \"9000:9000\"\n",
    "      - \"9001:9001\"\n",
    "    volumes:\n",
    "      - minio_data:/data\n",
    "\n",
    "  airflow-webserver:\n",
    "    build: ./airflow\n",
    "    container_name: airflow-webserver\n",
    "    ports:\n",
    "      - \"8080:8080\"\n",
    "    depends_on:\n",
    "      - postgres\n",
    "      - kafka\n",
    "      - mysql\n",
    "      - minio\n",
    "    environment:\n",
    "      AIRFLOW__CORE__EXECUTOR: LocalExecutor\n",
    "      AIRFLOW__CORE__LOAD_EXAMPLES: \"False\"\n",
    "    volumes:\n",
    "      - ./airflow/dags:/opt/airflow/dags\n",
    "      - ./airflow/logs:/opt/airflow/logs\n",
    "      - ./airflow/plugins:/opt/airflow/plugins\n",
    "      - ./great_expectations:/opt/airflow/great_expectations\n",
    "\n",
    "  spark:\n",
    "    build: ./spark\n",
    "    container_name: spark\n",
    "    depends_on:\n",
    "      - kafka\n",
    "      - mysql\n",
    "      - minio\n",
    "      - postgres\n",
    "    volumes:\n",
    "      - ./spark:/opt/spark_jobs\n",
    "      - ./great_expectations:/opt/spark_jobs/great_expectations\n",
    "\n",
    "volumes:\n",
    "  minio_data:\n",
    "'''\n",
    "\n",
    "print(docker_compose_yaml)"
   ]
  },
  {
   "cell_type": "code",
   "id": "airflow-dag-cell",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Airflow Batch Ingestion DAG (airflow/dags/batch_ingestion_dag.py)\n",
    "from airflow import DAG\n",
    "from airflow.operators.python_operator import PythonOperator\n",
    "from airflow.providers.mysql.hooks.mysql import MySqlHook\n",
    "from airflow.providers.postgres.hooks.postgres import PostgresHook\n",
    "from airflow.operators.bash_operator import BashOperator\n",
    "from datetime import datetime\n",
    "import pandas as pd\n",
    "import boto3\n",
    "import logging\n",
    "import great_expectations as ge\n",
    "\n",
    "default_args = {\n",
    "    'owner': 'airflow',\n",
    "    'start_date': datetime(2023, 1, 1),\n",
    "    'retries': 1\n",
    "}\n",
    "\n",
    "def extract_data_from_mysql(**kwargs):\n",
    "    logging.info(\"Extracting data from MySQL...\")\n",
    "    mysql_hook = MySqlHook(mysql_conn_id=\"mysql_default\")\n",
    "    df = mysql_hook.get_pandas_df(\"SELECT * FROM orders;\")\n",
    "    df.to_csv(\"/tmp/orders.csv\", index=False)\n",
    "    logging.info(f\"Extracted {len(df)} records.\")\n",
    "\n",
    "def validate_data_with_ge(**kwargs):\n",
    "    logging.info(\"Validating data with Great Expectations...\")\n",
    "    df = pd.read_csv(\"/tmp/orders.csv\")\n",
    "    ge_df = ge.from_pandas(df)\n",
    "    result = ge_df.expect_column_values_to_not_be_null(\"order_id\")\n",
    "    if not result[\"success\"]:\n",
    "        raise ValueError(\"Validation failed: 'order_id' contains null values\")\n",
    "    logging.info(\"Data validation passed.\")\n",
    "\n",
    "def load_to_minio(**kwargs):\n",
    "    logging.info(\"Uploading raw data to MinIO...\")\n",
    "    s3 = boto3.client(\n",
    "        's3',\n",
    "        endpoint_url='http://minio:9000',\n",
    "        aws_access_key_id='minio',\n",
    "        aws_secret_access_key='minio123',\n",
    "        region_name='us-east-1'\n",
    "    )\n",
    "    bucket_name = \"raw-data\"\n",
    "    try:\n",
    "        s3.create_bucket(Bucket=bucket_name)\n",
    "    except Exception as e:\n",
    "        logging.info(f\"Bucket may already exist: {e}\")\n",
    "    s3.upload_file(\"/tmp/orders.csv\", bucket_name, \"orders/orders.csv\")\n",
    "    logging.info(\"File uploaded to MinIO.\")\n",
    "\n",
    "def load_to_postgres(**kwargs):\n",
    "    logging.info(\"Loading transformed data to Postgres...\")\n",
    "    pg_hook = PostgresHook(postgres_conn_id=\"postgres_default\")\n",
    "    df = pd.read_csv(\"/tmp/transformed_orders.csv\")\n",
    "    pg_hook.run(\"TRUNCATE TABLE orders_transformed;\")\n",
    "    for _, row in df.iterrows():\n",
    "        insert_sql = \"\"\"\n",
    "        INSERT INTO orders_transformed(order_id, customer_id, amount, processed_timestamp)\n",
    "        VALUES (%s, %s, %s, %s)\n",
    "        \"\"\"\n",
    "        pg_hook.run(insert_sql, parameters=(row['order_id'], row['customer_id'], row['amount'], row['processed_timestamp']))\n",
    "    logging.info(\"Data loaded into Postgres.\")\n",
    "\n",
    "with DAG(\n",
    "    dag_id='batch_ingestion_dag',\n",
    "    default_args=default_args,\n",
    "    schedule_interval='@daily',\n",
    "    catchup=False\n",
    ") as dag:\n",
    "    extract_task = PythonOperator(\n",
    "        task_id='extract_mysql',\n",
    "        python_callable=extract_data_from_mysql\n",
    "    )\n",
    "\n",
    "    validate_task = PythonOperator(\n",
    "        task_id='validate_data',\n",
    "        python_callable=validate_data_with_ge\n",
    "    )\n",
    "\n",
    "    load_to_minio_task = PythonOperator(\n",
    "        task_id='load_to_minio',\n",
    "        python_callable=load_to_minio\n",
    "    )\n",
    "\n",
    "    spark_transform_task = BashOperator(\n",
    "        task_id='spark_transform',\n",
    "        bash_command='spark-submit --master local[2] /opt/spark_jobs/spark_batch_job.py'\n",
    "    )\n",
    "\n",
    "    load_postgres_task = PythonOperator(\n",
    "        task_id='load_to_postgres',\n",
    "        python_callable=load_to_postgres\n",
    "    )\n",
    "\n",
    "    extract_task >> validate_task >> load_to_minio_task >> spark_transform_task >> load_postgres_task\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "spark-batch-cell",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Spark Batch Job (spark/spark_batch_job.py)\n",
    "import os\n",
    "import logging\n",
    "import glob\n",
    "import shutil\n",
    "import sys\n",
    "import great_expectations as ge\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql.functions import current_timestamp, col, sum as spark_sum, count\n",
    "\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(asctime)s - %(levelname)s - %(message)s\")\n",
    "\n",
    "MINIO_ENDPOINT = \"http://minio:9000\"\n",
    "MINIO_ACCESS_KEY = \"minio\"\n",
    "MINIO_SECRET_KEY = \"minio123\"\n",
    "RAW_DATA_PATH = \"s3a://raw-data/orders/orders.csv\"\n",
    "PROCESSED_DATA_PATH = \"s3a://processed-data/orders_transformed.csv\"\n",
    "LOCAL_OUTPUT_DIR = \"/tmp/transformed_orders\"\n",
    "LOCAL_OUTPUT_FILE = \"/tmp/transformed_orders.csv\"\n",
    "\n",
    "def validate_schema(df):\n",
    "    logging.info(\"Validating schema using Great Expectations...\")\n",
    "    ge_df = ge.from_pandas(df.toPandas())\n",
    "    result_order_id = ge_df.expect_column_values_to_not_be_null(\"order_id\")\n",
    "    if not result_order_id.success:\n",
    "        raise ValueError(\"Validation failed: 'order_id' contains null values\")\n",
    "    result_customer = ge_df.expect_column_values_to_be_between(\"customer_id\", min_value=1)\n",
    "    if not result_customer.success:\n",
    "        raise ValueError(\"Validation failed: 'customer_id' is not positive\")\n",
    "    result_amount = ge_df.expect_column_values_to_be_between(\"amount\", min_value=0.01)\n",
    "    if not result_amount.success:\n",
    "        raise ValueError(\"Validation failed: 'amount' contains zero or negative values\")\n",
    "    logging.info(\"Schema validation passed.\")\n",
    "\n",
    "def main():\n",
    "    try:\n",
    "        spark = SparkSession.builder \\\n",
    "            .appName(\"BatchETL\") \\\n",
    "            .getOrCreate()\n",
    "\n",
    "        spark._jsc.hadoopConfiguration().set(\"fs.s3a.endpoint\", MINIO_ENDPOINT)\n",
    "        spark._jsc.hadoopConfiguration().set(\"fs.s3a.access.key\", MINIO_ACCESS_KEY)\n",
    "        spark._jsc.hadoopConfiguration().set(\"fs.s3a.secret.key\", MINIO_SECRET_KEY)\n",
    "        spark._jsc.hadoopConfiguration().set(\"fs.s3a.path.style.access\", \"true\")\n",
    "\n",
    "        logging.info(f\"Reading raw data from MinIO ({RAW_DATA_PATH})...\")\n",
    "        df = spark.read.option(\"header\", \"true\").csv(RAW_DATA_PATH)\n",
    "\n",
    "        if df.rdd.isEmpty():\n",
    "            raise Exception(\"No data found in the input file.\")\n",
    "\n",
    "        df = df.withColumn(\"order_id\", col(\"order_id\").cast(\"int\")) \\\n",
    "               .withColumn(\"customer_id\", col(\"customer_id\").cast(\"int\")) \\\n",
    "               .withColumn(\"amount\", col(\"amount\").cast(\"double\"))\n",
    "\n",
    "        validate_schema(df)\n",
    "\n",
    "        initial_record_count = df.count()\n",
    "        logging.info(f\"Initial record count: {initial_record_count}\")\n",
    "\n",
    "        df = df.dropDuplicates([\"order_id\"])\n",
    "        df = df.fillna({\"amount\": 0.0})\n",
    "        df_transformed = df.withColumn(\"processed_timestamp\", current_timestamp())\n",
    "\n",
    "        df_aggregated = df_transformed.groupBy(\"customer_id\").agg(\n",
    "            spark_sum(\"amount\").alias(\"total_spent\"),\n",
    "            count(\"*\").alias(\"order_count\")\n",
    "        )\n",
    "\n",
    "        final_record_count = df_transformed.count()\n",
    "        logging.info(f\"Final record count: {final_record_count}\")\n",
    "\n",
    "        logging.info(f\"Writing transformed data to MinIO ({PROCESSED_DATA_PATH})...\")\n",
    "        df_transformed.write.mode(\"overwrite\").option(\"header\", \"true\").csv(PROCESSED_DATA_PATH)\n",
    "\n",
    "        logging.info(f\"Writing transformed data to local directory ({LOCAL_OUTPUT_DIR})...\")\n",
    "        df_transformed.coalesce(1).write.mode(\"overwrite\").option(\"header\", \"true\").csv(LOCAL_OUTPUT_DIR)\n",
    "\n",
    "        csv_files = glob.glob(f\"{LOCAL_OUTPUT_DIR}/part-*.csv\")\n",
    "        if csv_files:\n",
    "            shutil.move(csv_files[0], LOCAL_OUTPUT_FILE)\n",
    "            logging.info(f\"Transformed file successfully written to {LOCAL_OUTPUT_FILE}\")\n",
    "        else:\n",
    "            logging.warning(\"No CSV files found after transformation.\")\n",
    "\n",
    "        spark.stop()\n",
    "        logging.info(\"Batch ETL process completed successfully.\")\n",
    "\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Batch processing failed: {str(e)}\")\n",
    "        sys.exit(1)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "spark-stream-cell",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Spark Streaming Job (spark/spark_streaming_job.py)\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql.functions import from_json, col\n",
    "from pyspark.sql.types import StructType, StructField, StringType, LongType, DoubleType\n",
    "import logging\n",
    "import great_expectations as ge\n",
    "\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(asctime)s - %(levelname)s - %(message)s\")\n",
    "\n",
    "KAFKA_BROKER = \"kafka:9092\"\n",
    "KAFKA_TOPIC = \"sensor_readings\"\n",
    "MINIO_ENDPOINT = \"http://minio:9000\"\n",
    "MINIO_ACCESS_KEY = \"minio\"\n",
    "MINIO_SECRET_KEY = \"minio123\"\n",
    "RAW_DATA_PATH = \"s3a://raw-data/streaming_raw/\"\n",
    "ANOMALY_DATA_PATH = \"s3a://processed-data/streaming_anomalies/\"\n",
    "\n",
    "POSTGRES_TABLE = \"anomalies_stream\"\n",
    "\n",
    "schema = StructType([\n",
    "    StructField(\"event_id\", StringType(), False),\n",
    "    StructField(\"timestamp\", LongType(), False),\n",
    "    StructField(\"device_id\", LongType(), False),\n",
    "    StructField(\"reading_value\", DoubleType(), False)\n",
    "])\n",
    "\n",
    "def validate_schema(df):\n",
    "    logging.info(\"Validating streaming data schema with Great Expectations...\")\n",
    "    ge_df = ge.from_pandas(df.toPandas())\n",
    "    result_event_id = ge_df.expect_column_values_to_not_be_null(\"event_id\")\n",
    "    if not result_event_id.success:\n",
    "        raise ValueError(\"Validation failed: 'event_id' contains null values\")\n",
    "    result_timestamp = ge_df.expect_column_values_to_be_between(\"timestamp\", min_value=1)\n",
    "    if not result_timestamp.success:\n",
    "        raise ValueError(\"Validation failed: 'timestamp' contains invalid values\")\n",
    "    result_device = ge_df.expect_column_values_to_be_between(\"device_id\", min_value=1)\n",
    "    if not result_device.success:\n",
    "        raise ValueError(\"Validation failed: 'device_id' contains invalid values\")\n",
    "    result_reading = ge_df.expect_column_values_to_be_between(\"reading_value\", min_value=0, max_value=100)\n",
    "    if not result_reading.success:\n",
    "        raise ValueError(\"Validation failed: 'reading_value' is out of expected range\")\n",
    "    logging.info(\"Streaming schema validation passed.\")\n",
    "\n",
    "def save_to_postgres(batch_df, batch_id):\n",
    "    logging.info(f\"Writing batch {batch_id} to PostgreSQL table {POSTGRES_TABLE}...\")\n",
    "    batch_df.write \\\n",
    "        .format(\"jdbc\") \\\n",
    "        .option(\"url\", \"jdbc:postgresql://postgres:5432/processed_db\") \\\n",
    "        .option(\"dbtable\", POSTGRES_TABLE) \\\n",
    "        .option(\"user\", \"user\") \\\n",
    "        .option(\"password\", \"pass\") \\\n",
    "        .option(\"driver\", \"org.postgresql.Driver\") \\\n",
    "        .mode(\"append\") \\\n",
    "        .save()\n",
    "    logging.info(f\"Batch {batch_id} written to PostgreSQL.\")\n",
    "\n",
    "def main():\n",
    "    spark = SparkSession.builder \\\n",
    "        .appName(\"StreamingETL\") \\\n",
    "        .config(\"spark.jars.packages\", \"org.postgresql:postgresql:42.2.18\") \\\n",
    "        .getOrCreate()\n",
    "\n",
    "    spark._jsc.hadoopConfiguration().set(\"fs.s3a.endpoint\", MINIO_ENDPOINT)\n",
    "    spark._jsc.hadoopConfiguration().set(\"fs.s3a.access.key\", MINIO_ACCESS_KEY)\n",
    "    spark._jsc.hadoopConfiguration().set(\"fs.s3a.secret.key\", MINIO_SECRET_KEY)\n",
    "    spark._jsc.hadoopConfiguration().set(\"fs.s3a.path.style.access\", \"true\")\n",
    "\n",
    "    logging.info(f\"Starting Spark Structured Streaming from Kafka topic: {KAFKA_TOPIC}\")\n",
    "    df_raw = spark.readStream \\\n",
    "        .format(\"kafka\") \\\n",
    "        .option(\"kafka.bootstrap.servers\", KAFKA_BROKER) \\\n",
    "        .option(\"subscribe\", KAFKA_TOPIC) \\\n",
    "        .option(\"startingOffsets\", \"latest\") \\\n",
    "        .load()\n",
    "\n",
    "    df_parsed = df_raw.select(\n",
    "        from_json(col(\"value\").cast(\"string\"), schema).alias(\"json_data\")\n",
    "    ).select(\"json_data.*\")\n",
    "\n",
    "    df_clean = df_parsed.filter(\n",
    "        col(\"event_id\").isNotNull() & \n",
    "        col(\"timestamp\").isNotNull() & \n",
    "        col(\"device_id\").isNotNull() & \n",
    "        col(\"reading_value\").isNotNull()\n",
    "    )\n",
    "\n",
    "    validate_schema(df_clean)\n",
    "\n",
    "    df_anomalies = df_clean.filter(col(\"reading_value\") > 70.0)\n",
    "\n",
    "    # Write raw data to MinIO\n",
    "    df_clean.writeStream \\\n",
    "        .format(\"parquet\") \\\n",
    "        .option(\"checkpointLocation\", \"/tmp/spark-checkpoints/raw\") \\\n",
    "        .option(\"path\", RAW_DATA_PATH) \\\n",
    "        .outputMode(\"append\") \\\n",
    "        .start()\n",
    "\n",
    "    # Write anomalies to MinIO\n",
    "    df_anomalies.writeStream \\\n",
    "        .format(\"parquet\") \\\n",
    "        .option(\"checkpointLocation\", \"/tmp/spark-checkpoints/anomalies\") \\\n",
    "        .option(\"path\", ANOMALY_DATA_PATH) \\\n",
    "        .outputMode(\"append\") \\\n",
    "        .start()\n",
    "\n",
    "    # Write anomalies to PostgreSQL using foreachBatch\n",
    "    df_anomalies.writeStream \\\n",
    "        .foreachBatch(lambda batch_df, batch_id: save_to_postgres(batch_df, batch_id)) \\\n",
    "        .outputMode(\"append\") \\\n",
    "        .start()\n",
    "\n",
    "    logging.info(\"Streaming job started. Processing events in real time.\")\n",
    "    spark.streams.awaitAnyTermination()\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "kafka-producer-cell",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Kafka Producer (kafka/producer.py)\n",
    "import json\n",
    "import time\n",
    "import random\n",
    "import uuid\n",
    "import logging\n",
    "from kafka import KafkaProducer\n",
    "import os\n",
    "\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(asctime)s - %(levelname)s - %(message)s\")\n",
    "\n",
    "KAFKA_BROKER = os.getenv(\"KAFKA_BROKER\", \"kafka:9092\")\n",
    "KAFKA_TOPIC = os.getenv(\"KAFKA_TOPIC\", \"sensor_readings\")\n",
    "MESSAGE_FREQUENCY = float(os.getenv(\"MESSAGE_FREQUENCY\", 1))\n",
    "BATCH_SIZE = int(os.getenv(\"BATCH_SIZE\", 10))\n",
    "ACKS_MODE = os.getenv(\"ACKS_MODE\", \"all\")\n",
    "\n",
    "producer = KafkaProducer(\n",
    "    bootstrap_servers=[KAFKA_BROKER],\n",
    "    value_serializer=lambda v: json.dumps(v).encode('utf-8'),\n",
    "    acks=ACKS_MODE\n",
    ")\n",
    "\n",
    "def create_kafka_topic(topic_name):\n",
    "    from kafka.admin import KafkaAdminClient, NewTopic\n",
    "    from kafka.errors import KafkaError\n",
    "    try:\n",
    "        admin_client = KafkaAdminClient(bootstrap_servers=KAFKA_BROKER)\n",
    "        existing_topics = admin_client.list_topics()\n",
    "        if topic_name not in existing_topics:\n",
    "            topic = NewTopic(name=topic_name, num_partitions=3, replication_factor=1)\n",
    "            admin_client.create_topics([topic])\n",
    "            logging.info(f\"Kafka topic '{topic_name}' created successfully.\")\n",
    "        else:\n",
    "            logging.info(f\"Kafka topic '{topic_name}' already exists.\")\n",
    "    except KafkaError as e:\n",
    "        logging.error(f\"Failed to create Kafka topic '{topic_name}': {e}\")\n",
    "\n",
    "def generate_event():\n",
    "    return {\n",
    "        \"event_id\": str(uuid.uuid4()),\n",
    "        \"timestamp\": int(time.time()),\n",
    "        \"device_id\": random.randint(1000, 9999),\n",
    "        \"reading_value\": round(random.uniform(20.0, 80.0), 2)\n",
    "    }\n",
    "\n",
    "def produce_messages():\n",
    "    logging.info(f\"Starting Kafka Producer. Sending messages to topic: {KAFKA_TOPIC}\")\n",
    "    create_kafka_topic(KAFKA_TOPIC)\n",
    "    batch = []\n",
    "    message_count = 0\n",
    "    try:\n",
    "        while True:\n",
    "            event = generate_event()\n",
    "            batch.append(event)\n",
    "            if len(batch) >= BATCH_SIZE:\n",
    "                for msg in batch:\n",
    "                    producer.send(KAFKA_TOPIC, msg)\n",
    "                producer.flush()\n",
    "                logging.info(f\"Sent batch of {len(batch)} messages to Kafka topic: {KAFKA_TOPIC}\")\n",
    "                message_count += len(batch)\n",
    "                batch.clear()\n",
    "            time.sleep(MESSAGE_FREQUENCY)\n",
    "    except KeyboardInterrupt:\n",
    "        logging.warning(\"Kafka Producer stopped manually.\")\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Kafka producer error: {e}\")\n",
    "    finally:\n",
    "        logging.info(f\"Kafka Producer shutting down. Total messages sent: {message_count}\")\n",
    "        producer.close()\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    produce_messages()\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "governance-cell",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Governance / Data Lineage Registration (governance/atlas_stub.py)\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "import requests\n",
    "\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(asctime)s - %(levelname)s - %(message)s\")\n",
    "\n",
    "ATLAS_API_URL = os.getenv(\"ATLAS_API_URL\", \"http://atlas:21000/api/atlas/v2/lineage\")\n",
    "ATLAS_USERNAME = os.getenv(\"ATLAS_USERNAME\", \"admin\")\n",
    "ATLAS_PASSWORD = os.getenv(\"ATLAS_PASSWORD\", \"admin\")\n",
    "\n",
    "HEADERS = { \"Content-Type\": \"application/json\" }\n",
    "\n",
    "def check_dataset_exists(dataset_name):\n",
    "    dataset_api_url = f\"{ATLAS_API_URL}/entities?type=Dataset&name={dataset_name}\"\n",
    "    try:\n",
    "        response = requests.get(dataset_api_url, auth=(ATLAS_USERNAME, ATLAS_PASSWORD), headers=HEADERS)\n",
    "        if response.status_code == 200:\n",
    "            data = response.json()\n",
    "            if \"entities\" in data and len(data[\"entities\"]) > 0:\n",
    "                logging.info(f\"Dataset '{dataset_name}' exists in Apache Atlas.\")\n",
    "                return True\n",
    "            else:\n",
    "                logging.warning(f\"Dataset '{dataset_name}' does not exist in Apache Atlas.\")\n",
    "                return False\n",
    "        else:\n",
    "            logging.error(f\"Failed to check dataset existence: {response.status_code} - {response.text}\")\n",
    "            return False\n",
    "    except requests.RequestException as e:\n",
    "        logging.error(f\"Error while checking dataset existence: {str(e)}\")\n",
    "        return False\n",
    "\n",
    "def register_dataset_lineage(source_name, target_name, extra_info=None):\n",
    "    if not check_dataset_exists(source_name) or not check_dataset_exists(target_name):\n",
    "        logging.error(\"Cannot register lineage: One or both datasets do not exist.\")\n",
    "        return\n",
    "    lineage_payload = {\n",
    "        \"guidEntityMap\": {},\n",
    "        \"relations\": [\n",
    "            {\n",
    "                \"typeName\": \"Process\",\n",
    "                \"fromEntityId\": source_name,\n",
    "                \"toEntityId\": target_name,\n",
    "                \"relationshipAttributes\": extra_info or {}\n",
    "            }\n",
    "        ]\n",
    "    }\n",
    "    try:\n",
    "        response = requests.post(\n",
    "            f\"{ATLAS_API_URL}/entities\",\n",
    "            auth=(ATLAS_USERNAME, ATLAS_PASSWORD),\n",
    "            headers=HEADERS,\n",
    "            data=json.dumps(lineage_payload)\n",
    "        )\n",
    "        if response.status_code in [200, 201]:\n",
    "            logging.info(f\"Successfully registered lineage from '{source_name}' to '{target_name}'\")\n",
    "        else:\n",
    "            logging.error(f\"Failed to register lineage: {response.status_code} - {response.text}\")\n",
    "    except requests.RequestException as e:\n",
    "        logging.error(f\"Error while registering dataset lineage: {str(e)}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    register_dataset_lineage(\n",
    "        \"mysql.orders\",\n",
    "        \"minio.raw-data.orders\",\n",
    "        {\"job\": \"batch_ingestion_dag\", \"transformation\": \"cleaning, enrichment\"}\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "monitoring-cell",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Monitoring Setup for Prometheus and Grafana (monitoring/monitoring.py)\n",
    "import os\n",
    "import json\n",
    "import subprocess\n",
    "import time\n",
    "import logging\n",
    "import requests\n",
    "\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(asctime)s - %(levelname)s - %(message)s\")\n",
    "\n",
    "PROMETHEUS_CONFIG_PATH = os.getenv(\"PROMETHEUS_CONFIG_PATH\", \"/etc/prometheus/prometheus.yml\")\n",
    "PROMETHEUS_PORT = os.getenv(\"PROMETHEUS_PORT\", \"9090\")\n",
    "GRAFANA_PORT = os.getenv(\"GRAFANA_PORT\", \"3000\")\n",
    "GRAFANA_API_URL = f\"http://localhost:{GRAFANA_PORT}/api\"\n",
    "GRAFANA_ADMIN_USER = os.getenv(\"GRAFANA_ADMIN_USER\", \"admin\")\n",
    "GRAFANA_ADMIN_PASS = os.getenv(\"GRAFANA_ADMIN_PASS\", \"admin\")\n",
    "DASHBOARDS_PATH = os.getenv(\"DASHBOARDS_PATH\", \"./monitoring/grafana_dashboards\")\n",
    "\n",
    "def start_prometheus():\n",
    "    if not os.path.exists(PROMETHEUS_CONFIG_PATH):\n",
    "        logging.error(f\"Prometheus config file not found: {PROMETHEUS_CONFIG_PATH}\")\n",
    "        return\n",
    "    logging.info(\"Starting Prometheus...\")\n",
    "    try:\n",
    "        subprocess.Popen([\"prometheus\", \"--config.file\", PROMETHEUS_CONFIG_PATH])\n",
    "        logging.info(f\"Prometheus started on port {PROMETHEUS_PORT}\")\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Failed to start Prometheus: {e}\")\n",
    "\n",
    "def start_grafana():\n",
    "    logging.info(\"Starting Grafana...\")\n",
    "    try:\n",
    "        subprocess.Popen([\"grafana-server\"])\n",
    "        logging.info(f\"Grafana started on port {GRAFANA_PORT}\")\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Failed to start Grafana: {e}\")\n",
    "\n",
    "def wait_for_grafana():\n",
    "    logging.info(\"Waiting for Grafana to be ready...\")\n",
    "    for _ in range(30):\n",
    "        try:\n",
    "            response = requests.get(f\"{GRAFANA_API_URL}/health\")\n",
    "            if response.status_code == 200:\n",
    "                logging.info(\"Grafana is ready.\")\n",
    "                return True\n",
    "        except requests.ConnectionError:\n",
    "            pass\n",
    "        time.sleep(1)\n",
    "    logging.error(\"Grafana did not start in time.\")\n",
    "    return False\n",
    "\n",
    "def create_grafana_datasource():\n",
    "    logging.info(\"Creating Grafana Prometheus datasource...\")\n",
    "    datasource_payload = {\n",
    "        \"name\": \"Prometheus\",\n",
    "        \"type\": \"prometheus\",\n",
    "        \"url\": f\"http://localhost:{PROMETHEUS_PORT}\",\n",
    "        \"access\": \"proxy\",\n",
    "        \"basicAuth\": False\n",
    "    }\n",
    "    response = requests.post(\n",
    "        f\"{GRAFANA_API_URL}/datasources\",\n",
    "        auth=(GRAFANA_ADMIN_USER, GRAFANA_ADMIN_PASS),\n",
    "        headers={\"Content-Type\": \"application/json\"},\n",
    "        data=json.dumps(datasource_payload)\n",
    "    )\n",
    "    if response.status_code in [200, 201]:\n",
    "        logging.info(\"Grafana Prometheus datasource created successfully.\")\n",
    "    else:\n",
    "        logging.error(f\"Failed to create Grafana datasource: {response.text}\")\n",
    "\n",
    "def import_grafana_dashboards():\n",
    "    if not os.path.exists(DASHBOARDS_PATH):\n",
    "        logging.warning(f\"Dashboard directory not found: {DASHBOARDS_PATH}\")\n",
    "        return\n",
    "    for dashboard_file in os.listdir(DASHBOARDS_PATH):\n",
    "        if dashboard_file.endswith(\".json\"):\n",
    "            dashboard_path = os.path.join(DASHBOARDS_PATH, dashboard_file)\n",
    "            with open(dashboard_path, \"r\") as f:\n",
    "                dashboard_data = json.load(f)\n",
    "                dashboard_payload = {\n",
    "                    \"dashboard\": dashboard_data,\n",
    "                    \"overwrite\": True\n",
    "                }\n",
    "                response = requests.post(\n",
    "                    f\"{GRAFANA_API_URL}/dashboards/db\",\n",
    "                    auth=(GRAFANA_ADMIN_USER, GRAFANA_ADMIN_PASS),\n",
    "                    headers={\"Content-Type\": \"application/json\"},\n",
    "                    data=json.dumps(dashboard_payload)\n",
    "                )\n",
    "                if response.status_code in [200, 201]:\n",
    "                    logging.info(f\"Imported Grafana dashboard: {dashboard_file}\")\n",
    "                else:\n",
    "                    logging.error(f\"Failed to import {dashboard_file}: {response.text}\")\n",
    "\n",
    "def main():\n",
    "    start_prometheus()\n",
    "    start_grafana()\n",
    "    if wait_for_grafana():\n",
    "        create_grafana_datasource()\n",
    "        import_grafana_dashboards()\n",
    "    logging.info(\"Monitoring setup complete.\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "mlflow-cell",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# MLflow Tracking (ml/mlflow_tracking.py)\n",
    "import mlflow\n",
    "import random\n",
    "\n",
    "def train_and_log_model():\n",
    "    mlflow.set_experiment(\"my_experiment\")\n",
    "    with mlflow.start_run():\n",
    "        param_value = random.randint(1, 100)\n",
    "        mlflow.log_param(\"param1\", param_value)\n",
    "        accuracy = round(random.uniform(0.8, 0.99), 4)\n",
    "        mlflow.log_metric(\"accuracy\", accuracy)\n",
    "        print(f\"Tracked model run with param1={param_value}, accuracy={accuracy}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    train_and_log_model()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "setup-md",
   "metadata": {},
   "source": [
    "## Setup Instructions\n",
    "\n",
    "1. **Clone the Repository**\n",
    "   ```bash\n",
    "   git clone https://github.com/your-repo/end-to-end-pipeline.git\n",
    "   cd end-to-end-pipeline\n",
    "   ```\n",
    "\n",
    "2. **Start the Pipeline Stack**\n",
    "   Use Docker Compose to launch all services:\n",
    "   ```bash\n",
    "   docker-compose up --build\n",
    "   ```\n",
    "   This starts MySQL, PostgreSQL, Kafka, MinIO, Airflow, Spark, Prometheus, and Grafana.\n",
    "\n",
    "3. **Access the Services**\n",
    "   - **Airflow UI:** [http://localhost:8080](http://localhost:8080) (set up connections as needed)\n",
    "   - **MinIO Console:** [http://localhost:9001](http://localhost:9001) (User: `minio`, Password: `minio123`)\n",
    "   - **Kafka:** Port `9092`\n",
    "   - **Prometheus:** [http://localhost:9090](http://localhost:9090)\n",
    "   - **Grafana:** [http://localhost:3000](http://localhost:3000) (Default login: `admin/admin`)\n",
    "\n",
    "4. **Run Batch Pipeline**\n",
    "   - In the Airflow UI, enable and trigger the `batch_ingestion_dag`.\n",
    "\n",
    "5. **Run Streaming Pipeline**\n",
    "   - Start the Kafka producer:\n",
    "     ```bash\n",
    "     docker-compose exec kafka python /opt/spark_jobs/../kafka/producer.py\n",
    "     ```\n",
    "   - Run the Spark streaming job:\n",
    "     ```bash\n",
    "     docker-compose exec spark spark-submit --master local[2] /opt/spark_jobs/spark_streaming_job.py\n",
    "     ```\n",
    "\n",
    "6. **Monitoring & Governance**\n",
    "   - Run the `monitoring.py` script to set up Prometheus and Grafana dashboards.\n",
    "   - Run the `governance/atlas_stub.py` script to register dataset lineage.\n",
    "\n",
    "7. **ML & Feature Store**\n",
    "   - Run `ml/mlflow_tracking.py` to simulate model training and tracking.\n",
    "   - Use `ml/feature_store_stub.py` (if implemented) for feature store integration.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "examples-md",
   "metadata": {},
   "source": [
    "## Example Applications\n",
    "\n",
    "- **E-Commerce & Retail:**\n",
    "  - Real-time recommendations and fraud detection using streaming data.\n",
    "  - Batch processing of historical sales data for demand forecasting.\n",
    "\n",
    "- **Financial Services:**\n",
    "  - Risk analysis and trade surveillance using aggregated transaction data.\n",
    "\n",
    "- **Healthcare:**\n",
    "  - Patient monitoring using real-time sensor data and predictive analytics.\n",
    "\n",
    "- **IoT & Manufacturing:**\n",
    "  - Predictive maintenance by monitoring equipment sensor data in real time.\n",
    "\n",
    "- **Media & Social Networks:**\n",
    "  - Sentiment analysis and ad fraud detection using streaming social media data.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "conclusion-md",
   "metadata": {},
   "source": [
    "## Conclusion & Further Steps\n",
    "\n",
    "This notebook has demonstrated an end-to-end data pipeline including batch and streaming processing, data quality checks, monitoring, governance, and ML tracking.\n",
    "\n",
    "### Next Steps:\n",
    "\n",
    "- Customize the transformation logic in Spark jobs.\n",
    "- Adjust Airflow DAG schedules and dependencies to match your business needs.\n",
    "- Scale the services in production using managed clusters or cloud services.\n",
    "- Integrate real-world governance and ML systems by replacing stub implementations with production APIs.\n",
    "\n",
    "Happy Coding!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "version": "3.x"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
