BigQuery Integration

Configure and run a BigQuery task

To run a parameterized BigQuery query as an asynchronous Flyte task, install the plugin and construct BigQueryConfig and BigQueryTask from the task module:

pip install flyteplugins-bigquery

from flyteplugins.bigquery.task import BigQueryConfig, BigQueryTask

config = BigQueryConfig(ProjectID="my-project", Location="US")
task = BigQueryTask(
    name="my_query",
    query_template="SELECT * FROM dataset.table WHERE id = {{ .user_id }}",
    plugin_config=config,
    inputs={"user_id": int},
)

BigQueryConfig.ProjectID is required. Location may be omitted, as in the repository's runnable example, and QueryJobConfig is optional. BigQueryTask is a TaskTemplate using the asynchronous task type bigquery_query_job_task; it accepts typed inputs, an optional dataframe output type, and an optional google_application_credentials secret name.

A dataframe-producing task can be invoked from another Flyte task. The output name created by BigQueryTask is results:

from flyte.io._dataframe.dataframe import DataFrame
from flyteplugins.bigquery.task import BigQueryConfig, BigQueryTask

config = BigQueryConfig(
    ProjectID="my-gcp-project",
    Location="US",
)

query_task = BigQueryTask(
    name="count_events",
    query_template="SELECT COUNT(*) AS total FROM `{ds}.events` WHERE date = @date",
    plugin_config=config,
    inputs={"date": str},
    output_dataframe_type=DataFrame[dict],
)

def run_query(date: str) -> DataFrame[dict]:
    return query_task(date=date)

Configure project, location, and job options

BigQueryConfig stores the following values:

Field	Value and behavior
ProjectID	Required Google Cloud project ID. It is serialized as ProjectID and used when creating, polling, and cancelling the BigQuery job.
Location	Optional location such as US or EU. It is serialized as Location; when omitted, the serialized value is None.
QueryJobConfig	Optional google.cloud.bigquery.QueryJobConfig. Its API representation's query fields are added to the task custom configuration.
google_application_credentials	Optional BigQueryTask constructor argument containing a Flyte secret name. It is serialized only when set.

For example, the task tests use QueryJobConfig to serialize query-cache, legacy-SQL, and byte-limit settings:

from google.cloud import bigquery
from flyteplugins.bigquery.task import BigQueryConfig, BigQueryTask

job_config = bigquery.QueryJobConfig(
    use_query_cache=False,
    use_legacy_sql=True,
    maximum_bytes_billed=1000000,
)
config = BigQueryConfig(
    ProjectID="test-project",
    Location="EU",
    QueryJobConfig=job_config,
)
task = BigQueryTask(
    name="test_task",
    query_template="SELECT * FROM table",
    plugin_config=config,
)

During serialization, BigQueryTask.custom_config() emits the basic settings and the serialization domain. With the values above and a serialization domain of test-domain, the tested fields are:

custom_config = task.custom_config(serialization_context)

assert custom_config["Location"] == "EU"
assert custom_config["ProjectID"] == "test-project"
assert custom_config["Domain"] == "test-domain"
assert custom_config["useQueryCache"] is False
assert custom_config["useLegacySql"] is True
assert custom_config["maximumBytesBilled"] == "1000000"

The task also supports a credentials secret:

task = BigQueryTask(
    name="test_task",
    query_template="SELECT * FROM table",
    plugin_config=BigQueryConfig(ProjectID="test-project", Location="US"),
    google_application_credentials="my-gcp-credentials-secret",
)

This produces secrets: {"google_application_credentials": "my-gcp-credentials-secret"} in custom configuration. If the argument is unset, secrets is omitted.

Use typed inputs in the query

Declare each runtime input in the task's inputs dictionary. BigQueryTask places those declarations in its NativeInterface. The connector later uses the declared Flyte types—not the runtime values—to select BigQuery scalar parameter types. The connector currently maps these Python types: list to ARRAY, bool to BOOL, bytes to BYTES, datetime.datetime to DATETIME, float to FLOAT64, int to INT64, and str to STRING.

At submission, an input such as {"user_id": 12345} becomes a bigquery.ScalarQueryParameter named user_id, with the mapped type INT64 and value 12345. Input keys therefore need to correspond to declared interface variables and the parameter names used by the SQL submitted to BigQuery.

