elaunira-airflow-providers-r2index 0.4.0__py3-none-any.whl

elaunira/__init__.py

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

elaunira/airflow/__init__.py

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

elaunira/airflow/providers/__init__.py

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

elaunira/airflow/providers/r2index/__init__.py

@@ -0,0 +1,31 @@
+"""Elaunira R2Index Airflow provider."""
+
+from importlib.metadata import version
+
+__version__ = version("elaunira-airflow-providers-r2index")
+
+
+def get_provider_info():
+    """Return provider metadata for Airflow."""
+    return {
+        "package-name": "elaunira-airflow-providers-r2index",
+        "name": "R2Index",
+        "description": "Airflow provider for R2Index connections",
+        "connection-types": [
+            {
+                "connection-type": "r2index",
+                "hook-class-name": "elaunira.airflow.providers.r2index.hooks.r2index.R2IndexHook",
+            }
+        ],
+        "task-decorators": [
+            {
+                "name": "r2index_upload",
+                "class-name": "elaunira.airflow.providers.r2index.decorators.r2index.r2index_upload",
+            },
+            {
+                "name": "r2index_download",
+                "class-name": "elaunira.airflow.providers.r2index.decorators.r2index.r2index_download",
+            },
+        ],
+        "versions": [__version__],
+    }

elaunira/airflow/providers/r2index/decorators/__init__.py

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

elaunira/airflow/providers/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.providers.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/providers/r2index/hooks/__init__.py

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

elaunira/airflow/providers/r2index/hooks/r2index.py

