From 1679c746a3ea249562ce57dd6a0b47d7611e6268 Mon Sep 17 00:00:00 2001
From: Leon Luithlen <leontimnaluithlen@gmail.com>
Date: Wed, 27 Nov 2024 14:15:03 +0100
Subject: [PATCH] Move class and functions to data.models

---
 cognee/modules/data/models/Data.py                  | 13 +++++++++++--
 .../modules/{ingestion => data}/models/Metadata.py  |  5 ++++-
 .../operations/delete_metadata.py                   |  0
 .../{ingestion => data}/operations/get_metadata.py  |  0
 .../operations/write_metadata.py                    |  0
 cognee/tasks/ingestion/ingest_data_with_metadata.py |  2 +-
 .../save_data_item_with_metadata_to_storage.py      |  2 +-
 7 files changed, 17 insertions(+), 5 deletions(-)
 rename cognee/modules/{ingestion => data}/models/Metadata.py (66%)
 rename cognee/modules/{ingestion => data}/operations/delete_metadata.py (100%)
 rename cognee/modules/{ingestion => data}/operations/get_metadata.py (100%)
 rename cognee/modules/{ingestion => data}/operations/write_metadata.py (100%)

diff --git a/cognee/modules/data/models/Data.py b/cognee/modules/data/models/Data.py
index 55991541d..7899b017f 100644
--- a/cognee/modules/data/models/Data.py
+++ b/cognee/modules/data/models/Data.py
@@ -8,7 +8,7 @@ from sqlalchemy.orm import Mapped, relationship
 from cognee.infrastructure.databases.relational import Base
 
 from .DatasetData import DatasetData
-
+from .Metadata import Metadata
 
 class Data(Base):
     __tablename__ = "data"
@@ -28,7 +28,7 @@ class Data(Base):
         DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc)
     )
 
-    datasets: Mapped[List["Dataset"]] = relationship(
+    datasets = relationship(
         "Dataset",
         secondary=DatasetData.__tablename__,
         back_populates="data",
@@ -36,6 +36,15 @@ class Data(Base):
         cascade="all, delete",
     )
 
+    metadata = relationship(
+        "Metadata",
+        secondary=Metadata.__tablename__,
+        back_populates="data",
+        lazy="noload",
+        cascade="all, delete",
+    )
+
+
     def to_json(self) -> dict:
         return {
             "id": str(self.id),
diff --git a/cognee/modules/ingestion/models/Metadata.py b/cognee/modules/data/models/Metadata.py
similarity index 66%
rename from cognee/modules/ingestion/models/Metadata.py
rename to cognee/modules/data/models/Metadata.py
index 5d6333bc3..adf254941 100644
--- a/cognee/modules/ingestion/models/Metadata.py
+++ b/cognee/modules/data/models/Metadata.py
@@ -1,7 +1,7 @@
 from datetime import datetime, timezone
 from uuid import uuid4
 
-from sqlalchemy import UUID, Column, DateTime, String
+from sqlalchemy import UUID, Column, DateTime, String, ForeignKey, relationship
 
 from cognee.infrastructure.databases.relational import Base
 
@@ -19,3 +19,6 @@ class Metadata(Base):
     updated_at = Column(
         DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc)
     )
+
+    dataset_id = Column(UUID, ForeignKey("datasets.id", ondelete="CASCADE"), primary_key = True)
+    data_id = Column(UUID, ForeignKey("data.id", ondelete="CASCADE"), primary_key = True)
\ No newline at end of file
diff --git a/cognee/modules/ingestion/operations/delete_metadata.py b/cognee/modules/data/operations/delete_metadata.py
similarity index 100%
rename from cognee/modules/ingestion/operations/delete_metadata.py
rename to cognee/modules/data/operations/delete_metadata.py
diff --git a/cognee/modules/ingestion/operations/get_metadata.py b/cognee/modules/data/operations/get_metadata.py
similarity index 100%
rename from cognee/modules/ingestion/operations/get_metadata.py
rename to cognee/modules/data/operations/get_metadata.py
diff --git a/cognee/modules/ingestion/operations/write_metadata.py b/cognee/modules/data/operations/write_metadata.py
similarity index 100%
rename from cognee/modules/ingestion/operations/write_metadata.py
rename to cognee/modules/data/operations/write_metadata.py
diff --git a/cognee/tasks/ingestion/ingest_data_with_metadata.py b/cognee/tasks/ingestion/ingest_data_with_metadata.py
index 07f5c8115..e7068a008 100644
--- a/cognee/tasks/ingestion/ingest_data_with_metadata.py
+++ b/cognee/tasks/ingestion/ingest_data_with_metadata.py
@@ -5,7 +5,7 @@ import dlt
 import cognee.modules.ingestion as ingestion
 from cognee.infrastructure.databases.relational import get_relational_engine
 from cognee.modules.data.methods import create_dataset
-from cognee.modules.ingestion.operations.delete_metadata import delete_metadata
+from cognee.modules.data.operations.delete_metadata import delete_metadata
 from cognee.modules.users.models import User
 from cognee.modules.users.permissions.methods import give_permission_on_document
 from cognee.shared.utils import send_telemetry
diff --git a/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py b/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py
index 6695153d0..11512a1c4 100644
--- a/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py
+++ b/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py
@@ -1,7 +1,7 @@
 from typing import Any, BinaryIO, Union
 
 from cognee.modules.ingestion import save_data_to_file
-from cognee.modules.ingestion.operations.write_metadata import write_metadata
+from cognee.modules.data.operations.write_metadata import write_metadata
 
 
 async def save_data_item_with_metadata_to_storage(