Move class and functions to data.models

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),

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)

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

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(