class TestFunctionHookNoDefaultProjectId(unittest.TestCase):
def setUp(self):
with mock.patch(
'airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__',
new=mock_base_gcp_hook_no_default_project_id):
self.gcf_function_hook_no_project_id = GcfHook(gcp_conn_id='test',
api_version='v1')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.functions.GcfHook._wait_for_operation_to_complete')
def test_create_new_function_missing_project_id(
self, wait_for_operation_to_complete, get_conn):
create_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.create
execute_method = create_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
with self.assertRaises(AirflowException) as cm:
self.gcf_function_hook_no_project_id.create_new_function(
location=GCF_LOCATION, body={})
create_method.assert_not_called()
execute_method.assert_not_called()
err = cm.exception
self.assertIn("The project id must be passed", str(err))
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.functions.GcfHook._wait_for_operation_to_complete')
def test_create_new_function_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
create_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.create
execute_method = create_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook_no_project_id.create_new_function(
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
location=GCF_LOCATION,
body={})
self.assertIsNone(res)
create_method.assert_called_once_with(
body={}, location='projects/example-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id')
@mock.patch('requests.put')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
def test_upload_function_zip_missing_project_id(self, get_conn,
requests_put):
mck = mock.mock_open()
with mock.patch('builtins.open', mck):
generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.generateUploadUrl
execute_method = generate_upload_url_method.return_value.execute
execute_method.return_value = {"uploadUrl": "http://uploadHere"}
requests_put.return_value = None
with self.assertRaises(AirflowException) as cm:
self.gcf_function_hook_no_project_id.upload_function_zip(
location=GCF_LOCATION, zip_path="/tmp/path.zip")
generate_upload_url_method.assert_not_called()
execute_method.assert_not_called()
mck.assert_not_called()
err = cm.exception
self.assertIn("The project id must be passed", str(err))
@mock.patch('requests.put')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
def test_upload_function_zip_overridden_project_id(self, get_conn,
requests_put):
mck, open_module = get_open_mock()
with mock.patch('{}.open'.format(open_module), mck):
generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.generateUploadUrl
execute_method = generate_upload_url_method.return_value.execute
execute_method.return_value = {"uploadUrl": "http://uploadHere"}
requests_put.return_value = None
res = self.gcf_function_hook_no_project_id.upload_function_zip(
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
location=GCF_LOCATION,
zip_path="/tmp/path.zip")
self.assertEqual("http://uploadHere", res)
generate_upload_url_method.assert_called_once_with(
parent='projects/example-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
requests_put.assert_called_once_with(
data=mock.ANY,
headers={
'Content-type': 'application/zip',
'x-goog-content-length-range': '0,104857600'
},
url='http://uploadHere')
class TestFunctionHookDefaultProjectId(unittest.TestCase):
def setUp(self):
with mock.patch(
'airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__',
new=mock_base_gcp_hook_default_project_id):
self.gcf_function_hook = GcfHook(gcp_conn_id='test',
api_version='v1')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.functions.GcfHook._wait_for_operation_to_complete')
def test_create_new_function(self, wait_for_operation_to_complete,
get_conn):
create_method = get_conn.return_value.projects.return_value.locations.\
return_value.functions.return_value.create
execute_method = create_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook.create_new_function(location=GCF_LOCATION,
body={})
self.assertIsNone(res)
create_method.assert_called_once_with(
body={}, location='projects/example-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.functions.GcfHook._wait_for_operation_to_complete')
def test_create_new_function_override_project_id(
self, wait_for_operation_to_complete, get_conn):
create_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.create
execute_method = create_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook.create_new_function(
project_id='new-project', location=GCF_LOCATION, body={})
self.assertIsNone(res)
create_method.assert_called_once_with(
body={}, location='projects/new-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
def test_get_function(self, get_conn):
get_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.get
execute_method = get_method.return_value.execute
execute_method.return_value = {"name": "function"}
res = self.gcf_function_hook.get_function(name=GCF_FUNCTION)
self.assertIsNotNone(res)
self.assertEqual('function', res['name'])
get_method.assert_called_once_with(name='function')
execute_method.assert_called_once_with(num_retries=5)
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.functions.GcfHook._wait_for_operation_to_complete')
def test_delete_function(self, wait_for_operation_to_complete, get_conn):
delete_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.delete
execute_method = delete_method.return_value.execute
wait_for_operation_to_complete.return_value = None
execute_method.return_value = {"name": "operation_id"}
res = self.gcf_function_hook.delete_function( # pylint: disable=assignment-from-no-return
name=GCF_FUNCTION)
self.assertIsNone(res)
delete_method.assert_called_once_with(name='function')
execute_method.assert_called_once_with(num_retries=5)
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.functions.GcfHook._wait_for_operation_to_complete')
def test_update_function(self, wait_for_operation_to_complete, get_conn):
patch_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.patch
execute_method = patch_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook.update_function( # pylint: disable=assignment-from-no-return
update_mask=['a', 'b', 'c'],
name=GCF_FUNCTION,
body={})
self.assertIsNone(res)
patch_method.assert_called_once_with(body={},
name='function',
updateMask='a,b,c')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id')
@mock.patch('requests.put')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
def test_upload_function_zip(self, get_conn, requests_put):
mck, open_module = get_open_mock()
with mock.patch('{}.open'.format(open_module), mck):
generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.generateUploadUrl
execute_method = generate_upload_url_method.return_value.execute
execute_method.return_value = {"uploadUrl": "http://uploadHere"}
requests_put.return_value = None
res = self.gcf_function_hook.upload_function_zip(
location=GCF_LOCATION, zip_path="/tmp/path.zip")
self.assertEqual("http://uploadHere", res)
generate_upload_url_method.assert_called_once_with(
parent='projects/example-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
requests_put.assert_called_once_with(
data=mock.ANY,
headers={
'Content-type': 'application/zip',
'x-goog-content-length-range': '0,104857600'
},
url='http://uploadHere')
@mock.patch('requests.put')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
def test_upload_function_zip_overridden_project_id(self, get_conn,
requests_put):
mck, open_module = get_open_mock()
with mock.patch('{}.open'.format(open_module), mck):
generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.generateUploadUrl
execute_method = generate_upload_url_method.return_value.execute
execute_method.return_value = {"uploadUrl": "http://uploadHere"}
requests_put.return_value = None
res = self.gcf_function_hook.upload_function_zip(
project_id='new-project',
location=GCF_LOCATION,
zip_path="/tmp/path.zip")
self.assertEqual("http://uploadHere", res)
generate_upload_url_method.assert_called_once_with(
parent='projects/new-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
requests_put.assert_called_once_with(
data=mock.ANY,
headers={
'Content-type': 'application/zip',
'x-goog-content-length-range': '0,104857600'
},
url='http://uploadHere')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
def test_call_function(self, mock_get_conn):
payload = {'executionId': 'wh41ppcyoa6l', 'result': 'Hello World!'}
call = mock_get_conn.return_value.projects.return_value.\
locations.return_value.functions.return_value.call
call.return_value.execute.return_value = payload
function_id = "function1234"
input_data = {'key': 'value'}
name = "projects/{project_id}/locations/{location}/functions/{function_id}".format(
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
location=GCF_LOCATION,
function_id=function_id)
result = self.gcf_function_hook.call_function(
function_id=function_id,
location=GCF_LOCATION,
input_data=input_data,
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST)
call.assert_called_once_with(body=input_data, name=name)
self.assertDictEqual(result, payload)
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
def test_call_function_error(self, mock_get_conn):
payload = {'error': 'Something very bad'}
call = mock_get_conn.return_value.projects.return_value. \
locations.return_value.functions.return_value.call
call.return_value.execute.return_value = payload
function_id = "function1234"
input_data = {'key': 'value'}
with self.assertRaises(AirflowException):
self.gcf_function_hook.call_function(
function_id=function_id,
location=GCF_LOCATION,
input_data=input_data,
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST)
class TestFunctionHookDefaultProjectId(unittest.TestCase):
def setUp(self):
with mock.patch(
'airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__',
new=mock_base_gcp_hook_default_project_id):
self.gcf_function_hook = GcfHook(gcp_conn_id='test',
api_version='v1')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.functions.GcfHook._wait_for_operation_to_complete')
def test_create_new_function(self, wait_for_operation_to_complete,
get_conn):
create_method = get_conn.return_value.projects.return_value.locations.\
return_value.functions.return_value.create
execute_method = create_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook.create_new_function(location=GCF_LOCATION,
body={})
self.assertIsNone(res)
create_method.assert_called_once_with(
body={}, location='projects/example-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.functions.GcfHook._wait_for_operation_to_complete')
def test_create_new_function_override_project_id(
self, wait_for_operation_to_complete, get_conn):
create_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.create
execute_method = create_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook.create_new_function(
project_id='new-project', location=GCF_LOCATION, body={})
self.assertIsNone(res)
create_method.assert_called_once_with(
body={}, location='projects/new-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
def test_get_function(self, get_conn):
get_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.get
execute_method = get_method.return_value.execute
execute_method.return_value = {"name": "function"}
res = self.gcf_function_hook.get_function(name=GCF_FUNCTION)
self.assertIsNotNone(res)
self.assertEqual('function', res['name'])
get_method.assert_called_once_with(name='function')
execute_method.assert_called_once_with(num_retries=5)
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.functions.GcfHook._wait_for_operation_to_complete')
def test_delete_function(self, wait_for_operation_to_complete, get_conn):
delete_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.delete
execute_method = delete_method.return_value.execute
wait_for_operation_to_complete.return_value = None
execute_method.return_value = {"name": "operation_id"}
res = self.gcf_function_hook.delete_function( # pylint: disable=assignment-from-no-return
name=GCF_FUNCTION)
self.assertIsNone(res)
delete_method.assert_called_once_with(name='function')
execute_method.assert_called_once_with(num_retries=5)
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.functions.GcfHook._wait_for_operation_to_complete')
def test_update_function(self, wait_for_operation_to_complete, get_conn):
patch_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.patch
execute_method = patch_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gcf_function_hook.update_function( # pylint: disable=assignment-from-no-return
update_mask=['a', 'b', 'c'],
name=GCF_FUNCTION,
body={})
self.assertIsNone(res)
patch_method.assert_called_once_with(body={},
name='function',
updateMask='a,b,c')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id')
@mock.patch('requests.put')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
def test_upload_function_zip(self, get_conn, requests_put):
mck, open_module = get_open_mock()
with mock.patch('{}.open'.format(open_module), mck):
generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.generateUploadUrl
execute_method = generate_upload_url_method.return_value.execute
execute_method.return_value = {"uploadUrl": "http://uploadHere"}
requests_put.return_value = None
res = self.gcf_function_hook.upload_function_zip(
location=GCF_LOCATION, zip_path="/tmp/path.zip")
self.assertEqual("http://uploadHere", res)
generate_upload_url_method.assert_called_once_with(
parent='projects/example-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
requests_put.assert_called_once_with(
data=mock.ANY,
headers={
'Content-type': 'application/zip',
'x-goog-content-length-range': '0,104857600'
},
url='http://uploadHere')
@mock.patch('requests.put')
@mock.patch('airflow.gcp.hooks.functions.GcfHook.get_conn')
def test_upload_function_zip_overridden_project_id(self, get_conn,
requests_put):
mck, open_module = get_open_mock()
with mock.patch('{}.open'.format(open_module), mck):
generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \
return_value.functions.return_value.generateUploadUrl
execute_method = generate_upload_url_method.return_value.execute
execute_method.return_value = {"uploadUrl": "http://uploadHere"}
requests_put.return_value = None
res = self.gcf_function_hook.upload_function_zip(
project_id='new-project',
location=GCF_LOCATION,
zip_path="/tmp/path.zip")
self.assertEqual("http://uploadHere", res)
generate_upload_url_method.assert_called_once_with(
parent='projects/new-project/locations/location')
execute_method.assert_called_once_with(num_retries=5)
requests_put.assert_called_once_with(
data=mock.ANY,
headers={
'Content-type': 'application/zip',
'x-goog-content-length-range': '0,104857600'
},
url='http://uploadHere')
class GcfFunctionDeployOperator(BaseOperator):
"""
Creates a function in Google Cloud Functions.
If a function with this name already exists, it will be updated.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:GcfFunctionDeployOperator`
:param location: Google Cloud Platform region where the function should be created.
:type location: str
:param body: Body of the Cloud Functions definition. The body must be a
Cloud Functions dictionary as described in:
https://cloud.google.com/functions/docs/reference/rest/v1/projects.locations.functions
. Different API versions require different variants of the Cloud Functions
dictionary.
:type body: dict or google.cloud.functions.v1.CloudFunction
:param project_id: (Optional) Google Cloud Platform project ID where the function
should be created.
:type project_id: str
:param gcp_conn_id: (Optional) The connection ID used to connect to Google Cloud
Platform - default 'google_cloud_default'.
:type gcp_conn_id: str
:param api_version: (Optional) API version used (for example v1 - default - or
v1beta1).
:type api_version: str
:param zip_path: Path to zip file containing source code of the function. If the path
is set, the sourceUploadUrl should not be specified in the body or it should
be empty. Then the zip file will be uploaded using the upload URL generated
via generateUploadUrl from the Cloud Functions API.
:type zip_path: str
:param validate_body: If set to False, body validation is not performed.
:type validate_body: bool
"""
# [START gcf_function_deploy_template_fields]
template_fields = ('project_id', 'location', 'gcp_conn_id', 'api_version')
# [END gcf_function_deploy_template_fields]
@apply_defaults
def __init__(self,
location,
body,
project_id=None,
gcp_conn_id='google_cloud_default',
api_version='v1',
zip_path=None,
validate_body=True,
*args,
**kwargs):
self.project_id = project_id
self.location = location
self.body = body
self.gcp_conn_id = gcp_conn_id
self.api_version = api_version
self.zip_path = zip_path
self.zip_path_preprocessor = ZipPathPreprocessor(body, zip_path)
self._field_validator = None
if validate_body:
self._field_validator = GcpBodyFieldValidator(
CLOUD_FUNCTION_VALIDATION, api_version=api_version)
self._hook = GcfHook(gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version)
self._validate_inputs()
super().__init__(*args, **kwargs)
def _validate_inputs(self):
if not self.location:
raise AirflowException(
"The required parameter 'location' is missing")
if not self.body:
raise AirflowException("The required parameter 'body' is missing")
self.zip_path_preprocessor.preprocess_body()
def _validate_all_body_fields(self):
if self._field_validator:
self._field_validator.validate(self.body)
def _create_new_function(self):
self._hook.create_new_function(project_id=self.project_id,
location=self.location,
body=self.body)
def _update_function(self):
self._hook.update_function(self.body['name'], self.body,
self.body.keys())
def _check_if_function_exists(self):
name = self.body.get('name')
if not name:
raise GcpFieldValidationException(
"The 'name' field should be present in "
"body: '{}'.".format(self.body))
try:
self._hook.get_function(name)
except HttpError as e:
status = e.resp.status
if status == 404:
return False
raise e
return True
def _upload_source_code(self):
return self._hook.upload_function_zip(project_id=self.project_id,
location=self.location,
zip_path=self.zip_path)
def _set_airflow_version_label(self):
if 'labels' not in self.body.keys():
self.body['labels'] = {}
self.body['labels'].update({
'airflow-version':
'v' + version.replace('.', '-').replace('+', '-')
})
def execute(self, context):
if self.zip_path_preprocessor.should_upload_function():
self.body[GCF_SOURCE_UPLOAD_URL] = self._upload_source_code()
self._validate_all_body_fields()
self._set_airflow_version_label()
if not self._check_if_function_exists():
self._create_new_function()
else:
self._update_function()