The task docstring specifies Flyte's Go templating format for query templates. The public README uses that form:

from flyteplugins.bigquery.task import BigQueryConfig, BigQueryTask

task = BigQueryTask(
    name="my_query",
    query_template="SELECT * FROM dataset.table WHERE id = {{ .user_id }}",
    plugin_config=BigQueryConfig(ProjectID="my-project", Location="US"),
    inputs={"user_id": int},
)

The package example also uses a named BigQuery parameter in SQL:

query_task = BigQueryTask(
    name="count_events",
    query_template="SELECT COUNT(*) AS total FROM `{ds}.events` WHERE date = @date",
    plugin_config=BigQueryConfig(ProjectID="my-gcp-project", Location="US"),
    inputs={"date": str},
    output_dataframe_type=DataFrame[dict],
)

Whichever form is used in the task definition, BigQueryTask.sql() serializes the normalized query as a Flyte ANSI tasks_pb2.Sql. BigQueryConnector.create() ultimately submits task_template.sql.statement to client.query().

Understand serialization and result shape

BigQueryTask normalizes its query in the constructor by replacing newlines and tabs and collapsing runs of whitespace to one space. For example, the tested multiline query becomes:

query_with_whitespace = """
SELECT
    col1, col2, col3
FROM table
WHERE condition = 1
"""

task = BigQueryTask(
    name="test_task",
    query_template=query_with_whitespace,
    plugin_config=BigQueryConfig(ProjectID="test-project"),
)

assert task.query_template == "SELECT col1, col2, col3 FROM table WHERE condition = 1"

sql() returns an ANSI SQL protobuf:

sql_proto = task.sql(serialization_context)
assert sql_proto.statement == "SELECT * FROM table"
assert sql_proto.dialect == 1

Passing output_dataframe_type=DataFrame or DataFrame[dict] adds an interface output named results. That declaration does not by itself create an output at completion: BigQueryConnector.get() only returns results when the completed BigQuery job reports a destination table.

Follow the asynchronous connector lifecycle

The plugin's connector is BigQueryConnector, an AsyncConnector registered with Flyte as soon as flyteplugins.bigquery.connector is imported. Its task_type_name is bigquery_query_job_task, matching BigQueryTask._TASK_TYPE.

Submit with create()

create() performs these operations:

1. If runtime inputs are present, it reads each variable from task_template.interface.inputs, converts its Flyte type with TypeEngine.guess_python_type(), and creates a bigquery.QueryJobConfig containing ScalarQueryParameter objects.
2. It reads ProjectID, Location, and Domain from serialized custom configuration.
3. It builds a user agent as Flyte/<runtime version> (GPN:Union;<domain>).
4. It obtains a cached BigQuery client and submits task_template.sql.statement.
5. It returns BigQueryMetadata, preserving the job ID, project, location, and user agent for subsequent lifecycle calls.

The connector-side test constructs a task template with an ANSI statement and an integer interface variable:

from flyteidl2.core.interface_pb2 import Variable, VariableEntry, VariableMap
from flyteidl2.core.tasks_pb2 import Sql, TaskTemplate
from flyteidl2.core.types_pb2 import LiteralType, SimpleType

template = TaskTemplate()
template.sql.CopyFrom(
    Sql(statement="SELECT * FROM table WHERE id = @user_id", dialect=Sql.Dialect.ANSI)
)
template.metadata.runtime.version = "1.0.0"

int_type = LiteralType()
int_type.simple = SimpleType.INTEGER

variables = VariableMap()
variables.variables.append(
    VariableEntry(key="user_id", value=Variable(type=int_type))
)
template.interface.inputs.CopyFrom(variables)

With serialized custom configuration containing ProjectID, Location, and Domain, calling await BigQueryConnector().create(template, inputs={"user_id": 12345}) submits the statement with a job configuration containing one user_id parameter. An unsupported declared Python type causes a lookup failure because create() indexes pythonTypeToBigQueryType directly.

The helper _get_bigquery_client() is decorated with functools.lru_cache. Its cache key includes project, location, user agent, and the credential string. When credentials are supplied at connector runtime, the helper parses the resolved string with json.loads() and creates service-account credentials with service_account.Credentials.from_service_account_info; without credentials it passes credentials=None to bigquery.Client.

