def poke(self, context): self.log.info('Sensor checks existence of objects: %s, %s', self.bucket, self.prefix) hook = GCSHook(google_cloud_storage_conn_id=self.google_cloud_conn_id, delegate_to=self.delegate_to) self._matches = hook.list(self.bucket, prefix=self.prefix) return bool(self._matches)
def execute(self, context): self.gcs_hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) self.gdrive_hook = GoogleDriveHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) if WILDCARD in self.source_object: total_wildcards = self.source_object.count(WILDCARD) if total_wildcards > 1: error_msg = ( "Only one wildcard '*' is allowed in source_object parameter. " "Found {} in {}.".format(total_wildcards, self.source_object)) raise AirflowException(error_msg) prefix, delimiter = self.source_object.split(WILDCARD, 1) objects = self.gcs_hook.list(self.source_bucket, prefix=prefix, delimiter=delimiter) for source_object in objects: if self.destination_object is None: destination_object = source_object else: destination_object = source_object.replace( prefix, self.destination_object, 1) self._copy_single_object(source_object=source_object, destination_object=destination_object) else: self._copy_single_object( source_object=self.source_object, destination_object=self.destination_object)
def execute(self, context): gcs_hook = GCSHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) sftp_hook = SFTPHook(self.sftp_conn_id) if WILDCARD in self.source_object: total_wildcards = self.source_object.count(WILDCARD) if total_wildcards > 1: raise AirflowException( "Only one wildcard '*' is allowed in source_object parameter. " "Found {} in {}.".format(total_wildcards, self.source_object)) prefix, delimiter = self.source_object.split(WILDCARD, 1) objects = gcs_hook.list(self.source_bucket, prefix=prefix, delimiter=delimiter) for source_object in objects: destination_path = os.path.join(self.destination_path, source_object) self._copy_single_object(gcs_hook, sftp_hook, source_object, destination_path) self.log.info("Done. Uploaded '%d' files to %s", len(objects), self.destination_path) else: destination_path = os.path.join(self.destination_path, self.source_object) self._copy_single_object(gcs_hook, sftp_hook, self.source_object, destination_path) self.log.info("Done. Uploaded '%s' file to %s", self.source_object, destination_path)
def poke(self, context): self.log.info('Sensor checks existence of : %s, %s', self.bucket, self.object) hook = GCSHook(google_cloud_storage_conn_id=self.google_cloud_conn_id, delegate_to=self.delegate_to) return hook.is_updated_after(self.bucket, self.object, self.ts_func(context))
def execute(self, context): gcs_hook = GCSHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) sftp_hook = SFTPHook(self.sftp_conn_id) if WILDCARD in self.source_path: total_wildcards = self.source_path.count(WILDCARD) if total_wildcards > 1: raise AirflowException( "Only one wildcard '*' is allowed in source_path parameter. " "Found {} in {}.".format(total_wildcards, self.source_path)) prefix, delimiter = self.source_path.split(WILDCARD, 1) base_path = os.path.dirname(prefix) files, _, _ = sftp_hook.get_tree_map(base_path, prefix=prefix, delimiter=delimiter) for file in files: destination_path = file.replace(base_path, self.destination_path, 1) self._copy_single_object(gcs_hook, sftp_hook, file, destination_path) else: destination_object = (self.destination_path if self.destination_path else self.source_path.rsplit("/", 1)[1]) self._copy_single_object(gcs_hook, sftp_hook, self.source_path, destination_object)
def execute(self, context): hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id) hook.insert_object_acl(bucket_name=self.bucket, object_name=self.object_name, entity=self.entity, role=self.role, generation=self.generation, user_project=self.user_project)
def _upload_to_gcs(self, files_to_upload: Dict[str, Any]): hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) for obj, tmp_file_handle in files_to_upload.items(): hook.upload(bucket_name=self.bucket, object_name=obj, filename=tmp_file_handle.name, mime_type='application/json', gzip=self.gzip)
def execute(self, context): hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) hook.create_bucket(bucket_name=self.bucket_name, resource=self.resource, storage_class=self.storage_class, location=self.location, project_id=self.project_id, labels=self.labels)
def apply_validate_fn(*args, **kwargs): prediction_path = kwargs["templates_dict"]["prediction_path"] scheme, bucket, obj, _, _ = urlsplit(prediction_path) if scheme != "gs" or not bucket or not obj: raise ValueError( "Wrong format prediction_path: {}".format(prediction_path)) summary = os.path.join(obj.strip("/"), "prediction.summary.json") gcs_hook = GCSHook() summary = json.loads(gcs_hook.download(bucket, summary)) return validate_fn(summary)
def execute(self, context): hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) hook.sync(source_bucket=self.source_bucket, destination_bucket=self.destination_bucket, source_object=self.source_object, destination_object=self.destination_object, recursive=self.recursive, delete_extra_files=self.delete_extra_files, allow_overwrite=self.allow_overwrite)
def execute(self, context): hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) self.log.info( 'Getting list of the files. Bucket: %s; Delimiter: %s; Prefix: %s', self.bucket, self.delimiter, self.prefix) return hook.list(bucket_name=self.bucket, prefix=self.prefix, delimiter=self.delimiter)
def hook(self): """ Returns GCS hook. """ remote_conn_id = conf.get('logging', 'REMOTE_LOG_CONN_ID') try: from airflow.gcp.hooks.gcs import GCSHook return GCSHook(google_cloud_storage_conn_id=remote_conn_id) except Exception as e: # pylint: disable=broad-except self.log.error( 'Could not create a GoogleCloudStorageHook with connection id ' '"%s". %s\n\nPlease make sure that airflow[gcp] is installed ' 'and the GCS connection exists.', remote_conn_id, str(e))
def _upload_to_gcs(self, files_to_upload): """ Upload all of the file splits (and optionally the schema .json file) to Google Cloud Storage. """ hook = GCSHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) for tmp_file in files_to_upload: hook.upload(self.bucket, tmp_file.get('file_name'), tmp_file.get('file_handle').name, mime_type=tmp_file.get('file_mime_type'), gzip=self.gzip if tmp_file.get('file_name') == self.schema_filename else False)
def execute(self, context): """ Uploads the file to Google Cloud Storage """ hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) hook.upload( bucket_name=self.bucket, object_name=self.dst, mime_type=self.mime_type, filename=self.src, gzip=self.gzip, )
def execute(self, context): hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) if self.objects: objects = self.objects else: objects = hook.list(bucket_name=self.bucket_name, prefix=self.prefix) self.log.info("Deleting %s objects from %s", len(objects), self.bucket_name) for object_name in objects: hook.delete(bucket_name=self.bucket_name, object_name=object_name)
def execute(self, context): hook = CloudTextToSpeechHook(gcp_conn_id=self.gcp_conn_id) result = hook.synthesize_speech( input_data=self.input_data, voice=self.voice, audio_config=self.audio_config, retry=self.retry, timeout=self.timeout, ) with NamedTemporaryFile() as temp_file: temp_file.write(result.audio_content) cloud_storage_hook = GCSHook( google_cloud_storage_conn_id=self.gcp_conn_id) cloud_storage_hook.upload(bucket_name=self.target_bucket_name, object_name=self.target_filename, filename=temp_file.name)
def execute(self, context: Dict): hook = GoogleDisplayVideo360Hook( gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to, api_version=self.api_version, ) gcs_hook = GCSHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to ) resource = hook.get_query(query_id=self.report_id) # Check if report is ready if resource["metadata"]["running"]: raise AirflowException('Report {} is still running'.format(self.report_id)) # If no custom report_name provided, use DV360 name file_url = resource["metadata"]["googleCloudStoragePathForLatestReport"] report_name = self.report_name or urlparse(file_url).path.split('/')[2] report_name = self._resolve_file_name(report_name) # Download the report self.log.info("Starting downloading report %s", self.report_id) with tempfile.NamedTemporaryFile(delete=False) as temp_file: with urllib.request.urlopen(file_url) as response: shutil.copyfileobj(response, temp_file, length=self.chunk_size) temp_file.flush() # Upload the local file to bucket gcs_hook.upload( bucket_name=self.bucket_name, object_name=report_name, gzip=self.gzip, filename=temp_file.name, mime_type="text/csv", ) self.log.info( "Report %s was saved in bucket %s as %s.", self.report_id, self.bucket_name, report_name, ) self.xcom_push(context, key='report_name', value=report_name)
def execute(self, context): self.log.info('Executing download: %s, %s, %s', self.bucket, self.object, self.filename) hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) if self.store_to_xcom_key: file_bytes = hook.download(bucket_name=self.bucket, object_name=self.object) if sys.getsizeof(file_bytes) < MAX_XCOM_SIZE: context['ti'].xcom_push(key=self.store_to_xcom_key, value=file_bytes) else: raise AirflowException( 'The size of the downloaded file is too large to push to XCom!' ) else: hook.download(bucket_name=self.bucket, object_name=self.object, filename=self.filename)
def execute(self, context): # use the super to list all files in an Google Cloud Storage bucket files = super().execute(context) s3_hook = S3Hook(aws_conn_id=self.dest_aws_conn_id, verify=self.dest_verify) if not self.replace: # if we are not replacing -> list all files in the S3 bucket # and only keep those files which are present in # Google Cloud Storage and not in S3 bucket_name, prefix = S3Hook.parse_s3_url(self.dest_s3_key) # look for the bucket and the prefix to avoid look into # parent directories/keys existing_files = s3_hook.list_keys(bucket_name, prefix=prefix) # in case that no files exists, return an empty array to avoid errors existing_files = existing_files if existing_files is not None else [] # remove the prefix for the existing files to allow the match existing_files = [ file.replace(prefix, '', 1) for file in existing_files ] files = list(set(files) - set(existing_files)) if files: hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) for file in files: file_bytes = hook.download(self.bucket, file) dest_key = self.dest_s3_key + file self.log.info("Saving file to %s", dest_key) s3_hook.load_bytes(file_bytes, key=dest_key, replace=self.replace) self.log.info("All done, uploaded %d files to S3", len(files)) else: self.log.info("In sync, no files needed to be uploaded to S3") return files
def execute(self, context: Dict): hook = GoogleSearchAdsHook( gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to, api_version=self.api_version, ) gcs_hook = GCSHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) # Resolve file name of the report report_name = self.report_name or self.report_id report_name = self._resolve_file_name(report_name) response = hook.get(report_id=self.report_id) if not response['isReportReady']: raise AirflowException('Report {} is not ready yet'.format( self.report_id)) # Resolve report fragments fragments_count = len(response["files"]) # Download chunks of report's data self.log.info("Downloading Search Ads report %s", self.report_id) with NamedTemporaryFile() as temp_file: for i in range(fragments_count): byte_content = hook.get_file(report_fragment=i, report_id=self.report_id) fragment = (byte_content if i == 0 else self._handle_report_fragment(byte_content)) temp_file.write(fragment) temp_file.flush() gcs_hook.upload( bucket_name=self.bucket_name, object_name=report_name, gzip=self.gzip, filename=temp_file.name, ) self.xcom_push(context, key="file_name", value=report_name)
def execute(self, context): # use the super to list all files in an Azure Data Lake path files = super().execute(context) g_hook = GCSHook( google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) if not self.replace: # if we are not replacing -> list all files in the ADLS path # and only keep those files which are present in # ADLS and not in Google Cloud Storage bucket_name, prefix = _parse_gcs_url(self.dest_gcs) existing_files = g_hook.list(bucket_name=bucket_name, prefix=prefix) files = set(files) - set(existing_files) if files: hook = AzureDataLakeHook( azure_data_lake_conn_id=self.azure_data_lake_conn_id ) for obj in files: with NamedTemporaryFile(mode='wb', delete=True) as f: hook.download_file(local_path=f.name, remote_path=obj) f.flush() dest_gcs_bucket, dest_gcs_prefix = _parse_gcs_url(self.dest_gcs) dest_path = os.path.join(dest_gcs_prefix, obj) self.log.info("Saving file to %s", dest_path) g_hook.upload( bucket_name=dest_gcs_bucket, object_name=dest_path, filename=f.name, gzip=self.gzip ) self.log.info("All done, uploaded %d files to GCS", len(files)) else: self.log.info("In sync, no files needed to be uploaded to GCS") return files
def execute(self, context: Dict): hook = GoogleCampaignManagerHook( gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to, api_version=self.api_version, ) gcs_hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) # Get name of the report report = hook.get_report(file_id=self.file_id, profile_id=self.profile_id, report_id=self.report_id) report_name = self.report_name or report.get("fileName", str(uuid.uuid4())) report_name = self._resolve_file_name(report_name) # Download the report self.log.info("Starting downloading report %s", self.report_id) request = hook.get_report_file(profile_id=self.profile_id, report_id=self.report_id, file_id=self.file_id) with tempfile.NamedTemporaryFile() as temp_file: downloader = http.MediaIoBaseDownload(fd=temp_file, request=request, chunksize=self.chunk_size) download_finished = False while not download_finished: _, download_finished = downloader.next_chunk() temp_file.flush() # Upload the local file to bucket gcs_hook.upload( bucket_name=self.bucket_name, object_name=report_name, gzip=self.gzip, filename=temp_file.name, mime_type="text/csv", ) self.xcom_push(context, key="report_name", value=report_name)
def execute(self, context: Dict): hook = GCSHook(gcp_conn_id=self.gcp_conn_id) with NamedTemporaryFile() as source_file, NamedTemporaryFile( ) as destination_file: self.log.info("Downloading file from %s", self.source_bucket) hook.download(bucket_name=self.source_bucket, object_name=self.source_object, filename=source_file.name) self.log.info("Starting the transformation") cmd = [self.transform_script] if isinstance( self.transform_script, str) else self.transform_script cmd += [source_file.name, destination_file.name] process = subprocess.Popen(args=cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, close_fds=True) self.log.info("Process output:") for line in iter(process.stdout.readline, b''): self.log.info(line.decode(self.output_encoding).rstrip()) process.wait() if process.returncode > 0: raise AirflowException("Transform script failed: {0}".format( process.returncode)) self.log.info( "Transformation succeeded. Output temporarily located at %s", destination_file.name) self.log.info("Uploading file to %s as %s", self.destination_bucket, self.destination_object) hook.upload(bucket_name=self.destination_bucket, object_name=self.destination_object, filename=destination_file.name)
def execute(self, context): bq_hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id, delegate_to=self.delegate_to, location=self.location) if not self.schema_fields: if self.schema_object and self.source_format != 'DATASTORE_BACKUP': gcs_hook = GCSHook(google_cloud_storage_conn_id=self. google_cloud_storage_conn_id, delegate_to=self.delegate_to) schema_fields = json.loads( gcs_hook.download(self.bucket, self.schema_object).decode("utf-8")) elif self.schema_object is None and self.autodetect is False: raise AirflowException( 'At least one of `schema_fields`, ' '`schema_object`, or `autodetect` must be passed.') else: schema_fields = None else: schema_fields = self.schema_fields source_uris = [ 'gs://{}/{}'.format(self.bucket, source_object) for source_object in self.source_objects ] conn = bq_hook.get_conn() cursor = conn.cursor() if self.external_table: cursor.create_external_table( external_project_dataset_table=self. destination_project_dataset_table, schema_fields=schema_fields, source_uris=source_uris, source_format=self.source_format, compression=self.compression, skip_leading_rows=self.skip_leading_rows, field_delimiter=self.field_delimiter, max_bad_records=self.max_bad_records, quote_character=self.quote_character, ignore_unknown_values=self.ignore_unknown_values, allow_quoted_newlines=self.allow_quoted_newlines, allow_jagged_rows=self.allow_jagged_rows, encoding=self.encoding, src_fmt_configs=self.src_fmt_configs, encryption_configuration=self.encryption_configuration) else: cursor.run_load( destination_project_dataset_table=self. destination_project_dataset_table, schema_fields=schema_fields, source_uris=source_uris, source_format=self.source_format, autodetect=self.autodetect, create_disposition=self.create_disposition, skip_leading_rows=self.skip_leading_rows, write_disposition=self.write_disposition, field_delimiter=self.field_delimiter, max_bad_records=self.max_bad_records, quote_character=self.quote_character, ignore_unknown_values=self.ignore_unknown_values, allow_quoted_newlines=self.allow_quoted_newlines, allow_jagged_rows=self.allow_jagged_rows, encoding=self.encoding, schema_update_options=self.schema_update_options, src_fmt_configs=self.src_fmt_configs, time_partitioning=self.time_partitioning, cluster_fields=self.cluster_fields, encryption_configuration=self.encryption_configuration) if self.max_id_key: cursor.execute('SELECT MAX({}) FROM {}'.format( self.max_id_key, self.destination_project_dataset_table)) row = cursor.fetchone() max_id = row[0] if row[0] else 0 self.log.info('Loaded BQ data with max %s.%s=%s', self.destination_project_dataset_table, self.max_id_key, max_id)
def __init__(self, gcp_conn_id: str = 'google_cloud_default', delegate_to: Optional[str] = None) -> None: self._gcs_hook = GCSHook(gcp_conn_id, delegate_to)
def execute(self, context): # use the super method to list all the files in an S3 bucket/key files = super().execute(context) gcs_hook = GCSHook(google_cloud_storage_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) if not self.replace: # if we are not replacing -> list all files in the GCS bucket # and only keep those files which are present in # S3 and not in Google Cloud Storage bucket_name, object_prefix = _parse_gcs_url(self.dest_gcs) existing_files_prefixed = gcs_hook.list(bucket_name, prefix=object_prefix) existing_files = [] if existing_files_prefixed: # Remove the object prefix itself, an empty directory was found if object_prefix in existing_files_prefixed: existing_files_prefixed.remove(object_prefix) # Remove the object prefix from all object string paths for f in existing_files_prefixed: if f.startswith(object_prefix): existing_files.append(f[len(object_prefix):]) else: existing_files.append(f) files = list(set(files) - set(existing_files)) if len(files) > 0: self.log.info('%s files are going to be synced: %s.', len(files), files) else: self.log.info( 'There are no new files to sync. Have a nice day!') if files: hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify) for file in files: # GCS hook builds its own in-memory file so we have to create # and pass the path file_object = hook.get_key(file, self.bucket) with NamedTemporaryFile(mode='wb', delete=True) as f: file_object.download_fileobj(f) f.flush() dest_gcs_bucket, dest_gcs_object_prefix = _parse_gcs_url( self.dest_gcs) # There will always be a '/' before file because it is # enforced at instantiation time dest_gcs_object = dest_gcs_object_prefix + file # Sync is sequential and the hook already logs too much # so skip this for now # self.log.info( # 'Saving file {0} from S3 bucket {1} in GCS bucket {2}' # ' as object {3}'.format(file, self.bucket, # dest_gcs_bucket, # dest_gcs_object)) gcs_hook.upload(dest_gcs_bucket, dest_gcs_object, f.name, gzip=self.gzip) self.log.info( "All done, uploaded %d files to Google Cloud Storage", len(files)) else: self.log.info( 'In sync, no files needed to be uploaded to Google Cloud' 'Storage') return files
def poke(self, context): hook = GCSHook() return self.is_bucket_updated(len(hook.list(self.bucket, prefix=self.prefix)))