Update DatabaseSessionService db_engine construction (#195)

Co-authored-by: Hangfei Lin <hangfei@google.com>
This commit is contained in:
DeanChensj 2025-04-14 17:55:15 -07:00 committed by GitHub
parent 9a44831a08
commit 3e85df2641
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -30,6 +30,7 @@ from sqlalchemy.dialects import postgresql
from sqlalchemy.engine import create_engine from sqlalchemy.engine import create_engine
from sqlalchemy.engine import Engine from sqlalchemy.engine import Engine
from sqlalchemy.ext.mutable import MutableDict from sqlalchemy.ext.mutable import MutableDict
from sqlalchemy.exc import ArgumentError
from sqlalchemy.inspection import inspect from sqlalchemy.inspection import inspect
from sqlalchemy.orm import DeclarativeBase from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped from sqlalchemy.orm import Mapped
@ -189,13 +190,20 @@ class DatabaseSessionService(BaseSessionService):
# 2. Create all tables based on schema # 2. Create all tables based on schema
# 3. Initialize all properies # 3. Initialize all properies
supported_dialects = ["postgresql", "mysql", "sqlite"] try:
dialect = db_url.split("://")[0]
if dialect in supported_dialects:
db_engine = create_engine(db_url) db_engine = create_engine(db_url)
else: except Exception as e:
raise ValueError(f"Unsupported database URL: {db_url}") if isinstance(e, ArgumentError):
raise ValueError(
f"Invalid database URL format or argument '{db_url}'."
) from e
if isinstance(e, ImportError):
raise ValueError(
f"Database related module not found for URL '{db_url}'."
) from e
raise ValueError(
f"Failed to create database engine for URL '{db_url}'"
) from e
# Get the local timezone # Get the local timezone
local_timezone = get_localzone() local_timezone = get_localzone()