wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-12-18 18:43:05 +08:00
Code Issues Actions 10 Packages Projects Releases Wiki Activity

fixed new PgSQL bug

Browse Source
This commit is contained in:
bigcat88 2025-09-14 10:12:33 +03:00
parent 975650060f
commit 37b81e6658
No known key found for this signature in database
GPG Key ID: 1F0BF0EC3CF22721
2 changed files with 63 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
app/api/assets_routes.py
View File

@ -1,4 +1,5 @@
import contextlib import contextlib
import logging
import os import os
import urllib.parse import urllib.parse
import uuid import uuid
@ -13,7 +14,8 @@ from .. import assets_manager, assets_scanner, user_manager
from . import schemas_in, schemas_out from . import schemas_in, schemas_out
ROUTES = web.RouteTableDef() ROUTES = web.RouteTableDef()
UserManager: Optional[user_manager.UserManager] = None USER_MANAGER: Optional[user_manager.UserManager] = None
LOGGER = logging.getLogger(__name__)
# UUID regex (canonical hyphenated form, case-insensitive) # UUID regex (canonical hyphenated form, case-insensitive)
UUID_RE = r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}" UUID_RE = r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}"
@ -58,7 +60,7 @@ async def list_assets(request: web.Request) -> web.Response:
offset=q.offset, offset=q.offset,
sort=q.sort, sort=q.sort,
order=q.order, order=q.order,
owner_id=UserManager.get_request_user_id(request), owner_id=USER_MANAGER.get_request_user_id(request),
) )
return web.json_response(payload.model_dump(mode="json")) return web.json_response(payload.model_dump(mode="json"))
@ -72,7 +74,7 @@ async def download_asset_content(request: web.Request) -> web.Response:
try: try:
abs_path, content_type, filename = await assets_manager.resolve_asset_content_for_download( abs_path, content_type, filename = await assets_manager.resolve_asset_content_for_download(
asset_info_id=str(uuid.UUID(request.match_info["id"])), asset_info_id=str(uuid.UUID(request.match_info["id"])),
owner_id=UserManager.get_request_user_id(request), owner_id=USER_MANAGER.get_request_user_id(request),
) )
except ValueError as ve: except ValueError as ve:
return _error_response(404, "ASSET_NOT_FOUND", str(ve)) return _error_response(404, "ASSET_NOT_FOUND", str(ve))
@ -105,7 +107,7 @@ async def create_asset_from_hash(request: web.Request) -> web.Response:
name=body.name, name=body.name,
tags=body.tags, tags=body.tags,
user_metadata=body.user_metadata, user_metadata=body.user_metadata,
owner_id=UserManager.get_request_user_id(request), owner_id=USER_MANAGER.get_request_user_id(request),
) )
if result is None: if result is None:
return _error_response(404, "ASSET_NOT_FOUND", f"Asset content {body.hash} does not exist") return _error_response(404, "ASSET_NOT_FOUND", f"Asset content {body.hash} does not exist")
@ -234,7 +236,7 @@ async def upload_asset(request: web.Request) -> web.Response:
400, "INVALID_BODY", f"unknown models category '{spec.tags[1] if len(spec.tags) >= 2 else ''}'" 400, "INVALID_BODY", f"unknown models category '{spec.tags[1] if len(spec.tags) >= 2 else ''}'"
) )
owner_id = UserManager.get_request_user_id(request) owner_id = USER_MANAGER.get_request_user_id(request)
# Fast path: if a valid provided hash exists, create AssetInfo without writing anything # Fast path: if a valid provided hash exists, create AssetInfo without writing anything
if spec.hash and provided_hash_exists is True: if spec.hash and provided_hash_exists is True:
@ -247,6 +249,7 @@ async def upload_asset(request: web.Request) -> web.Response:
owner_id=owner_id, owner_id=owner_id,
) )
except Exception: except Exception:
LOGGER.exception("create_asset_from_hash failed for hash=%s, owner_id=%s", spec.hash, owner_id)
return _error_response(500, "INTERNAL", "Unexpected server error.") return _error_response(500, "INTERNAL", "Unexpected server error.")
if result is None: if result is None:
@ -289,6 +292,7 @@ async def upload_asset(request: web.Request) -> web.Response:
except Exception: except Exception:
if tmp_path and os.path.exists(tmp_path): if tmp_path and os.path.exists(tmp_path):
os.remove(tmp_path) os.remove(tmp_path)
LOGGER.exception("upload_asset_from_temp_path failed for tmp_path=%s, owner_id=%s", tmp_path, owner_id)
return _error_response(500, "INTERNAL", "Unexpected server error.") return _error_response(500, "INTERNAL", "Unexpected server error.")
@ -298,11 +302,16 @@ async def get_asset(request: web.Request) -> web.Response:
try: try:
result = await assets_manager.get_asset( result = await assets_manager.get_asset(
asset_info_id=asset_info_id, asset_info_id=asset_info_id,
owner_id=UserManager.get_request_user_id(request), owner_id=USER_MANAGER.get_request_user_id(request),
) )
except ValueError as ve: except ValueError as ve:
return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id}) return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id})
except Exception: except Exception:
LOGGER.exception(
"get_asset failed for asset_info_id=%s, owner_id=%s",
asset_info_id,
USER_MANAGER.get_request_user_id(request),
)
return _error_response(500, "INTERNAL", "Unexpected server error.") return _error_response(500, "INTERNAL", "Unexpected server error.")
return web.json_response(result.model_dump(mode="json"), status=200) return web.json_response(result.model_dump(mode="json"), status=200)
@ -323,11 +332,16 @@ async def update_asset(request: web.Request) -> web.Response:
name=body.name, name=body.name,
tags=body.tags, tags=body.tags,
user_metadata=body.user_metadata, user_metadata=body.user_metadata,
owner_id=UserManager.get_request_user_id(request), owner_id=USER_MANAGER.get_request_user_id(request),
) )
except (ValueError, PermissionError) as ve: except (ValueError, PermissionError) as ve:
return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id}) return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id})
except Exception: except Exception:
LOGGER.exception(
"update_asset failed for asset_info_id=%s, owner_id=%s",
asset_info_id,
USER_MANAGER.get_request_user_id(request),
)
return _error_response(500, "INTERNAL", "Unexpected server error.") return _error_response(500, "INTERNAL", "Unexpected server error.")
return web.json_response(result.model_dump(mode="json"), status=200) return web.json_response(result.model_dump(mode="json"), status=200)
@ -346,11 +360,16 @@ async def set_asset_preview(request: web.Request) -> web.Response:
result = await assets_manager.set_asset_preview( result = await assets_manager.set_asset_preview(
asset_info_id=asset_info_id, asset_info_id=asset_info_id,
preview_asset_id=body.preview_id, preview_asset_id=body.preview_id,
owner_id=UserManager.get_request_user_id(request), owner_id=USER_MANAGER.get_request_user_id(request),
) )
except (PermissionError, ValueError) as ve: except (PermissionError, ValueError) as ve:
return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id}) return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id})
except Exception: except Exception:
LOGGER.exception(
"set_asset_preview failed for asset_info_id=%s, owner_id=%s",
asset_info_id,
USER_MANAGER.get_request_user_id(request),
)
return _error_response(500, "INTERNAL", "Unexpected server error.") return _error_response(500, "INTERNAL", "Unexpected server error.")
return web.json_response(result.model_dump(mode="json"), status=200) return web.json_response(result.model_dump(mode="json"), status=200)
@ -364,10 +383,15 @@ async def delete_asset(request: web.Request) -> web.Response:
try: try:
deleted = await assets_manager.delete_asset_reference( deleted = await assets_manager.delete_asset_reference(
asset_info_id=asset_info_id, asset_info_id=asset_info_id,
owner_id=UserManager.get_request_user_id(request), owner_id=USER_MANAGER.get_request_user_id(request),
delete_content_if_orphan=delete_content, delete_content_if_orphan=delete_content,
) )
except Exception: except Exception:
LOGGER.exception(
"delete_asset_reference failed for asset_info_id=%s, owner_id=%s",
asset_info_id,
USER_MANAGER.get_request_user_id(request),
)
return _error_response(500, "INTERNAL", "Unexpected server error.") return _error_response(500, "INTERNAL", "Unexpected server error.")
if not deleted: if not deleted:
@ -393,7 +417,7 @@ async def get_tags(request: web.Request) -> web.Response:
offset=query.offset, offset=query.offset,
order=query.order, order=query.order,
include_zero=query.include_zero, include_zero=query.include_zero,
owner_id=UserManager.get_request_user_id(request), owner_id=USER_MANAGER.get_request_user_id(request),
) )
return web.json_response(result.model_dump(mode="json")) return web.json_response(result.model_dump(mode="json"))
@ -414,11 +438,16 @@ async def add_asset_tags(request: web.Request) -> web.Response:
asset_info_id=asset_info_id, asset_info_id=asset_info_id,
tags=data.tags, tags=data.tags,
origin="manual", origin="manual",
owner_id=UserManager.get_request_user_id(request), owner_id=USER_MANAGER.get_request_user_id(request),
) )
except (ValueError, PermissionError) as ve: except (ValueError, PermissionError) as ve:
return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id}) return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id})
except Exception: except Exception:
LOGGER.exception(
"add_tags_to_asset failed for asset_info_id=%s, owner_id=%s",
asset_info_id,
USER_MANAGER.get_request_user_id(request),
)
return _error_response(500, "INTERNAL", "Unexpected server error.") return _error_response(500, "INTERNAL", "Unexpected server error.")
return web.json_response(result.model_dump(mode="json"), status=200) return web.json_response(result.model_dump(mode="json"), status=200)
@ -439,11 +468,16 @@ async def delete_asset_tags(request: web.Request) -> web.Response:
result = await assets_manager.remove_tags_from_asset( result = await assets_manager.remove_tags_from_asset(
asset_info_id=asset_info_id, asset_info_id=asset_info_id,
tags=data.tags, tags=data.tags,
owner_id=UserManager.get_request_user_id(request), owner_id=USER_MANAGER.get_request_user_id(request),
) )
except ValueError as ve: except ValueError as ve:
return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id}) return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id})
except Exception: except Exception:
LOGGER.exception(
"remove_tags_from_asset failed for asset_info_id=%s, owner_id=%s",
asset_info_id,
USER_MANAGER.get_request_user_id(request),
)
return _error_response(500, "INTERNAL", "Unexpected server error.") return _error_response(500, "INTERNAL", "Unexpected server error.")
return web.json_response(result.model_dump(mode="json"), status=200) return web.json_response(result.model_dump(mode="json"), status=200)
@ -476,8 +510,8 @@ async def get_asset_scan_status(request: web.Request) -> web.Response:
def register_assets_system(app: web.Application, user_manager_instance: user_manager.UserManager) -> None: def register_assets_system(app: web.Application, user_manager_instance: user_manager.UserManager) -> None:
global UserManager global USER_MANAGER
UserManager = user_manager_instance USER_MANAGER = user_manager_instance
app.add_routes(ROUTES) app.add_routes(ROUTES)

