examples/openlineage/notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7fb27b941602401d91542211134fc71a",
   "metadata": {},
   "source": [
    "Licensed to the Apache Software Foundation (ASF) under one\nor more contributor license agreements.  See the NOTICE file\ndistributed with this work for additional information\nregarding copyright ownership.  The ASF licenses this file\nto you under the Apache License, Version 2.0 (the\n\"License\"); you may not use this file except in compliance\nwith the License.  You may obtain a copy of the License at\n\n  http://www.apache.org/licenses/LICENSE-2.0\n\nUnless required by applicable law or agreed to in writing,\nsoftware distributed under the License is distributed on an\n\"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\nKIND, either express or implied.  See the License for the\nspecific language governing permissions and limitations\nunder the License."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc0e298d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this cell to install dependencies\n",
    "%pip install apache-hamilton[openlineage,visualization] sqlalchemy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f937d835",
   "metadata": {},
   "source": [
    "# OpenLineage example pipeline [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dagworks-inc/hamilton/blob/main/examples/openlineage/notebook.ipynb) [![GitHub badge](https://img.shields.io/badge/github-view_source-2b3137?logo=github)](https://github.com/apache/hamilton/blob/main/examples/openlineage/notebook.ipynb)\n",
    "\n",
    "This pipeline reads a query joining `orders` and `customers` from a sales database, aggregates daily revenue in Python, and writes `daily_revenue` to a warehouse database. The pipeline does not import OpenLineage and does not need to know about it: the built-in SQL materializers record the datasource they used, and the OpenLineage adapter names the physical tables from that metadata.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab02bddb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext hamilton.plugins.jupyter_magic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ccc7699",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-06T17:30:38.628295Z",
     "start_time": "2024-09-06T17:30:38.182382Z"
    }
   },
   "outputs": [],
   "source": [
    "%%cell_to_module pipeline --display\n",
    "\n",
    "\"\"\"Revenue reporting with built-in SQL materializers.\n",
    "\n",
    "- order_lines reads a query joining orders and customers from the sales database.\n",
    "- daily_revenue aggregates in Python.\n",
    "- revenue_report writes daily_revenue into the warehouse database.\n",
    "\n",
    "No custom loader and no hand-written lineage metadata: the SQL materializers record the\n",
    "datasource they used, and the OpenLineage adapter names the physical tables from it.\n",
    "\"\"\"\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "from hamilton.function_modifiers import load_from, save_to, source, value\n",
    "\n",
    "REVENUE_QUERY = \"\"\"\n",
    "-- paid order lines with the customer's country\n",
    "WITH paid AS (SELECT * FROM orders WHERE status = 'paid')\n",
    "SELECT p.order_date, c.country, p.amount\n",
    "FROM paid p\n",
    "JOIN customers c ON p.customer_id = c.id\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "@load_from.sql(query_or_table=value(REVENUE_QUERY), db_connection=source(\"sales_db\"))\n",
    "def order_lines(df: pd.DataFrame) -> pd.DataFrame:\n",
    "    return df\n",
    "\n",
    "\n",
    "def daily_revenue(order_lines: pd.DataFrame) -> pd.DataFrame:\n",
    "    return order_lines.groupby([\"order_date\", \"country\"], as_index=False)[\"amount\"].sum()\n",
    "\n",
    "\n",
    "@save_to.sql(\n",
    "    table_name=value(\"daily_revenue\"),\n",
    "    db_connection=source(\"warehouse_db\"),\n",
    "    if_exists=value(\"replace\"),\n",
    "    index=value(False),\n",
    "    output_name_=\"saved_revenue\",\n",
    ")\n",
    "def revenue_report(daily_revenue: pd.DataFrame) -> pd.DataFrame:\n",
    "    return daily_revenue\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50a7f746",
   "metadata": {},
   "source": [
    "# Seed a small sales database and create the OpenLineage client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7dcecc8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-06T17:31:52.088986Z",
     "start_time": "2024-09-06T17:31:51.984490Z"
    }
   },
   "outputs": [],
   "source": [
    "import sqlite3\n",
    "from pathlib import Path\n",
    "\n",
    "import pandas as pd\n",
    "from openlineage.client import OpenLineageClient\n",
    "from openlineage.client.transport.file import FileConfig, FileTransport\n",
    "from sqlalchemy import create_engine\n",
    "\n",
    "sales_db = sqlite3.connect(\"sales.db\")\n",
    "pd.DataFrame(\n",
    "    {\n",
    "        \"customer_id\": [1, 1, 2, 2],\n",
    "        \"order_date\": [\"2026-09-01\", \"2026-09-01\", \"2026-09-01\", \"2026-09-02\"],\n",
    "        \"amount\": [10.0, 5.0, 7.5, 3.0],\n",
    "        \"status\": [\"paid\", \"paid\", \"paid\", \"open\"],\n",
    "    }\n",
    ").to_sql(\"orders\", sales_db, index=False, if_exists=\"replace\")\n",
    "pd.DataFrame({\"id\": [1, 2], \"country\": [\"NL\", \"DE\"]}).to_sql(\n",
    "    \"customers\", sales_db, index=False, if_exists=\"replace\"\n",
    ")\n",
    "warehouse_db = create_engine(f\"sqlite:///{Path('warehouse.db').resolve()}\")\n",
    "\n",
    "# without a running OpenLineage server, the FileTransport writes events to a file\n",
    "Path(\"pipeline.json\").unlink(missing_ok=True)\n",
    "client = OpenLineageClient(\n",
    "    transport=FileTransport(FileConfig(log_file_path=\"pipeline.json\", append=True))\n",
    ")\n",
    "# with a running server, e.g. marquez: client = OpenLineageClient(url=\"http://localhost:5000\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a49f34a",
   "metadata": {},
   "source": [
    "# Run the dataflow with the OpenLineage adapter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6db87906",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-06T17:32:31.294469Z",
     "start_time": "2024-09-06T17:32:30.852466Z"
    }
   },
   "outputs": [],
   "source": [
    "from hamilton import driver\n",
    "from hamilton.plugins import h_openlineage\n",
    "\n",
    "adapter = h_openlineage.OpenLineageAdapter(\n",
    "    client, \"demo_namespace\", \"revenue_job\", sql_dataset_identity=\"datasource\"\n",
    ")\n",
    "dr = driver.Builder().with_modules(pipeline).with_adapters(adapter).build()\n",
    "result = dr.execute([\"saved_revenue\"], inputs={\"sales_db\": sales_db, \"warehouse_db\": warehouse_db})\n",
    "sales_db.close()\n",
    "result[\"saved_revenue\"][\"sql_metadata\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec295725",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-06T17:32:31.534973Z",
     "start_time": "2024-09-06T17:32:31.475171Z"
    }
   },
   "outputs": [],
   "source": [
    "# The datasets are named after the datasource (the SQLite files), not the job namespace\n",
    "import json\n",
    "\n",
    "for line in Path(\"pipeline.json\").read_text().splitlines():\n",
    "    event = json.loads(line)\n",
    "    for kind in (\"inputs\", \"outputs\"):\n",
    "        for dataset in event.get(kind) or []:\n",
    "            print(f\"{event['eventType']} {kind[:-1]}: {dataset['namespace']}  {dataset['name']}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}