From 11634cb58d495cd8ead5caabc27beab6958ff1dd Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Fri, 13 Dec 2024 16:54:53 +0100
Subject: [PATCH] feat: Add unauth access error to getting data

Raise unauth access error when trying to read data without access

Feature COG-656
---
 cognee/modules/data/exceptions/__init__.py   | 1 +
 cognee/modules/data/exceptions/exceptions.py | 9 +++++++++
 cognee/modules/data/methods/get_data.py      | 3 ++-
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/cognee/modules/data/exceptions/__init__.py b/cognee/modules/data/exceptions/__init__.py
index fa8468c88..6f74c627e 100644
--- a/cognee/modules/data/exceptions/__init__.py
+++ b/cognee/modules/data/exceptions/__init__.py
@@ -6,4 +6,5 @@ This module defines a set of exceptions for handling various data errors
 
 from .exceptions import (
     UnstructuredLibraryImportError,
+    UnauthorizedDataAccessError,
 )
\ No newline at end of file
diff --git a/cognee/modules/data/exceptions/exceptions.py b/cognee/modules/data/exceptions/exceptions.py
index 3b1aac52c..7d7bbd632 100644
--- a/cognee/modules/data/exceptions/exceptions.py
+++ b/cognee/modules/data/exceptions/exceptions.py
@@ -7,5 +7,14 @@ class UnstructuredLibraryImportError(CogneeApiError):
             message: str = "Import error. Unstructured library is not installed.",
             name: str = "UnstructuredModuleImportError",
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+    ):
+        super().__init__(message, name, status_code)
+
+class UnauthorizedDataAccessError(CogneeApiError):
+    def __init__(
+            self,
+            message: str = "Usesr does not have permission to access this data.",
+            name: str = "UnauthorizedDataAccessError",
+            status_code=status.HTTP_401_UNAUTHORIZED,
     ):
         super().__init__(message, name, status_code)
\ No newline at end of file
diff --git a/cognee/modules/data/methods/get_data.py b/cognee/modules/data/methods/get_data.py
index 9161e3df2..d7daff29b 100644
--- a/cognee/modules/data/methods/get_data.py
+++ b/cognee/modules/data/methods/get_data.py
@@ -1,6 +1,7 @@
 from uuid import UUID
 from typing import Optional
 from cognee.infrastructure.databases.relational import get_relational_engine
+from ..exceptions import UnauthorizedDataAccessError
 from ..models import Data
 
 async def get_data(user_id: UUID, data_id: UUID) -> Optional[Data]:
@@ -19,6 +20,6 @@ async def get_data(user_id: UUID, data_id: UUID) -> Optional[Data]:
         data = await session.get(Data, data_id)
 
         if data and data.owner_id != user_id:
-            return None
+            raise UnauthorizedDataAccessError(message=f"User {user_id} is not authorized to access data {data_id}")
 
         return data
\ No newline at end of file