|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +# |
| 3 | +# Licensed to the Apache Software Foundation (ASF) under one |
| 4 | +# or more contributor license agreements. See the NOTICE file |
| 5 | +# distributed with this work for additional information |
| 6 | +# regarding copyright ownership. The ASF licenses this file |
| 7 | +# to you under the Apache License, Version 2.0 (the |
| 8 | +# "License"); you may not use this file except in compliance |
| 9 | +# with the License. You may obtain a copy of the License at |
| 10 | +# |
| 11 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 12 | +# |
| 13 | +# Unless required by applicable law or agreed to in writing, |
| 14 | +# software distributed under the License is distributed on an |
| 15 | +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 16 | +# KIND, either express or implied. See the License for the |
| 17 | +# specific language governing permissions and limitations |
| 18 | +# under the License. |
| 19 | + |
| 20 | +import time |
| 21 | +import requests |
| 22 | +from googleapiclient.discovery import build |
| 23 | + |
| 24 | +from airflow import AirflowException |
| 25 | +from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook |
| 26 | + |
| 27 | +# Number of retries - used by googleapiclient method calls to perform retries |
| 28 | +# For requests that are "retriable" |
| 29 | +NUM_RETRIES = 5 |
| 30 | + |
| 31 | +# Time to sleep between active checks of the operation results |
| 32 | +TIME_TO_SLEEP_IN_SECONDS = 1 |
| 33 | + |
| 34 | + |
| 35 | +# noinspection PyAbstractClass |
| 36 | +class GcfHook(GoogleCloudBaseHook): |
| 37 | + """ |
| 38 | + Hook for Google Cloud Functions APIs. |
| 39 | + """ |
| 40 | + _conn = None |
| 41 | + |
| 42 | + def __init__(self, |
| 43 | + api_version, |
| 44 | + gcp_conn_id='google_cloud_default', |
| 45 | + delegate_to=None): |
| 46 | + super(GcfHook, self).__init__(gcp_conn_id, delegate_to) |
| 47 | + self.api_version = api_version |
| 48 | + |
| 49 | + def get_conn(self): |
| 50 | + """ |
| 51 | + Retrieves connection to cloud functions. |
| 52 | +
|
| 53 | + :return: Google Cloud Function services object |
| 54 | + :rtype: dict |
| 55 | + """ |
| 56 | + if not self._conn: |
| 57 | + http_authorized = self._authorize() |
| 58 | + self._conn = build('cloudfunctions', self.api_version, |
| 59 | + http=http_authorized, cache_discovery=False) |
| 60 | + return self._conn |
| 61 | + |
| 62 | + def get_function(self, name): |
| 63 | + """ |
| 64 | + Returns the function with a given name. |
| 65 | +
|
| 66 | + :param name: name of the function |
| 67 | + :type name: str |
| 68 | + :return: a CloudFunction object representing the function |
| 69 | + :rtype: dict |
| 70 | + """ |
| 71 | + return self.get_conn().projects().locations().functions().get( |
| 72 | + name=name).execute(num_retries=NUM_RETRIES) |
| 73 | + |
| 74 | + def list_functions(self, full_location): |
| 75 | + """ |
| 76 | + Lists all functions created in the location. |
| 77 | +
|
| 78 | + :param full_location: full location including project. On the form |
| 79 | + of /projects/<PROJECT>/location/<LOCATION> |
| 80 | + :type full_location: str |
| 81 | + :return: array of CloudFunction objects - representing functions in the location |
| 82 | + :rtype: [dict] |
| 83 | + """ |
| 84 | + list_response = self.get_conn().projects().locations().functions().list( |
| 85 | + parent=full_location).execute(num_retries=NUM_RETRIES) |
| 86 | + return list_response.get("functions", []) |
| 87 | + |
| 88 | + def create_new_function(self, full_location, body): |
| 89 | + """ |
| 90 | + Creates new cloud function in location given with body specified. |
| 91 | +
|
| 92 | + :param full_location: full location including project. On the form |
| 93 | + of /projects/<PROJECT>/location/<LOCATION> |
| 94 | + :type full_location: str |
| 95 | + :param body: body required by the cloud function insert API |
| 96 | + :type body: dict |
| 97 | + :return: response returned by the operation |
| 98 | + :rtype: dict |
| 99 | + """ |
| 100 | + response = self.get_conn().projects().locations().functions().create( |
| 101 | + location=full_location, |
| 102 | + body=body |
| 103 | + ).execute(num_retries=NUM_RETRIES) |
| 104 | + operation_name = response["name"] |
| 105 | + return self._wait_for_operation_to_complete(operation_name) |
| 106 | + |
| 107 | + def update_function(self, name, body, update_mask): |
| 108 | + """ |
| 109 | + Updates cloud function according to the update mask specified. |
| 110 | +
|
| 111 | + :param name: name of the function |
| 112 | + :type name: str |
| 113 | + :param body: body required by the cloud function patch API |
| 114 | + :type body: str |
| 115 | + :param update_mask: update mask - array of fields that should be patched |
| 116 | + :type update_mask: [str] |
| 117 | + :return: response returned by the operation |
| 118 | + :rtype: dict |
| 119 | + """ |
| 120 | + response = self.get_conn().projects().locations().functions().patch( |
| 121 | + updateMask=",".join(update_mask), |
| 122 | + name=name, |
| 123 | + body=body |
| 124 | + ).execute(num_retries=NUM_RETRIES) |
| 125 | + operation_name = response["name"] |
| 126 | + return self._wait_for_operation_to_complete(operation_name) |
| 127 | + |
| 128 | + def upload_function_zip(self, parent, zip_path): |
| 129 | + """ |
| 130 | + Uploads zip file with sources. |
| 131 | +
|
| 132 | + :param parent: project and location in which signed upload URL should be generated |
| 133 | + in the form of /projects/<PROJECT>/location/<LOCATION> |
| 134 | + :type parent: str |
| 135 | + :param zip_path: path of the file to upload (should point to valid .zip file) |
| 136 | + :type zip_path: str |
| 137 | + :return: Upload URL that was returned by generateUploadUrl method |
| 138 | + """ |
| 139 | + response = self.get_conn().projects().locations().functions().generateUploadUrl( |
| 140 | + parent=parent |
| 141 | + ).execute(num_retries=NUM_RETRIES) |
| 142 | + upload_url = response.get('uploadUrl') |
| 143 | + with open(zip_path, 'rb') as fp: |
| 144 | + requests.put( |
| 145 | + url=upload_url, |
| 146 | + data=fp.read(), |
| 147 | + # Those two headers needs to be specified according to: |
| 148 | + # https://cloud.google.com/functions/docs/reference/rest/v1/projects.locations.functions/generateUploadUrl |
| 149 | + # nopep8 |
| 150 | + headers={ |
| 151 | + 'Content-type': 'application/zip', |
| 152 | + 'x-goog-content-length-range': '0,104857600', |
| 153 | + } |
| 154 | + ) |
| 155 | + return upload_url |
| 156 | + |
| 157 | + def delete_function(self, name): |
| 158 | + """ |
| 159 | + Deletes cloud function specified by name. |
| 160 | +
|
| 161 | + :param name: name of the function |
| 162 | + :type name: str |
| 163 | + :return: response returned by the operation |
| 164 | + :rtype: dict |
| 165 | + """ |
| 166 | + response = self.get_conn().projects().locations().functions().delete( |
| 167 | + name=name).execute(num_retries=NUM_RETRIES) |
| 168 | + operation_name = response["name"] |
| 169 | + return self._wait_for_operation_to_complete(operation_name) |
| 170 | + |
| 171 | + def _wait_for_operation_to_complete(self, operation_name): |
| 172 | + """ |
| 173 | + Waits for the named operation to complete - checks status of the |
| 174 | + asynchronous call. |
| 175 | +
|
| 176 | + :param operation_name: name of the operation |
| 177 | + :type operation_name: str |
| 178 | + :return: response returned by the operation |
| 179 | + :rtype: dict |
| 180 | + :exception: AirflowException in case error is returned |
| 181 | + """ |
| 182 | + service = self.get_conn() |
| 183 | + while True: |
| 184 | + operation_response = service.operations().get( |
| 185 | + name=operation_name, |
| 186 | + ).execute(num_retries=NUM_RETRIES) |
| 187 | + if operation_response.get("done"): |
| 188 | + response = operation_response.get("response") |
| 189 | + error = operation_response.get("error") |
| 190 | + # Note, according to documentation always either response or error is |
| 191 | + # set when "done" == True |
| 192 | + if error: |
| 193 | + raise AirflowException(str(error)) |
| 194 | + return response |
| 195 | + time.sleep(TIME_TO_SLEEP_IN_SECONDS) |
0 commit comments