From f47d7b95fe68fe4c0c9db6503ddbf2ed13ea43dd Mon Sep 17 00:00:00 2001
From: Niko <65743084+o-nikolas@users.noreply.github.com>
Date: Fri, 22 Oct 2021 17:30:33 -0700
Subject: [PATCH] Allow specifying extras when using breeze
 initialize_local_virtualenv (#19178)

---
 LOCAL_VIRTUALENV.rst                    | 7 +++++++
 breeze                                  | 5 ++++-
 scripts/ci/libraries/_initialization.sh | 6 ++++++
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/LOCAL_VIRTUALENV.rst b/LOCAL_VIRTUALENV.rst
index e96fd6a068170..bf982f20ed40a 100644
--- a/LOCAL_VIRTUALENV.rst
+++ b/LOCAL_VIRTUALENV.rst
@@ -204,6 +204,13 @@ Activate your virtualenv, e.g. by using ``workon``, and once you are in it, run:
 
   ./breeze initialize-local-virtualenv
 
+By default Breeze installs the ``devel`` extra only. You can optionally control which extras are installed by exporting ``VIRTUALENV_EXTRAS`` before calling Breeze:
+
+.. code-block:: bash
+
+  export VIRTUALENV_EXTRAS="devel,google,postgres"
+  ./breeze initialize-local-virtualenv
+
 5. (optionally) run yarn build if you plan to run the webserver
 
 .. code-block:: bash
diff --git a/breeze b/breeze
index d85b17c20fb4c..d332271ea0feb 100755
--- a/breeze
+++ b/breeze
@@ -230,6 +230,7 @@ function breeze::setup_default_breeze_constants() {
 #    PYTHON_MAJOR_MINOR_VERSION
 #    AIRFLOW_HOME_DIR
 #    AIRFLOW_SOURCES
+#    VIRTUALENV_EXTRAS
 #    DEFAULT_CONSTRAINTS_BRANCH
 #    OSTYPE
 #
@@ -252,13 +253,15 @@ function breeze::initialize_virtualenv() {
         echo
         echo "Initializing the virtualenv: $(command -v python)!"
         echo
+        echo "Extras to be installed: ${VIRTUALENV_EXTRAS}"
+        echo
         echo "This will wipe out ${AIRFLOW_HOME_DIR} and reset all the databases!"
         echo
         "${AIRFLOW_SOURCES}/confirm" "Proceeding with the initialization"
         echo
         pushd "${AIRFLOW_SOURCES}" >/dev/null 2>&1 || exit 1
         set +e
-        pip install -e ".[devel]" \
+        pip install -e ".[${VIRTUALENV_EXTRAS}]" \
             --constraint "https://raw.githubusercontent.com/${CONSTRAINTS_GITHUB_REPOSITORY}/${DEFAULT_CONSTRAINTS_BRANCH}/constraints-source-providers-${PYTHON_MAJOR_MINOR_VERSION}.txt"
         res=$?
         set -e
diff --git a/scripts/ci/libraries/_initialization.sh b/scripts/ci/libraries/_initialization.sh
index cad80ffb3d130..0a14af7bc2447 100644
--- a/scripts/ci/libraries/_initialization.sh
+++ b/scripts/ci/libraries/_initialization.sh
@@ -564,6 +564,11 @@ function initialization::initialize_kubernetes_variables() {
     readonly API_SERVER_PORT
 }
 
+function initialization::initialize_virtualenv_variables() {
+    # The extras to install when initializing a virtual env with breeze
+    export VIRTUALENV_EXTRAS=${VIRTUALENV_EXTRAS:="devel"}
+}
+
 function initialization::initialize_git_variables() {
     # SHA of the commit for the current sources
     COMMIT_SHA="$(git rev-parse HEAD 2>/dev/null || echo "Unknown")"
@@ -638,6 +643,7 @@ function initialization::initialize_common_environment() {
     initialization::initialize_image_build_variables
     initialization::initialize_provider_package_building
     initialization::initialize_kubernetes_variables
+    initialization::initialize_virtualenv_variables
     initialization::initialize_git_variables
     initialization::initialize_github_variables
     initialization::initialize_test_variables