Optimized DAGs parse time (#258)

* replace `Variable.get()` when field is templatable

* streamline imports

* remove `Variable.get()` dependency from non templatable field

* only set `config_file` when in default namespace

Author: lucaszanotelli

dags/asset_#_pipeline_dag.py

@@ -1,5 +1,5 @@
 import datetime
-import json
+from json import loads
 
 from airflow import DAG
 from airflow.models.variable import Variable
@@ -16,7 +16,7 @@
     description="This DAG runs dbt to calculate asset # based on stablecoin and XLM trades",
     schedule_interval="0 2 * * *",  # daily at 2am
     params={},
-    user_defined_filters={"fromjson": lambda s: json.loads(s)},
+    user_defined_filters={"fromjson": lambda s: loads(s)},
     user_defined_macros={
         "subtract_data_interval": macros.subtract_data_interval,
         "batch_run_date_as_datetime_string": macros.batch_run_date_as_datetime_string,

dags/audit_log_dag.py

@@ -1,8 +1,8 @@
 """
 This DAG runs an audit log SQL to update the audit log dashboard.
 """
-import datetime
-import json
+from datetime import datetime
+from json import loads
 
 from airflow import DAG
 from airflow.models import Variable
@@ -16,12 +16,12 @@
     "audit_log_dag",
     default_args=get_default_dag_args(),
     description="This DAG runs periodically to update the audit log dashboard.",
-    start_date=datetime.datetime(2023, 1, 1, 0, 0),
+    start_date=datetime(2023, 1, 1, 0, 0),
     schedule_interval="10 9 * * *",
     params={
         "alias": "audit-log",
     },
-    user_defined_filters={"fromjson": lambda s: json.loads(s)},
+    user_defined_filters={"fromjson": lambda s: loads(s)},
     catchup=False,
 )
 

dags/bucket_list_dag.py

@@ -4,9 +4,9 @@
 to stop exporting. This end ledger  is determined by when the Airflow DAG is run. This DAG should be triggered manually
 when initializing the tables in order to catch up to the current state in the network, but should not be scheduled to run constantly.
 """
-import ast
-import datetime
-import json
+from ast import literal_eval
+from datetime import datetime
+from json import loads
 
 from airflow import DAG
 from airflow.models import Variable
@@ -23,29 +23,31 @@
 dag = DAG(
     "bucket_list_export",
     default_args=get_default_dag_args(),
-    start_date=datetime.datetime(2021, 10, 15),
-    end_date=datetime.datetime(2021, 10, 15),
+    start_date=datetime(2021, 10, 15),
+    end_date=datetime(2021, 10, 15),
     description="This DAG loads a point forward view of state tables. Caution: Does not capture historical changes!",
     schedule_interval="@daily",
     params={
         "alias": "bucket",
     },
-    user_defined_filters={"fromjson": lambda s: json.loads(s)},
+    user_defined_filters={
+        "fromjson": lambda s: loads(s),
+        "literal_eval": lambda e: literal_eval(e),
+    },
     user_defined_macros={
         "subtract_data_interval": macros.subtract_data_interval,
         "batch_run_date_as_datetime_string": macros.batch_run_date_as_datetime_string,
     },
 )
 
-file_names = Variable.get("output_file_names", deserialize_json=True)
 table_names = Variable.get("table_ids", deserialize_json=True)
-internal_project = Variable.get("bq_project")
-internal_dataset = Variable.get("bq_dataset")
-public_project = Variable.get("public_project")
-public_dataset = Variable.get("public_dataset")
-use_testnet = ast.literal_eval(Variable.get("use_testnet"))
-use_futurenet = ast.literal_eval(Variable.get("use_futurenet"))
-
+internal_project = "{{ var.value.bq_project }}"
+internal_dataset = "{{ var.value.bq_dataset }}"
+public_project = "{{ var.value.public_project }}"
+public_dataset = "{{ var.value.public_dataset }}"
+public_dataset_new = "{{ var.value.public_dataset_new }}"
+use_testnet = "{{ var.value.use_testnet | literal_eval }}"
+use_futurenet = "{{ var.value.use_futurenet | literal_eval }}"
 """
 The time task reads in the execution time of the current run, as well as the next
 execution time. It converts these two times into ledger ranges.
@@ -60,7 +62,7 @@
     dag,
     "bucket",
     "export_accounts",
-    file_names["accounts"],
+    "{{ var.json.output_file_names.accounts }}",
     use_testnet=use_testnet,
     use_futurenet=use_futurenet,
     use_gcs=True,
@@ -69,7 +71,7 @@
     dag,
     "bucket",
     "export_claimable_balances",
-    file_names["claimable_balances"],
+    "{{ var.json.output_file_names.claimable_balances }}",
     use_testnet=use_testnet,
     use_futurenet=use_futurenet,
     use_gcs=True,
@@ -78,7 +80,7 @@
     dag,
     "bucket",
     "export_offers",
-    file_names["offers"],
+    "{{ var.json.output_file_names.offers }}",
     use_testnet=use_testnet,
     use_futurenet=use_futurenet,
     use_gcs=True,
@@ -87,7 +89,7 @@
     dag,
     "bucket",
     "export_pools",
-    file_names["liquidity_pools"],
+    "{{ var.json.output_file_names.liquidity_pools }}",
     use_testnet=use_testnet,
     use_futurenet=use_futurenet,
     use_gcs=True,
@@ -96,7 +98,7 @@
     dag,
     "bucket",
     "export_signers",
-    file_names["signers"],
+    "{{ var.json.output_file_names.signers }}",
     use_testnet=use_testnet,
     use_futurenet=use_futurenet,
     use_gcs=True,
@@ -105,7 +107,7 @@
     dag,
     "bucket",
     "export_trustlines",
-    file_names["trustlines"],
+    "{{ var.json.output_file_names.trustlines }}",
     use_testnet=use_testnet,
     use_futurenet=use_futurenet,
     use_gcs=True,
@@ -131,37 +133,37 @@
     dag, internal_project, internal_dataset, table_names["accounts"]
 )
 delete_acc_pub_task = build_delete_data_task(
-    dag, public_project, public_dataset, table_names["accounts"]
+    dag, public_project, public_dataset, table_names["accounts"], "pub"
 )
 delete_bal_task = build_delete_data_task(
     dag, internal_project, internal_dataset, table_names["claimable_balances"]
 )
 delete_bal_pub_task = build_delete_data_task(
-    dag, public_project, public_dataset, table_names["claimable_balances"]
+    dag, public_project, public_dataset, table_names["claimable_balances"], "pub"
 )
 delete_off_task = build_delete_data_task(
     dag, internal_project, internal_dataset, table_names["offers"]
 )
 delete_off_pub_task = build_delete_data_task(
-    dag, public_project, public_dataset, table_names["offers"]
+    dag, public_project, public_dataset, table_names["offers"], "pub"
 )
 delete_pool_task = build_delete_data_task(
     dag, internal_project, internal_dataset, table_names["liquidity_pools"]
 )
 delete_pool_pub_task = build_delete_data_task(
-    dag, public_project, public_dataset, table_names["liquidity_pools"]
+    dag, public_project, public_dataset, table_names["liquidity_pools"], "pub"
 )
 delete_sign_task = build_delete_data_task(
     dag, internal_project, internal_dataset, table_names["signers"]
 )
 delete_sign_pub_task = build_delete_data_task(
-    dag, public_project, public_dataset, table_names["signers"]
+    dag, public_project, public_dataset, table_names["signers"], "pub"
 )
 delete_trust_task = build_delete_data_task(
     dag, internal_project, internal_dataset, table_names["trustlines"]
 )
 delete_trust_pub_task = build_delete_data_task(
-    dag, public_project, public_dataset, table_names["trustlines"]
+    dag, public_project, public_dataset, table_names["trustlines"], "pub"
 )
 
 """
@@ -244,6 +246,7 @@
     "",
     partition=True,
     cluster=True,
+    dataset_type="pub",
 )
 send_bal_to_pub_task = build_gcs_to_bq_task(
     dag,
@@ -254,6 +257,7 @@
     "",
     partition=True,
     cluster=True,
+    dataset_type="pub",
 )
 send_off_to_pub_task = build_gcs_to_bq_task(
     dag,
@@ -264,6 +268,7 @@
     "",
     partition=True,
     cluster=True,
+    dataset_type="pub",
 )
 send_pool_to_pub_task = build_gcs_to_bq_task(
     dag,
@@ -274,6 +279,7 @@
     "",
     partition=True,
     cluster=True,
+    dataset_type="pub",
 )
 send_sign_to_pub_task = build_gcs_to_bq_task(
     dag,
@@ -284,6 +290,7 @@
     "",
     partition=True,
     cluster=True,
+    dataset_type="pub",
 )
 send_trust_to_pub_task = build_gcs_to_bq_task(
     dag,
@@ -294,6 +301,7 @@
     "",
     partition=True,
     cluster=True,
+    dataset_type="pub",
 )
 
 (

dags/daily_euro_ohlc_dag.py

@@ -2,8 +2,8 @@
 The daily_euro_ohlc_dag DAG updates the currency table in Bigquey every day.
 """
 
