Hannah

.dockerignore

@@ -0,0 +1,8 @@
+astro
+.git
+.env
+airflow_settings.yaml
+logs/
+.venv
+airflow.db
+airflow.cfg

dags/exampledag.py

@@ -0,0 +1,98 @@
+"""
+## Astronaut ETL example DAG
+
+This DAG queries the list of astronauts currently in space from the
+Open Notify API and prints each astronaut's name and flying craft.
+
+There are two tasks, one to get the data from the API and save the results,
+and another to print the results. Both tasks are written in Python using
+Airflow's TaskFlow API, which allows you to easily turn Python functions into
+Airflow tasks, and automatically infer dependencies and pass data.
+
+The second task uses dynamic task mapping to create a copy of the task for
+each Astronaut in the list retrieved from the API. This list will change
+depending on how many Astronauts are in space, and the DAG will adjust
+accordingly each time it runs.
+
+For more explanation and getting started instructions, see our Write your
+first DAG tutorial: https://www.astronomer.io/docs/learn/get-started-with-airflow
+
+![Picture of the ISS](https://www.esa.int/var/esa/storage/images/esa_multimedia/images/2010/02/space_station_over_earth/10293696-3-eng-GB/Space_Station_over_Earth_card_full.jpg)
+"""
+
+from airflow.sdk import Asset, dag, task
+from pendulum import datetime
+import requests
+
+
+# Define the basic parameters of the DAG, like schedule and start_date
+@dag(
+    start_date=datetime(2025, 4, 22),
+    schedule="@daily",
+    doc_md=__doc__,
+    default_args={"owner": "Astro", "retries": 3},
+    tags=["example"],
+)
+def example_astronauts():
+    # Define tasks
+    @task(
+        # Define an asset outlet for the task. This can be used to schedule downstream DAGs when this task has run.
+        outlets=[Asset("current_astronauts")]
+    )  # Define that this task updates the `current_astronauts` Asset
+    def get_astronauts(**context) -> list[dict]:
+        """
+        This task uses the requests library to retrieve a list of Astronauts
+        currently in space. The results are pushed to XCom with a specific key
+        so they can be used in a downstream pipeline. The task returns a list
+        of Astronauts to be used in the next task.
+        """
+        try:
+            r = requests.get("http://api.open-notify.org/astros.json")
+            r.raise_for_status()
+            number_of_people_in_space = r.json()["number"]
+            list_of_people_in_space = r.json()["people"]
+        except Exception:
+            print("API currently not available, using hardcoded data instead.")
+            number_of_people_in_space = 12
+            list_of_people_in_space = [
+                {"craft": "ISS", "name": "Oleg Kononenko"},
+                {"craft": "ISS", "name": "Nikolai Chub"},
+                {"craft": "ISS", "name": "Tracy Caldwell Dyson"},
+                {"craft": "ISS", "name": "Matthew Dominick"},
+                {"craft": "ISS", "name": "Michael Barratt"},
+                {"craft": "ISS", "name": "Jeanette Epps"},
+                {"craft": "ISS", "name": "Alexander Grebenkin"},
+                {"craft": "ISS", "name": "Butch Wilmore"},
+                {"craft": "ISS", "name": "Sunita Williams"},
+                {"craft": "Tiangong", "name": "Li Guangsu"},
+                {"craft": "Tiangong", "name": "Li Cong"},
+                {"craft": "Tiangong", "name": "Ye Guangfu"},
+            ]
+
+        context["ti"].xcom_push(
+            key="number_of_people_in_space", value=number_of_people_in_space
+        )
+        return list_of_people_in_space
+
+    @task
+    def print_astronaut_craft(greeting: str, person_in_space: dict) -> None:
+        """
+        This task creates a print statement with the name of an
+        Astronaut in space and the craft they are flying on from
+        the API request results of the previous task, along with a
+        greeting which is hard-coded in this example.
+        """
+        craft = person_in_space["craft"]
+        name = person_in_space["name"]
+
+        print(f"{name} is currently in space flying on the {craft}! {greeting}")
+
+    # Use dynamic task mapping to run the print_astronaut_craft task for each
+    # Astronaut in space
+    print_astronaut_craft.partial(greeting="Hello! :)").expand(
+        person_in_space=get_astronauts()  # Define dependencies using TaskFlow API syntax
+    )
+
+
+# Instantiate the DAG
+example_astronauts()

