def __init__(self): """ Attempt to create hook with airflow[gcloud] (and set use_gcloud = True), otherwise uses airflow[gcp_api] """ remote_conn_id = configuration.get('core', 'REMOTE_LOG_CONN_ID') self.use_gcloud = False try: from airflow.contrib.hooks import GCSHook self.hook = GCSHook(remote_conn_id) self.use_gcloud = True except: try: from airflow.contrib.hooks import GoogleCloudStorageHook self.hook = GoogleCloudStorageHook( scope= 'https://www.googleapis.com/auth/devstorage.read_write', google_cloud_storage_conn_id=remote_conn_id) except: self.hook = None logging.error( 'Could not create a GCSHook with connection id "{}". ' 'Please make sure that either airflow[gcloud] or ' 'airflow[gcp_api] is installed and the GCS connection ' 'exists.'.format(remote_conn_id))
def __init__(self): """ Attempt to create hook with airflow[gcp_api]. """ remote_conn_id = configuration.get('core', 'REMOTE_LOG_CONN_ID') self.hook = None try: from airflow.contrib.hooks import GoogleCloudStorageHook self.hook = GoogleCloudStorageHook( google_cloud_storage_conn_id=remote_conn_id) except: logging.error( 'Could not create a GoogleCloudStorageHook with connection id ' '"{}". Please make sure that airflow[gcp_api] is installed ' 'and the GCS connection exists.'.format(remote_conn_id))
def do_list_predictions_files(**kwargs): """ Retrieves all the predictions files that should be loaded to BigQuery. Can not do a GoogleCloudStorageToBigQueryOperator directly due to the possible multiple files. """ # List all relevant files # TODO Add when Composer is on Airflow 2.0 # predictions_files = gcs_list_operator.GoogleCloudStorageListOperator( # task_id='predictions_files', # bucket=COMPOSER_BUCKET_NAME, # prefix='predictions/output/prediction.results-' # ) # TODO Remove when Composer on Airflow 2.0 gcs = GoogleCloudStorageHook() predictions_files = gcs.list( bucket=COMPOSER_BUCKET_NAME, prefix='predictions/output/prediction.results-') logging.info("Predictions files are: {}".format(predictions_files)) # Create a variable that can be used in the next task kwargs['ti'].xcom_push(key='predictions_files', value=predictions_files)
def output_manager(self, file_name): """ Takes output and uploads to corresponding destination. """ if self.destination.lower() == 's3': s3 = S3Hook(self.dest_conn_id) s3.load_file(filename=file_name, key=self.key, bucket_name=self.bucket, replace=True) elif self.destination.lower() == 'gcs': print("Uploading File!") gcs = GoogleCloudStorageHook(self.dest_conn_id) gcs.upload( bucket=self.bucket, object=self.key, filename=file_name, ) print("Uploaded file to {0}/{1}".format(self.bucket, self.key)) os.remove(file_name)
def execute(self, context): hook = HttpHook(self.t1, self.t2) with open("launches.json", "w+t") as f: f.write(json.dumps(hook.get_results())) GoogleCloudStorageHook().upload(bucket="launchbucket",object="launches.json",filename="launches.json")