-import datetime
-import json
+from datetime import datetime
+from json import loads
 
 from airflow import DAG
 from airflow.decorators import dag
@@ -18,13 +18,13 @@
 
 with DAG(
     dag_id="daily_euro_ohlc_dag",
-    start_date=datetime.datetime(2023, 1, 1, 0, 0),
+    start_date=datetime(2023, 1, 1, 0, 0),
     description="This DAG updates the currency tables in Bigquey every day",
     schedule_interval="35 0 * * *",
     params={
         "alias": "euro",
     },
-    user_defined_filters={"fromjson": lambda s: json.loads(s)},
+    user_defined_filters={"fromjson": lambda s: loads(s)},
     catchup=False,
 ) as dag:
     currency_ohlc = Variable.get("currency_ohlc", deserialize_json=True)

dags/dataset_reset_dag.py

@@ -1,9 +1,9 @@
 """
 When the Test net server is reset, the dataset reset DAG deletes all the datasets in the test Hubble.
 """
-import ast
-import datetime
-import json
+from ast import literal_eval
+from datetime import datetime
+from json import loads
 
 from airflow import DAG
 from airflow.models import Variable
@@ -20,13 +20,13 @@
     "testnet_data_reset",
     default_args=get_default_dag_args(),
     description="This DAG runs after the Testnet data reset that occurs periodically.",
