EricChea / LAI-BigQuery-Component

A Lightning component to run queries on BigQuery.
______________________________________________________________________

About

This component lets you run queries against a BigQuery warehouse.

Use the component

To Run a query

import pickle
import lightning as L
from lightning_bigquery import BigQuery


class YourComponent(L.LightningFlow):
    def __init__(self):
        super().__init__()
        self.bq_client = BigQuery(
            project="<YOUR GOOGLE PROJECT ID>",
            location="<LOCATION OF DATASET>",  # Region where the dataset is located.
            credential="<SERVICE ACCOUNT KEY>",  # Service account credentials
        )
        self.reader = ReadResults()

    def run(self):
        sqlquery = "select * from <BigQuery table>"
        self.bq_client.query(sqlquery, to_dataframe=True)

        if self.bq_client.has_succeeded:
            # The results from the query are saved as a pickled file.
            self.reader.run(self.bq_client.result_path)


class ReadResults(L.LightningWork):
    def run(self, result_filepath):
        with open(result_filepath, "rb") as _file:
            data = pickle.load(_file)

        # Print top results from the dataframe
        print(data.head())


app = L.LightningApp(YourComponent())

Install

Run the following to install:

git clone https://github.com/PyTorchLightning/LAI-bigquery
cd LAI-bigquery
pip install -r requirements.txt
pip install -e .

Tests

To run unit tests locally:

# From the root level of the package (LAI-bigquery)
pip install -r tests/requirements.txt
pytest