fix(config): Allow DB_EXTRAS to set search_path via options (#19560)

Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
-LAN- 2025-05-13 10:10:18 +08:00 committed by GitHub
parent a8b82d2b67
commit 0fed5c1193
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 88 additions and 48 deletions

View File

@ -1,6 +1,6 @@
import os
from typing import Any, Literal, Optional
from urllib.parse import quote_plus
from urllib.parse import parse_qsl, quote_plus
from pydantic import Field, NonNegativeInt, PositiveFloat, PositiveInt, computed_field
from pydantic_settings import BaseSettings
@ -176,14 +176,28 @@ class DatabaseConfig(BaseSettings):
default=os.cpu_count() or 1,
)
@computed_field
@computed_field # type: ignore[misc]
@property
def SQLALCHEMY_ENGINE_OPTIONS(self) -> dict[str, Any]:
# Parse DB_EXTRAS for 'options'
db_extras_dict = dict(parse_qsl(self.DB_EXTRAS))
options = db_extras_dict.get("options", "")
# Always include timezone
timezone_opt = "-c timezone=UTC"
if options:
# Merge user options and timezone
merged_options = f"{options} {timezone_opt}"
else:
merged_options = timezone_opt
connect_args = {"options": merged_options}
return {
"pool_size": self.SQLALCHEMY_POOL_SIZE,
"max_overflow": self.SQLALCHEMY_MAX_OVERFLOW,
"pool_recycle": self.SQLALCHEMY_POOL_RECYCLE,
"pool_pre_ping": self.SQLALCHEMY_POOL_PRE_PING,
"connect_args": {"options": "-c timezone=UTC"},
"connect_args": connect_args,
}

View File

@ -1,49 +1,28 @@
import os
from textwrap import dedent
import pytest
from flask import Flask
from yarl import URL
from configs.app_config import DifyConfig
EXAMPLE_ENV_FILENAME = ".env"
@pytest.fixture
def example_env_file(tmp_path, monkeypatch) -> str:
monkeypatch.chdir(tmp_path)
file_path = tmp_path.joinpath(EXAMPLE_ENV_FILENAME)
file_path.write_text(
dedent(
"""
CONSOLE_API_URL=https://example.com
CONSOLE_WEB_URL=https://example.com
HTTP_REQUEST_MAX_WRITE_TIMEOUT=30
"""
)
)
return str(file_path)
def test_dify_config_undefined_entry(example_env_file):
# NOTE: See https://github.com/microsoft/pylance-release/issues/6099 for more details about this type error.
# load dotenv file with pydantic-settings
config = DifyConfig(_env_file=example_env_file)
# entries not defined in app settings
with pytest.raises(TypeError):
# TypeError: 'AppSettings' object is not subscriptable
assert config["LOG_LEVEL"] == "INFO"
# NOTE: If there is a `.env` file in your Workspace, this test might not succeed as expected.
# This is due to `pymilvus` loading all the variables from the `.env` file into `os.environ`.
def test_dify_config(example_env_file):
def test_dify_config(monkeypatch):
# clear system environment variables
os.environ.clear()
# Set environment variables using monkeypatch
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
monkeypatch.setenv("HTTP_REQUEST_MAX_WRITE_TIMEOUT", "30")
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("HTTP_REQUEST_MAX_READ_TIMEOUT", "600")
# load dotenv file with pydantic-settings
config = DifyConfig(_env_file=example_env_file)
config = DifyConfig()
# constant values
assert config.COMMIT_SHA == ""
@ -54,7 +33,7 @@ def test_dify_config(example_env_file):
assert config.SENTRY_TRACES_SAMPLE_RATE == 1.0
# annotated field with default value
assert config.HTTP_REQUEST_MAX_READ_TIMEOUT == 60
assert config.HTTP_REQUEST_MAX_READ_TIMEOUT == 600
# annotated field with configured value
assert config.HTTP_REQUEST_MAX_WRITE_TIMEOUT == 30
@ -64,11 +43,24 @@ def test_dify_config(example_env_file):
# NOTE: If there is a `.env` file in your Workspace, this test might not succeed as expected.
# This is due to `pymilvus` loading all the variables from the `.env` file into `os.environ`.
def test_flask_configs(example_env_file):
def test_flask_configs(monkeypatch):
flask_app = Flask("app")
# clear system environment variables
os.environ.clear()
flask_app.config.from_mapping(DifyConfig(_env_file=example_env_file).model_dump()) # pyright: ignore
# Set environment variables using monkeypatch
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
monkeypatch.setenv("HTTP_REQUEST_MAX_WRITE_TIMEOUT", "30")
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("WEB_API_CORS_ALLOW_ORIGINS", "http://127.0.0.1:3000,*")
monkeypatch.setenv("CODE_EXECUTION_ENDPOINT", "http://127.0.0.1:8194/")
flask_app.config.from_mapping(DifyConfig().model_dump()) # pyright: ignore
config = flask_app.config
# configs read from pydantic-settings
@ -83,7 +75,7 @@ def test_flask_configs(example_env_file):
# fallback to alias choices value as CONSOLE_API_URL
assert config["FILES_URL"] == "https://example.com"
assert config["SQLALCHEMY_DATABASE_URI"] == "postgresql://postgres:@localhost:5432/dify"
assert config["SQLALCHEMY_DATABASE_URI"] == "postgresql://postgres:postgres@localhost:5432/dify"
assert config["SQLALCHEMY_ENGINE_OPTIONS"] == {
"connect_args": {
"options": "-c timezone=UTC",
@ -96,13 +88,47 @@ def test_flask_configs(example_env_file):
assert config["CONSOLE_WEB_URL"] == "https://example.com"
assert config["CONSOLE_CORS_ALLOW_ORIGINS"] == ["https://example.com"]
assert config["WEB_API_CORS_ALLOW_ORIGINS"] == ["*"]
assert config["WEB_API_CORS_ALLOW_ORIGINS"] == ["http://127.0.0.1:3000", "*"]
assert str(config["CODE_EXECUTION_ENDPOINT"]) == "http://sandbox:8194/"
assert str(URL(str(config["CODE_EXECUTION_ENDPOINT"])) / "v1") == "http://sandbox:8194/v1"
assert str(config["CODE_EXECUTION_ENDPOINT"]) == "http://127.0.0.1:8194/"
assert str(URL(str(config["CODE_EXECUTION_ENDPOINT"])) / "v1") == "http://127.0.0.1:8194/v1"
def test_inner_api_config_exist(example_env_file):
config = DifyConfig(_env_file=example_env_file)
def test_inner_api_config_exist(monkeypatch):
# Set environment variables using monkeypatch
monkeypatch.setenv("CONSOLE_API_URL", "https://example.com")
monkeypatch.setenv("CONSOLE_WEB_URL", "https://example.com")
monkeypatch.setenv("HTTP_REQUEST_MAX_WRITE_TIMEOUT", "30")
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("INNER_API_KEY", "test-inner-api-key")
config = DifyConfig()
assert config.INNER_API is False
assert config.INNER_API_KEY is None
assert isinstance(config.INNER_API_KEY, str)
assert len(config.INNER_API_KEY) > 0
def test_db_extras_options_merging(monkeypatch):
"""Test that DB_EXTRAS options are properly merged with default timezone setting"""
# Set environment variables
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")
# 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"]
assert "search_path=myschema" in options
assert "timezone=UTC" in options