elaunira-airflow 0.1.0__py3-none-any.whl

elaunira/__init__.py

@@ -0,0 +1 @@
+"""Elaunira namespace package."""

elaunira/airflow/__init__.py

@@ -0,0 +1,5 @@
+"""Elaunira Airflow operators and hooks for R2Index."""
+
+from importlib.metadata import version
+
+__version__ = version("elaunira-airflow")

elaunira/airflow/decorators/__init__.py

@@ -0,0 +1,5 @@
+"""Elaunira Airflow TaskFlow decorators."""
+
+from elaunira.airflow.decorators.r2index import r2index_download, r2index_upload
+
+__all__ = ["r2index_upload", "r2index_download"]

elaunira/airflow/decorators/r2index.py

@@ -0,0 +1,139 @@
+"""R2Index TaskFlow decorators."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
+
+from airflow.decorators.base import task_decorator_factory
+
+from elaunira.airflow.operators.r2index import (
+    DownloadItem,
+    R2IndexDownloadOperator,
+    R2IndexUploadOperator,
+    UploadItem,
+)
+
+if TYPE_CHECKING:
+    from airflow.decorators.base 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/hooks/__init__.py

@@ -0,0 +1,5 @@
+"""Elaunira Airflow hooks."""
+
+from elaunira.airflow.provider.r2index.hooks import R2IndexHook
+
+__all__ = ["R2IndexHook"]

elaunira/airflow/links/__init__.py

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

elaunira/airflow/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/operators/__init__.py

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

elaunira/airflow/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.hooks.r2index import R2IndexHook
+from elaunira.airflow.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-0.1.0.dist-info/METADATA

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: elaunira-airflow
+Version: 0.1.0
+Summary: Elaunira Airflow operators and hooks for R2Index
+Project-URL: Repository, https://github.com/elaunira/elaunira-airflow
+License-Expression: MIT
+Requires-Python: >=3.12
+Requires-Dist: apache-airflow-providers-edge3
+Requires-Dist: apache-airflow-providers-git
+Requires-Dist: apache-airflow>=3.0.0
+Requires-Dist: cryptography
+Requires-Dist: elaunira-airflow-provider-r2index
+Requires-Dist: elaunira-r2index
+Requires-Dist: pyjwt

elaunira_airflow-0.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+elaunira/__init__.py,sha256=qaXVGBU6uIJyveNTEbWux5EcfVSM186PvDwjyxiXLw4,34
+elaunira/airflow/__init__.py,sha256=7pO55pTaG_1iZrfvDW8A26V_yBgwK7MbxiD7g22s2_4,139
+elaunira/airflow/decorators/__init__.py,sha256=YP5PKFqMAAN0F-T-0imzyix2Ld_FWQISilc33XNmeE0,176
+elaunira/airflow/decorators/r2index.py,sha256=tXAyiJpg3t271xYXuab0A6E87mBd9UsezY7sHgAj0fs,4294
+elaunira/airflow/hooks/__init__.py,sha256=N17rG-6soMLH4WBvx-495-n7UXF3RDdLuEX7icxqjxI,122
+elaunira/airflow/links/__init__.py,sha256=PYO9UsgLrLMO6tyOjNDfFK8ZvA5WejIpYVVaqCIe4nQ,136
+elaunira/airflow/links/r2index.py,sha256=YSYrTZVxyQ-6bXFTKJcZOPlnQFhoZJZ_ZDlqMiCXnD8,990
+elaunira/airflow/operators/__init__.py,sha256=xZuu18w7Vh5Azmxq9lYz_oNGjH7E26n4ogdfeVPrPBw,289
+elaunira/airflow/operators/r2index.py,sha256=s1RONuoEk1xnOtlLVCjyw1bXCuFHqRREXOxIicp8lso,8670
+elaunira_airflow-0.1.0.dist-info/METADATA,sha256=Q0Xiqxq3b0Cnx_YDxdqz0c8oFT17wCpedTz3ofX011I,493
+elaunira_airflow-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+elaunira_airflow-0.1.0.dist-info/entry_points.txt,sha256=A8fFcVohG1m1_0pzCIJcqoZJIFoexzjZdrpvgIEn_JY,166
+elaunira_airflow-0.1.0.dist-info/RECORD,,

elaunira_airflow-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

elaunira_airflow-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[airflow.task_decorators]
+r2index_download = elaunira.airflow.decorators.r2index:r2index_download
+r2index_upload = elaunira.airflow.decorators.r2index:r2index_upload