elaunira-airflow-provider-r2index 0.2.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

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/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/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/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/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, key="return_value")
+        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/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/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]

{elaunira_airflow_provider_r2index-0.2.2.dist-info → elaunira_airflow_provider_r2index-0.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: elaunira-airflow-provider-r2index
-Version: 0.2.2
+Version: 0.3.1
 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.3.1.dist-info/RECORD

@@ -0,0 +1,17 @@
+elaunira/__init__.py,sha256=qaXVGBU6uIJyveNTEbWux5EcfVSM186PvDwjyxiXLw4,34
+elaunira/airflow/__init__.py,sha256=joy3454dE9nG4_TL4Mh2gw-m2pm-3PijpWpfpTR6nmQ,42
+elaunira/airflow/provider/__init__.py,sha256=zapSJV1rJS7wCpkwklwtIgm-H65A93dulwXGgmArN3M,51
+elaunira/airflow/provider/r2index/__init__.py,sha256=FfEbVB2LUfEx8AIvN-uk2pOXskcMyT2-f5Rlzk_VmgY,1026
+elaunira/airflow/provider/r2index/decorators/__init__.py,sha256=FHlxBfYIDB33GVfQDBllraKdpSmTfYPkufkZI4USMHs,197
+elaunira/airflow/provider/r2index/decorators/r2index.py,sha256=umjV4NQdkBjU-kBkCs-tufe9B_iY5uLpQQC3u8luxq8,4319
+elaunira/airflow/provider/r2index/hooks/__init__.py,sha256=fLCw13d7Z61tf2qhKdNiJN3KdhwE4PDF3SyEZaf2PfU,130
+elaunira/airflow/provider/r2index/hooks/r2index.py,sha256=MxwNCTgPSkuDRJ6qmX8TLtcrsnfij7_RcK8h8kuO0bE,12489
+elaunira/airflow/provider/r2index/links/__init__.py,sha256=0Uv3lwDd_DWz1ko1JcaEQI_2TDIMnbPLuKaBSdNjDWg,144
+elaunira/airflow/provider/r2index/links/r2index.py,sha256=n8L1RumKlo1C_t7NaMfdsWrqzqzsBtHTny6ddoiE5Js,1010
+elaunira/airflow/provider/r2index/operators/__init__.py,sha256=GT13c3srTKvVD9dbRTcjvihQox7pjj5mMVOsV-qDenM,309
+elaunira/airflow/provider/r2index/operators/r2index.py,sha256=fRsqxg1ftwOsg61hdxqmHTPtncQSD7irNmVjYi7AiRk,8696
+elaunira_airflow_provider_r2index-0.3.1.dist-info/METADATA,sha256=TCWMpMzBogagKaTf81oE_DPj4jnGJU8Ca1bNfq6-MXo,412
+elaunira_airflow_provider_r2index-0.3.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+elaunira_airflow_provider_r2index-0.3.1.dist-info/entry_points.txt,sha256=LHYlCOo_WbImlAdr1JvsogvEQsKPJwRe68cCA_Ne9A0,94
+elaunira_airflow_provider_r2index-0.3.1.dist-info/licenses/LICENSE,sha256=lEpO3lmi57nowVFEMLFKL7oMBgLPDi5xB4eVIcCf_KQ,1065
+elaunira_airflow_provider_r2index-0.3.1.dist-info/RECORD,,

elaunira_airflow_provider_r2index-0.2.2.dist-info/RECORD

@@ -1,11 +0,0 @@
-elaunira/__init__.py,sha256=qaXVGBU6uIJyveNTEbWux5EcfVSM186PvDwjyxiXLw4,34
-elaunira/airflow/__init__.py,sha256=joy3454dE9nG4_TL4Mh2gw-m2pm-3PijpWpfpTR6nmQ,42
-elaunira/airflow/provider/__init__.py,sha256=zapSJV1rJS7wCpkwklwtIgm-H65A93dulwXGgmArN3M,51
-elaunira/airflow/provider/r2index/__init__.py,sha256=9JI0sdPmX6Z1pgNM03O_dWFsglzYxbJfR5NCKiupvg0,638
-elaunira/airflow/provider/r2index/hooks/__init__.py,sha256=fLCw13d7Z61tf2qhKdNiJN3KdhwE4PDF3SyEZaf2PfU,130
-elaunira/airflow/provider/r2index/hooks/r2index.py,sha256=MxwNCTgPSkuDRJ6qmX8TLtcrsnfij7_RcK8h8kuO0bE,12489
-elaunira_airflow_provider_r2index-0.2.2.dist-info/METADATA,sha256=cVxeI27z0pIDiLt_ht_GZ8k1cFKYMTH11kNarcaPFmk,380
-elaunira_airflow_provider_r2index-0.2.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-elaunira_airflow_provider_r2index-0.2.2.dist-info/entry_points.txt,sha256=LHYlCOo_WbImlAdr1JvsogvEQsKPJwRe68cCA_Ne9A0,94
-elaunira_airflow_provider_r2index-0.2.2.dist-info/licenses/LICENSE,sha256=lEpO3lmi57nowVFEMLFKL7oMBgLPDi5xB4eVIcCf_KQ,1065
-elaunira_airflow_provider_r2index-0.2.2.dist-info/RECORD,,