feat: enhance database queries with error handling and improve SQL statement readability
Co-authored-by: Copilot <copilot@github.com>
This commit is contained in:
@@ -129,15 +129,13 @@ async def generate_image(
|
||||
user_id = current_user.get("id") or current_user.get("sub")
|
||||
now = datetime.now(timezone.utc).replace(tzinfo=None)
|
||||
stored: list[ImageResult] = []
|
||||
sql_insert = "INSERT INTO generated_images (user_id, model_id, prompt, image_data, created_at) VALUES (?, ?, ?, ?, ?) RETURNING id"
|
||||
async with get_write_lock():
|
||||
conn = get_conn()
|
||||
for img in images:
|
||||
if img.url:
|
||||
row = conn.execute(
|
||||
"""INSERT INTO generated_images (user_id, model_id, prompt, image_data, created_at)
|
||||
VALUES (?, ?, ?, ?, ?) RETURNING id""",
|
||||
[user_id, body.model, body.prompt, img.url, now],
|
||||
).fetchone()
|
||||
sql_insert, [user_id, body.model, body.prompt, img.url, now],).fetchone()
|
||||
image_id = str(row[0]) if row else None
|
||||
else:
|
||||
image_id = None
|
||||
@@ -167,13 +165,8 @@ async def list_generated_images(
|
||||
"""Return all generated images for the current user, newest first."""
|
||||
user_id = current_user.get("id") or current_user.get("sub")
|
||||
conn = get_conn()
|
||||
rows = conn.execute(
|
||||
"""SELECT id, model_id, prompt, image_data, created_at
|
||||
FROM generated_images
|
||||
WHERE user_id = ?
|
||||
ORDER BY created_at DESC""",
|
||||
[user_id],
|
||||
).fetchall()
|
||||
sql_fetch = "SELECT id, model_id, prompt, image_data, created_at FROM generated_images WHERE user_id = ? ORDER BY created_at DESC"
|
||||
rows = conn.execute(sql_fetch, [user_id]).fetchall()
|
||||
return [
|
||||
{
|
||||
"id": str(r[0]),
|
||||
@@ -245,12 +238,12 @@ async def generate_video(
|
||||
conn = get_conn()
|
||||
conn.execute(
|
||||
"""INSERT INTO generated_videos (user_id, job_id, model_id, prompt, polling_url, status, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
[user_id, job_id, body.model, body.prompt,
|
||||
polling_url, job_status, now, now],
|
||||
polling_url, job_status, now, now],
|
||||
)
|
||||
|
||||
urls = result.get("unsigned_urls") or result.get("video_urls")
|
||||
urls = result.get("unsigned_urls") or result.get("video_urls")
|
||||
return VideoResponse(
|
||||
id=job_id,
|
||||
model=body.model,
|
||||
@@ -298,22 +291,22 @@ async def generate_video_from_image(
|
||||
conn = get_conn()
|
||||
conn.execute(
|
||||
"""INSERT INTO generated_videos (user_id, job_id, model_id, prompt, polling_url, status, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
[user_id, job_id, body.model, body.prompt,
|
||||
polling_url, job_status, now, now],
|
||||
)
|
||||
|
||||
urls = result.get("unsigned_urls") or result.get("video_urls")
|
||||
return VideoResponse(
|
||||
id=job_id,
|
||||
model=body.model,
|
||||
status=job_status,
|
||||
polling_url=polling_url,
|
||||
video_urls=urls,
|
||||
video_url=(urls or [None])[0],
|
||||
error=result.get("error"),
|
||||
metadata=result.get("metadata"),
|
||||
)
|
||||
urls = result.get("unsigned_urls") or result.get("video_urls")
|
||||
return VideoResponse(
|
||||
id=job_id,
|
||||
model=body.model,
|
||||
status=job_status,
|
||||
polling_url=polling_url,
|
||||
video_urls=urls,
|
||||
video_url=(urls or [None])[0],
|
||||
error=result.get("error"),
|
||||
metadata=result.get("metadata"),
|
||||
)
|
||||
|
||||
|
||||
@router.get("/video/status", response_model=VideoResponse)
|
||||
@@ -339,8 +332,8 @@ async def poll_video_status(
|
||||
conn = get_conn()
|
||||
conn.execute(
|
||||
"""UPDATE generated_videos
|
||||
SET status = ?, video_url = ?, updated_at = ?
|
||||
WHERE job_id = ?""",
|
||||
SET status = ?, video_url = ?, updated_at = ?
|
||||
WHERE job_id = ?""",
|
||||
[job_status, video_url, now, result.get("id", "")],
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user