elaunira-airflow-provider-r2index 0.2.1__tar.gz → 0.3.0__tar.gz

{elaunira_airflow_provider_r2index-0.2.1 → elaunira_airflow_provider_r2index-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: elaunira-airflow-provider-r2index
-Version: 0.2.1
+Version: 0.3.0
 Summary: Airflow provider for Elaunira R2Index connections
 Project-URL: Repository, https://github.com/elaunira/elaunira-airflow-provider-r2index
 License-Expression: MIT
@@ -8,3 +8,4 @@ License-File: LICENSE
 Requires-Python: >=3.12
 Requires-Dist: apache-airflow-providers-hashicorp
 Requires-Dist: apache-airflow>=3.0.0
+Requires-Dist: elaunira-r2index

{elaunira_airflow_provider_r2index-0.2.1 → elaunira_airflow_provider_r2index-0.3.0}/pyproject.toml

@@ -4,13 +4,14 @@ build-backend = "hatchling.build"
 
 [project]
 name = "elaunira-airflow-provider-r2index"
-version = "0.2.1"
+version = "0.3.0"
 description = "Airflow provider for Elaunira R2Index connections"
 requires-python = ">=3.12"
 license = "MIT"
 dependencies = [
     "apache-airflow>=3.0.0",
     "apache-airflow-providers-hashicorp",
+    "elaunira-r2index",
 ]
 
 [project.urls]

{elaunira_airflow_provider_r2index-0.2.1 → elaunira_airflow_provider_r2index-0.3.0}/src/elaunira/airflow/provider/r2index/__init__.py

@@ -17,5 +17,15 @@ def get_provider_info():
                 "hook-class-name": "elaunira.airflow.provider.r2index.hooks.r2index.R2IndexHook",
             }
         ],
+        "task-decorators": [
+            {
+                "name": "r2index_upload",
+                "class-name": "elaunira.airflow.provider.r2index.decorators.r2index.r2index_upload",
+            },
+            {
+                "name": "r2index_download",
+                "class-name": "elaunira.airflow.provider.r2index.decorators.r2index.r2index_download",
+            },
+        ],
         "versions": [__version__],
     }

elaunira_airflow_provider_r2index-0.3.0/src/elaunira/airflow/provider/r2index/decorators/__init__.py

@@ -0,0 +1,8 @@
+"""R2Index TaskFlow decorators."""
+
+from elaunira.airflow.provider.r2index.decorators.r2index import (
+    r2index_download,
+    r2index_upload,
+)
+
+__all__ = ["r2index_download", "r2index_upload"]

elaunira_airflow_provider_r2index-0.3.0/src/elaunira/airflow/provider/r2index/decorators/r2index.py

@@ -0,0 +1,139 @@
+"""R2Index TaskFlow decorators."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
+
+from airflow.sdk.bases.decorator import task_decorator_factory
+
+from elaunira.airflow.provider.r2index.operators.r2index import (
+    DownloadItem,
+    R2IndexDownloadOperator,
+    R2IndexUploadOperator,
+    UploadItem,
+)
+
+if TYPE_CHECKING:
+    from airflow.sdk.bases.decorator import TaskDecorator
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def r2index_upload(
+    bucket: str,
+    r2index_conn_id: str = "r2index_default",
+    **kwargs: Any,
+) -> TaskDecorator:
+    """
+    Decorator to upload files to R2Index.
+
+    The decorated function should return an UploadItem or list of UploadItems.
+
+    Example:
+        @task.r2index_upload(bucket="my-bucket")
+        def prepare_upload() -> UploadItem:
+            return UploadItem(
+                source="/tmp/data.csv",
+                category="acme",
+                entity="acme-data",
+                extension="csv",
+                media_type="text/csv",
+                destination_path="acme/data",
+                destination_filename="data.csv",
+                destination_version="{{ ds }}",
+            )
+
+    :param bucket: R2 bucket name.
+    :param r2index_conn_id: Airflow connection ID for R2Index.
+    """
+    return task_decorator_factory(
+        decorated_operator_class=_R2IndexUploadDecoratedOperator,
+        bucket=bucket,
+        r2index_conn_id=r2index_conn_id,
+        **kwargs,
+    )
+
+
+def r2index_download(
+    bucket: str,
+    r2index_conn_id: str = "r2index_default",
+    **kwargs: Any,
+) -> TaskDecorator:
+    """
+    Decorator to download files from R2Index.
+
+    The decorated function should return a DownloadItem or list of DownloadItems.
+
+    Example:
+        @task.r2index_download(bucket="my-bucket")
+        def prepare_download() -> DownloadItem:
+            return DownloadItem(
+                source_path="acme/data",
+                source_filename="data.csv",
+                source_version="{{ ds }}",
+                destination="/tmp/downloaded.csv",
+            )
+
+    :param bucket: R2 bucket name.
+    :param r2index_conn_id: Airflow connection ID for R2Index.
+    """
+    return task_decorator_factory(
+        decorated_operator_class=_R2IndexDownloadDecoratedOperator,
+        bucket=bucket,
+        r2index_conn_id=r2index_conn_id,
+        **kwargs,
+    )
+
+
+class _R2IndexUploadDecoratedOperator(R2IndexUploadOperator):
+    """Decorated operator for R2Index uploads."""
+
+    custom_operator_name = "@task.r2index_upload"
+
+    def __init__(
+        self,
+        *,
+        python_callable: Callable[..., UploadItem | list[UploadItem]],
+        op_args: list[Any] | None = None,
+        op_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self.python_callable = python_callable
+        self.op_args = op_args or []
+        self.op_kwargs = op_kwargs or {}
+        # items will be set in execute()
+        kwargs["items"] = []
+        super().__init__(**kwargs)
+
+    def execute(self, context: Any) -> list[dict[str, Any]]:
+        """Execute the decorated function and upload the result."""
+        items = self.python_callable(*self.op_args, **self.op_kwargs)
+        self.items = [items] if isinstance(items, UploadItem) else items
+        return super().execute(context)
+
+
+class _R2IndexDownloadDecoratedOperator(R2IndexDownloadOperator):
+    """Decorated operator for R2Index downloads."""
+
+    custom_operator_name = "@task.r2index_download"
+
+    def __init__(
+        self,
+        *,
+        python_callable: Callable[..., DownloadItem | list[DownloadItem]],
+        op_args: list[Any] | None = None,
+        op_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self.python_callable = python_callable
+        self.op_args = op_args or []
+        self.op_kwargs = op_kwargs or {}
+        # items will be set in execute()
+        kwargs["items"] = []
+        super().__init__(**kwargs)
+
+    def execute(self, context: Any) -> list[dict[str, Any]]:
+        """Execute the decorated function and download the result."""
+        items = self.python_callable(*self.op_args, **self.op_kwargs)
+        self.items = [items] if isinstance(items, DownloadItem) else items
+        return super().execute(context)

{elaunira_airflow_provider_r2index-0.2.1 → elaunira_airflow_provider_r2index-0.3.0}/src/elaunira/airflow/provider/r2index/hooks/r2index.py

@@ -25,7 +25,7 @@ class R2IndexHook(BaseHook):
         {
             "vault_conn_id": "openbao_default",
             "vault_namespace": "elaunira/production",
-            "vault_secrets": {
+            "vault_secrets_mapping": {
                 "r2index_api_url": "cloudflare/r2index#api-url",
                 "r2index_api_token": "cloudflare/r2index#api-token",
                 "r2_access_key_id": "cloudflare/r2/airflow#access-key-id",
@@ -37,7 +37,7 @@ class R2IndexHook(BaseHook):
         The vault_conn_id references an Airflow HashiCorp Vault connection
         configured with AppRole or other auth method.
 
-        vault_secrets format: "path#key" or "path" (uses config key as secret key)
+        vault_secrets_mapping format: "path#key" or "path" (uses config key as secret key)
         Required keys:
         - r2index_api_url
         - r2index_api_token
@@ -87,7 +87,7 @@ class R2IndexHook(BaseHook):
                 "password": "API token for direct connection",
                 "vault_conn_id": "openbao-myservice",
                 "vault_namespace": "myservice/production",
-                "vault_secrets": '{"r2index_api_url": "cloudflare/r2index#api-url", ...}',
+                "vault_secrets_mapping": '{"r2index_api_url": "cloudflare/r2index#api-url", ...}',
                 "r2_access_key_id": "Direct mode: R2 access key ID",
                 "r2_secret_access_key": "Direct mode: R2 secret access key",
                 "r2_endpoint_url": "https://account.r2.cloudflarestorage.com",
@@ -112,7 +112,7 @@ class R2IndexHook(BaseHook):
                 widget=BS3TextFieldWidget(),
                 description="OpenBao namespace (e.g., elaunira/production)",
             ),
-            "vault_secrets": StringField(
+            "vault_secrets_mapping": StringField(
                 lazy_gettext("Vault Secrets (JSON)"),
                 widget=BS3TextFieldWidget(),
                 description="JSON mapping of config keys to secret paths",
@@ -213,7 +213,7 @@ class R2IndexHook(BaseHook):
 
             vault_conn_id = extra.get("vault_conn_id")
             if vault_conn_id:
-                secrets_raw = extra.get("vault_secrets")
+                secrets_raw = extra.get("vault_secrets_mapping")
                 if not secrets_raw:
                     return None
                 if isinstance(secrets_raw, str):

elaunira_airflow_provider_r2index-0.3.0/src/elaunira/airflow/provider/r2index/links/__init__.py

@@ -0,0 +1,5 @@
+"""R2Index operator extra links."""
+
+from elaunira.airflow.provider.r2index.links.r2index import R2IndexFileLink
+
+__all__ = ["R2IndexFileLink"]

elaunira_airflow_provider_r2index-0.3.0/src/elaunira/airflow/provider/r2index/links/r2index.py

@@ -0,0 +1,37 @@
+"""R2Index operator extra links."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from airflow.models import BaseOperatorLink
+
+if TYPE_CHECKING:
+    from airflow.models.taskinstancekey import TaskInstanceKey
+
+
+class R2IndexFileLink(BaseOperatorLink):
+    """
+    Link to the R2Index file details.
+
+    This link extracts the file ID from the operator's XCom return value
+    and constructs a URL to view the file in the R2Index UI.
+    """
+
+    name = "R2Index File"
+
+    def get_link(
+        self,
+        operator: Any,
+        *,
+        ti_key: TaskInstanceKey,
+    ) -> str:
+        """Get the link to the R2Index file."""
+        from airflow.models import XCom
+
+        result = XCom.get_value(ti_key=ti_key)
+        if result and isinstance(result, dict):
+            file_id = result.get("id") or result.get("file_record", {}).get("id")
+            if file_id:
+                return f"https://r2index.elaunira.com/files/{file_id}"
+        return ""

elaunira_airflow_provider_r2index-0.3.0/src/elaunira/airflow/provider/r2index/operators/__init__.py

@@ -0,0 +1,15 @@
+"""R2Index operators for Airflow."""
+
+from elaunira.airflow.provider.r2index.operators.r2index import (
+    DownloadItem,
+    R2IndexDownloadOperator,
+    R2IndexUploadOperator,
+    UploadItem,
+)
+
+__all__ = [
+    "DownloadItem",
+    "R2IndexDownloadOperator",
+    "R2IndexUploadOperator",
+    "UploadItem",
+]

elaunira_airflow_provider_r2index-0.3.0/src/elaunira/airflow/provider/r2index/operators/r2index.py

@@ -0,0 +1,232 @@
+"""R2Index operators for Airflow."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+
+from elaunira.airflow.provider.r2index.hooks import R2IndexHook
+from elaunira.airflow.provider.r2index.links.r2index import R2IndexFileLink
+from elaunira.r2index import AsyncR2IndexClient
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+@dataclass
+class UploadItem:
+    """Defines a single file to upload."""
+
+    source: str
+    category: str
+    entity: str
+    extension: str
+    media_type: str
+    destination_path: str
+    destination_filename: str
+    destination_version: str
+    bucket: str | None = None
+    name: str | None = None
+    tags: list[str] | None = None
+    extra: dict[str, Any] | None = None
+    create_checksum_files: bool = False
+    r2index_conn_id: str | None = None
+
+
+@dataclass
+class DownloadItem:
+    """Defines a single file to download."""
+
+    source_path: str
+    source_filename: str
+    source_version: str
+    destination: str
+    bucket: str | None = None
+    verify_checksum: bool = True
+    r2index_conn_id: str | None = None
+
+
+class R2IndexUploadOperator(BaseOperator):
+    """
+    Upload one or more files to R2 in parallel.
+
+    When multiple items are provided, all uploads run concurrently using
+    asyncio within the worker process.
+
+    :param bucket: Default R2 bucket name (can be overridden per item).
+    :param items: List of UploadItem or single UploadItem defining files to upload.
+    :param r2index_conn_id: Default Airflow connection ID (can be overridden per item).
+    """
+
+    template_fields: Sequence[str] = ("items",)
+    template_ext: Sequence[str] = ()
+    ui_color = "#e4f0e8"
+    operator_extra_links = (R2IndexFileLink(),)
+
+    def __init__(
+        self,
+        *,
+        bucket: str,
+        items: list[UploadItem] | UploadItem,
+        r2index_conn_id: str = "r2index_default",
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.bucket = bucket
+        self.items = [items] if isinstance(items, UploadItem) else items
+        self.r2index_conn_id = r2index_conn_id
+
+    def _get_client_config(self, conn_id: str) -> dict[str, Any]:
+        """Get client configuration using the hook's priority chain."""
+        hook = R2IndexHook(r2index_conn_id=conn_id)
+        config = hook._get_config_from_connection()
+        if config is None or not config.get("index_api_url"):
+            config = hook._get_config_from_env()
+        return config or {}
+
+    def execute(self, context: Context) -> list[dict[str, Any]]:
+        """Execute the uploads in parallel."""
+        # Group items by connection ID for efficient client reuse
+        conn_configs: dict[str, dict[str, Any]] = {}
+
+        async def get_or_create_config(conn_id: str) -> dict[str, Any]:
+            if conn_id not in conn_configs:
+                conn_configs[conn_id] = self._get_client_config(conn_id)
+            return conn_configs[conn_id]
+
+        async def upload_one(item: UploadItem) -> dict[str, Any]:
+            conn_id = item.r2index_conn_id or self.r2index_conn_id
+            bucket = item.bucket or self.bucket
+            config = await get_or_create_config(conn_id)
+
+            try:
+                async with AsyncR2IndexClient(**config) as client:
+                    file_record = await client.upload(
+                        bucket=bucket,
+                        source=item.source,
+                        category=item.category,
+                        entity=item.entity,
+                        extension=item.extension,
+                        media_type=item.media_type,
+                        destination_path=item.destination_path,
+                        destination_filename=item.destination_filename,
+                        destination_version=item.destination_version,
+                        name=item.name,
+                        tags=item.tags,
+                        extra=item.extra,
+                        create_checksum_files=item.create_checksum_files,
+                    )
+                    return {"status": "success", "file_record": file_record.model_dump()}
+            except Exception as e:
+                return {"status": "error", "message": str(e), "source": item.source}
+
+        async def upload_all() -> list[dict[str, Any]]:
+            tasks = [upload_one(item) for item in self.items]
+            return await asyncio.gather(*tasks)
+
+        results = asyncio.run(upload_all())
+
+        # Check for errors
+        errors = [r for r in results if r.get("status") == "error"]
+        if errors:
+            self.log.error("Failed uploads: %s", errors)
+            raise AirflowException(f"{len(errors)}/{len(results)} uploads failed")
+
+        self.log.info("Uploaded %d files in parallel", len(results))
+        return [r["file_record"] for r in results]
+
+
+class R2IndexDownloadOperator(BaseOperator):
+    """
+    Download one or more files from R2 in parallel.
+
+    When multiple items are provided, all downloads run concurrently using
+    asyncio within the worker process.
+
+    :param bucket: Default R2 bucket name (can be overridden per item).
+    :param items: List of DownloadItem or single DownloadItem defining files to download.
+    :param r2index_conn_id: Default Airflow connection ID (can be overridden per item).
+    """
+
+    template_fields: Sequence[str] = ("items",)
+    template_ext: Sequence[str] = ()
+    ui_color = "#f0e4e8"
+    operator_extra_links = (R2IndexFileLink(),)
+
+    def __init__(
+        self,
+        *,
+        bucket: str,
+        items: list[DownloadItem] | DownloadItem,
+        r2index_conn_id: str = "r2index_default",
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.bucket = bucket
+        self.items = [items] if isinstance(items, DownloadItem) else items
+        self.r2index_conn_id = r2index_conn_id
+
+    def _get_client_config(self, conn_id: str) -> dict[str, Any]:
+        """Get client configuration using the hook's priority chain."""
+        hook = R2IndexHook(r2index_conn_id=conn_id)
+        config = hook._get_config_from_connection()
+        if config is None or not config.get("index_api_url"):
+            config = hook._get_config_from_env()
+        return config or {}
+
+    def execute(self, context: Context) -> list[dict[str, Any]]:
+        """Execute the downloads in parallel."""
+        # Group items by connection ID for efficient client reuse
+        conn_configs: dict[str, dict[str, Any]] = {}
+
+        async def get_or_create_config(conn_id: str) -> dict[str, Any]:
+            if conn_id not in conn_configs:
+                conn_configs[conn_id] = self._get_client_config(conn_id)
+            return conn_configs[conn_id]
+
+        async def download_one(item: DownloadItem) -> dict[str, Any]:
+            conn_id = item.r2index_conn_id or self.r2index_conn_id
+            bucket = item.bucket or self.bucket
+            config = await get_or_create_config(conn_id)
+
+            try:
+                async with AsyncR2IndexClient(**config) as client:
+                    downloaded_path, file_record = await client.download(
+                        bucket=bucket,
+                        source_path=item.source_path,
+                        source_filename=item.source_filename,
+                        source_version=item.source_version,
+                        destination=item.destination,
+                        verify_checksum=item.verify_checksum,
+                    )
+                    return {
+                        "status": "success",
+                        "path": str(downloaded_path),
+                        "file_record": file_record.model_dump(),
+                    }
+            except Exception as e:
+                return {
+                    "status": "error",
+                    "message": str(e),
+                    "source_path": item.source_path,
+                }
+
+        async def download_all() -> list[dict[str, Any]]:
+            tasks = [download_one(item) for item in self.items]
+            return await asyncio.gather(*tasks)
+
+        results = asyncio.run(download_all())
+
+        # Check for errors
+        errors = [r for r in results if r.get("status") == "error"]
+        if errors:
+            self.log.error("Failed downloads: %s", errors)
+            raise AirflowException(f"{len(errors)}/{len(results)} downloads failed")
+
+        self.log.info("Downloaded %d files in parallel", len(results))
+        return [{"path": r["path"], "file_record": r["file_record"]} for r in results]