18
app/database/services/content.py
View File

@ -514,24 +514,36 @@ async def ingest_fs_asset(
d_sqlite.insert(Asset) d_sqlite.insert(Asset)
.values(**vals) .values(**vals)
.on_conflict_do_nothing(index_elements=[Asset.hash]) .on_conflict_do_nothing(index_elements=[Asset.hash])
.returning(Asset.id)
) )
elif dialect == "postgresql": elif dialect == "postgresql":
ins = ( ins = (
d_pg.insert(Asset) d_pg.insert(Asset)
.values(**vals) .values(**vals)
.on_conflict_do_nothing(index_elements=[Asset.hash]) .on_conflict_do_nothing(
index_elements=[Asset.hash],
index_where=Asset.__table__.c.hash.isnot(None),
)
.returning(Asset.id)
) )
else: else:
raise NotImplementedError(f"Unsupported database dialect: {dialect}") raise NotImplementedError(f"Unsupported database dialect: {dialect}")
res = await session.execute(ins) res = await session.execute(ins)
rowcount = int(res.rowcount or 0) inserted_id = res.scalar_one_or_none()
asset = ( asset = (
await session.execute(select(Asset).where(Asset.hash == asset_hash).limit(1)) await session.execute(select(Asset).where(Asset.hash == asset_hash).limit(1))
).scalars().first() ).scalars().first()
if not asset: if not asset:
raise RuntimeError("Asset row not found after upsert.") raise RuntimeError("Asset row not found after upsert.")
if rowcount > 0: if inserted_id:
out["asset_created"] = True out["asset_created"] = True
asset = await session.get(Asset, inserted_id)
else:
asset = (
await session.execute(select(Asset).where(Asset.hash == asset_hash).limit(1))
).scalars().first()
if not asset:
raise RuntimeError("Asset row not found after upsert.")
else: else:
changed = False changed = False
if asset.size_bytes != int(size_bytes) and int(size_bytes) > 0: if asset.size_bytes != int(size_bytes) and int(size_bytes) > 0: