flyteplugins-connectors 2.0.0b29__py3-none-any.whl

flyteplugins/connectors/bigquery/__init__.py

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

flyteplugins/connectors/bigquery/connector.py

@@ -0,0 +1,140 @@
+import datetime
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any, Dict, Optional
+
+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
+
+
+@lru_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,
+        **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 = _get_bigquery_client(
+            project=project,
+            location=location,
+            user_agent=user_agent,
+            google_application_credentials=custom.get("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 = _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",
+        )
+
+        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 = _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/connectors/bigquery/task.py

@@ -0,0 +1,83 @@
+import re
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Type
+
+from flyte._task import TaskTemplate
+from flyte.connectors._connector import AsyncConnectorExecutorMixin
+from flyte.io import DataFrame
+from flyte.models import NativeInterface, SerializationContext
+from flyteidl2.core import tasks_pb2
+from google.cloud import bigquery
+from google.protobuf import json_format
+from google.protobuf.struct_pb2 import Struct
+
+
+@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,
+            **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}
+        s = Struct()
+        s.update(config)
+        return json_format.MessageToDict(s)
+
+    def sql(self, sctx: SerializationContext) -> Optional[str]:
+        sql = tasks_pb2.Sql(statement=self.query_template, dialect=tasks_pb2.Sql.Dialect.ANSI)
+        return sql

flyteplugins_connectors-2.0.0b29.dist-info/METADATA