dags/taxi_pipeline.py

@@ -10,10 +10,15 @@
 autograder fails while any NotImplementedError remains.
 """
 
+import io
 import os
 from datetime import datetime
 from pathlib import Path
 
+import pandas as pd
+import requests
+from airflow.operators.bash import BashOperator
+from airflow.providers.postgres.hooks.postgres import PostgresHook
 from airflow.sdk import dag, task
 
 # Your per-student schema. AIRFLOW_STUDENT is set in .env for local Astro dev;
@@ -35,30 +40,132 @@ def find_dbt_dir() -> str:
 
 
 DBT_DIR = find_dbt_dir()
+DBT_ENV = {
+    "PG_HOST": "{{ conn.azure_pg.host }}",
+    "PG_USER": "{{ conn.azure_pg.login }}",
+    "PG_PASSWORD": "{{ conn.azure_pg.password }}",
+    "PG_DBNAME": "{{ conn.azure_pg.schema }}",
+    "PG_SCHEMA": SCHEMA,
+}
+
+DBT = (
+    "uvx --python 3.11 "
+    "--from 'dbt-core==1.10.*' "
+    "--with 'dbt-postgres==1.10.*' "
+    "dbt"
+)
 
 
 @dag(
-    # Task 1 (see README): configure the decorator — schedule, start_date,
-    # catchup=False, max_active_runs=1, default_args retries, tags.
+    dag_id="hannahwn_taxi_pipeline",
+    schedule="@monthly",
     start_date=datetime(2024, 1, 1),
+    catchup=False,
+    max_active_runs=1,
+    tags=["week12", "taxi", "student:hannahwn"],
+    default_args={"retries": 2, "retry_delay": 300},
+    # retry transient failures twice
+
 )
 def taxi_pipeline():
     @task
     def ingest_taxi_month() -> int:
-        """Download one month of TLC green-taxi data and load it into
-        ``{SCHEMA}.raw_trips`` idempotently. Return the number of rows.
-
-        Task 2 and Task 3 (see README): derive the partition from the
-        logical date, DELETE-then-append that month, and filter the
-        parquet to the logical month before write (Gotcha #4).
-        """
-        raise NotImplementedError
+        ds = _partition_date()
+        year_month = ds[:7]  # YYYY-MM  
+
+        print(f"Processing partition {year_month} for schema {SCHEMA}")
+
+        url = f"{TLC_BASE}/green_tripdata_{year_month}.parquet"
+
+
+         #download parquet
+        response = requests.get(
+            url,
+            timeout=60
+        )
+
+        response.raise_for_status()
+
+        #parquet to dataframe
+        df = pd.read_parquet(
+            io.BytesIO(response.content)
+        )
+
+
+        hook = PostgresHook(
+            postgres_conn_id="azure_pg"
+        )
+
+
+        engine = hook.get_sqlalchemy_engine()
+        #create schema
+        with hook.get_conn() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    f'CREATE SCHEMA IF NOT EXISTS "{SCHEMA}"'
+                )
+        #create table if missing
+        df.head(0).to_sql(
+            "raw_trips",
+            engine,
+            schema=SCHEMA,
+            if_exists="append",
+            index=False,
+        )
+
+        # idempotency remove old data for same month
+        with hook.get_conn() as conn:
+            with conn.cursor() as cur:
+
+                cur.execute(
+                    f"""
+                    DELETE FROM "{SCHEMA}".raw_trips
+                    WHERE to_char(
+                        lpep_pickup_datetime,
+                        'YYYY-MM'
+                    ) = %s
+                    """,
+                    (year_month,),
+                )
+        #insert fresh data        
+        df.to_sql(
+            "raw_trips",
+            engine,
+            schema=SCHEMA,
+            if_exists="append",
+            index=False,
+        )
+
+
+        return len(df)
+
+
+
+
+    dbt_run = BashOperator(
+        task_id="dbt_run",
+        bash_command=(
+            f"{DBT} deps --project-dir {DBT_DIR} --profiles-dir {DBT_DIR} && "
+            f"{DBT} run --project-dir {DBT_DIR} --profiles-dir {DBT_DIR}"
+        ),
+        env=DBT_ENV,
+        append_env=True,
+    )
+
+    dbt_test = BashOperator(
+        task_id="dbt_test",
+        bash_command=(
+            f"{DBT} test --project-dir {DBT_DIR} --profiles-dir {DBT_DIR}"
+        ),
+        env=DBT_ENV,
+        append_env=True,
+    )
+
+
+    ingest_taxi_month() >> dbt_run >> dbt_test
+
 
-    # Task 2 (see README): add the two transform tasks, wire the full
-    # chain, and run the transform through the Chapter 4 command so it works
-    # on the image's Python. Task 4: add retry behaviour.
+taxi_pipeline()
 
-    ingest_taxi_month()
 
 
-taxi_pipeline()

tests/dags/test_dag_example.py

@@ -0,0 +1,91 @@
+"""Example DAGs test. This test ensures that all Dags have tags, retries set to two, and no import errors. This is an example pytest and may not be fit the context of your DAGs. Feel free to add and remove tests."""
+
+import os
+import logging
+from contextlib import contextmanager
+import pytest
+from airflow.models import DagBag
+
+
+def _make_dag_bag():
+    """Build a DagBag across Airflow versions with differing signatures."""
+    try:
+        return DagBag(include_examples=False)
+    except TypeError:
+        return DagBag()
+
+
+@contextmanager
+def suppress_logging(namespace):
+    logger = logging.getLogger(namespace)
+    old_value = logger.disabled
+    logger.disabled = True
+    try:
+        yield
+    finally:
+        logger.disabled = old_value
+
+
+def get_import_errors():
+    """
+    Generate a tuple for import errors in the dag bag
+    """
+    with suppress_logging("airflow"):
+        dag_bag = _make_dag_bag()
+
+        def strip_path_prefix(path):
+            return os.path.relpath(path, os.environ.get("AIRFLOW_HOME"))
+
+        # prepend "(None,None)" to ensure that a test object is always created even if it's a no op.
+        return [(None, None)] + [
+            (strip_path_prefix(k), v.strip()) for k, v in dag_bag.import_errors.items()
+        ]
+
+
+def get_dags():
+    """
+    Generate a tuple of dag_id, <DAG objects> in the DagBag
+    """
+    with suppress_logging("airflow"):
+        dag_bag = _make_dag_bag()
+
+    def strip_path_prefix(path):
+        return os.path.relpath(path, os.environ.get("AIRFLOW_HOME"))
+
+    return [(k, v, strip_path_prefix(v.fileloc)) for k, v in dag_bag.dags.items()]
+
+
+@pytest.mark.parametrize(
+    "rel_path,rv", get_import_errors(), ids=[x[0] for x in get_import_errors()]
+)
+def test_file_imports(rel_path, rv):
+    """Test for import errors on a file"""
+    if rel_path and rv:
+        raise Exception(f"{rel_path} failed to import with message \n {rv}")
+
+
+APPROVED_TAGS = {}
+
+
+@pytest.mark.parametrize(
+    "dag_id,dag,fileloc", get_dags(), ids=[x[2] for x in get_dags()]
+)
+def test_dag_tags(dag_id, dag, fileloc):
+    """
+    test if a DAG is tagged and if those TAGs are in the approved list
+    """
+    assert dag.tags, f"{dag_id} in {fileloc} has no tags"
+    if APPROVED_TAGS:
+        assert not set(dag.tags) - APPROVED_TAGS
+
+
+@pytest.mark.parametrize(
+    "dag_id,dag, fileloc", get_dags(), ids=[x[2] for x in get_dags()]
+)
+def test_dag_retries(dag_id, dag, fileloc):
+    """
+    test if a DAG has retries set
+    """
+    assert (
+        dag.default_args.get("retries", None) >= 2
+    ), f"{dag_id} in {fileloc} must have task retries >= 2."