-    start_date=datetime.datetime(2023, 1, 1, 0, 0),
+    start_date=datetime(2023, 1, 1, 0, 0),
     schedule_interval="10 9 * * *",
-    is_paused_upon_creation=ast.literal_eval(Variable.get("use_testnet")),
+    is_paused_upon_creation=literal_eval(Variable.get("use_testnet")),
     params={
         "alias": "testnet-reset",
     },
-    user_defined_filters={"fromjson": lambda s: json.loads(s)},
+    user_defined_filters={"fromjson": lambda s: loads(s)},
 )
 
 internal_project = "test-hubble-319619"

dags/enriched_tables_dag.py

@@ -1,7 +1,6 @@
-import datetime
+from datetime import datetime
 
 from airflow import DAG
-from airflow.models.variable import Variable
 from stellar_etl_airflow.build_cross_dependency_task import build_cross_deps
 from stellar_etl_airflow.build_dbt_task import build_dbt_task
 from stellar_etl_airflow.default import get_default_dag_args, init_sentry
@@ -11,7 +10,7 @@
 dag = DAG(
     "enriched_tables",
     default_args=get_default_dag_args(),
-    start_date=datetime.datetime(2023, 4, 12, 0, 0),
+    start_date=datetime(2023, 4, 12, 0, 0),
     description="This DAG runs dbt to create the tables for the models in marts/enriched/.",
     schedule_interval="*/30 * * * *",  # Runs every 30 mins
     params={},