@@ -0,0 +1,156 @@
+Metadata-Version: 2.4
+Name: flyteplugins-connectors
+Version: 2.0.0b29
+Summary: Connector plugin for flyte
+Author-email: Kevin Su <pingsutw@users.noreply.github.com>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: flyte
+Requires-Dist: flyteidl
+Provides-Extra: bigquery
+Requires-Dist: google-cloud-bigquery; extra == "bigquery"
+Requires-Dist: google-cloud-bigquery-storage; extra == "bigquery"
+
+<div align="center">
+
+# 🔗 Flyte Connectors Plugin
+
+[![PyPI version](https://badge.fury.io/py/flyteplugins-connectors.svg)](https://badge.fury.io/py/flyteplugins-connectors)
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![Tests](https://github.com/flyteorg/flyte/workflows/tests/badge.svg)](https://github.com/flyteorg/flyte/actions)
+[![Documentation](https://img.shields.io/badge/docs-latest-brightgreen.svg)](https://docs.flyte.org)
+
+**🚀 Seamlessly connect Flyte workflows to external data sources and services**
+
+*Build powerful data pipelines with native integrations to popular cloud services*
+
+</div>
+
+## 🚀 Quick Start
+
+### Installation
+
+
+```bash
+pip install --pre flyteplugins-connectors[bigquery]  # Install BigQuery connector
+```
+
+
+
+### BigQuery Integration
+
+Execute SQL queries on BigQuery and seamlessly integrate results into your Flyte workflows:
+
+```python
+from flyteplugins.connectors.bigquery.task import BigQueryConfig, BigQueryTask
+import flyte
+from flyte.io import DataFrame
+
+# Configure your BigQuery connection
+config = BigQueryConfig(
+    ProjectID="your-gcp-project",
+    Location="US"  # Optional: specify region
+)
+
+# Create a task environment
+env = flyte.TaskEnvironment(name="analytics_env")
+
+# Define your BigQuery task
+analytics_task = BigQueryTask(
+    name="user_analytics",
+    inputs={
+        "user_id": int,
+        "start_date": str,
+        "end_date": str
+    },
+    output_dataframe_type=DataFrame,
+    plugin_config=config,
+    query_template="""
+        SELECT
+            user_id,
+            COUNT(*) as event_count,
+            MAX(timestamp) as last_activity
+        FROM events
+        WHERE user_id = {{ .user_id }}
+          AND DATE(timestamp) BETWEEN '{{ .start_date }}' AND '{{ .end_date }}'
+        GROUP BY user_id
+    """
+)
+
+env.from_task(analytics_task)
+
+# Run your workflow
+if __name__ == "__main__":
+    flyte.init_from_config()
+    result = flyte.with_runcontext(mode="remote").run(
+        analytics_task,
+        user_id=12345,
+        start_date="2024-01-01",
+        end_date="2024-01-31"
+    )
+    print(f"Workflow URL: {result.url}")
+```
+
+## 📚 Available Connectors
+
+| Connector | Status | Description | Use Cases |
+|-----------|--------|-------------|-----------|
+| 🔷 **BigQuery** | ✅ Stable | Google Cloud data warehouse | Analytics, ML training, reporting |
+| 🔗 **More Coming Soon** | 🚧 | Additional connectors in development | - |
+
+## 🧪 Testing
+
+Run the test suite to ensure everything works correctly:
+
+```bash
+# Run all connector tests
+pytest plugins/connectors/tests/ -v
+
+# Run specific connector tests
+pytest plugins/connectors/tests/test_bigquery.py -v
+
+# Run with coverage
+pytest plugins/connectors/tests/ --cov=flyteplugins.connectors --cov-report=html
+```
+
+## 🤝 Contributing
+
+We welcome contributions! Here's how to get started:
+
+1. **Fork the repository**
+2. **Create a feature branch**: `git checkout -b feature/amazing-connector`
+3. **Write tests** for your changes
+4. **Ensure tests pass**: `pytest plugins/connectors/tests/`
+5. **Submit a pull request**
+
+### Adding a New Connector
+
+1. Create your connector module in `src/flyteplugins/connectors/`
+2. Implement the `TaskTemplate` interface
+3. Add comprehensive tests in `tests/`
+4. Update this README with examples
+5. Add example usage in `examples/connectors/`
+
+## 🔧 Requirements
+
+- **Python**: 3.10+
+- **Flyte**: Latest version
+- **Dependencies**: See `pyproject.toml` for full requirements
+
+
+## 🆘 Support
+
+- **📬 Community**: [Flyte Slack](https://slack.flyte.org/)
+- **🐛 Issues**: [GitHub Issues](https://github.com/flyteorg/flyte-sdk/issues)
+- **💬 Discussions**: [GitHub Discussions](https://github.com/flyteorg/flyte/discussions)
+- **📚 Documentation**: [docs.flyte.org](https://docs.flyte.org)
+
+---
+
+<div align="center">
+
+**Made with ❤️ by the Flyte Community**
+
+[⭐ Star us on GitHub](https://github.com/flyteorg/flyte) • [🐦 Follow us on Twitter](https://twitter.com/flyteorg) • [💼 LinkedIn](https://linkedin.com/company/flyte-org)
+
+</div>

flyteplugins_connectors-2.0.0b29.dist-info/RECORD

@@ -0,0 +1,8 @@
+flyteplugins/connectors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flyteplugins/connectors/bigquery/__init__.py,sha256=TzgxgF_qb5RL4Ajb36mJAP6oRoMWOmJdUL24LpxFjRI,219
+flyteplugins/connectors/bigquery/connector.py,sha256=KkBb5IInWsu6RdFRs-eKfK9OW5P3tl2LbFpnIWd3qGo,5175
+flyteplugins/connectors/bigquery/task.py,sha256=XCiZElNlRa0PjI5hVNxVAL8tRwsEXY6EhDo8tWUZU8Q,3330
+flyteplugins_connectors-2.0.0b29.dist-info/METADATA,sha256=Z27UVvTDk338Fa8xlz1U9NOsKCWFHp9as24pW0HPJdk,4533
+flyteplugins_connectors-2.0.0b29.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+flyteplugins_connectors-2.0.0b29.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_connectors-2.0.0b29.dist-info/RECORD,,

flyteplugins_connectors-2.0.0b29.dist-info/WHEEL

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

flyteplugins_connectors-2.0.0b29.dist-info/top_level.txt

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