From 121c8c51b14aaaeaa4dfaf1c28c0095f5e71b99a Mon Sep 17 00:00:00 2001
From: arkohut <39525455+arkohut@users.noreply.github.com>
Date: Sat, 1 Jun 2024 13:03:53 +0800
Subject: [PATCH] feat: add entity

---
 memos/crud.py    | 15 +++++++++++++--
 memos/models.py  |  2 +-
 memos/schemas.py |  5 ++---
 memos/server.py  | 12 ++++--------
 4 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/memos/crud.py b/memos/crud.py
index 8274fc1..a0d378e 100644
--- a/memos/crud.py
+++ b/memos/crud.py
@@ -1,6 +1,6 @@
 from sqlalchemy.orm import Session
-from .schemas import Library, NewLibraryParam, Folder
-from .models import LibraryModel, FolderModel
+from .schemas import Library, NewLibraryParam, Folder, NewEntityParam, Entity
+from .models import LibraryModel, FolderModel, EntityModel, EntityModel
 
 
 def get_library_by_id(library_id: int, db: Session) -> Library | None:
@@ -24,3 +24,14 @@ def create_library(library: NewLibraryParam, db: Session) -> Library:
         folders=[Folder(id=db_folder.id, name=db_folder.path) for db_folder in db_library.folders],
         plugins=[]
     )
+
+
+def create_entity(library_id: int, entity: NewEntityParam, db: Session) -> Entity:
+    db_entity = EntityModel(
+        **entity.model_dump(),
+        library_id=library_id
+    )
+    db.add(db_entity)
+    db.commit()
+    db.refresh(db_entity)
+    return db_entity
diff --git a/memos/models.py b/memos/models.py
index 0d9af0a..32e4489 100644
--- a/memos/models.py
+++ b/memos/models.py
@@ -95,5 +95,5 @@ class LibraryPluginModel(Base):
 
 
 # Create the database engine with the path from config
-engine = create_engine(f"sqlite:///:memory:", echo=True)
+engine = create_engine(f"sqlite:///{get_database_path()}", echo=True)
 Base.metadata.create_all(engine)
diff --git a/memos/schemas.py b/memos/schemas.py
index a7853fc..1c4b9c1 100644
--- a/memos/schemas.py
+++ b/memos/schemas.py
@@ -26,9 +26,8 @@ class NewFolderParam(BaseModel):
 
 class NewEntityParam(BaseModel):
     filename: str
-    path: str
+    filepath: str
     size: int
-    created_at: datetime
     file_created_at: datetime
     file_last_modified_at: datetime
     file_type: str
@@ -92,7 +91,7 @@ class Library(BaseModel):
 
 class Entity(BaseModel):
     id: int
-    path: str
+    filepath: str
     filename: str
     size: int
     file_created_at: datetime
diff --git a/memos/server.py b/memos/server.py
index 0178de4..5a6ed63 100644
--- a/memos/server.py
+++ b/memos/server.py
@@ -6,7 +6,7 @@ from sqlalchemy.orm import sessionmaker
 from typing import List
 
 from .config import get_database_path
-from .crud import get_library_by_id, create_library
+from .crud import get_library_by_id, create_library, create_entity
 from .schemas import (
     Library,
     Folder,
@@ -57,14 +57,10 @@ def new_folder(
 
 @app.post("/libraries/{library_id}/entities", response_model=Entity)
 def new_entity(
-    entity: NewEntityParam, library_id: int, db: Session = Depends(get_db)
+    new_entity: NewEntityParam, library_id: int, db: Session = Depends(get_db)
 ):
-
-    db_entity = Entity(**entity.model_dump(), library_id=library_id)
-    db.add(db_entity)
-    db.commit()
-    db.refresh(db_entity)
-    return db_entity
+    entity = create_entity(library_id, new_entity, db)
+    return entity
 
 
 @app.post("/plugins", response_model=Plugin)