1085 Ecocounter Open Data Schema+DAG+QC plots

dags/ecocounter_open_data.py

@@ -3,7 +3,7 @@
 """
 import sys
 import os
-from datetime import timedelta
+from datetime import timedelta, datetime
 import logging
 import pendulum
 from functools import partial
@@ -12,15 +12,14 @@
 from airflow.models import Variable 
 from airflow.hooks.base_hook import BaseHook
 from airflow.providers.postgres.operators.postgres import PostgresOperator
-from airflow.sensors.external_task import ExternalTaskSensor
 from airflow.sensors.date_time import DateTimeSensor
 from airflow.macros import ds_format
-
+from airflow.operators.python import get_current_context
 
 try:
     repo_path = os.path.abspath(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
     sys.path.insert(0, repo_path)
-    from dags.dag_functions import task_fail_slack_alert, send_slack_msg
+    from dags.dag_functions import task_fail_slack_alert, send_slack_msg, get_readme_docmd
     from dags.custom_operators import SQLCheckOperatorWithReturnValue
 except:
     raise ImportError("Cannot import DAG helper functions.")
@@ -31,11 +30,14 @@
 DAG_NAME = 'ecocounter_open_data'
 DAG_OWNERS = Variable.get('dag_owners', deserialize_json=True).get(DAG_NAME, ["Unknown"])
 
+README_PATH = os.path.join(repo_path, 'volumes/ecocounter/readme.md')
+DOC_MD = get_readme_docmd(README_PATH, DAG_NAME)
+
 default_args = {
     'owner': ','.join(DAG_OWNERS),
     'depends_on_past':False,
     #set earlier start_date + catchup when ready?
-    'start_date': pendulum.datetime(2024, 10, 1, tz="America/Toronto"),
+    'start_date': pendulum.datetime(2024, 1, 1, tz="America/Toronto"),
     'email_on_failure': False,
     'email_on_success': False,
     'retries': 0,
@@ -46,26 +48,14 @@
 @dag(
     dag_id=DAG_NAME,
     default_args=default_args,
-    schedule='0 11 1 * *', # 10am, 1st day of each month
-    catchup=True,
+    schedule='0 12 1 * *', # 12pm, 1st day of each month
+    catchup=False,
     max_active_runs=1,
     tags=["ecocounter", "open_data"],
-    doc_md=__doc__
+    doc_md=DOC_MD
 )
 def ecocounter_open_data_dag():
 
-    t_upstream_done = ExternalTaskSensor(
-        task_id="starting_point",
-        external_dag_id="ecocounter_pull",
-        external_task_id="done",
-        poke_interval=3600, #retry hourly
-        mode="reschedule",
-        doc_md="Wait for last day of month to run before running monthly DAG.",
-        timeout=86400, #one day
-        #wait for the 1st of the following month
-        execution_date_fn=lambda dt: dt + pendulum.duration(months=1, hours=-1) #ecocounter_pull scheduled at '0 10 * * *'
-    )
-
     check_data_availability = SQLCheckOperatorWithReturnValue(
         task_id="check_data_availability",
         sql="""WITH daily_volumes AS (
@@ -106,46 +96,93 @@ def reminder_message(ds = None, **context):
     wait_till_10th.doc_md = """
     Wait until the 10th day of the month to export data. Alternatively mark task as success to proceed immediately.
     """
+
+    @task()
+    def get_years(ds=None):
+        mnth = pendulum.from_format(ds, 'YYYY-MM-DD')
+        prev_mnth = mnth.subtract(months=1)
+        yrs = [str(mnth.year), str(prev_mnth.year)]
+        return list(set(yrs)) #unique
 
     @task_group()
-    def insert_and_download_data():
-        insert_daily = PostgresOperator(
-            sql="SELECT ecocounter.open_data_daily_counts_insert({{ macros.ds_format(ds, '%Y-%m-%d', '%Y') }})",
-            task_id='insert_daily_open_data',
-            postgres_conn_id='ecocounter_bot',
-            autocommit=True,
-            retries = 0
+    def insert_and_download_data(yr):
+        @task(map_index_template="{{ yr }}")
+        def insert_daily(yr):
+            context = get_current_context()
+            context["yr"] = yr
+            t = PostgresOperator(
+                sql=f"SELECT ecocounter.open_data_daily_counts_insert({yr}::int)",
+                task_id='insert_daily_open_data',
+                postgres_conn_id='ecocounter_bot',
+                autocommit=True,
+                retries = 0
+            )
+            return t.execute(context=context)
+
+        @task(map_index_template="{{ yr }}")
+        def insert_15min(yr):
+            context = get_current_context()
+            context["yr"] = yr
+            t = PostgresOperator(
+                sql=f"SELECT ecocounter.open_data_15min_counts_insert({yr}::int)",
+                task_id='insert_15min_open_data',
+                postgres_conn_id='ecocounter_bot',
+                autocommit=True,
+                retries = 0
+            )
+            return t.execute(context=context)
+
+        @task.bash(
+            map_index_template="{{ yr }}",
+            env={
+                "HOST":  BaseHook.get_connection("ecocounter_bot").host,
+                "USER" :  BaseHook.get_connection("ecocounter_bot").login,
+                "PGPASSWORD": BaseHook.get_connection("ecocounter_bot").password
+            }
         )
-        insert_raw = PostgresOperator(
-            sql="SELECT ecocounter.open_data_raw_counts_insert({{ macros.ds_format(ds, '%Y-%m-%d', '%Y') }})",
-            task_id='insert_raw_open_data',
-            postgres_conn_id='ecocounter_bot',
-            autocommit=True,
-            retries = 0
-        )
-
-        @task.bash(env={
-            "HOST":  BaseHook.get_connection("ecocounter_bot").host,
-            "USER" :  BaseHook.get_connection("ecocounter_bot").login,
-            "PGPASSWORD": BaseHook.get_connection("ecocounter_bot").password
-        })
-        def download_daily_open_data()->str:
-            return '''psql -h $HOST -U $USER -d bigdata -c \
-                "SELECT site_description, direction, dt, daily_volume FROM ecocounter.open_data_daily_counts WHERE dt >= date_trunc('year'::text, '{{ ds }}'::date) LIMIT 100" \
-                --csv -o /data/open_data/permanent-bike-counters/ecocounter_daily_counts_{{ macros.ds_format(ds, '%Y-%m-%d', '%Y') }}.csv'''
+        def download_daily_open_data(yr)->str:
+            context = get_current_context()
+            context["yr"] = yr
+            return f'''/usr/bin/psql -h $HOST -U $USER -d bigdata -c \
+                "SELECT site_description, direction, dt, daily_volume
+                FROM ecocounter.open_data_daily_counts
+                WHERE
+                    dt >= to_date({yr}::text, 'yyyy')
+                    AND dt < LEAST(date_trunc('month', now()), to_date(({yr}::int+1)::text, 'yyyy'));" \
+                --csv -o "/data/open_data/permanent-bike-counters/ecocounter_daily_counts_{yr}.csv"'''
 
-        @task.bash(env={
-            "HOST":  BaseHook.get_connection("ecocounter_bot").host,
-            "USER" :  BaseHook.get_connection("ecocounter_bot").login,
-            "PGPASSWORD": BaseHook.get_connection("ecocounter_bot").password
-        })
-        def download_raw_open_data()->str:
-            return '''psql -h $HOST -U $USER -d bigdata -c \
-                "SELECT site_description, direction, datetime_bin, bin_volume FROM ecocounter.open_data_raw_counts WHERE datetime_bin >= date_trunc('year'::text, '{{ ds }}'::date) LIMIT 100" \
-                --csv -o /data/open_data/permanent-bike-counters/ecocounter_raw_counts_{{ macros.ds_format(ds, '%Y-%m-%d', '%Y') }}.csv'''
+        @task.bash(
+            map_index_template="{{ yr }}",
+            env={
+                "HOST":  BaseHook.get_connection("ecocounter_bot").host,
+                "USER" :  BaseHook.get_connection("ecocounter_bot").login,
+                "PGPASSWORD": BaseHook.get_connection("ecocounter_bot").password
+            }
+        )
+        def download_15min_open_data(yr)->str:
+            context = get_current_context()
+            context["yr"] = yr
+            return f'''/usr/bin/psql -h $HOST -U $USER -d bigdata -c \
+                "SELECT site_description, direction, datetime_bin, bin_volume
+                FROM ecocounter.open_data_15min_counts
+                WHERE
+                    datetime_bin >= to_date({yr}::text, 'yyyy')
+                    AND datetime_bin < LEAST(date_trunc('month', now()), to_date(({yr}+1)::text, 'yyyy'));" \
+                --csv -o "/data/open_data/permanent-bike-counters/ecocounter_15min_counts_{yr}.csv"'''
 
-        insert_daily >> download_daily_open_data()
-        insert_raw >> download_raw_open_data()
+        insert_daily(yr) >> download_daily_open_data(yr)
+        insert_15min(yr) >> download_15min_open_data(yr)
+
+    @task.bash(env={
+        "HOST":  BaseHook.get_connection("ecocounter_bot").host,
+        "USER" :  BaseHook.get_connection("ecocounter_bot").login,
+        "PGPASSWORD": BaseHook.get_connection("ecocounter_bot").password
+    })
+    def download_locations_open_data()->str:
+        return '''/usr/bin/psql -h $HOST -U $USER -d bigdata -c \
+            "SELECT location_name, direction, linear_name_full, side_street, lng, lat, centreline_id, bin_size, latest_calibration_study, first_active, last_active, date_decommissioned, technology
+                FROM ecocounter.open_data_locations" \
+            --csv -o /data/open_data/permanent-bike-counters/locations.csv'''
 
     @task(
         retries=0,
@@ -159,14 +196,14 @@ def status_message(ds = None, **context):
             msg=f"Ecocounter :open_data_to: DAG ran successfully for {mnth} :white_check_mark:",
             use_proxy=True
         )
-
+
+    yrs = get_years()
     (
-        t_upstream_done >>
         check_data_availability >>
         reminder_message() >>
         wait_till_10th >> 
-        insert_and_download_data() >>
+        [insert_and_download_data.expand(yr = yrs), download_locations_open_data()] >>
         status_message()
     )
-
+        
 ecocounter_open_data_dag()

volumes/ecocounter/open_data/create-function-open_data_raw_counts_insert.sql

@@ -1,4 +1,4 @@
-CREATE OR REPLACE FUNCTION ecocounter.open_data_raw_counts_insert(
+CREATE OR REPLACE FUNCTION ecocounter.open_data_15min_counts_insert(
     yyyy integer
 )
 RETURNS void
@@ -30,7 +30,7 @@ AS $BODY$
         HAVING COUNT(DISTINCT cc.datetime_bin) = (3600*24 / EXTRACT(epoch FROM bin_size))
     )
 
-    INSERT INTO ecocounter.open_data_raw_counts (
+    INSERT INTO ecocounter.open_data_15min_counts (
         site_id, site_description, direction, datetime_bin, bin_volume
     )
     SELECT  
@@ -63,12 +63,12 @@ AS $BODY$
 
 $BODY$;
 
-ALTER FUNCTION ecocounter.open_data_raw_counts_insert(integer) OWNER TO ecocounter_admins;
+ALTER FUNCTION ecocounter.open_data_15min_counts_insert(integer) OWNER TO ecocounter_admins;
 
-GRANT EXECUTE ON FUNCTION ecocounter.open_data_raw_counts_insert(integer) TO ecocounter_admins;
-GRANT EXECUTE ON FUNCTION ecocounter.open_data_raw_counts_insert(integer) TO ecocounter_bot;
-REVOKE EXECUTE ON FUNCTION ecocounter.open_data_raw_counts_insert(integer) FROM bdit_humans;
+GRANT EXECUTE ON FUNCTION ecocounter.open_data_15min_counts_insert(integer) TO ecocounter_admins;
+GRANT EXECUTE ON FUNCTION ecocounter.open_data_15min_counts_insert(integer) TO ecocounter_bot;
+REVOKE EXECUTE ON FUNCTION ecocounter.open_data_15min_counts_insert(integer) FROM bdit_humans;
 
-COMMENT ON FUNCTION ecocounter.open_data_raw_counts_insert(integer) IS
-'Function to insert disaggregate data for a year into `ecocounter.open_data_raw_counts`. '
+COMMENT ON FUNCTION ecocounter.open_data_15min_counts_insert(integer) IS
+'Function to insert disaggregate data for a year into `ecocounter.open_data_15min_counts`. '
 'Does not overwrite existing data (eg. if sensitivity was retroactively updated).';

volumes/ecocounter/open_data/create-table-open_data_raw_counts.sql

@@ -1,8 +1,8 @@
--- Table: ecocounter.open_data_raw_counts
+-- Table: ecocounter.open_data_15min_counts
 
--- DROP TABLE IF EXISTS ecocounter.open_data_raw_counts;
+-- DROP TABLE IF EXISTS ecocounter.open_data_15min_counts;
 
-CREATE TABLE IF NOT EXISTS ecocounter.open_data_raw_counts
+CREATE TABLE IF NOT EXISTS ecocounter.open_data_15min_counts
 (
     site_id numeric,
     site_description text COLLATE pg_catalog."default",
@@ -14,12 +14,12 @@ CREATE TABLE IF NOT EXISTS ecocounter.open_data_raw_counts
 
 TABLESPACE pg_default;
 
-ALTER TABLE IF EXISTS ecocounter.open_data_raw_counts OWNER TO ecocounter_admins;
+ALTER TABLE IF EXISTS ecocounter.open_data_15min_counts OWNER TO ecocounter_admins;
 
-REVOKE ALL ON TABLE ecocounter.open_data_raw_counts FROM bdit_humans;
-GRANT SELECT ON TABLE ecocounter.open_data_raw_counts TO bdit_humans;
+REVOKE ALL ON TABLE ecocounter.open_data_15min_counts FROM bdit_humans;
+GRANT SELECT ON TABLE ecocounter.open_data_15min_counts TO bdit_humans;
 
-GRANT SELECT, INSERT ON TABLE ecocounter.open_data_raw_counts TO ecocounter_bot;
+GRANT SELECT, INSERT ON TABLE ecocounter.open_data_15min_counts TO ecocounter_bot;
 
-COMMENT ON TABLE ecocounter.open_data_raw_counts IS
+COMMENT ON TABLE ecocounter.open_data_15min_counts IS
 'Disaggregate Ecocounter data by site and direction.';

volumes/ecocounter/open_data/readme.md

@@ -28,7 +28,7 @@ ecocounter.open_data_locations
 | dt               | date      | 06/26/1994                                    | Date of count.                           |
 | daily_volume     | integer   | 939                                           | Count of users on date `dt`.             |
 
-## ecocounter.open_data_raw_counts
+## ecocounter.open_data_15min_counts
 
 | column_name      | data_type                   | sample                                        | explanation                                                                                  |
 |------------------|-----------------------------|-----------------------------------------------|----------------------------------------------------------------------------------------------|

volumes/ecocounter/readme.md

@@ -1,5 +1,3 @@
-# Ecocounter <!-- omit in toc -->
-
 <!-- TOC -->
 
 - [Bicycle loop detectors](#bicycle-loop-detectors)
@@ -12,6 +10,7 @@
   - [Historical data](#historical-data)
   - [`ecocounter_pull` DAG](#ecocounter_pull-dag)
   - [`ecocounter_check` DAG](#ecocounter_check-dag)
+  - [`ecocounter_check` DAG](#ecocounter_check-dag-1)
 - [SQL Tables](#sql-tables)
   - [Main Tables](#main-tables)
     - [`ecocounter.sites_unfiltered`](#ecocountersites_unfiltered)
@@ -143,6 +142,25 @@ The `ecocounter_check` DAG runs daily at 4am following completion of `ecocounter
 - `check_unvalidated_sites` runs a `SQLCheckOperatorWithReturnValue` to check for unvalidated sites or flows with non-zero volumes this week and send a slack notification with their details. 
 <!-- ecocounter_check_doc_md -->
 
+<!-- ecocounter_open_data_doc_md -->
+
+## `ecocounter_open_data` DAG
+The `ecocounter_open_data` DAG runs monthly on the 1st of the month to perform insert/download of open data extracts for the month. 
+
+- `check_data_availability` A SQLCheckOperatorWithReturnValue to check if there is data for every day of the previous month before proceeding and report if not. 
+- `reminder_message`: A slack message to indicate data should be verified and any necessary anomalous ranges added to exclude irregular data. 
+- `wait_till_10th`: A DateTimeSensor which waits until the 10th day of the month to proceed with exporting data. Can also be marked as "Success" manually to proceed earlier. 
+- `get_years`: Identifies this month and last month's years in case of additional data added for the previous month.
+- `insert_and_download_data`: TaskGroup for tasks which update and download data. Dynamically mapped over output of `get_years`. 
+  - `insert_daily_open_data`: Inserts daily data into open data table.
+  - `insert_15min_open_data`: Inserts 15min data into open data table.
+  - `download_daily_open_data`: Downloads daily data to Open Data mounted drive on Morbius.
+  - `download_15min_open_data`: Downloads 15min data to Open Data mounted drive on Morbius.
+- `download_locations_open_data`: Downloads locations table to Open Data mounted drive on Morbius.
+- `status_message`: A slack message indicating successful pipeline completion.
+
+<!-- ecocounter_open_data_doc_md -->
+
 # SQL Tables
 
 ## Main Tables