Merge pull request #12 from michaelconan/feat/notion-habits

Feat/notion habits

Author: michaelconan

.devcontainer/devcontainer.json

@@ -10,7 +10,8 @@
 			"extensions": [
 				"ms-python.python",
 				"GitHub.copilot-chat",
-				"ms-python.black-formatter"
+				"ms-python.black-formatter",
+				"bierner.markdown-mermaid"
 			]
 		}
 	},

.devcontainer/docker-compose.yml

@@ -3,14 +3,14 @@ version: '3.8'
 services:
   app:
     environment:
+      AIRFLOW_HOME: "/workspaces/personal-reporting"
       AIRFLOW__DATABASE__SQL_ALCHEMY_CONN: postgresql+psycopg2://postgres:postgres@db/postgres
       AIRFLOW__DATABASE__LOAD_DEFAULT_CONNECTIONS: "false"
       AIRFLOW__CORE__EXECUTOR: LocalExecutor
       AIRFLOW__CORE__FERNET_KEY: ""
       AIRFLOW__CORE__LOAD_EXAMPLES: "false"
       AIRFLOW__WEBSERVER__INSTANCE_NAME: "Local Instance"
       AIRFLOW__WEBSERVER__EXPOSE_CONFIG: "true"
-      AIRFLOW__CORE__DAGS_FOLDER: "/workspaces/personal-reporting/dags"
     build:
       context: .
       dockerfile: Dockerfile

.gitignore

@@ -1,7 +1,20 @@
+# Cache data
 __pycache__
+.pytest_cache
+.coverage
+
+# Logs
 logs
+*.log
+
+# Data files
 *.csv
 
 # Secrets and variables
 .env
-config/*
+config/*
+
+# DBT compiled files
+dbt/*/target
+dbt/*/dbt_modules
+dbt/*/dbt_packages

.vscode/settings.json

@@ -2,12 +2,15 @@
     "python.testing.pytestArgs": [
         "tests",
         // Coverage of source code
-        "--cov=src/",
+        "--cov=dags/",
         // INFO level pytest logging
         "--log-cli-level=INFO",
     ],
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true,
     "editor.defaultFormatter": "ms-python.black-formatter",
-    "editor.formatOnSave": true
+    "editor.formatOnSave": true,
+    "[json]": {
+        "editor.defaultFormatter": "vscode.json-language-features"
+    }
 }

README.md

@@ -14,6 +14,47 @@ The following data will be ingested from my personal systems into a BigQuery war
 2. HubSpot
 3. Google Contacts
 