@@ -0,0 +1,350 @@
+"""R2Index hook for Airflow."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import TYPE_CHECKING, Any
+
+from airflow.sdk.bases.hook import BaseHook
+
+if TYPE_CHECKING:
+    from elaunira.r2index import R2IndexClient
+
+
+class R2IndexHook(BaseHook):
+    """
+    Hook for interacting with R2Index API and R2 storage.
+
+    Configuration priority:
+        1. Airflow connection with Vault/OpenBao reference (fetches secrets dynamically)
+        2. Airflow connection with direct credentials
+        3. Environment variables (fallback)
+
+    Airflow connection with Vault/OpenBao reference (extra JSON):
+        {
+            "vault_conn_id": "openbao_default",
+            "vault_namespace": "elaunira/production",
+            "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",
+                "r2_secret_access_key": "cloudflare/r2/airflow#secret-access-key",
+                "r2_endpoint_url": "cloudflare/r2/airflow#endpoint-url"
+            }
+        }
+
+        The vault_conn_id references an Airflow HashiCorp Vault connection
+        configured with AppRole or other auth method.
+
+        vault_secrets_mapping format: "path#key" or "path" (uses config key as secret key)
+        Required keys:
+        - r2index_api_url
+        - r2index_api_token
+        - r2_access_key_id
+        - r2_secret_access_key
+        - r2_endpoint_url
+
+    Airflow connection with direct credentials:
+        - host: R2Index API URL
+        - password: R2Index API token
+        - extra.r2_access_key_id: R2 access key ID
+        - extra.r2_secret_access_key: R2 secret access key
+        - extra.r2_endpoint_url: R2 endpoint URL
+
+    Environment variables (fallback):
+        - R2INDEX_API_URL
+        - R2INDEX_API_TOKEN
+        - R2_ACCESS_KEY_ID
+        - R2_SECRET_ACCESS_KEY
+        - R2_ENDPOINT_URL
+    """
+
+    conn_name_attr = "r2index_conn_id"
+    default_conn_name = "r2index_default"
+    conn_type = "r2index"
+    hook_name = "R2Index"
+
+    CONFIG_KEYS = [
+        "r2index_api_url",
+        "r2index_api_token",
+        "r2_access_key_id",
+        "r2_secret_access_key",
+        "r2_endpoint_url",
+    ]
+
+    @classmethod
+    def get_ui_field_behaviour(cls) -> dict[str, Any]:
+        """Customize connection form UI."""
+        return {
+            "hidden_fields": ["port", "schema", "login", "extra"],
+            "relabeling": {
+                "host": "R2Index API URL (direct mode only)",
+                "password": "R2Index API Token (direct mode only)",
+            },
+            "placeholders": {
+                "host": "https://r2index.example.com",
+                "password": "API token for direct connection",
+                "vault_conn_id": "openbao-myservice",
+                "vault_namespace": "myservice/production",
+                "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",
+            },
+        }
+
+    @classmethod
+    def get_connection_form_widgets(cls) -> dict[str, Any]:
+        """Define custom connection form widgets."""
+        from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget
+        from flask_babel import lazy_gettext
+        from wtforms import PasswordField, StringField
+
+        return {
+            "vault_conn_id": StringField(
+                lazy_gettext("Vault Connection ID"),
+                widget=BS3TextFieldWidget(),
+                description="Airflow Vault connection ID (e.g., openbao-elaunira)",
+            ),
+            "vault_namespace": StringField(
+                lazy_gettext("Vault Namespace"),
+                widget=BS3TextFieldWidget(),
+                description="OpenBao namespace (e.g., elaunira/production)",
+            ),
+            "vault_secrets_mapping": StringField(
+                lazy_gettext("Vault Secrets (JSON)"),
+                widget=BS3TextFieldWidget(),
+                description="JSON mapping of config keys to secret paths",
+            ),
+            "r2_access_key_id": StringField(
+                lazy_gettext("R2 Access Key ID"),
+                widget=BS3TextFieldWidget(),
+                description="Direct mode: Cloudflare R2 access key ID",
+            ),
+            "r2_secret_access_key": PasswordField(
+                lazy_gettext("R2 Secret Access Key"),
+                widget=BS3PasswordFieldWidget(),
+                description="Direct mode: Cloudflare R2 secret access key",
+            ),
+            "r2_endpoint_url": StringField(
+                lazy_gettext("R2 Endpoint URL"),
+                widget=BS3TextFieldWidget(),
+                description="Direct mode: Cloudflare R2 endpoint URL",
+            ),
+        }
+
+    def __init__(self, r2index_conn_id: str = default_conn_name) -> None:
+        super().__init__()
+        self.r2index_conn_id = r2index_conn_id
+        self._client: R2IndexClient | None = None
+
+    def _parse_secret_ref(self, secret_ref: str, default_key: str) -> tuple[str, str]:
+        """Parse a secret reference into (path, key).
+
+        Format: "path#key" or just "path" (uses default_key).
+        """
+        if "#" in secret_ref:
+            path, key = secret_ref.rsplit("#", 1)
+            return path, key
+        return secret_ref, default_key
+
+    def _get_config_from_env(self) -> dict[str, str | None]:
+        """Get configuration from environment variables."""
+        return {
+            "index_api_url": os.environ.get("R2INDEX_API_URL"),
+            "index_api_token": os.environ.get("R2INDEX_API_TOKEN"),
+            "r2_access_key_id": os.environ.get("R2_ACCESS_KEY_ID"),
+            "r2_secret_access_key": os.environ.get("R2_SECRET_ACCESS_KEY"),
+            "r2_endpoint_url": os.environ.get("R2_ENDPOINT_URL"),
+        }
+
+    def _get_config_from_vault(
+        self,
+        vault_conn_id: str,
+        secrets: dict[str, str],
+        namespace: str | None = None,
+    ) -> dict[str, str | None] | None:
+        """Get configuration from Vault/OpenBao using Airflow's VaultHook.
+
+        Args:
+            vault_conn_id: Airflow connection ID for Vault/OpenBao
+            secrets: Mapping of config key to secret reference (path#key format)
+            namespace: OpenBao namespace to use
+        """
+        from airflow.providers.hashicorp.hooks.vault import VaultHook
+
+        try:
+            # Note: namespace should be configured in the Vault connection's extra field
+            vault_hook = VaultHook(vault_conn_id=vault_conn_id)
+            secret_cache: dict[str, dict[str, Any]] = {}
+
+            def get_secret_value(config_key: str) -> str | None:
+                secret_ref = secrets.get(config_key)
+                if not secret_ref:
+                    return None
+
+                path, key = self._parse_secret_ref(secret_ref, config_key)
+
+                if path not in secret_cache:
+                    data = vault_hook.get_secret(secret_path=path, secret_version=None)
+                    secret_cache[path] = data or {}
+
+                return secret_cache[path].get(key)
+
+            config = {
+                "index_api_url": get_secret_value("r2index_api_url"),
+                "index_api_token": get_secret_value("r2index_api_token"),
+                "r2_access_key_id": get_secret_value("r2_access_key_id"),
+                "r2_secret_access_key": get_secret_value("r2_secret_access_key"),
+                "r2_endpoint_url": get_secret_value("r2_endpoint_url"),
+            }
+            # Log which values are missing
+            missing = [k for k, v in config.items() if v is None]
+            if missing:
+                self.log.warning("Missing Vault secrets: %s", missing)
+            return config
+        except Exception as e:
+            self.log.error("Failed to get config from Vault: %s", e)
+            return None
+
+    def _get_config_from_connection(self) -> dict[str, str | None] | None:
+        """Get configuration from Airflow connection.
+
+        If connection has vault_conn_id, fetches from Vault/OpenBao.
+        Otherwise uses direct credentials from connection fields.
+        """
+        try:
+            conn = self.get_connection(self.r2index_conn_id)
+            extra = conn.extra_dejson
+
+            vault_conn_id = extra.get("vault_conn_id")
+            if vault_conn_id:
+                secrets_raw = extra.get("vault_secrets_mapping")
+                if not secrets_raw:
+                    return None
+                if isinstance(secrets_raw, str):
+                    secrets = json.loads(secrets_raw)
+                else:
+                    secrets = secrets_raw
+                return self._get_config_from_vault(
+                    vault_conn_id=vault_conn_id,
+                    secrets=secrets,
+                    namespace=extra.get("vault_namespace"),
+                )
+
+            return {
+                "index_api_url": conn.host,
+                "index_api_token": conn.password,
+                "r2_access_key_id": extra.get("r2_access_key_id"),
+                "r2_secret_access_key": extra.get("r2_secret_access_key"),
+                "r2_endpoint_url": extra.get("r2_endpoint_url"),
+            }
+        except Exception as e:
+            self.log.error("Failed to get config from connection: %s", e)
+            return None
+
+    def get_conn(self) -> R2IndexClient:
+        """Get the R2IndexClient."""
+        if self._client is not None:
+            return self._client
+
+        from elaunira.r2index import R2IndexClient
+
+        config = self._get_config_from_connection()
+        if config is None or not config.get("index_api_url"):
+            config = self._get_config_from_env()
+
+        self._client = R2IndexClient(
+            index_api_url=config["index_api_url"],
+            index_api_token=config["index_api_token"],
+            r2_access_key_id=config["r2_access_key_id"],
+            r2_secret_access_key=config["r2_secret_access_key"],
+            r2_endpoint_url=config["r2_endpoint_url"],
+        )
+        return self._client
+
+    def upload(
+        self,
+        bucket: str,
+        source: str,
+        category: str,
+        entity: str,
+        extension: str,
+        media_type: str,
+        destination_path: str,
+        destination_filename: str,
+        destination_version: str,
+        name: str | None = None,
+        tags: list[str] | None = None,
+        extra: dict[str, Any] | None = None,
+        create_checksum_files: bool = False,
+    ) -> dict[str, Any]:
+        """Upload a file to R2 and register it with R2Index."""
+        client = self.get_conn()
+        file_record = client.upload(
+            bucket=bucket,
+            source=source,
+            category=category,
+            entity=entity,
+            extension=extension,
+            media_type=media_type,
+            destination_path=destination_path,
+            destination_filename=destination_filename,
+            destination_version=destination_version,
+            name=name,
+            tags=tags,
+            extra=extra,
+            create_checksum_files=create_checksum_files,
+        )
+        return file_record.model_dump()
+
+    def download(
+        self,
+        bucket: str,
+        source_path: str,
+        source_filename: str,
+        source_version: str,
+        destination: str,
+        verify_checksum: bool = True,
+    ) -> dict[str, Any]:
+        """Download a file from R2."""
+        client = self.get_conn()
+        downloaded_path, file_record = client.download(
+            bucket=bucket,
+            source_path=source_path,
+            source_filename=source_filename,
+            source_version=source_version,
+            destination=destination,
+            verify_checksum=verify_checksum,
+        )
+        return {
+            "path": str(downloaded_path),
+            "file_record": file_record.model_dump(),
+        }
+
+    def get_file(self, file_id: str) -> dict[str, Any]:
+        """Get a file record by ID."""
+        client = self.get_conn()
+        return client.get(file_id).model_dump()
+
+    def list_files(
+        self,
+        bucket: str | None = None,
+        category: str | None = None,
+        entity: str | None = None,
+        extension: str | None = None,
+        tags: list[str] | None = None,
+        limit: int | None = None,
+    ) -> dict[str, Any]:
+        """List files with optional filters."""
+        client = self.get_conn()
+        response = client.list_files(
+            bucket=bucket,
+            category=category,
+            entity=entity,
+            extension=extension,
+            tags=tags,
+            limit=limit,
+        )
+        return response.model_dump()

elaunira/airflow/providers/r2index/links/__init__.py

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

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

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

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

@@ -0,0 +1,11 @@
+Metadata-Version: 2.4
+Name: elaunira-airflow-providers-r2index
+Version: 0.4.0
+Summary: Airflow provider for Elaunira R2Index connections
+Project-URL: Repository, https://github.com/elaunira/elaunira-airflow-providers-r2index
+License-Expression: MIT
+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_providers_r2index-0.4.0.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/providers/__init__.py,sha256=h-aj5DTpHI6WreTrM0NecDgqTnakxp2j6xaa36vCZUs,52
+elaunira/airflow/providers/r2index/__init__.py,sha256=HKjGhB-tnvzbMfpNbyQ4oXYQcK02etALq88ACbwkXOk,1031
+elaunira/airflow/providers/r2index/decorators/__init__.py,sha256=kEXhGWHF8jM-Cbv-w4WVFlnI1m1LJrgSnVRJ-Drfi_s,198
+elaunira/airflow/providers/r2index/decorators/r2index.py,sha256=kQV1xZSWg__nwCWAj-7NlNF2MbSTzY7ypuIvzjsk9sg,4320
+elaunira/airflow/providers/r2index/hooks/__init__.py,sha256=RKgxvpwyOe_uV0JgZAgscAXmZb1pQ24ncxLQrDpfOPw,131
+elaunira/airflow/providers/r2index/hooks/r2index.py,sha256=Bf5y3sFm9X4LJvIvMJWFi0rJz4vgAZG7EuJt4fTIadU,12936
+elaunira/airflow/providers/r2index/links/__init__.py,sha256=MqB1EPonN17M0djF0j3PS9oZ88Vr5JZ5YlHDZkq4G-o,145
+elaunira/airflow/providers/r2index/links/r2index.py,sha256=n8L1RumKlo1C_t7NaMfdsWrqzqzsBtHTny6ddoiE5Js,1010
+elaunira/airflow/providers/r2index/operators/__init__.py,sha256=Tc_b6N1daRIEYgMVON_Z_wnbgfodJXJusnygvEgq0Aw,310
+elaunira/airflow/providers/r2index/operators/r2index.py,sha256=h--O7knmwK3FnjLuF-DjpfRhRwiGlrLTtculti5VQlM,8698
+elaunira_airflow_providers_r2index-0.4.0.dist-info/METADATA,sha256=v5TSDj3FWXXfcd_JjJdZd3d9gcrbvw_ajSY8ftqDtd8,414
+elaunira_airflow_providers_r2index-0.4.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+elaunira_airflow_providers_r2index-0.4.0.dist-info/entry_points.txt,sha256=T04D_FeKrQcuUh7I0ijx4F0mrJquFEYU9KbzljNpHFA,95
+elaunira_airflow_providers_r2index-0.4.0.dist-info/licenses/LICENSE,sha256=lEpO3lmi57nowVFEMLFKL7oMBgLPDi5xB4eVIcCf_KQ,1065
+elaunira_airflow_providers_r2index-0.4.0.dist-info/RECORD,,

elaunira_airflow_providers_r2index-0.4.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_providers_r2index-0.4.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[apache_airflow_provider]
+provider_info = elaunira.airflow.providers.r2index:get_provider_info

elaunira_airflow_providers_r2index-0.4.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Elaunira
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.