laipeng/reflex
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Ping db connections before returning them from pool

Browse Source 
Move connect engine kwargs to a separate function
This commit is contained in:
Masen Furer 2024-12-06 10:04:33 -08:00
parent b63110e2cb
commit 904df85da1
No known key found for this signature in database
GPG Key ID: B0008AD22B3B3A95
2 changed files with 30 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
reflex/config.py
View File

@ -512,6 +512,9 @@ class EnvironmentVariables:
# Whether to print the SQL queries if the log level is INFO or lower.
SQLALCHEMY_ECHO: EnvVar[bool] = env_var(False)
# Whether to check db connections before using them.
SQLALCHEMY_POOL_PRE_PING: EnvVar[bool] = env_var(True)
# Whether to ignore the redis config error. Some redis servers only allow out-of-band configuration.
REFLEX_IGNORE_REDIS_CONFIG_ERROR: EnvVar[bool] = env_var(False)

38
reflex/model.py
View File

@ -27,6 +27,29 @@ _ASYNC_ENGINE: dict[str, sqlalchemy.engine.Engine] = {}
_AsyncSessionLocal: dict[str | None, sqlalchemy.orm.sessionmaker] = {}
def get_engine_args(url: str | None = None) -> dict[str, Any]:
"""Get the database engine arguments.
Args:
url: The database url.
Returns:
The database engine arguments as a dict.
"""
kwargs: dict[str, Any] = dict(
# Check connections before returning them.
pool_pre_ping=environment.SQLALCHEMY_POOL_PRE_PING.get(),
# Print the SQL queries if the log level is INFO or lower.
echo_db_query=environment.SQLALCHEMY_ECHO.get(),
)
conf = get_config()
url = url or conf.db_url
if url is not None and url.startswith("sqlite"):
# Needed for the admin dash on sqlite.
kwargs["connect_args"] = {"check_same_thread": False}
return kwargs
def get_engine(url: str | None = None) -> sqlalchemy.engine.Engine:
"""Get the database engine.
@ -52,13 +75,9 @@ def get_engine(url: str | None = None) -> sqlalchemy.engine.Engine:
console.warn(
"Database is not initialized, run [bold]reflex db init[/bold] first."
)
# Print the SQL queries if the log level is INFO or lower.
echo_db_query = environment.SQLALCHEMY_ECHO.get()
# Needed for the admin dash on sqlite.
connect_args = {"check_same_thread": False} if url.startswith("sqlite") else {}
_ENGINE[url] = sqlmodel.create_engine(
url, echo=echo_db_query, connect_args=connect_args
url,
**get_engine_args(url),
)
return _ENGINE[url]
@ -88,12 +107,9 @@ def get_async_engine(url: str | None) -> sqlalchemy.engine.Engine:
console.warn(
"Database is not initialized, run [bold]reflex db init[/bold] first."
)
# Print the SQL queries if the log level is INFO or lower.
echo_db_query = environment.SQLALCHEMY_ECHO.get()
# Needed for the admin dash on sqlite.
connect_args = {"check_same_thread": False} if url.startswith("sqlite") else {}
_ASYNC_ENGINE[url] = sqlmodel.create_engine(
url, echo=echo_db_query, connect_args=connect_args
url,
**get_engine_args(url),
)
return _ASYNC_ENGINE[url]