Set env value based on airflow version

astronomer · Feb 8, 2023 · e03b907 · e03b907
1 parent 80a0f20
commit e03b907
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/python-sdk/dev/docker-compose.yaml b/python-sdk/dev/docker-compose.yaml
@@ -18,7 +18,7 @@ x-airflow-common:
     AIRFLOW__WEBSERVER__EXPOSE_CONFIG: "True"
     AIRFLOW__SCHEDULER__DAG_DIR_LIST_INTERVAL: "5"
     ASTRONOMER_ENVIRONMENT: local
-    AIRFLOW__CORE__ALLOWED_DESERIALIZATION_CLASSES: airflow\.* astro\.*
+    # AIRFLOW__CORE__ALLOWED_DESERIALIZATION_CLASSES: airflow\.* astro\.*
     AIRFLOW__LINEAGE__BACKEND: openlineage.lineage_backend.OpenLineageBackend
     OPENLINEAGE_URL: http://host.docker.internal:5050/
     OPENLINEAGE_NAMESPACE: "astro"

diff --git a/python-sdk/example_dags/example_dataframe_api.py b/python-sdk/example_dags/example_dataframe_api.py
@@ -3,6 +3,7 @@
 from io import StringIO
 from typing import List
 
+import airflow
 import pandas as pd
 import requests
 from airflow import DAG
@@ -17,9 +18,13 @@
 @task or @aql.transform tasks.
 """
 
-ALLOWED_DESERIALIZATION_CLASSES = os.getenv(
-    "AIRFLOW__CORE__ALLOWED_DESERIALIZATION_CLASSES", default="airflow\\.* astro\\.*"
-)
+# We need deserialization classes for airflow 2.5 with the value where as 2.2.5 we don't need it.
+if airflow.__version__ == "2.2.5":
+    ALLOWED_DESERIALIZATION_CLASSES = os.getenv("AIRFLOW__CORE__ALLOWED_DESERIALIZATION_CLASSES")
+else:
+    ALLOWED_DESERIALIZATION_CLASSES = os.getenv(
+        "AIRFLOW__CORE__ALLOWED_DESERIALIZATION_CLASSES", default="airflow\\.* astro\\.*"
+    )
 
 
 def _load_covid_data():