Ejemplo n.º 1
0
class GcfFunctionDeleteOperator(BaseOperator):
    """
    Deletes the specified function from Google Cloud Functions.

    .. seealso::
        For more information on how to use this operator, take a look at the guide:
        :ref:`howto/operator:GcfFunctionDeleteOperator`

    :param name: A fully-qualified function name, matching
        the pattern: `^projects/[^/]+/locations/[^/]+/functions/[^/]+$`
    :type name: str
    :param gcp_conn_id: The connection ID to use to connect to Google Cloud Platform.
    :type gcp_conn_id: str
    :param api_version: API version used (for example v1 or v1beta1).
    :type api_version: str
    """
    # [START gcf_function_delete_template_fields]
    template_fields = ('name', 'gcp_conn_id', 'api_version')
    # [END gcf_function_delete_template_fields]

    @apply_defaults
    def __init__(self,
                 name,
                 gcp_conn_id='google_cloud_default',
                 api_version='v1',
                 *args,
                 **kwargs):
        self.name = name
        self.gcp_conn_id = gcp_conn_id
        self.api_version = api_version
        self._validate_inputs()
        self.hook = GcfHook(gcp_conn_id=self.gcp_conn_id,
                            api_version=self.api_version)
        super().__init__(*args, **kwargs)

    def _validate_inputs(self):
        if not self.name:
            raise AttributeError('Empty parameter: name')
        else:
            pattern = FUNCTION_NAME_COMPILED_PATTERN
            if not pattern.match(self.name):
                raise AttributeError(
                    'Parameter name must match pattern: {}'.format(
                        FUNCTION_NAME_PATTERN))

    def execute(self, context):
        try:
            return self.hook.delete_function(self.name)
        except HttpError as e:
            status = e.resp.status
            if status == 404:
                self.log.info('The function does not exist in this project')
            else:
                self.log.error('An error occurred. Exiting.')
                raise e
Ejemplo n.º 2
0
 def execute(self, context):
     hook = GcfHook(gcp_conn_id=self.gcp_conn_id, api_version=self.api_version)
     try:
         return hook.delete_function(self.name)
     except HttpError as e:
         status = e.resp.status
         if status == 404:
             self.log.info('The function does not exist in this project')
         else:
             self.log.error('An error occurred. Exiting.')
             raise e
Ejemplo n.º 3
0
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')
Ejemplo n.º 4
0
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)