feat(sqlalchemy): use session factory instead of context manager

This commit is contained in:
Jonathan Rohde
2024-06-24 13:06:15 +02:00
parent eb01e8d275
commit da403f3e3c
15 changed files with 640 additions and 759 deletions

View File

@@ -3,9 +3,8 @@ from typing import List, Optional
import time
from sqlalchemy import String, Column, BigInteger
from sqlalchemy.orm import Session
from apps.webui.internal.db import Base, get_session
from apps.webui.internal.db import Base, Session
import json
@@ -50,65 +49,59 @@ class PromptsTable:
def insert_new_prompt(
self, user_id: str, form_data: PromptForm
) -> Optional[PromptModel]:
with get_session() as db:
prompt = PromptModel(
**{
"user_id": user_id,
"command": form_data.command,
"title": form_data.title,
"content": form_data.content,
"timestamp": int(time.time()),
}
)
prompt = PromptModel(
**{
"user_id": user_id,
"command": form_data.command,
"title": form_data.title,
"content": form_data.content,
"timestamp": int(time.time()),
}
)
try:
result = Prompt(**prompt.dict())
db.add(result)
db.commit()
db.refresh(result)
if result:
return PromptModel.model_validate(result)
else:
return None
except Exception as e:
try:
result = Prompt(**prompt.dict())
Session.add(result)
Session.commit()
Session.refresh(result)
if result:
return PromptModel.model_validate(result)
else:
return None
except Exception as e:
return None
def get_prompt_by_command(self, command: str) -> Optional[PromptModel]:
with get_session() as db:
try:
prompt = db.query(Prompt).filter_by(command=command).first()
return PromptModel.model_validate(prompt)
except:
return None
try:
prompt = Session.query(Prompt).filter_by(command=command).first()
return PromptModel.model_validate(prompt)
except:
return None
def get_prompts(self) -> List[PromptModel]:
with get_session() as db:
return [
PromptModel.model_validate(prompt) for prompt in db.query(Prompt).all()
]
return [
PromptModel.model_validate(prompt) for prompt in Session.query(Prompt).all()
]
def update_prompt_by_command(
self, command: str, form_data: PromptForm
) -> Optional[PromptModel]:
with get_session() as db:
try:
prompt = db.query(Prompt).filter_by(command=command).first()
prompt.title = form_data.title
prompt.content = form_data.content
prompt.timestamp = int(time.time())
db.commit()
return prompt
# return self.get_prompt_by_command(command)
except:
return None
try:
prompt = Session.query(Prompt).filter_by(command=command).first()
prompt.title = form_data.title
prompt.content = form_data.content
prompt.timestamp = int(time.time())
Session.commit()
return PromptModel.model_validate(prompt)
except:
return None
def delete_prompt_by_command(self, command: str) -> bool:
with get_session() as db:
try:
db.query(Prompt).filter_by(command=command).delete()
return True
except:
return False
try:
Session.query(Prompt).filter_by(command=command).delete()
return True
except:
return False
Prompts = PromptsTable()