Google BigQuery agent
Installation
To install the BigQuery agent, run the following command:
pip install flytekitplugins-bigquery
Example usage
For an example query, see BigQuery agent example.
Local testing
To test the BigQuery agent locally, create a class for the agent task that inherits from AsyncAgentExecutorMixin. This mixin allows flytekit to mimic FlytePropeller's behavior in calling the agent.
To test the BigQuery agent, copy the following code to a file called bigquery_task.py
, modifying as needed.
Note
When testing the BigQuery integration locally, you will need to set the GOOGLE_APPLICATION_CREDENTIALS
environment variable.
Add AsyncAgentExecutorMixin
to this class to tell flytekit to use the agent to run the task:
class BigQueryTask(AsyncAgentExecutorMixin, SQLTask[BigQueryConfig]):
def __init__(self, name: str, **kwargs):
...
Flytekit will automatically use the agent to run the task in the local execution.
bigquery_doge_coin = BigQueryTask(
name=f"bigquery.doge_coin",
inputs=kwtypes(version=int),
query_template="SELECT * FROM `bigquery-public-data.crypto_dogecoin.transactions` WHERE version = @version LIMIT 10;",
output_structured_dataset_type=StructuredDataset,
task_config=BigQueryConfig(ProjectID="flyte-test-340607")
)
You can run the above example task locally and test the agent with the following command:
$ pyflyte run bigquery_task.py bigquery_doge_coin --version 10
Union cluster deployment
After you have finished testing the agent locally, contact the Union team to enable it in your cluster.