add airflow setup

adchia · adchia · commit 4d65e4a82555 · 2022-09-29T14:35:56.000-05:00
Signed-off-by: Danny Chiao &lt;danny@tecton.ai&gt;
diff --git a/.gitignore b/.gitignore
@@ -6,4 +6,5 @@ terraform.tfstate
 terraform.tfstate.backup
 .terraform*
 *.iml
-**/feast-postgres-data/*
+**/feast-postgres-data/*
+**/airflow_demo/airflow_home/*
diff --git a/module_1/airflow_demo/dag.py b/module_1/airflow_demo/dag.py
@@ -1,16 +1,21 @@
+import os
 from airflow.decorators import dag, task
 from datetime import datetime
-from feast import RepoConfig, RegistryConfig, FeatureStore, RedisOnlineStoreConfig
+from feast import RepoConfig, FeatureStore
+from feast.repo_config import RegistryConfig
+from feast.infra.online_stores.redis import RedisOnlineStoreConfig
+import pendulum
 
 
 @dag(
     schedule="@daily",
     catchup=False,
+    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
     tags=["feast"],
 )
 def materialize_dag():
     @task()
-    def materialize():
+    def materialize(data_interval_start=None, data_interval_end=None):
         repo_config = RepoConfig(
             registry=RegistryConfig(
                 registry_type="sql",
@@ -20,11 +25,15 @@ def materialize():
             provider="local",
             offline_store="file",
             online_store=RedisOnlineStoreConfig(connection_string="localhost:6379"),
+            entity_key_serialization_version=2,
         )
+        # Needed for Mac OS users because of a seg fault in requests for standalone Airflow (not needed in prod)
+        os.environ["NO_PROXY"] = "*"
         store = FeatureStore(config=repo_config)
-        store.materialize_incremental(datetime.now())
+        # Add 1 hr overlap to account for late data
+        store.materialize(data_interval_start.subtract(hours=1), data_interval_end)
 
     materialize()
 
 
-materialize_dag()
+materialization_dag = materialize_dag()
diff --git a/module_1/airflow_demo/setup_airflow.sh b/module_1/airflow_demo/setup_airflow.sh
@@ -3,6 +3,9 @@
 export AIRFLOW_HOME=$(pwd)/airflow_home
 export AIRFLOW__CORE__LOAD_EXAMPLES=False
 
+# Cleanup previous state, if it exists
+rm -rf $AIRFLOW_HOME
+
 # Install Airflow using the constraints file
 AIRFLOW_VERSION=2.4.0
 PYTHON_VERSION="$(python --version | cut -d " " -f 2 | cut -d "." -f 1-2)"
@@ -11,13 +14,13 @@ CONSTRAINT_URL="https://raw.githubusercontent.com/apache/airflow/constraints-${A
 # For example: https://raw.githubusercontent.com/apache/airflow/constraints-2.4.0/constraints-3.7.txt
 pip install "apache-airflow==${AIRFLOW_VERSION}" --constraint "${CONSTRAINT_URL}"
 
+# Setup Feast dags
+mkdir -p $AIRFLOW_HOME/dags
+cp dag.py $AIRFLOW_HOME/dags
+
 # The Standalone command will initialise the database, make a user,
 # and start all components for you.
 airflow standalone
 
 # Visit localhost:8080 in the browser and use the admin account details
 # shown on the terminal to login.
-
-# Setup dags
-mkdir $AIRFLOW_HOME/dags
-cp dag.py $AIRFLOW_HOME/dags
diff --git a/module_1/feature_repo/data_sources.py b/module_1/feature_repo/data_sources.py
@@ -1,3 +1,4 @@
+import os
 from feast import (
     FileSource,
     PushSource,
@@ -6,7 +7,7 @@
 # Feast also supports pulling data from data warehouses like BigQuery, Snowflake, Redshift and data lakes (e.g. via Redshift Spectrum, Trino, Spark)
 driver_stats = FileSource(
     name="driver_stats_source",
-    path="data/driver_stats.parquet",
+    path=f"{os.getcwd()}/data/driver_stats.parquet",
     timestamp_field="event_timestamp",
     created_timestamp_column="created",
     description="A table describing the stats of a driver based on hourly logs",
@@ -15,5 +16,6 @@
 
 # A push source is useful if you have upstream systems that transform features (e.g. stream processing jobs)
 driver_stats_push_source = PushSource(
-    name="driver_stats_push_source", batch_source=driver_stats,
+    name="driver_stats_push_source",
+    batch_source=driver_stats,
 )
