add motherduck support for duckdb plugin

plugins/flytekit-duckdb/flytekitplugins/duckdb/__init__.py

@@ -8,4 +8,4 @@
    DuckDBQuery
 """
 
-from .task import DuckDBQuery
+from .task import DuckDBProvider, DuckDBQuery

plugins/flytekit-duckdb/flytekitplugins/duckdb/task.py

@@ -1,7 +1,9 @@
 import json
+from enum import Enum
+from functools import partial
 from typing import Dict, List, NamedTuple, Optional, Union
 
-from flytekit import PythonInstanceTask, lazy_module
+from flytekit import PythonInstanceTask, Secret, current_context, lazy_module
 from flytekit.extend import Interface
 from flytekit.types.structured.structured_dataset import StructuredDataset
 
@@ -10,6 +12,26 @@
 pa = lazy_module("pyarrow")
 
 
+def connect_local(hosted_secret: Optional[Secret]):
+    """Connect to local DuckDB."""
+    return duckdb.connect(":memory:")
+
+
+def connect_motherduck(hosted_secret: Secret):
+    """Connect to MotherDuck."""
+    motherduck_token = current_context().secrets.get(
+        group=hosted_secret.group,
+        key=hosted_secret.key,
+        group_version=hosted_secret.group_version,
+    )
+    return duckdb.connect("md:", config={"motherduck_token": motherduck_token})
+
+
+class DuckDBProvider(Enum):
+    LOCAL = partial(connect_local)
+    MOTHERDUCK = partial(connect_motherduck)
+
+
 class QueryOutput(NamedTuple):
     counter: int = -1
     output: Optional[str] = None
@@ -21,8 +43,9 @@ class DuckDBQuery(PythonInstanceTask):
     def __init__(
         self,
         name: str,
-        query: Union[str, List[str]],
+        query: Optional[Union[str, List[str]]] = None,
         inputs: Optional[Dict[str, Union[StructuredDataset, list]]] = None,
+        provider: DuckDBProvider = DuckDBProvider.LOCAL,
         **kwargs,
     ):
         """
@@ -32,8 +55,13 @@ def __init__(
             name: Name of the task
             query: DuckDB query to execute
             inputs: The query parameters to be used while executing the query
+            provider: DuckDB provider (e.g., LOCAL, MOTHERDUCK, ANOTHERPRODUCT)
         """
         self._query = query
+        self._provider = provider
+        secret_requests: Optional[list[Secret]] = kwargs.get("secret_requests", None)
+        self._hosted_secret = secret_requests[0] if secret_requests else None
+
         outputs = {"result": StructuredDataset}
 
         super(DuckDBQuery, self).__init__(
@@ -44,6 +72,22 @@ def __init__(
             **kwargs,
         )
 
+    def _connect_to_duckdb(self):
+        """
+        Handles the connection to DuckDB based on the provider.
+
+        Returns:
+            A DuckDB connection object.
+        """
+
+        if self._provider not in DuckDBProvider:
+            raise ValueError(f"Unknown DuckDB provider: {self._provider}")
+
+        if self._provider != DuckDBProvider.LOCAL and self._hosted_secret is None:
+            raise ValueError(f"A secret is required for the {self._provider} provider.")
+
+        return self._provider.value(self._hosted_secret)
+
     def _execute_query(
         self, con: duckdb.DuckDBPyConnection, params: list, query: str, counter: int, multiple_params: bool
     ):
@@ -76,14 +120,15 @@ def _execute_query(
 
     def execute(self, **kwargs) -> StructuredDataset:
         # TODO: Enable iterative download after adding the functionality to structured dataset code.
-
-        # create an in-memory database that's non-persistent
-        con = duckdb.connect(":memory:")
+        con = self._connect_to_duckdb()
 
         params = None
         for key in self.python_interface.inputs.keys():
             val = kwargs.get(key)
-            if isinstance(val, StructuredDataset):
+            if key == "query" and val is not None:
+                # Execution query takes priority
+                self._query = val
+            elif isinstance(val, StructuredDataset):
                 # register structured dataset
                 con.register(key, val.open(pa.Table).all())
             elif isinstance(val, (pd.DataFrame, pa.Table)):
@@ -98,6 +143,9 @@ def execute(self, **kwargs) -> StructuredDataset:
             else:
                 raise ValueError(f"Expected inputs of type StructuredDataset, str or list, received {type(val)}")
 
+        if self._query is None:
+            raise ValueError("A query must be specified when defining or executing a DuckDBQuery.")
+
         final_query = self._query
         query_output = QueryOutput()
         # set flag to indicate the presence of params for multiple queries