flyteplugins-bigquery 2.0.0b54__py3-none-any.whl

flyteplugins/bigquery/__init__.py

@@ -0,0 +1,4 @@
+from flyteplugins.bigquery.connector import BigQueryConnector
+from flyteplugins.bigquery.task import BigQueryConfig, BigQueryTask
+
+__all__ = ["BigQueryConfig", "BigQueryConnector", "BigQueryTask"]

flyteplugins/bigquery/connector.py

@@ -0,0 +1,143 @@
+import datetime
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from async_lru import alru_cache
+from flyte import logger
+from flyte.connectors import (
+    AsyncConnector,
+    ConnectorRegistry,
+    Resource,
+    ResourceMeta,
+)
+from flyte.connectors.utils import convert_to_flyte_phase
+from flyte.io import DataFrame
+from flyte.types import TypeEngine
+from flyteidl2.core.execution_pb2 import TaskExecution, TaskLog
+from flyteidl2.core.tasks_pb2 import TaskTemplate
+from google.api_core.client_info import ClientInfo
+from google.cloud import bigquery
+from google.oauth2 import service_account
+from google.protobuf import json_format
+
+pythonTypeToBigQueryType: Dict[type, str] = {
+    # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#data_type_sizes
+    list: "ARRAY",
+    bool: "BOOL",
+    bytes: "BYTES",
+    datetime.datetime: "DATETIME",
+    float: "FLOAT64",
+    int: "INT64",
+    str: "STRING",
+}
+
+
+@dataclass
+class BigQueryMetadata(ResourceMeta):
+    job_id: str
+    project: str
+    location: str
+    user_agent: str
+
+
+@alru_cache
+async def _get_bigquery_client(
+    project: str, location: str, user_agent: str, google_application_credentials: str
+) -> bigquery.Client:
+    if google_application_credentials is not None:
+        credentials = service_account.Credentials.from_service_account_info(google_application_credentials)
+    else:
+        credentials = None
+    cinfo = ClientInfo(user_agent=user_agent)
+    return bigquery.Client(project=project, location=location, client_info=cinfo, credentials=credentials)
+
+
+class BigQueryConnector(AsyncConnector):
+    name = "Bigquery Connector"
+    task_type_name = "bigquery_query_job_task"
+    metadata_type = BigQueryMetadata
+
+    async def create(
+        self,
+        task_template: TaskTemplate,
+        inputs: Optional[Dict[str, Any]] = None,
+        google_application_credentials: Optional[str] = None,
+        **kwargs,
+    ) -> BigQueryMetadata:
+        job_config = None
+        if inputs:
+            python_interface_inputs = {
+                name: TypeEngine.guess_python_type(lt.type)
+                for name, lt in task_template.interface.inputs.variables.items()
+            }
+            job_config = bigquery.QueryJobConfig(
+                query_parameters=[
+                    bigquery.ScalarQueryParameter(name, pythonTypeToBigQueryType[python_interface_inputs[name]], val)
+                    for name, val in inputs.items()
+                ]
+            )
+
+        custom = json_format.MessageToDict(task_template.custom) if task_template.custom else None
+
+        domain = custom.get("Domain")
+        sdk_version = task_template.metadata.runtime.version
+
+        user_agent = f"Flyte/{sdk_version} (GPN:Union;{domain or ''})"
+        project = custom["ProjectID"]
+        location = custom["Location"]
+
+        client = await _get_bigquery_client(
+            project=project,
+            location=location,
+            user_agent=user_agent,
+            google_application_credentials=google_application_credentials,
+        )
+        query_job = client.query(task_template.sql.statement, job_config=job_config)
+
+        return BigQueryMetadata(job_id=str(query_job.job_id), location=location, project=project, user_agent=user_agent)
+
+    async def get(
+        self, resource_meta: BigQueryMetadata, google_application_credentials: Optional[str] = None, **kwargs
+    ) -> Resource:
+        client = await _get_bigquery_client(
+            project=resource_meta.project,
+            location=resource_meta.location,
+            user_agent=resource_meta.user_agent,
+            google_application_credentials=google_application_credentials,
+        )
+        log_link = TaskLog(
+            uri=f"https://console.cloud.google.com/bigquery?project={resource_meta.project}&j=bq:{resource_meta.location}:{resource_meta.job_id}&page=queryresults",
+            name="BigQuery Console",
+            ready=True,
+            link_type=TaskLog.DASHBOARD,
+        )
+
+        job = client.get_job(resource_meta.job_id, resource_meta.project, resource_meta.location)
+        if job.errors:
+            logger.error("failed to run BigQuery job with error:", job.errors.__str__())
+            return Resource(phase=TaskExecution.FAILED, message=job.errors.__str__(), log_links=[log_link])
+
+        cur_phase = convert_to_flyte_phase(str(job.state))
+        res = None
+
+        if cur_phase == TaskExecution.SUCCEEDED:
+            dst = job.destination
+            if dst:
+                output_location = f"bq://{dst.project}:{dst.dataset_id}.{dst.table_id}"
+                res = {"results": DataFrame(uri=output_location)}
+
+        return Resource(phase=cur_phase, message=str(job.state), log_links=[log_link], outputs=res)
+
+    async def delete(
+        self, resource_meta: BigQueryMetadata, google_application_credentials: Optional[str] = None, **kwargs
+    ):
+        client = await _get_bigquery_client(
+            project=resource_meta.project,
+            location=resource_meta.location,
+            user_agent=resource_meta.user_agent,
+            google_application_credentials=google_application_credentials,
+        )
+        client.cancel_job(resource_meta.job_id, resource_meta.project, resource_meta.location)
+
+
+ConnectorRegistry.register(BigQueryConnector())

