fix(composer): use separate airflow db file per session (GoogleCloudPlatform#5515)

Takashi Matsuo · leahecole · web-flow · commit ca9cd09c0fbb · 2021-03-16T13:56:17.000-07:00
* fix(composer): use separate AIRFLOW_HOME per session fixes GoogleCloudPlatform#5482 * add a commit for triggering builds in blog and workflow subdirs. * set AIRFLOW_HOME in noxfile_config.py * fix errors in noxfile_config.py Co-authored-by: Leah E. Cole <6719667+leahecole@users.noreply.github.com>
diff --git a/composer/blog/gcp-tech-blog/unit-test-dags-cloud-build/dags/example_dag_test.py b/composer/blog/gcp-tech-blog/unit-test-dags-cloud-build/dags/example_dag_test.py
@@ -21,6 +21,7 @@
 
 
 @pytest.fixture(autouse=True, scope="function")
+# The fixture `airflow_database` lives in composer/conftest.py.
 def set_variables(airflow_database):
     models.Variable.set('gcp_project', PROJECT_ID)
     yield
diff --git a/composer/blog/gcp-tech-blog/unit-test-dags-cloud-build/noxfile_config.py b/composer/blog/gcp-tech-blog/unit-test-dags-cloud-build/noxfile_config.py
@@ -20,6 +20,15 @@
 # The source of truth:
 # https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/noxfile_config.py
 
+import tempfile
+
+
+# Airflow creates a config file at the installation, so we want to set
+# `AIRFLOW_HOME` envvar before running pytest.
+
+_tmpdir = tempfile.TemporaryDirectory()
+
+
 TEST_CONFIG_OVERRIDE = {
     # You can opt out from the test for specific Python versions.
     # Skipping for Python 3.9 due to numpy compilation failure.
@@ -35,5 +44,5 @@
     # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT',
     # A dictionary you want to inject into your test. Don't put any
     # secrets here. These values will override predefined values.
-    "envs": {},
+    "envs": {"AIRFLOW_HOME": _tmpdir.name},
 }
diff --git a/composer/conftest.py b/composer/conftest.py
@@ -13,6 +13,9 @@
 # limitations under the License.
 
 
+import os
+import tempfile
+
 import pytest
 
 
@@ -21,5 +24,17 @@
 @pytest.fixture(scope="session")
 def airflow_database():
     import airflow.utils.db
+
+    # We use separate directory for local db path per session
+    # by setting AIRFLOW_HOME env var, which is done in noxfile_config.py.
+
+    assert('AIRFLOW_HOME' in os.environ)
+
+    airflow_home = os.environ["AIRFLOW_HOME"]
+    airflow_db = f"{airflow_home}/airflow.db"
+
     # reset both resets and initializes a new database
     airflow.utils.db.resetdb(rbac=None)  # this command will change in Airflow 2.0
+
+    # Making sure we are using a data file there.
+    assert(os.path.isfile(airflow_db))
diff --git a/composer/workflows/noxfile_config.py b/composer/workflows/noxfile_config.py
@@ -20,6 +20,14 @@
 # The source of truth:
 # https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/noxfile_config.py
 
+import tempfile
+
+
+# Airflow creates a config file at the installation, so we want to set
+# `AIRFLOW_HOME` envvar before running pytest.
+
+_tmpdir = tempfile.TemporaryDirectory()
+
 TEST_CONFIG_OVERRIDE = {
     # You can opt out from the test for specific Python versions.
     # Skipping for Python 3.9 due to numpy compilation failure.
@@ -35,5 +43,5 @@
     # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT',
     # A dictionary you want to inject into your test. Don't put any
     # secrets here. These values will override predefined values.
-    "envs": {},
+    "envs": {"AIRFLOW_HOME": _tmpdir.name},
 }
diff --git a/composer/workflows/quickstart_test.py b/composer/workflows/quickstart_test.py
@@ -19,6 +19,7 @@
 
 
 @pytest.fixture(autouse=True, scope="function")
+# The fixture `airflow_database` lives in composer/conftest.py.
 def set_variables(airflow_database):
 
     models.Variable.set('gcs_bucket', 'example_bucket')
