"""
**File:** ``storage_account.py``
**Region:** ``ds_provider_azure_py_lib/linked_service/storage_account``
Azure Linked Service
This module implements a linked service for Azure Storage services (Blob and Table).
Authentication supports either an account access key (shared key) or an Azure AD
service principal (client credentials).
Example (access key):
>>> linked_service = AzureLinkedService(
... settings=AzureLinkedServiceSettings(
... account_name="account name",
... access_key="account key",
... ),
... id=uuid.uuid4(),
... name="testazurepackage",
... version="0.0.1",
... description="testazurepackage"
... )
Example (client credentials):
>>> linked_service = AzureLinkedService(
... settings=AzureLinkedServiceSettings(
... account_name="account name",
... tenant_id="tenant id",
... client_id="client id",
... client_secret="client secret",
... ),
... id=uuid.uuid4(),
... name="testazurepackage",
... version="0.0.1",
... description="testazurepackage"
... )
"""
from dataclasses import dataclass, field
from typing import Generic, TypeVar
from azure.core.credentials import AzureNamedKeyCredential
from azure.data.tables import TableServiceClient
from azure.identity import ClientSecretCredential
from azure.storage.blob import BlobServiceClient
from ds_common_logger_py_lib import Logger
from ds_resource_plugin_py_lib.common.resource.linked_service import LinkedService, LinkedServiceSettings
from ds_resource_plugin_py_lib.common.resource.linked_service.errors import (
AuthenticationError,
ConnectionError,
)
from ..enums import ResourceType
logger = Logger.get_logger(__name__, package=True)
_AUTH_REQUIRED_MESSAGE = (
"Access Key, or Tenant ID, Client ID and Client Secret, are required for Azure authentication."
)
[docs]
@dataclass(kw_only=True)
class AzureLinkedServiceSettings(LinkedServiceSettings):
"""
The object containing the Azure linked service settings.
Either ``access_key`` or the full set of ``tenant_id``, ``client_id`` and
``client_secret`` must be provided.
"""
account_name: str
access_key: str | None = field(default=None, metadata={"mask": True})
tenant_id: str | None = None
client_id: str | None = None
client_secret: str | None = field(default=None, metadata={"mask": True})
AzureLinkedServiceSettingsType = TypeVar(
"AzureLinkedServiceSettingsType",
bound=AzureLinkedServiceSettings,
)
[docs]
@dataclass(kw_only=True)
class AzureLinkedServiceConnection:
"""
The object containing the Azure linked service connection clients.
"""
blob_service_client: BlobServiceClient
table_service_client: TableServiceClient
[docs]
@dataclass(kw_only=True)
class AzureLinkedService(LinkedService[AzureLinkedServiceSettingsType], Generic[AzureLinkedServiceSettingsType]):
"""
Linked service for connecting to AzureLinkedService.
"""
settings: AzureLinkedServiceSettingsType
_blob_service_client: BlobServiceClient | None = None
_table_service_client: TableServiceClient | None = None
_credential: AzureNamedKeyCredential | ClientSecretCredential | None = field(
default=None, metadata={"serialize": False}
)
[docs]
def check_settings_is_set(self) -> None:
"""
Check if settings are set correctly.
Either ``access_key`` or the full set of ``tenant_id``, ``client_id`` and
``client_secret`` must be provided.
Returns:
None
Raises:
AttributeError: If settings are not set correctly.
AuthenticationError: If the account name, or a full set of credentials
(access key, or tenant/client id and secret), is missing.
"""
if not isinstance(self.settings, AzureLinkedServiceSettings):
raise AttributeError("settings not set.")
if not self.settings.account_name:
raise AuthenticationError("Account Name is required for Azure authentication.")
has_access_key = bool(self.settings.access_key)
has_client_credential = bool(
self.settings.tenant_id and self.settings.client_id and self.settings.client_secret
)
if not has_access_key and not has_client_credential:
raise AuthenticationError(_AUTH_REQUIRED_MESSAGE)
@property
def type(self) -> ResourceType:
"""
Get the type of the linked service.
Returns:
ResourceType
"""
return ResourceType.STORAGE_ACCOUNT
@property
def connection(self) -> AzureLinkedServiceConnection:
"""
Get the connection object for Azure StorageAccount.
Returns:
AzureLinkedServiceConnection
"""
if self._table_service_client is None or self._blob_service_client is None:
raise ConnectionError(
message="Connection has not been established. Call connect() first.",
details={"provider": self.type.value},
)
return AzureLinkedServiceConnection(
blob_service_client=self._blob_service_client, table_service_client=self._table_service_client
)
@property
def blob_service_client(self) -> BlobServiceClient:
"""
Get the BlobServiceClient instance.
Returns:
BlobServiceClient
Raises:
ConnectionError: If blob service client is not connected.
"""
if not self._blob_service_client:
raise ConnectionError("Blob service client is not connected. Call connect() first.")
return self._blob_service_client
@property
def table_service_client(self) -> TableServiceClient:
"""
Get the TableServiceClient instance.
Returns:
TableServiceClient
Raises:
ConnectionError: If table service client is not connected.
"""
if not self._table_service_client:
raise ConnectionError("Table service client is not connected. Call connect() first.")
return self._table_service_client
[docs]
def get_blob_service(self) -> BlobServiceClient:
"""
Connect to Azure Blob StorageAccount.
Returns:
BlobServiceClient
"""
account_url = f"https://{self.settings.account_name}.blob.core.windows.net/"
return BlobServiceClient(
account_url=account_url,
credential=self._credential,
)
[docs]
def get_table_service(self) -> TableServiceClient:
"""
Connect to Azure Table StorageAccount.
Returns:
TableServiceClient
"""
account_url = f"https://{self.settings.account_name}.table.core.windows.net/"
return TableServiceClient(
endpoint=account_url,
credential=self._credential,
)
[docs]
def get_credential(self) -> AzureNamedKeyCredential | ClientSecretCredential:
"""
Build the credential used to authenticate against Azure Storage.
Prefers the account access key when set, otherwise falls back to
client credential (service principal) authentication.
Returns:
AzureNamedKeyCredential | ClientSecretCredential
Raises:
AuthenticationError: If the account name, or a full set of credentials
(access key, or tenant/client id and secret), is missing.
"""
if not self.settings.account_name:
raise AuthenticationError("Account Name is required for Azure authentication.")
if self.settings.access_key:
return AzureNamedKeyCredential(
name=self.settings.account_name,
key=self.settings.access_key,
)
if not (self.settings.tenant_id and self.settings.client_id and self.settings.client_secret):
raise AuthenticationError(_AUTH_REQUIRED_MESSAGE)
return ClientSecretCredential(
tenant_id=self.settings.tenant_id,
client_id=self.settings.client_id,
client_secret=self.settings.client_secret,
)
[docs]
def connect(self) -> None:
"""
Connect to Azure Storage (Blob and Table), ensuring both service clients are initialized.
Returns:
None
"""
self.check_settings_is_set()
self._credential = self.get_credential()
self._blob_service_client = self.get_blob_service()
self._table_service_client = self.get_table_service()
logger.debug("Connected to Azure StorageAccount.")
[docs]
def test_connection(self) -> tuple[bool, str]:
"""
Test the connection to Azure Storage (Blob or Table).
Returns:
tuple[bool, str]
"""
try:
if not self._blob_service_client or not self._table_service_client:
self.connect()
_ = self.blob_service_client.list_containers()
_ = self.table_service_client.list_tables()
logger.debug("Tested connection to Azure StorageAccount successfully.")
return True, "Connection successful."
except Exception as exc:
logger.error(f"Failed to test connection: {exc}", exc_info=True)
return False, str(exc)
[docs]
def close(self) -> None:
"""
No need to close the linked service. Just to comply with the interface.
Returns:
None
"""
pass
[docs]
def __enter__(self) -> "AzureLinkedService[AzureLinkedServiceSettingsType]":
"""
Enter context manager.
Returns:
AzureLinkedService: Returns self for use in with statement.
"""
return self
[docs]
def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
"""
Exit context manager and close the connection.
Args:
exc_type: Exception type if an exception occurred.
exc_val: Exception value if an exception occurred.
exc_tb: Exception traceback if an exception occurred.
Returns:
None
"""
self.close()