Allow downloading requirements file from GCS in `BeamRunPythonPipelin…

…eOperator` (#31645)

* Update beam python run operator to download requirements if needed

---------

Co-authored-by: Hussein Awala <hussein@awala.fr>
Co-authored-by: eladkal <45845474+eladkal@users.noreply.github.com>
Co-authored-by: Peng Yu <peng.montreal@gmail.com>

airflow/providers/apache/beam/operators/beam.py

@@ -309,10 +309,15 @@ def execute(self, context: Context):
 
     def execute_sync(self, context: Context):
         with ExitStack() as exit_stack:
+            gcs_hook = GCSHook(gcp_conn_id=self.gcp_conn_id)
             if self.py_file.lower().startswith("gs://"):
-                gcs_hook = GCSHook(gcp_conn_id=self.gcp_conn_id)
                 tmp_gcs_file = exit_stack.enter_context(gcs_hook.provide_file(object_url=self.py_file))
                 self.py_file = tmp_gcs_file.name
+            if self.snake_case_pipeline_options.get("requirements_file", "").startswith("gs://"):
+                tmp_req_file = exit_stack.enter_context(
+                    gcs_hook.provide_file(object_url=self.snake_case_pipeline_options["requirements_file"])
+                )
+                self.snake_case_pipeline_options["requirements_file"] = tmp_req_file.name
 
             if self.is_dataflow and self.dataflow_hook:
                 with self.dataflow_hook.provide_authorized_gcloud():

tests/providers/apache/beam/operators/test_beam.py

@@ -39,6 +39,7 @@
 JAR_FILE = "gs://my-bucket/example/test.jar"
 JOB_CLASS = "com.test.NotMain"
 PY_FILE = "gs://my-bucket/my-object.py"
+REQURIEMENTS_FILE = "gs://my-bucket/my-requirements.txt"
 PY_INTERPRETER = "python3"
 PY_OPTIONS = ["-m"]
 GO_FILE = "gs://my-bucket/example/main.go"
@@ -48,6 +49,10 @@
     "project": "test",
     "stagingLocation": "gs://test/staging",
 }
+PY_DEFAULT_OPTIONS = {
+    **DEFAULT_OPTIONS,
+    "requirements_file": REQURIEMENTS_FILE,
+}
 ADDITIONAL_OPTIONS = {"output": "gs://test/output", "labels": {"foo": "bar"}}
 TEST_VERSION = f"v{version.replace('.', '-').replace('+', '-')}"
 EXPECTED_ADDITIONAL_OPTIONS = {
@@ -63,7 +68,7 @@ def setup_method(self):
             task_id=TASK_ID,
             py_file=PY_FILE,
             py_options=PY_OPTIONS,
-            default_pipeline_options=DEFAULT_OPTIONS,
+            default_pipeline_options=PY_DEFAULT_OPTIONS,
             pipeline_options=ADDITIONAL_OPTIONS,
         )
 
@@ -74,7 +79,7 @@ def test_init(self):
         assert self.operator.runner == DEFAULT_RUNNER
         assert self.operator.py_options == PY_OPTIONS
         assert self.operator.py_interpreter == PY_INTERPRETER
-        assert self.operator.default_pipeline_options == DEFAULT_OPTIONS
+        assert self.operator.default_pipeline_options == PY_DEFAULT_OPTIONS
         assert self.operator.pipeline_options == EXPECTED_ADDITIONAL_OPTIONS
 
     @mock.patch("airflow.providers.apache.beam.operators.beam.BeamHook")
@@ -92,8 +97,10 @@ def test_exec_direct_runner(self, gcs_hook, beam_hook_mock):
             "staging_location": "gs://test/staging",
             "output": "gs://test/output",
             "labels": {"foo": "bar", "airflow-version": TEST_VERSION},
+            "requirements_file": gcs_provide_file.return_value.__enter__.return_value.name,
         }
-        gcs_provide_file.assert_called_once_with(object_url=PY_FILE)
+        gcs_provide_file.assert_any_call(object_url=PY_FILE)
+        gcs_provide_file.assert_any_call(object_url=REQURIEMENTS_FILE)
         start_python_hook.assert_called_once_with(
             variables=expected_options,
             py_file=gcs_provide_file.return_value.__enter__.return_value.name,
@@ -134,8 +141,10 @@ def test_exec_dataflow_runner(self, gcs_hook, dataflow_hook_mock, beam_hook_mock
             "labels": {"foo": "bar", "airflow-version": TEST_VERSION},
             "region": "us-central1",
             "impersonate_service_account": TEST_IMPERSONATION_ACCOUNT,
+            "requirements_file": gcs_provide_file.return_value.__enter__.return_value.name,
         }
-        gcs_provide_file.assert_called_once_with(object_url=PY_FILE)
+        gcs_provide_file.assert_any_call(object_url=PY_FILE)
+        gcs_provide_file.assert_any_call(object_url=REQURIEMENTS_FILE)
         persist_link_mock.assert_called_once_with(
             self.operator,
             None,