+### Warehouse Data Flow
+
+```mermaid
+graph TB
+
+    %% Soruces
+    S1[Notion]
+    S2[HubSpot]
+
+    subgraph staging
+        direction TB
+        L1[Daily Habits]
+        L2[Weekly Habits]
+        L3[Contacts]
+        L4[Companies]
+        L5[Engagements]
+    end
+
+    %% Source to Staging Flows
+    S1 --> L1
+    S1 --> L2
+    S2 --> L3
+    S2 --> L4
+    S2 --> L5
+
+    subgraph cleansed
+        C1[Habits]
+    end
+
+    %% Staging to Cleansed Flows
+    L1 --> C1
+    L2 --> C1
+```
+
+## Frameworks
+
+1. [Alembic](https://alembic.sqlalchemy.org/en/latest/) database migrations for raw tables which should generally match the schema of the source system, run via [Airflow provider package](https://pypi.org/project/airflow-provider-alembic/)
+2. [Airflow](https://airflow.apache.org/) to orchestrate data loading scripts and additional automated workflows
+3. [DBT core](https://docs.getdbt.com/) to define data models and transformations, again orchestrated by Airflow (via CLI / `BashOperator`) 
+
+
 ## Setup
 
 ### Airflow Setup

dags/.airflowignore

@@ -0,0 +1,4 @@
+# Modules to ignore for DAG discovery using REGEX
+
+# Common tasks to import into dags
+reporting/common/.*

dags/michael/__init__.py

@@ -0,0 +1,36 @@
+"""
+bigquery.py
+
+Common functions to interact with Google BigQuery in DAG tasks.
+"""
+
+# Airflow hooks and operators
+from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+
+# Additional libraries
+from google.cloud import bigquery
+
+
+def load_file_to_bq(conn_id: str, file_path: str, table_id: str) -> str:
+    """Insert newline-delimited JSON file (JSONL) into BigQuery table
+
+    Args:
+        conn_id (str): Airflow connection ID for BigQuery.
+        file_path (str): Path to the file to load.
+        table_name (str): Table name to load the data into.
+
+    Returns:
+        str: State of completed insert job.
+    """
+    # Load the data into BigQuery
+    connection = BigQueryHook(gcp_conn_id=conn_id)
+    client = connection.get_client()
+
+    job_config = bigquery.LoadJobConfig(
+        source_format=bigquery.SourceFormat.NEWLINE_DELIMITED_JSON,
+    )
+    with open(file_path, "rb") as source_file:
+        job = client.load_table_from_file(source_file, table_id, job_config=job_config)
+    job.result()
+
+    return job.state

dags/michael/common/__init__.py

@@ -0,0 +1,81 @@
+# Base imports
+import os
+import pendulum
+import datetime
+
+# PyPI imports
+import yaml
+
+# Airflow imports
+from airflow.decorators import dag, task
+from airflow.operators.bash import BashOperator
+from airflow.hooks.base import BaseHook
+
+# Dataset name in BigQuery for DBT
+DBT_DATASET = "reporting"
+
+
+@dag(
+    schedule="@daily",
+    catchup=False,
+    start_date=pendulum.datetime(2025, 1, 1),
+    dagrun_timeout=datetime.timedelta(minutes=20),
+)
+def DBT():
+
+    # File paths for service account key and dbt profile
+    PROFILES_DIR = "/tmp/.dbt"
+    KEYFILE_PATH = os.path.join(PROFILES_DIR, "bq-service-account.json")
+    PROFILE_PATH = os.path.join(PROFILES_DIR, "dbt_profile.yml")
+
+    @task(
+        task_id="generate_dbt_profile",
+    )
+    def generate_dbt_profile():
+        # Get BigQuery connection details
+        conn = BaseHook.get_connection("bigquery_reporting")
+
+        # Write keyfile to temporary file
+        os.makedirs(os.path.dirname(KEYFILE_PATH), exist_ok=True)
+        with open(KEYFILE_PATH, "w") as f:
+            f.write(conn.extra_dejson.get("keyfile_dict"))
+
+        # Generate profile with BigQuery details
+        profile = {
+            "michael": {
+                "outputs": {
+                    "dev": {
+                        "type": "bigquery",
+                        "method": "service-account",
+                        "keyfile": KEYFILE_PATH,
+                        "dataset": DBT_DATASET,
+                        "project": conn.extra_dejson.get("project"),
+                        "location": conn.extra_dejson.get("location"),
+                        "priority": "interactive",
+                        "job_execution_timeout_seconds": 300,
+                        "job_retries": 1,
+                        "threads": 1,
+                    },
+                },
+                "target": "dev",
+            }
+        }
+        # Create profile file for dbt run
+        with open(PROFILE_PATH, "w") as f:
+            yaml.dump(profile, f)
+
+    dbt_run = BashOperator(
+        task_id="dbt_run",
+        bash_command=f"dbt run --profiles-dir {PROFILES_DIR}",
+        env={"DBT_PROFILES_DIR": PROFILES_DIR},
+    )
+
+    @task(
+        task_id="cleanup_files",
+    )
+    def cleanup_files():
+        # Remove temporary files
+        os.remove(PROFILES_DIR)
+
+    # Define DAG workflow
+    generate_dbt_profile() >> dbt_run >> cleanup_files()

dags/michael/common/bigquery.py

@@ -0,0 +1,47 @@
+# Base imports
+import os
+from datetime import datetime
+
+# Airflow imports
+from airflow.models import DAG
+from airflow.decorators import task
+from airflow.models.param import Param
+
+# Airflow hooks and operators
+from airflow.hooks.base import BaseHook
+from airflow_provider_alembic.operators.alembic import AlembicOperator
+
+# Get migration folder relative to DAG
+migration_folder = os.path.join(os.path.dirname(__file__), "migrations")
+
+with DAG(
+    "raw_migrations",
+    schedule="@once",  # also consider "None"
+    start_date=datetime(1970, 1, 1),
+    params={"command": Param("upgrade"), "revision": Param("head")},
+) as dag:
+
+    KEYFILE_PATH = "/tmp/bq-service-account.json"
+    BIGQUERY_CONN_ID = "bigquery_reporting"
+    # Set keyfile as application default credentials
+    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = KEYFILE_PATH
+
+    @task(task_id="bigquery_keyfile")
+    def get_keyfile():
+        # Get BigQuery connection details
+        conn = BaseHook.get_connection(BIGQUERY_CONN_ID)
+
+        # Write keyfile to temporary file
+        with open(KEYFILE_PATH, "w") as f:
+            f.write(conn.extra_dejson.get("keyfile_dict"))
+
+    # Run migrations
+    alembic_op = AlembicOperator(
+        task_id="alembic_op",
+        conn_id=BIGQUERY_CONN_ID,
+        command="{{ params.command }}",
+        revision="{{ params.revision }}",
+        script_location=migration_folder,
+    )
+
+    get_keyfile() >> alembic_op

dags/michael/dbt.py

@@ -0,0 +1 @@
+Generic single-database configuration.