chore: allow disabling app-level PostgreSQL timezone injection (#35129)

This commit is contained in:
非法操作 2026-04-14 13:57:27 +08:00 committed by GitHub
parent 2ef9a8a769
commit 2333d75c56
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 36 additions and 12 deletions

View File

@ -160,6 +160,16 @@ class DatabaseConfig(BaseSettings):
default="",
)
DB_SESSION_TIMEZONE_OVERRIDE: str = Field(
description=(
"PostgreSQL session timezone override injected via startup options."
" Default is 'UTC' for out-of-the-box consistency."
" Set to empty string to disable app-level timezone injection, for example when using RDS Proxy"
" together with a database-side default timezone."
),
default="UTC",
)
@computed_field # type: ignore[prop-decorator]
@property
def SQLALCHEMY_DATABASE_URI_SCHEME(self) -> str:
@ -227,12 +237,13 @@ class DatabaseConfig(BaseSettings):
connect_args: dict[str, str] = {}
# Use the dynamic SQLALCHEMY_DATABASE_URI_SCHEME property
if self.SQLALCHEMY_DATABASE_URI_SCHEME.startswith("postgresql"):
timezone_opt = "-c timezone=UTC"
if options:
merged_options = f"{options} {timezone_opt}"
else:
merged_options = timezone_opt
connect_args = {"options": merged_options}
merged_options = options.strip()
session_timezone_override = self.DB_SESSION_TIMEZONE_OVERRIDE.strip()
if session_timezone_override:
timezone_opt = f"-c timezone={session_timezone_override}"
merged_options = f"{merged_options} {timezone_opt}".strip() if merged_options else timezone_opt
if merged_options:
connect_args = {"options": merged_options}
result: SQLAlchemyEngineOptionsDict = {
"pool_size": self.SQLALCHEMY_POOL_SIZE,

View File

@ -145,7 +145,7 @@ def test_inner_api_config_exist(monkeypatch: pytest.MonkeyPatch):
def test_db_extras_options_merging(monkeypatch: pytest.MonkeyPatch):
"""Test that DB_EXTRAS options are properly merged with default timezone setting"""
"""Test that DB_EXTRAS options are merged with the default timezone startup option."""
# Set environment variables
monkeypatch.setenv("DB_TYPE", "postgresql")
monkeypatch.setenv("DB_USERNAME", "postgres")
@ -158,15 +158,28 @@ def test_db_extras_options_merging(monkeypatch: pytest.MonkeyPatch):
# Create config
config = DifyConfig()
# Get engine options
engine_options = config.SQLALCHEMY_ENGINE_OPTIONS
# Verify options contains both search_path and timezone
options = engine_options["connect_args"]["options"]
options = config.SQLALCHEMY_ENGINE_OPTIONS["connect_args"]["options"]
assert "search_path=myschema" in options
assert "timezone=UTC" in options
def test_db_session_timezone_override_can_disable_app_level_timezone_injection(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("DB_TYPE", "postgresql")
monkeypatch.setenv("DB_USERNAME", "postgres")
monkeypatch.setenv("DB_PASSWORD", "postgres")
monkeypatch.setenv("DB_HOST", "localhost")
monkeypatch.setenv("DB_PORT", "5432")
monkeypatch.setenv("DB_DATABASE", "dify")
monkeypatch.setenv("DB_EXTRAS", "options=-c search_path=myschema")
monkeypatch.setenv("DB_SESSION_TIMEZONE_OVERRIDE", "")
config = DifyConfig()
assert config.SQLALCHEMY_ENGINE_OPTIONS["connect_args"] == {
"options": "-c search_path=myschema",
}
def test_pubsub_redis_url_default(monkeypatch: pytest.MonkeyPatch):
os.environ.clear()