From 742540929c66d5169b96d7d1b2a25dcfd64e9be7 Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Tue, 5 Nov 2024 22:17:28 +0100
Subject: [PATCH] docs: Documentation update for getting and deleting Data

Updated docs to functions related to getting and deleting Data model objects

Docs #COG-334
---
 cognee/api/v1/datasets/routers/get_datasets_router.py |  2 +-
 cognee/modules/data/methods/delete_data.py            | 11 +++++++++++
 cognee/modules/data/methods/get_data.py               | 10 +++++++++-
 cognee/modules/data/methods/get_dataset.py            |  3 ++-
 4 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/cognee/api/v1/datasets/routers/get_datasets_router.py b/cognee/api/v1/datasets/routers/get_datasets_router.py
index 698e6bb9e..c1cb26e85 100644
--- a/cognee/api/v1/datasets/routers/get_datasets_router.py
+++ b/cognee/api/v1/datasets/routers/get_datasets_router.py
@@ -77,7 +77,7 @@ def get_datasets_router() -> APIRouter:
                 detail=f"Dataset ({dataset_id}) not found."
             )
 
-        data = await get_data(user.id, data_id)
+        data = await get_data(data_id)
 
         if data is None:
             raise HTTPException(
diff --git a/cognee/modules/data/methods/delete_data.py b/cognee/modules/data/methods/delete_data.py
index 4fe437d92..7560762e1 100644
--- a/cognee/modules/data/methods/delete_data.py
+++ b/cognee/modules/data/methods/delete_data.py
@@ -3,6 +3,17 @@ from cognee.infrastructure.databases.relational import get_relational_engine
 
 
 async def delete_data(data: Data):
+    """Delete a data record from the database.
+
+       Args:
+           data (Data): The data object to be deleted.
+
+       Raises:
+           ValueError: If the data object is invalid.
+    """
+    if not hasattr(data, '__tablename__'):
+        raise ValueError("The provided data object is missing the required '__tablename__' attribute.")
+
     db_engine = get_relational_engine()
 
     return await db_engine.delete_data_by_id(data.__tablename__, data.id)
diff --git a/cognee/modules/data/methods/get_data.py b/cognee/modules/data/methods/get_data.py
index dac3eb202..b07401463 100644
--- a/cognee/modules/data/methods/get_data.py
+++ b/cognee/modules/data/methods/get_data.py
@@ -3,7 +3,15 @@ from typing import Optional
 from cognee.infrastructure.databases.relational import get_relational_engine
 from ..models import Data
 
-async def get_data(user_id: UUID, data_id: UUID) -> Optional[Data]:
+async def get_data(data_id: UUID) -> Optional[Data]:
+    """Retrieve data by ID.
+
+        Args:
+            data_id (UUID): ID of the data to retrieve
+
+        Returns:
+            Optional[Data]: The requested data object if found, None otherwise
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:
diff --git a/cognee/modules/data/methods/get_dataset.py b/cognee/modules/data/methods/get_dataset.py
index f66b707a1..9f46fa223 100644
--- a/cognee/modules/data/methods/get_dataset.py
+++ b/cognee/modules/data/methods/get_dataset.py
@@ -1,8 +1,9 @@
+from typing import Optional
 from uuid import UUID
 from cognee.infrastructure.databases.relational import get_relational_engine
 from ..models import Dataset
 
-async def get_dataset(user_id: UUID, dataset_id: UUID) -> Dataset:
+async def get_dataset(user_id: UUID, dataset_id: UUID) -> Optional[Dataset]:
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session: