Support YAML input for CloudBuildCreateOperator (#8808)

apache · May 23, 2020 · cf5cf45 · cf5cf45
1 parent db70da2
commit cf5cf45
Show file tree

Hide file tree

Showing 6 changed files with 98 additions and 9 deletions.
diff --git a/airflow/providers/google/cloud/example_dags/example_cloud_build.py b/airflow/providers/google/cloud/example_dags/example_cloud_build.py
@@ -29,6 +29,7 @@
 """
 
 import os
+from pathlib import Path
 
 from future.backports.urllib.parse import urlparse
 
@@ -45,6 +46,8 @@
 GCP_SOURCE_ARCHIVE_URL_PARTS = urlparse(GCP_SOURCE_ARCHIVE_URL)
 GCP_SOURCE_BUCKET_NAME = GCP_SOURCE_ARCHIVE_URL_PARTS.netloc
 
+CURRENT_FOLDER = Path(__file__).parent
+
 # [START howto_operator_gcp_create_build_from_storage_body]
 create_build_from_storage_body = {
  "source": {"storageSource": GCP_SOURCE_ARCHIVE_URL},
@@ -99,6 +102,13 @@
  task_id="create_build_from_repo_result",
  )
 
+ # [START howto_operator_gcp_create_build_from_yaml_body]
+ create_build_from_file = CloudBuildCreateOperator(
+ task_id="create_build_from_file", project_id=GCP_PROJECT_ID,
+ body=str(CURRENT_FOLDER.joinpath('example_cloud_build.yaml')),
+ params={'name': 'Airflow'}
+ )
+ # [END howto_operator_gcp_create_build_from_yaml_body]
  create_build_from_storage >> create_build_from_storage_result # pylint: disable=pointless-statement
 
  create_build_from_repo >> create_build_from_repo_result # pylint: disable=pointless-statement
diff --git a/airflow/providers/google/cloud/example_dags/example_cloud_build.yaml b/airflow/providers/google/cloud/example_dags/example_cloud_build.yaml
@@ -0,0 +1,20 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+steps:
+ - name: 'ubuntu'
+ args: ['echo', 'Hello {{ params.name}}']
diff --git a/airflow/providers/google/cloud/hooks/cloud_build.py b/airflow/providers/google/cloud/hooks/cloud_build.py
@@ -75,7 +75,7 @@ def create_build(self, body: Dict, project_id: str) -> Dict:
  Starts a build with the specified configuration.
 
  :param body: The request body.
- See: https://cloud.google.com/cloud-build/docs/api/reference/rest/Shared.Types/Build
+ See: https://cloud.google.com/cloud-build/docs/api/reference/rest/v1/projects.builds
  :type body: dict
  :param project_id: Optional, Google Cloud Project project_id where the function belongs.
  If set to None or missing, the default project_id from the GCP connection is used.

diff --git a/airflow/providers/google/cloud/operators/cloud_build.py b/airflow/providers/google/cloud/operators/cloud_build.py
@@ -16,11 +16,14 @@
 # specific language governing permissions and limitations
 # under the License.
 """Operators that integrat with Google Cloud Build service."""
+import json
 import re
 from copy import deepcopy
-from typing import Any, Dict, Iterable, Optional
+from typing import Any, Dict, Iterable, Optional, Union
 from urllib.parse import unquote, urlparse
 
+import yaml
+
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.cloud_build import CloudBuildHook
@@ -39,9 +42,10 @@ class BuildProcessor:
  * It is possible to provide the source as the URL address instead dict.
 
  :param body: The request body.
- See: https://cloud.google.com/cloud-build/docs/api/reference/rest/Shared.Types/Build
+ See: https://cloud.google.com/cloud-build/docs/api/reference/rest/v1/projects.builds
  :type body: dict
  """
+
  def __init__(self, body: Dict) -> None:
  self.body = deepcopy(body)
 
@@ -90,8 +94,9 @@ def process_body(self):
  :return: the body.
  :type: dict
  """
- self._verify_source()
- self._reformat_source()
+ if 'source' in self.body:
+ self._verify_source()
+ self._reformat_source()
  return self.body
 
  @staticmethod
@@ -162,9 +167,10 @@ class CloudBuildCreateOperator(BaseOperator):
  For more information on how to use this operator, take a look at the guide:
  :ref:`howto/operator:CloudBuildCreateOperator`
 
- :param body: The request body.
- See: https://cloud.google.com/cloud-build/docs/api/reference/rest/Shared.Types/Build
- :type body: dict
+ :param body: The build config with instructions to perform with CloudBuild.
+ Can be a dictionary or path to a file type like YAML or JSON.
+ See: https://cloud.google.com/cloud-build/docs/api/reference/rest/v1/projects.builds
+ :type body: dict or string
  :param project_id: ID of the Google Cloud project if None then
  default project_id is used.
  :type project_id: str
@@ -175,21 +181,34 @@ class CloudBuildCreateOperator(BaseOperator):
  """
 
  template_fields = ("body", "gcp_conn_id", "api_version") # type: Iterable[str]
+ template_ext = ['.yml', '.yaml', '.json']
 
  @apply_defaults
  def __init__(self,
- body: dict,
+ body: Union[dict, str],
  project_id: Optional[str] = None,
  gcp_conn_id: str = "google_cloud_default",
  api_version: str = "v1",
  *args, **kwargs) -> None:
  super().__init__(*args, **kwargs)
  self.body = body
+ # Not template fields to keep original value
+ self.body_raw = body
  self.project_id = project_id
  self.gcp_conn_id = gcp_conn_id
  self.api_version = api_version
  self._validate_inputs()
 
+ def prepare_template(self) -> None:
+ # if no file is specified, skip
+ if not isinstance(self.body_raw, str):
+ return
+ with open(self.body_raw, 'r') as file:
+ if any(self.body_raw.endswith(ext) for ext in ['.yaml', '.yml']):
+ self.body = yaml.load(file.read(), Loader=yaml.FullLoader)
+ if self.body_raw.endswith('.json'):
+ self.body = json.loads(file.read())
+
  def _validate_inputs(self):
  if not self.body:
  raise AirflowException("The required parameter 'body' is missing")

diff --git a/docs/howto/operator/gcp/cloud_build.rst b/docs/howto/operator/gcp/cloud_build.rst
@@ -42,6 +42,7 @@ Build configuration overview
 
 In order to trigger a build, it is necessary to pass the build configuration.
 
+
 .. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_build.py
  :language: python
  :dedent: 0
@@ -80,6 +81,14 @@ It is also possible to specify it using the URL:
  :start-after: [START howto_operator_gcp_cloud_build_source_repo_url]
  :end-before: [END howto_operator_gcp_cloud_build_source_repo_url]
 
+It is also possible to specify it using a YAML or JSON format.
+
+.. exampleinclude:: ../../../../airflow/providers/google/cloud/example_dags/example_cloud_build.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_operator_gcp_create_build_from_yaml_body]
+ :end-before: [END howto_operator_gcp_create_build_from_yaml_body]
+
 Read `Build Configuration Overview <https://cloud.google.com/cloud-build/docs/build-config>`__ to understand all the fields you can include in a build config file.
 
 

