Source Code: Save Query Output To Google Drive From BigQuery Using Python

Buy Me a Coffee? Your support is much appreciated!

demo.py

import time
from google.cloud import bigquery
from gdrive import GDrive

# Step 1. Retrieve query results from a SQL Statement
client_file = 'client-secret.json'
client = bigquery.Client()

sql_statement = 'SELECT * FROM `sql-for-bigquery.JJ_Datasets.demo_product`'
query_job = client.query('SELECT * FROM `sql-for-bigquery.JJ_Datasets.demo_product`')

while query_job.state != 'DONE':
    time.sleep(2)
    query_job.reload()

response = query_job.result()    
df = response.to_dataframe()

# Step 2. Push resultset to Google Drive
drive = GDrive(client_file)
drive.init_service()

folder_id = ['<folder id1>', '<folder id2>']
file_uploaded = drive.upload_file_bytes(df.to_csv().encode(), 'demo.csv', 'text/csv', folder_id)
print(file_uploaded)

gdrive.py

import os
import io
from typing import List
import mimetypes
from googleapiclient.http import MediaFileUpload, MediaIoBaseUpload
from google_apis import create_service

class GDrive:
    API_NAME = 'drive'
    API_VERSION = 'v3'
    SCOPES = ['https://www.googleapis.com/auth/drive']

    def __init__(self, client_file):
        self.client_file = client_file
        self.service = None

    def init_service(self):
        try:
            self.service = create_service(self.client_file, self.API_NAME, self.API_VERSION, self.SCOPES)
            if self.service is None:
                print('Failed')
        except Exception as e:
            self.service = None
            print(e)

    def upload_file_bytes(self, byte_strings, file_name, file_mime_type, parent_folder: List[str]=None):
        f = io.BytesIO(byte_strings)
        media_content = MediaIoBaseUpload(f, mimetype=file_mime_type)
        if parent_folder is None:
            parent_folder = []
        file_metadata = {
            'name': file_name,
            'parents': parent_folder
        }

        if self.service is None:
            print('Please initialize Google Drive service first')
            return

        file = self.service.files().create(
            body=file_metadata,
            media_body=media_content
        ).execute()    
        return file

    def upload_file(self, file_path, parent_folder: List[str]=None):
        if parent_folder is None:
            parent_folder = []
        file_mime_type = mimetypes.guess_type(file_path)[0]
        media_content = MediaFileUpload(file_path, mimetype=file_mime_type)
        file_base_name = os.path.basename(file_path)
        file_metadata = {
            'name': file_base_name,
            'parents': parent_folder
        }

        if self.service is None:
            print('Please initialize Google Drive service first')
            return
            
        file = self.service.files().create(
            body=file_metadata,
            media_body=media_content
        ).execute()
        return file