Add --database-isolation as option to run breeze with (apache#35140)

Not yet fully functional, but we should be able to start airflow in `database-isolation` mode in order to start testing AIP-44. There are quite many options in the "breeze" command alredy that's why this PR also regroups them to avoid a wall of unrlated options listed together.
pemako · Oct 23, 2023 · f457228 · f457228
1 parent 8531396
commit f457228
Show file tree

Hide file tree

Showing 15 changed files with 298 additions and 185 deletions.
diff --git a/Dockerfile.ci b/Dockerfile.ci
@@ -766,6 +766,13 @@ if [[ ${SKIP_ENVIRONMENT_INITIALIZATION=} != "true" ]]; then
         echo
         export AIRFLOW__SCHEDULER__STANDALONE_DAG_PROCESSOR=True
     fi
+    if [[ ${DATABASE_ISOLATION=} == "true" ]]; then
+        echo "${COLOR_BLUE}Force database isolation configuration:${COLOR_RESET}"
+        export AIRFLOW__CORE__DATABASE_ACCESS_ISOLATION=True
+        export AIRFLOW__CORE__INTERNAL_API_URL=http://localhost:8080
+        export AIRFLOW__WEBSERVER_RUN_INTERNAL_API=True
+    fi
+
     RUN_TESTS=${RUN_TESTS:="false"}
     CI=${CI:="false"}
     USE_AIRFLOW_VERSION="${USE_AIRFLOW_VERSION:=""}"

diff --git a/dev/breeze/src/airflow_breeze/commands/developer_commands.py b/dev/breeze/src/airflow_breeze/commands/developer_commands.py
@@ -50,6 +50,7 @@
     option_builder,
     option_celery_broker,
     option_celery_flower,
+    option_database_isolation,
     option_db_reset,
     option_downgrade_sqlalchemy,
     option_dry_run,
@@ -175,6 +176,7 @@ def run(self):
 @option_celery_broker
 @option_celery_flower
 @option_standalone_dag_processor
+@option_database_isolation
 @click.argument("extra-args", nargs=-1, type=click.UNPROCESSED)
 def shell(
     python: str,
@@ -206,6 +208,7 @@ def shell(
     upgrade_boto: bool,
     downgrade_sqlalchemy: bool,
     standalone_dag_processor: bool,
+    database_isolation: bool,
 ):
     """Enter breeze environment. this is the default command use when no other is selected."""
     if get_verbose() or get_dry_run():
@@ -246,6 +249,7 @@ def shell(
         upgrade_boto=upgrade_boto,
         downgrade_sqlalchemy=downgrade_sqlalchemy,
         standalone_dag_processor=standalone_dag_processor,
+        database_isolation=database_isolation,
     )
     sys.exit(result.returncode)
 
@@ -292,6 +296,7 @@ def shell(
 @option_celery_broker
 @option_celery_flower
 @option_standalone_dag_processor
+@option_database_isolation
 def start_airflow(
     python: str,
     backend: str,
@@ -321,6 +326,7 @@ def start_airflow(
     celery_broker: str,
     celery_flower: bool,
     standalone_dag_processor: bool,
+    database_isolation: bool,
 ):
     """
     Enter breeze environment and starts all Airflow components in the tmux session.
@@ -366,6 +372,7 @@ def start_airflow(
         celery_broker=celery_broker,
         celery_flower=celery_flower,
         standalone_dag_processor=standalone_dag_processor,
+        database_isolation=database_isolation,
     )
     sys.exit(result.returncode)
 

diff --git a/dev/breeze/src/airflow_breeze/commands/developer_commands_config.py b/dev/breeze/src/airflow_breeze/commands/developer_commands_config.py
@@ -32,54 +32,69 @@
 DEVELOPER_PARAMETERS: dict[str, list[dict[str, str | list[str]]]] = {
     "breeze": [
         {
-            "name": "Basic flags",
+            "name": "Execution mode",
             "options": [
                 "--python",
+                "--integration",
+                "--standalone-dag-processor",
+                "--database-isolation",
+            ],
+        },
+        {
+            "name": "Database",
+            "options": [
                 "--backend",
                 "--postgres-version",
                 "--mysql-version",
                 "--mssql-version",
-                "--integration",
-                "--standalone-dag-processor",
-                "--forward-credentials",
                 "--db-reset",
-                "--max-time",
+            ],
+        },
+        {
+            "name": "Build CI image (before entering shell)",
+            "options": [
                 "--github-repository",
                 "--builder",
             ],
         },
+        {
+            "name": "Other options",
+            "options": [
+                "--forward-credentials",
+                "--max-time",
+            ],
+        },
     ],
     "breeze shell": [
         {
-            "name": "Basic flags",
+            "name": "Execution mode",
             "options": [
                 "--python",
+                "--integration",
+                "--standalone-dag-processor",
+                "--database-isolation",
+            ],
+        },
+        {
+            "name": "Database",
+            "options": [
                 "--backend",
                 "--postgres-version",
                 "--mysql-version",
                 "--mssql-version",
-                "--integration",
-                "--forward-credentials",
                 "--db-reset",
-                "--max-time",
             ],
         },
         {
-            "name": "Choosing executor",
+            "name": "Choose executor",
             "options": [
                 "--executor",
                 "--celery-broker",
                 "--celery-flower",
             ],
         },
         {
-            "name": "Airflow execution modes",
-            "options": [
-                "--standalone-dag-processor",
-            ],
-        },
-        {
-            "name": "Building image before entering shell",
+            "name": "Build CI image (before entering shell)",
             "options": [
                 "--force-build",
                 "--platform",
@@ -113,6 +128,13 @@
                 "--downgrade-sqlalchemy",
             ],
         },
+        {
+            "name": "Other options",
+            "options": [
+                "--forward-credentials",
+                "--max-time",
+            ],
+        },
     ],
     "breeze compile-www-assets": [
         {
@@ -124,18 +146,24 @@
     ],
     "breeze start-airflow": [
         {
-            "name": "Basic flags",
+            "name": "Execution mode",
             "options": [
                 "--python",
+                "--platform",
+                "--integration",
+                "--standalone-dag-processor",
+                "--database-isolation",
                 "--load-example-dags",
                 "--load-default-connections",
+            ],
+        },
+        {
+            "name": "Database",
+            "options": [
                 "--backend",
-                "--platform",
                 "--postgres-version",
                 "--mysql-version",
                 "--mssql-version",
-                "--integration",
-                "--forward-credentials",
                 "--db-reset",
             ],
         },
@@ -147,12 +175,6 @@
                 "--celery-flower",
             ],
         },
-        {
-            "name": "Airflow execution modes",
-            "options": [
-                "--standalone-dag-processor",
-            ],
-        },
         {
             "name": "Asset compilation options",
             "options": [
@@ -161,7 +183,7 @@
             ],
         },
         {
-            "name": "Building image before entering shell",
+            "name": "Build CI image (before entering shell)",
             "options": [
                 "--force-build",
                 "--image-tag",
@@ -185,6 +207,12 @@
                 "--package-format",
             ],
         },
+        {
+            "name": "Other options",
+            "options": [
+                "--forward-credentials",
+            ],
+        },
     ],
     "breeze exec": [
         {"name": "Drops in the interactive shell of active airflow container"},

diff --git a/dev/breeze/src/airflow_breeze/commands/main_command.py b/dev/breeze/src/airflow_breeze/commands/main_command.py
@@ -31,6 +31,7 @@
     option_answer,
     option_backend,
     option_builder,
+    option_database_isolation,
     option_db_reset,
     option_dry_run,
     option_forward_credentials,
@@ -108,6 +109,7 @@ def get_command(self, ctx: Context, cmd_name: str):
 @option_mysql_version
 @option_mssql_version
 @option_integration
+@option_database_isolation
 @option_standalone_dag_processor
 @option_forward_credentials
 @option_db_reset

diff --git a/dev/breeze/src/airflow_breeze/params/shell_params.py b/dev/breeze/src/airflow_breeze/params/shell_params.py
@@ -128,6 +128,7 @@ class ShellParams:
     regenerate_missing_docs: bool = False
     skip_provider_dependencies_check: bool = False
     standalone_dag_processor: bool = False
+    database_isolation: bool = False
 
     def clone_with_test(self, test_type: str) -> ShellParams:
         new_params = deepcopy(self)

diff --git a/dev/breeze/src/airflow_breeze/utils/common_options.py b/dev/breeze/src/airflow_breeze/utils/common_options.py
@@ -581,6 +581,12 @@ def _set_default_from_parent(ctx: click.core.Context, option: click.core.Option,
     is_flag=True,
     envvar="STANDALONE_DAG_PROCESSOR",
 )
+option_database_isolation = click.option(
+    "--database-isolation",
+    help="Run airflow in database isolation mode.",
+    is_flag=True,
+    envvar="DATABASE_ISOLATION",
+)
 option_install_selected_providers = click.option(
     "--install-selected-providers",
     help="Comma-separated list of providers selected to be installed (implies --use-packages-from-dist).",

diff --git a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py
@@ -632,6 +632,7 @@ def update_expected_environment_variables(env: dict[str, str]) -> None:
     "BACKEND": "backend",
     "BASE_BRANCH": "base_branch",
     "COMPOSE_FILE": "compose_file",
+    "DATABASE_ISOLATION": "database_isolation",
     "DB_RESET": "db_reset",
     "DEV_MODE": "dev_mode",
     "DEFAULT_CONSTRAINTS_BRANCH": "default_constraints_branch",

diff --git a/images/breeze/output-commands-hash.txt b/images/breeze/output-commands-hash.txt
@@ -1,7 +1,7 @@
 # This file is automatically generated by pre-commit. If you have a conflict with this file
 # Please do not solve it but run `breeze setup regenerate-command-images`.
 # This command should fix the conflict and regenerate help images that you have conflict with.
-main:4098f8e87f8aab37816dfe7b91f1b675
+main:c1af89b83fe9c261801630807cad77ce
 build-docs:1b1848521361fc7fa711409cf9132108
 ci:find-backtracking-candidates:17fe56b867a745e5032a08dfcd3f73ee
 ci:fix-ownership:3e5a73533cc96045e72cb258783cfc96
@@ -62,8 +62,8 @@ setup:regenerate-command-images:a228071206e00b6004ede2e69b8f4a6c
 setup:self-upgrade:4af905a147fcd6670a0e33d3d369a94b
 setup:version:be116d90a21c2afe01087f7609774e1e
 setup:c569719a4fc7be726556a266e447c871
-shell:676e7d054d7ce079301018a98f406545
-start-airflow:9a5668007f0acf1b954456c12d8b4051
+shell:6840d0c0e939dcd03d127d5a040bf2f5
+start-airflow:7981958d896eebbab19ec32147657eb6
 static-checks:19926b8fcea5784b28d4a0d99865363c
 testing:docker-compose-tests:fd154a058082fcfda12eb877a9a89338
 testing:helm-tests:0669be17b744ba057adbf38681bd8e68