diff --git a/tests/providers/google/cloud/operators/test_cloud_build.py b/tests/providers/google/cloud/operators/test_cloud_build.py
@@ -17,13 +17,16 @@
 # under the License.
 """Tests for Google Cloud Build operators """
 
+import tempfile
 from copy import deepcopy
+from datetime import datetime
 from unittest import TestCase
 
 import mock
 from parameterized import parameterized
 
 from airflow.exceptions import AirflowException
+from airflow.models import DAG, TaskInstance
 from airflow.providers.google.cloud.operators.cloud_build import BuildProcessor, CloudBuildCreateOperator
 
 TEST_CREATE_BODY = {
@@ -34,6 +37,7 @@
  "images": ["gcr.io/$PROJECT_ID/my-image"],
 }
 TEST_PROJECT_ID = "example-id"
+TEST_DEFAULT_DATE = datetime(year=2020, month=1, day=1)
 
 
 class TestBuildProcessor(TestCase):
@@ -180,6 +184,7 @@ def test_repo_source_replace(self, hook_mock):
  operator = CloudBuildCreateOperator(
  body=current_body, project_id=TEST_PROJECT_ID, task_id="task-id"
  )
+
  return_value = operator.execute({})
  expected_body = {
  # [START howto_operator_gcp_cloud_build_source_repo_dict]
@@ -203,3 +208,29 @@ def test_repo_source_replace(self, hook_mock):
  body=expected_body, project_id=TEST_PROJECT_ID
  )
  self.assertEqual(return_value, TEST_CREATE_BODY)
+
+ def test_load_templated_yaml(self):
+ dag = DAG(dag_id='example_cloudbuild_operator', start_date=TEST_DEFAULT_DATE)
+ with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w+t') as build:
+ build.writelines("""
+ steps:
+ - name: 'ubuntu'
+ args: ['echo', 'Hello {{ params.name }}!']
+ """)
+ build.seek(0)
+ body_path = build.name
+ operator = CloudBuildCreateOperator(
+ body=body_path,
+ task_id="task-id", dag=dag,
+ params={'name': 'airflow'}
+ )
+ operator.prepare_template()
+ ti = TaskInstance(operator, TEST_DEFAULT_DATE)
+ ti.render_templates()
+ expected_body = {'steps': [
+ {'name': 'ubuntu',
+ 'args': ['echo', 'Hello airflow!']
+ }
+ ]
+ }
+ self.assertEqual(expected_body, operator.body)