Poll with get()

get() retrieves the job using the metadata's job ID, project, and location. Every returned resource includes a ready BigQuery Console dashboard link of the form used by the connector:

https://console.cloud.google.com/bigquery?project=<project>&j=bq:<location>:<job_id>&page=queryresults

If job.errors is truthy, the connector returns a failed Flyte resource with the stringified errors and no outputs, before converting the BigQuery state. Otherwise, it passes str(job.state) to convert_to_flyte_phase() and returns that phase with the state as the resource message.

For a successful job with a destination table, the output is a Flyte DataFrame whose URI is built from the destination reference:

output_location = "bq://output-project:output_dataset.output_table"

The resulting resource has outputs["results"] pointing to DataFrame(uri=output_location). A successful job with destination is None returns outputs=None, even when the task declared a dataframe output.

Cancel with delete()

When Flyte calls BigQueryConnector.delete(), it obtains the same client and calls:

client.cancel_job(resource_meta.job_id, resource_meta.project, resource_meta.location)

The method is asynchronous at the connector interface, while the Google client calls inside create(), get(), and delete() are synchronous.

Run the repository example locally

The repository's connector example registers the task in a TaskEnvironment and runs it with a local run context:

from flyteplugins.bigquery.task import BigQueryConfig, BigQueryTask

import flyte
from flyte.io._dataframe.dataframe import DataFrame

bigquery_task = BigQueryTask(
    name="bigquery",
    inputs={"version": int},
    output_dataframe_type=DataFrame,
    plugin_config=BigQueryConfig(ProjectID="dogfood-gcp-dataplane"),
    query_template="SELECT * from dataset.flyte_table3;",
)

flyte.TaskEnvironment.from_task("bigquery_env", bigquery_task)

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.with_runcontext(mode="local").run(bigquery_task, 123)
    print(run.url)

For a connector application, the example image installs flyteplugins-bigquery and deploys a flyte.app.ConnectorEnvironment; the app does not provide connector implementation code separately:

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: Parsed source module 'flyte.app' has no member 'app' 'flyte.Resources' does not accept keyword(s): cpu, memory

import flyte.app
from flyte import Image, Resources

image = Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
    "flyteplugins-bigquery", "flyteplugins-databricks", "flyteplugins-snowflake"
)

app_env = flyte.app.ConnectorEnvironment(
    name="bigquery-connector-app",
    image=image,
    resources=Resources(cpu="1", memory="1Gi"),
)

The package declares the flyte.connectors entry point named bigquery, pointing at BigQueryConnector.

Troubleshoot common submission and completion issues

• Omitted location: BigQueryTask.custom_config() serializes "Location": None, and BigQueryConnector.create() indexes that value directly and passes it to the BigQuery client. If the underlying client or deployment does not accept None, set an explicit location such as US or EU.
• Unsupported input type: Automatic parameters are limited to list, bool, bytes, datetime.datetime, float, int, and str. Use one of those declared Python types; the connector does not infer a replacement from the runtime value.
• Name mismatch: Runtime input names are looked up against the serialized interface variables and become parameter names. Keep the inputs keys, SQL parameter names, and invocation keyword arguments aligned.
• No dataframe result: A declared output_dataframe_type produces an interface output named results, but get() populates it only if BigQuery supplies a destination table. A successful query with no destination has no outputs.
• Credentials format: The task argument is a secret name, while the connector argument is the resolved secret value. The resolved value must be service-account JSON accepted by json.loads() and from_service_account_info(); a secret name or malformed JSON will fail during client creation.
• Empty query: Empty and whitespace-only query templates are accepted by the constructor and normalize to ""; there is no task-level validation before submission.
• Advanced configuration and runtime inputs: custom_config() serializes QueryJobConfig fields, but BigQueryConnector.create() constructs a new QueryJobConfig when runtime inputs are supplied and does not merge the serialized advanced fields into that object. Verify the resulting submission when combining advanced job settings with runtime parameters.
• Cached clients: _get_bigquery_client() reuses clients for identical cache-key values. A changed credential string creates a different key; reusing the same project, location, user agent, and credential string reuses the cached client.