flyteplugins/bigquery/task.py

@@ -0,0 +1,80 @@
+import re
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Type
+
+from flyte.connectors import AsyncConnectorExecutorMixin
+from flyte.extend import TaskTemplate
+from flyte.io import DataFrame
+from flyte.models import NativeInterface, SerializationContext
+from flyteidl2.core import tasks_pb2
+from google.cloud import bigquery
+
+
+@dataclass
+class BigQueryConfig(object):
+    """
+    BigQueryConfig should be used to configure a BigQuery Task.
+    """
+
+    ProjectID: str
+    Location: Optional[str] = None
+    QueryJobConfig: Optional[bigquery.QueryJobConfig] = None
+
+
+class BigQueryTask(AsyncConnectorExecutorMixin, TaskTemplate):
+    _TASK_TYPE = "bigquery_query_job_task"
+
+    def __init__(
+        self,
+        name: str,
+        query_template: str,
+        plugin_config: BigQueryConfig,
+        inputs: Optional[Dict[str, Type]] = None,
+        output_dataframe_type: Optional[Type[DataFrame]] = None,
+        google_application_credentials: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        To be used to query BigQuery Tables.
+
+        :param name: The Name of this task, should be unique in the project
+        :param query_template: The actual query to run. We use Flyte's Golang templating format for Query templating.
+         Refer to the templating documentation
+        :param plugin_config: BigQueryConfig object
+        :param inputs: Name and type of inputs specified as an ordered dictionary
+        :param output_dataframe_type: If some data is produced by this query, then you can specify the
+         output dataframe type.
+         :param google_application_credentials: The name of the secret containing the Google Application Credentials.
+        """
+        outputs = None
+        if output_dataframe_type is not None:
+            outputs = {
+                "results": output_dataframe_type,
+            }
+        super().__init__(
+            name=name,
+            interface=NativeInterface({k: (v, None) for k, v in inputs.items()} if inputs else {}, outputs or {}),
+            task_type=self._TASK_TYPE,
+            image=None,
+            **kwargs,
+        )
+        self.output_dataframe_type = output_dataframe_type
+        self.plugin_config = plugin_config
+        self.query_template = re.sub(r"\s+", " ", query_template.replace("\n", " ").replace("\t", " ")).strip()
+        self.google_application_credentials = google_application_credentials
+
+    def custom_config(self, sctx: SerializationContext) -> Optional[Dict[str, Any]]:
+        config = {
+            "Location": self.plugin_config.Location,
+            "ProjectID": self.plugin_config.ProjectID,
+            "Domain": sctx.domain,
+        }
+        if self.plugin_config.QueryJobConfig is not None:
+            config.update(self.plugin_config.QueryJobConfig.to_api_repr()["query"])
+        if self.google_application_credentials is not None:
+            config["secrets"] = {"google_application_credentials:": self.google_application_credentials}
+        return config
+
+    def sql(self, sctx: SerializationContext) -> Optional[str]:
+        sql = tasks_pb2.Sql(statement=self.query_template, dialect=tasks_pb2.Sql.Dialect.ANSI)
+        return sql

flyteplugins_bigquery-2.0.0b54.dist-info/METADATA

@@ -0,0 +1,34 @@
+Metadata-Version: 2.4
+Name: flyteplugins-bigquery
+Version: 2.0.0b54
+Summary: BigQuery plugin for flyte
+Author-email: Kevin Su <pingsutw@users.noreply.github.com>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: flyte[connector]
+Requires-Dist: google-cloud-bigquery
+Requires-Dist: google-cloud-bigquery-storage
+
+# BigQuery Plugin for Flyte
+
+This plugin provides BigQuery integration for Flyte, enabling you to run BigQuery queries as Flyte tasks.
+
+## Installation
+
+```bash
+pip install flyteplugins-bigquery
+```
+
+## Usage
+
+```python
+from flyteplugins.bigquery 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},
+)
+```

flyteplugins_bigquery-2.0.0b54.dist-info/RECORD

@@ -0,0 +1,9 @@
+flyteplugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flyteplugins/bigquery/__init__.py,sha256=aXxeQ2_aLVqZej1ERmXPXtpsJkYcPKg448W_4nXWKe0,197
+flyteplugins/bigquery/connector.py,sha256=Ds5MFW0Pndv8F-IG73TfA34iQsDvnbtwFvMX5Qq99jE,5308
+flyteplugins/bigquery/task.py,sha256=VphsNEb0C7pTTwN5SmHM8IjkKfcOB-MODofSX_1OB2Q,3190
+flyteplugins_bigquery-2.0.0b54.dist-info/METADATA,sha256=qVkMPznvvR0WyzjXqL_TQqYmso0CubcuQnluIDI6qJA,867
+flyteplugins_bigquery-2.0.0b54.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+flyteplugins_bigquery-2.0.0b54.dist-info/entry_points.txt,sha256=oBu4yigPOb5w59qt1pJpqLVXltD2IAxUdU1FbYpPeLE,80
+flyteplugins_bigquery-2.0.0b54.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_bigquery-2.0.0b54.dist-info/RECORD,,

flyteplugins_bigquery-2.0.0b54.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.10.2)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flyteplugins_bigquery-2.0.0b54.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[flyte.connectors]
+bigquery = flyteplugins.bigquery.connector:BigQueryConnector

flyteplugins_bigquery-2.0.0b54.dist-info/top_level.txt

@@ -0,0 +1 @@
+flyteplugins