hexdag 0.5.0.dev1__py3-none-any.whl

hexdag_plugins/azure/azure_blob_adapter.py

@@ -0,0 +1,450 @@
+"""Azure Blob Storage adapter for hexDAG framework.
+
+Provides file storage and retrieval for pipelines and agents.
+"""
+
+import os
+import time
+from typing import Any
+
+from hexdag.core.ports.healthcheck import HealthStatus
+
+
+class AzureBlobAdapter:
+    """Azure Blob Storage adapter for file operations.
+
+    Provides scalable file storage for documents, artifacts, and pipeline
+    outputs using Azure Blob Storage.
+
+    Parameters
+    ----------
+    connection_string : str
+        Azure Storage connection string (auto-resolved from AZURE_STORAGE_CONNECTION_STRING)
+    container_name : str
+        Blob container name (default: "hexdag")
+    account_url : str, optional
+        Storage account URL (for managed identity auth)
+    use_managed_identity : bool
+        Use Managed Identity instead of connection string (default: False)
+
+    Examples
+    --------
+    YAML configuration::
+
+        spec:
+          ports:
+            storage:
+              adapter: hexdag_plugins.azure.AzureBlobAdapter
+              config:
+                container_name: "pipeline-artifacts"
+
+    With managed identity::
+
+        spec:
+          ports:
+            storage:
+              adapter: hexdag_plugins.azure.AzureBlobAdapter
+              config:
+                account_url: "https://mystorageaccount.blob.core.windows.net"
+                container_name: "pipeline-artifacts"
+                use_managed_identity: true
+
+    Python usage::
+
+        from hexdag_plugins.azure import AzureBlobAdapter
+
+        adapter = AzureBlobAdapter(
+            connection_string="...",  # or auto-resolved
+            container_name="pipeline-artifacts"
+        )
+
+        # Upload file
+        await adapter.aupload("reports/output.json", json_bytes)
+
+        # Download file
+        content = await adapter.adownload("reports/output.json")
+
+        # List files
+        files = await adapter.alist("reports/")
+    """
+
+    def __init__(
+        self,
+        connection_string: str | None = None,
+        container_name: str = "hexdag",
+        account_url: str | None = None,
+        use_managed_identity: bool = False,
+    ):
+        """Initialize Azure Blob Storage adapter.
+
+        Args
+        ----
+            connection_string: Azure Storage connection string
+                (auto-resolved from AZURE_STORAGE_CONNECTION_STRING)
+            container_name: Blob container name (default: "hexdag")
+            account_url: Storage account URL (for managed identity)
+            use_managed_identity: Use Managed Identity (default: False)
+        """
+        self.connection_string = connection_string or os.getenv("AZURE_STORAGE_CONNECTION_STRING")
+        self.container_name = container_name
+        self.account_url = account_url
+        self.use_managed_identity = use_managed_identity
+
+        self._container_client = None
+
+    async def _get_container(self):
+        """Get or create blob container client."""
+        if self._container_client is None:
+            try:
+                from azure.storage.blob.aio import BlobServiceClient
+            except ImportError as e:
+                raise ImportError(
+                    "Azure Storage SDK not installed. Install with: pip install azure-storage-blob"
+                ) from e
+
+            if self.use_managed_identity:
+                if not self.account_url:
+                    raise ValueError("account_url is required when use_managed_identity=True")
+                try:
+                    from azure.identity.aio import DefaultAzureCredential
+
+                    credential = DefaultAzureCredential()
+                    service_client = BlobServiceClient(
+                        account_url=self.account_url, credential=credential
+                    )
+                except ImportError as e:
+                    raise ImportError(
+                        "Azure Identity SDK not installed. Install with: pip install azure-identity"
+                    ) from e
+            else:
+                if not self.connection_string:
+                    raise ValueError(
+                        "connection_string is required when use_managed_identity=False"
+                    )
+                service_client = BlobServiceClient.from_connection_string(self.connection_string)
+
+            # Create container if not exists
+            self._container_client = service_client.get_container_client(self.container_name)
+            try:
+                await self._container_client.create_container()
+            except Exception:  # noqa: SIM105 - contextlib.suppress doesn't work with async
+                # Container might already exist
+                pass
+
+        return self._container_client
+
+    async def aupload(
+        self,
+        blob_name: str,
+        data: bytes | str,
+        content_type: str | None = None,
+        metadata: dict[str, str] | None = None,
+        overwrite: bool = True,
+    ) -> str:
+        """Upload data to blob storage.
+
+        Args
+        ----
+            blob_name: Name/path of the blob
+            data: Data to upload (bytes or string)
+            content_type: MIME type of the content
+            metadata: Optional blob metadata
+            overwrite: Overwrite existing blob (default: True)
+
+        Returns
+        -------
+            URL of the uploaded blob
+        """
+        container = await self._get_container()
+
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+            if content_type is None:
+                content_type = "text/plain"
+
+        blob_client = container.get_blob_client(blob_name)
+
+        await blob_client.upload_blob(
+            data,
+            content_type=content_type,
+            metadata=metadata,
+            overwrite=overwrite,
+        )
+
+        return blob_client.url
+
+    async def adownload(self, blob_name: str) -> bytes:
+        """Download blob content.
+
+        Args
+        ----
+            blob_name: Name/path of the blob
+
+        Returns
+        -------
+            Blob content as bytes
+
+        Raises
+        ------
+            FileNotFoundError: If blob doesn't exist
+        """
+        container = await self._get_container()
+        blob_client = container.get_blob_client(blob_name)
+
+        try:
+            stream = await blob_client.download_blob()
+            return await stream.readall()
+        except Exception as e:
+            if "BlobNotFound" in str(e):
+                raise FileNotFoundError(f"Blob '{blob_name}' not found") from e
+            raise
+
+    async def adownload_text(self, blob_name: str, encoding: str = "utf-8") -> str:
+        """Download blob content as text.
+
+        Args
+        ----
+            blob_name: Name/path of the blob
+            encoding: Text encoding (default: "utf-8")
+
+        Returns
+        -------
+            Blob content as string
+        """
+        content = await self.adownload(blob_name)
+        return content.decode(encoding)
+
+    async def adelete(self, blob_name: str) -> bool:
+        """Delete a blob.
+
+        Args
+        ----
+            blob_name: Name/path of the blob
+
+        Returns
+        -------
+            True if deleted, False if not found
+        """
+        container = await self._get_container()
+        blob_client = container.get_blob_client(blob_name)
+
+        try:
+            await blob_client.delete_blob()
+            return True
+        except Exception:
+            return False
+
+    async def aexists(self, blob_name: str) -> bool:
+        """Check if blob exists.
+
+        Args
+        ----
+            blob_name: Name/path of the blob
+
+        Returns
+        -------
+            True if exists, False otherwise
+        """
+        container = await self._get_container()
+        blob_client = container.get_blob_client(blob_name)
+        return await blob_client.exists()
+
+    async def alist(
+        self,
+        prefix: str | None = None,
+        max_results: int | None = None,
+    ) -> list[dict[str, Any]]:
+        """List blobs in container.
+
+        Args
+        ----
+            prefix: Optional prefix to filter blobs
+            max_results: Maximum number of results
+
+        Returns
+        -------
+            List of blob info dicts with name, size, last_modified, etc.
+        """
+        container = await self._get_container()
+
+        results = []
+        count = 0
+
+        async for blob in container.list_blobs(name_starts_with=prefix):
+            results.append(
+                {
+                    "name": blob.name,
+                    "size": blob.size,
+                    "last_modified": blob.last_modified,
+                    "content_type": blob.content_settings.content_type
+                    if blob.content_settings
+                    else None,
+                    "metadata": blob.metadata,
+                }
+            )
+            count += 1
+            if max_results and count >= max_results:
+                break
+
+        return results
+
+    async def acopy(self, source_blob: str, dest_blob: str) -> str:
+        """Copy blob within container.
+
+        Args
+        ----
+            source_blob: Source blob name
+            dest_blob: Destination blob name
+
+        Returns
+        -------
+            URL of the copied blob
+        """
+        container = await self._get_container()
+        source_client = container.get_blob_client(source_blob)
+        dest_client = container.get_blob_client(dest_blob)
+
+        await dest_client.start_copy_from_url(source_client.url)
+        return dest_client.url
+
+    async def aget_url(self, blob_name: str, expiry_hours: int = 1) -> str:
+        """Get a SAS URL for blob access.
+
+        Args
+        ----
+            blob_name: Name/path of the blob
+            expiry_hours: URL expiry time in hours
+
+        Returns
+        -------
+            SAS URL for blob access
+        """
+        from datetime import datetime, timedelta
+
+        try:
+            from azure.storage.blob import BlobSasPermissions, generate_blob_sas
+        except ImportError as e:
+            raise ImportError(
+                "Azure Storage SDK not installed. Install with: pip install azure-storage-blob"
+            ) from e
+
+        container = await self._get_container()
+        blob_client = container.get_blob_client(blob_name)
+
+        # Generate SAS token
+        sas_token = generate_blob_sas(
+            account_name=blob_client.account_name,
+            container_name=self.container_name,
+            blob_name=blob_name,
+            account_key=self._get_account_key(),
+            permission=BlobSasPermissions(read=True),
+            expiry=datetime.utcnow() + timedelta(hours=expiry_hours),
+        )
+
+        return f"{blob_client.url}?{sas_token}"
+
+    def _get_account_key(self) -> str | None:
+        """Extract account key from connection string."""
+        if not self.connection_string:
+            return None
+
+        for part in self.connection_string.split(";"):
+            if part.startswith("AccountKey="):
+                return part[11:]
+        return None
+
+    async def aupload_json(
+        self,
+        blob_name: str,
+        data: dict | list,
+        metadata: dict[str, str] | None = None,
+    ) -> str:
+        """Upload JSON data to blob storage.
+
+        Args
+        ----
+            blob_name: Name/path of the blob
+            data: JSON-serializable data
+            metadata: Optional blob metadata
+
+        Returns
+        -------
+            URL of the uploaded blob
+        """
+        import json as json_module
+
+        json_str = json_module.dumps(data, indent=2, default=str)
+        return await self.aupload(
+            blob_name,
+            json_str.encode("utf-8"),
+            content_type="application/json",
+            metadata=metadata,
+        )
+
+    async def adownload_json(self, blob_name: str) -> dict | list:
+        """Download and parse JSON blob.
+
+        Args
+        ----
+            blob_name: Name/path of the blob
+
+        Returns
+        -------
+            Parsed JSON data
+        """
+        import json as json_module
+
+        content = await self.adownload_text(blob_name)
+        return json_module.loads(content)
+
+    async def ahealth_check(self) -> HealthStatus:
+        """Check Azure Blob Storage connectivity.
+
+        Returns
+        -------
+            HealthStatus with connectivity details
+        """
+        try:
+            start_time = time.time()
+            container = await self._get_container()
+
+            # Count blobs to verify access
+            count = 0
+            async for _ in container.list_blobs():
+                count += 1
+                if count >= 10:  # Sample only
+                    break
+
+            latency_ms = (time.time() - start_time) * 1000
+
+            return HealthStatus(
+                status="healthy",
+                adapter_name="AzureBlob",
+                latency_ms=latency_ms,
+                details={
+                    "container": self.container_name,
+                    "sample_blob_count": count,
+                },
+            )
+
+        except Exception as e:
+            return HealthStatus(
+                status="unhealthy",
+                adapter_name="AzureBlob",
+                latency_ms=0.0,
+                details={"error": str(e)},
+            )
+
+    async def aclose(self) -> None:
+        """Close the blob storage client."""
+        if self._container_client:
+            await self._container_client.close()
+            self._container_client = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize adapter configuration (excluding secrets)."""
+        return {
+            "container_name": self.container_name,
+            "account_url": self.account_url,
+            "use_managed_identity": self.use_managed_identity,
+        }