diff --git a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py index fe6e2c5af..8a8fa4adf 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py @@ -7,6 +7,17 @@ import logging from gooddata_sdk._version import __version__ +from gooddata_sdk.catalog.ai_lake.model import ( + CatalogDatabaseInstance, + CatalogFailedOperation, + CatalogOperation, + CatalogOperationError, + CatalogPendingOperation, + CatalogProvisionDatabaseInstanceRequest, + CatalogSucceededOperation, + OperationKind, +) +from gooddata_sdk.catalog.ai_lake.service import AiLakeService from gooddata_sdk.catalog.data_source.action_model.requests.ldm_request import ( CatalogGenerateLdmRequest, CatalogPdmLdmRequest, diff --git a/packages/gooddata-sdk/src/gooddata_sdk/catalog/__init__.py b/packages/gooddata-sdk/src/gooddata_sdk/catalog/__init__.py index 67106a19b..b3c6f730a 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/catalog/__init__.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/catalog/__init__.py @@ -1 +1,12 @@ # (C) 2022 GoodData Corporation +from gooddata_sdk.catalog.ai_lake.model import ( + CatalogDatabaseInstance, + CatalogFailedOperation, + CatalogOperation, + CatalogOperationError, + CatalogPendingOperation, + CatalogProvisionDatabaseInstanceRequest, + CatalogSucceededOperation, + OperationKind, +) +from gooddata_sdk.catalog.ai_lake.service import AiLakeService diff --git a/packages/gooddata-sdk/src/gooddata_sdk/catalog/ai_lake/__init__.py b/packages/gooddata-sdk/src/gooddata_sdk/catalog/ai_lake/__init__.py new file mode 100644 index 000000000..0a896878f --- /dev/null +++ b/packages/gooddata-sdk/src/gooddata_sdk/catalog/ai_lake/__init__.py @@ -0,0 +1,24 @@ +# (C) 2025 GoodData Corporation +from gooddata_sdk.catalog.ai_lake.model import ( + CatalogDatabaseInstance, + CatalogFailedOperation, + CatalogOperation, + CatalogOperationError, + CatalogPendingOperation, + CatalogProvisionDatabaseInstanceRequest, + CatalogSucceededOperation, + OperationKind, +) +from gooddata_sdk.catalog.ai_lake.service import AiLakeService + +__all__ = [ + "AiLakeService", + "CatalogDatabaseInstance", + "CatalogFailedOperation", + "CatalogOperation", + "CatalogOperationError", + "CatalogPendingOperation", + "CatalogProvisionDatabaseInstanceRequest", + "CatalogSucceededOperation", + "OperationKind", +] diff --git a/packages/gooddata-sdk/src/gooddata_sdk/catalog/ai_lake/model.py b/packages/gooddata-sdk/src/gooddata_sdk/catalog/ai_lake/model.py new file mode 100644 index 000000000..ca9fa0423 --- /dev/null +++ b/packages/gooddata-sdk/src/gooddata_sdk/catalog/ai_lake/model.py @@ -0,0 +1,105 @@ +# (C) 2025 GoodData Corporation +from __future__ import annotations + +from typing import Any, Dict, List, Literal, Optional, Set + +from attr import define +from gooddata_api_client.model.provision_database_instance_request import ProvisionDatabaseInstanceRequest + +# TypeAlias for OperationKind values matching the API's allowed_values +OperationKind = Literal["provision-database", "deprovision-database"] + + +@define(kw_only=True) +class CatalogOperationError: + title: str + status: int + detail: str + + @classmethod + def from_api(cls, api_obj: Any) -> "CatalogOperationError": + return cls( + title=api_obj.title, + status=api_obj.status, + detail=api_obj.detail, + ) + + +@define(kw_only=True) +class CatalogOperation: + id: str + kind: str + status: str + + +@define(kw_only=True) +class CatalogPendingOperation(CatalogOperation): + status: str = "pending" + + @classmethod + def from_api(cls, api_obj: Any) -> "CatalogPendingOperation": + return cls( + id=api_obj.id, + kind=api_obj.kind, + status="pending", + ) + + +@define(kw_only=True) +class CatalogSucceededOperation(CatalogOperation): + status: str = "succeeded" + result: Optional[Dict[str, Any]] = None + + @classmethod + def from_api(cls, api_obj: Any) -> "CatalogSucceededOperation": + result = None + if hasattr(api_obj, "result"): + result = api_obj.result + return cls( + id=api_obj.id, + kind=api_obj.kind, + status="succeeded", + result=result, + ) + + +@define(kw_only=True) +class CatalogFailedOperation(CatalogOperation): + status: str = "failed" + error: CatalogOperationError = None # type: ignore[assignment] + + @classmethod + def from_api(cls, api_obj: Any) -> "CatalogFailedOperation": + return cls( + id=api_obj.id, + kind=api_obj.kind, + status="failed", + error=CatalogOperationError.from_api(api_obj.error), + ) + + +@define(kw_only=True) +class CatalogProvisionDatabaseInstanceRequest: + name: str + storage_ids: Set[str] + + def as_api_model(self) -> ProvisionDatabaseInstanceRequest: + return ProvisionDatabaseInstanceRequest( + name=self.name, + storage_ids=list(self.storage_ids), + ) + + +@define(kw_only=True) +class CatalogDatabaseInstance: + id: str + name: str + storage_ids: List[str] + + @classmethod + def from_api(cls, api_obj: Any) -> "CatalogDatabaseInstance": + return cls( + id=api_obj.id, + name=api_obj.name, + storage_ids=list(api_obj.storage_ids), + ) diff --git a/packages/gooddata-sdk/src/gooddata_sdk/catalog/ai_lake/service.py b/packages/gooddata-sdk/src/gooddata_sdk/catalog/ai_lake/service.py new file mode 100644 index 000000000..9a7cf4926 --- /dev/null +++ b/packages/gooddata-sdk/src/gooddata_sdk/catalog/ai_lake/service.py @@ -0,0 +1,113 @@ +# (C) 2025 GoodData Corporation +from __future__ import annotations + +from typing import Optional, Union + +from gooddata_api_client.api.ai_lake_api import AILakeApi +from gooddata_api_client.model.failed_operation import FailedOperation as ApiFailedOperation +from gooddata_api_client.model.pending_operation import PendingOperation as ApiPendingOperation +from gooddata_api_client.model.succeeded_operation import SucceededOperation as ApiSucceededOperation + +from gooddata_sdk.catalog.ai_lake.model import ( + CatalogDatabaseInstance, + CatalogFailedOperation, + CatalogOperation, + CatalogPendingOperation, + CatalogProvisionDatabaseInstanceRequest, + CatalogSucceededOperation, +) +from gooddata_sdk.client import GoodDataApiClient + + +def _convert_operation( + api_op: Union[ApiPendingOperation, ApiSucceededOperation, ApiFailedOperation], +) -> CatalogOperation: + if isinstance(api_op, ApiSucceededOperation): + return CatalogSucceededOperation.from_api(api_op) + elif isinstance(api_op, ApiFailedOperation): + return CatalogFailedOperation.from_api(api_op) + elif isinstance(api_op, ApiPendingOperation): + return CatalogPendingOperation.from_api(api_op) + else: + raise ValueError(f"Unknown operation type: {type(api_op)}") + + +class AiLakeService: + def __init__(self, api_client: GoodDataApiClient) -> None: + self._client = api_client + self._api = AILakeApi(api_client._api_client) + + def provision_database_instance( + self, + request: CatalogProvisionDatabaseInstanceRequest, + operation_id: Optional[str] = None, + ) -> None: + """(BETA) Provision a new AI Lake database instance. + + Args: + request (CatalogProvisionDatabaseInstanceRequest): + Request containing name and storage IDs for the new database instance. + operation_id (Optional[str]): + Optional idempotency key for the operation. + + Returns: + None + """ + kwargs: dict = {"provision_database_instance_request": request.as_api_model()} + if operation_id is not None: + kwargs["operation_id"] = operation_id + self._api.provision_ai_lake_database_instance(**kwargs) + + def deprovision_database_instance( + self, + instance_id: str, + operation_id: Optional[str] = None, + ) -> None: + """(BETA) Delete an existing AI Lake database instance. + + Args: + instance_id (str): + ID of the database instance to delete. + operation_id (Optional[str]): + Optional idempotency key for the operation. + + Returns: + None + """ + kwargs: dict = {} + if operation_id is not None: + kwargs["operation_id"] = operation_id + self._api.deprovision_ai_lake_database_instance(instance_id, **kwargs) + + def get_database_instance( + self, + instance_id: str, + ) -> CatalogDatabaseInstance: + """(BETA) Get details of an AI Lake database instance. + + Args: + instance_id (str): + ID of the database instance to retrieve. + + Returns: + CatalogDatabaseInstance: Details of the database instance. + """ + result = self._api.get_ai_lake_database_instance(instance_id) + return CatalogDatabaseInstance.from_api(result) + + def get_operation( + self, + operation_id: str, + ) -> CatalogOperation: + """(BETA) Get the status of a long-running AI Lake operation. + + Args: + operation_id (str): + ID of the operation to retrieve. + + Returns: + CatalogOperation: The operation status, one of CatalogPendingOperation, + CatalogSucceededOperation, or CatalogFailedOperation. + """ + result = self._api.get_ai_lake_operation(operation_id) + return _convert_operation(result) diff --git a/packages/gooddata-sdk/src/gooddata_sdk/sdk.py b/packages/gooddata-sdk/src/gooddata_sdk/sdk.py index 5e7587464..8f3535e23 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/sdk.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/sdk.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Optional +from gooddata_sdk.catalog.ai_lake.service import AiLakeService from gooddata_sdk.catalog.data_source.service import CatalogDataSourceService from gooddata_sdk.catalog.export.service import ExportService from gooddata_sdk.catalog.organization.service import CatalogOrganizationService @@ -77,6 +78,7 @@ def __init__(self, client: GoodDataApiClient) -> None: """ self._client = client + self._ai_lake = AiLakeService(self._client) self._catalog_workspace = CatalogWorkspaceService(self._client) self._catalog_workspace_content = CatalogWorkspaceContentService(self._client) self._catalog_data_source = CatalogDataSourceService(self._client) @@ -89,6 +91,10 @@ def __init__(self, client: GoodDataApiClient) -> None: self._catalog_permission = CatalogPermissionService(self._client) self._export = ExportService(self._client) + @property + def ai_lake(self) -> AiLakeService: + return self._ai_lake + @property def catalog_workspace(self) -> CatalogWorkspaceService: return self._catalog_workspace diff --git a/packages/gooddata-sdk/tests/catalog/test_ai_lake.py b/packages/gooddata-sdk/tests/catalog/test_ai_lake.py new file mode 100644 index 000000000..c02e7daf5 --- /dev/null +++ b/packages/gooddata-sdk/tests/catalog/test_ai_lake.py @@ -0,0 +1,232 @@ +# (C) 2025 GoodData Corporation +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from gooddata_sdk import ( + AiLakeService, + CatalogDatabaseInstance, + CatalogFailedOperation, + CatalogOperation, + CatalogPendingOperation, + CatalogProvisionDatabaseInstanceRequest, + CatalogSucceededOperation, + GoodDataApiClient, +) +from gooddata_sdk.catalog.ai_lake.model import CatalogOperationError + + +def _make_mock_api_client() -> MagicMock: + """Create a mock GoodDataApiClient.""" + mock_client = MagicMock(spec=GoodDataApiClient) + mock_client._api_client = MagicMock() + return mock_client + + +class TestCatalogProvisionDatabaseInstanceRequest: + def test_as_api_model_round_trip(self) -> None: + request = CatalogProvisionDatabaseInstanceRequest( + name="test-db", + storage_ids={"storage-1", "storage-2"}, + ) + api_model = request.as_api_model() + assert api_model.name == "test-db" + assert set(api_model.storage_ids) == {"storage-1", "storage-2"} + + def test_as_api_model_single_storage(self) -> None: + request = CatalogProvisionDatabaseInstanceRequest( + name="my-database", + storage_ids={"only-storage"}, + ) + api_model = request.as_api_model() + assert api_model.name == "my-database" + assert api_model.storage_ids == ["only-storage"] + + +class TestAiLakeServiceProvision: + def test_provision_database_instance_no_operation_id(self) -> None: + mock_client = _make_mock_api_client() + with patch("gooddata_sdk.catalog.ai_lake.service.AILakeApi") as mock_api_class: + mock_api = MagicMock() + mock_api_class.return_value = mock_api + + service = AiLakeService(mock_client) + request = CatalogProvisionDatabaseInstanceRequest( + name="test-db", + storage_ids={"storage-1"}, + ) + service.provision_database_instance(request) + + mock_api.provision_ai_lake_database_instance.assert_called_once() + call_kwargs = mock_api.provision_ai_lake_database_instance.call_args[1] + assert "operation_id" not in call_kwargs + assert call_kwargs["provision_database_instance_request"].name == "test-db" + + def test_provision_database_instance_with_operation_id(self) -> None: + mock_client = _make_mock_api_client() + with patch("gooddata_sdk.catalog.ai_lake.service.AILakeApi") as mock_api_class: + mock_api = MagicMock() + mock_api_class.return_value = mock_api + + service = AiLakeService(mock_client) + request = CatalogProvisionDatabaseInstanceRequest( + name="test-db", + storage_ids={"storage-1"}, + ) + service.provision_database_instance(request, operation_id="op-uuid-123") + + mock_api.provision_ai_lake_database_instance.assert_called_once() + call_kwargs = mock_api.provision_ai_lake_database_instance.call_args[1] + assert call_kwargs["operation_id"] == "op-uuid-123" + + +class TestAiLakeServiceDeprovision: + def test_deprovision_database_instance_no_operation_id(self) -> None: + mock_client = _make_mock_api_client() + with patch("gooddata_sdk.catalog.ai_lake.service.AILakeApi") as mock_api_class: + mock_api = MagicMock() + mock_api_class.return_value = mock_api + + service = AiLakeService(mock_client) + service.deprovision_database_instance("instance-abc") + + mock_api.deprovision_ai_lake_database_instance.assert_called_once_with("instance-abc") + + def test_deprovision_database_instance_with_operation_id(self) -> None: + mock_client = _make_mock_api_client() + with patch("gooddata_sdk.catalog.ai_lake.service.AILakeApi") as mock_api_class: + mock_api = MagicMock() + mock_api_class.return_value = mock_api + + service = AiLakeService(mock_client) + service.deprovision_database_instance("instance-abc", operation_id="op-uuid-456") + + mock_api.deprovision_ai_lake_database_instance.assert_called_once_with( + "instance-abc", operation_id="op-uuid-456" + ) + + +class TestAiLakeServiceGetDatabaseInstance: + def test_get_database_instance(self) -> None: + mock_client = _make_mock_api_client() + with patch("gooddata_sdk.catalog.ai_lake.service.AILakeApi") as mock_api_class: + mock_api = MagicMock() + mock_api_class.return_value = mock_api + + api_result = MagicMock() + api_result.id = "instance-abc" + api_result.name = "test-db" + api_result.storage_ids = ["storage-1", "storage-2"] + mock_api.get_ai_lake_database_instance.return_value = api_result + + service = AiLakeService(mock_client) + result = service.get_database_instance("instance-abc") + + assert isinstance(result, CatalogDatabaseInstance) + assert result.id == "instance-abc" + assert result.name == "test-db" + assert result.storage_ids == ["storage-1", "storage-2"] + mock_api.get_ai_lake_database_instance.assert_called_once_with("instance-abc") + + +class TestAiLakeServiceGetOperation: + def test_get_operation_pending(self) -> None: + from gooddata_api_client.model.pending_operation import PendingOperation as ApiPendingOperation + + mock_client = _make_mock_api_client() + with patch("gooddata_sdk.catalog.ai_lake.service.AILakeApi") as mock_api_class: + mock_api = MagicMock() + mock_api_class.return_value = mock_api + + api_result = MagicMock(spec=ApiPendingOperation) + api_result.id = "op-123" + api_result.kind = "provision-database" + mock_api.get_ai_lake_operation.return_value = api_result + + service = AiLakeService(mock_client) + result = service.get_operation("op-123") + + assert isinstance(result, CatalogPendingOperation) + assert result.id == "op-123" + assert result.kind == "provision-database" + assert result.status == "pending" + + def test_get_operation_succeeded(self) -> None: + from gooddata_api_client.model.succeeded_operation import SucceededOperation as ApiSucceededOperation + + mock_client = _make_mock_api_client() + with patch("gooddata_sdk.catalog.ai_lake.service.AILakeApi") as mock_api_class: + mock_api = MagicMock() + mock_api_class.return_value = mock_api + + api_result = MagicMock(spec=ApiSucceededOperation) + api_result.id = "op-456" + api_result.kind = "provision-database" + api_result.result = {"instanceId": "inst-789"} + mock_api.get_ai_lake_operation.return_value = api_result + + service = AiLakeService(mock_client) + result = service.get_operation("op-456") + + assert isinstance(result, CatalogSucceededOperation) + assert result.id == "op-456" + assert result.kind == "provision-database" + assert result.status == "succeeded" + assert result.result == {"instanceId": "inst-789"} + + def test_get_operation_succeeded_null_result(self) -> None: + """SucceededOperation.result is nullable (e.g., for deprovision operations).""" + from gooddata_api_client.model.succeeded_operation import SucceededOperation as ApiSucceededOperation + + mock_client = _make_mock_api_client() + with patch("gooddata_sdk.catalog.ai_lake.service.AILakeApi") as mock_api_class: + mock_api = MagicMock() + mock_api_class.return_value = mock_api + + api_result = MagicMock(spec=ApiSucceededOperation) + api_result.id = "op-deprov" + api_result.kind = "deprovision-database" + api_result.result = None + mock_api.get_ai_lake_operation.return_value = api_result + + service = AiLakeService(mock_client) + result = service.get_operation("op-deprov") + + assert isinstance(result, CatalogSucceededOperation) + assert result.id == "op-deprov" + assert result.kind == "deprovision-database" + assert result.status == "succeeded" + assert result.result is None + + def test_get_operation_failed(self) -> None: + from gooddata_api_client.model.failed_operation import FailedOperation as ApiFailedOperation + + mock_client = _make_mock_api_client() + with patch("gooddata_sdk.catalog.ai_lake.service.AILakeApi") as mock_api_class: + mock_api = MagicMock() + mock_api_class.return_value = mock_api + + mock_error = MagicMock() + mock_error.title = "Provisioning failed" + mock_error.status = 500 + mock_error.detail = "Database provisioning encountered an error" + + api_result = MagicMock(spec=ApiFailedOperation) + api_result.id = "op-789" + api_result.kind = "provision-database" + api_result.error = mock_error + mock_api.get_ai_lake_operation.return_value = api_result + + service = AiLakeService(mock_client) + result = service.get_operation("op-789") + + assert isinstance(result, CatalogFailedOperation) + assert result.id == "op-789" + assert result.kind == "provision-database" + assert result.status == "failed" + assert isinstance(result.error, CatalogOperationError) + assert result.error.title == "Provisioning failed" + assert result.error.status == 500 + assert result.error.detail == "Database provisioning encountered an error"