How to run a BigQuery query and then send the output CSV to Google Cloud Storage in Apache Airflow?

Viewed 9886

I need to run a bigquery script in python, which needs to output as a CSV in google cloud storage. Currently, my script triggers the big query code and saves to my PC directly.

However, I need to get this running in Airflow so I can't have any local dependencies.

My current script saves the output to my local machine and then I have to move it into GCS. Looked online and I can't figure it out. (ps im very new to python so im sorry in advance if this has been asked before!)

import pandas as pd
from googleapiclient import discovery
from oauth2client.client import GoogleCredentials

def run_script():

    df = pd.read_gbq('SELECT * FROM `table/veiw` LIMIT 15000',
                 project_id='PROJECT',
                 dialect='standard'
                 )

    df.to_csv('XXX.csv', index=False)

def copy_to_gcs(filename, bucket, destination_filename):

    credentials = GoogleCredentials.get_application_default()
    service = discovery.build('storage', 'v1', credentials=credentials)

    body = {'name': destination_filename}
    req = service.objects().insert(bucket=bucket,body=body, media_body=filename)
    resp = req.execute()

current_date = datetime.date.today()
filename = (r"C:\Users\LOCALDRIVE\ETC\ETC\ETC.csv")
bucket = 'My GCS BUCKET'

str_prefix_datetime = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
destfile = 'XXX' + str_prefix_datetime + '.csv'
print('')

    ```
1 Answers

Airflow provides several operators for working with BigQuery.

You can see an example of running a query, followed by exporting the results to a CSV in the Cloud Composer code samples.

# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Query recent StackOverflow questions.

bq_recent_questions_query = bigquery_operator.BigQueryOperator(
    task_id='bq_recent_questions_query',
    sql="""
    SELECT owner_display_name, title, view_count
    FROM `bigquery-public-data.stackoverflow.posts_questions`
    WHERE creation_date < CAST('{max_date}' AS TIMESTAMP)
        AND creation_date >= CAST('{min_date}' AS TIMESTAMP)
    ORDER BY view_count DESC
    LIMIT 100
    """.format(max_date=max_query_date, min_date=min_query_date),
    use_legacy_sql=False,
    destination_dataset_table=bq_recent_questions_table_id)

# Export query result to Cloud Storage.
export_questions_to_gcs = bigquery_to_gcs.BigQueryToCloudStorageOperator(
    task_id='export_recent_questions_to_gcs',
    source_project_dataset_table=bq_recent_questions_table_id,
    destination_cloud_storage_uris=[output_file],
    export_format='CSV')
Related