mirror of
https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git
synced 2025-08-12 17:29:04 +08:00
fix(config): Allow DB_EXTRAS to set search_path via options (#19560)
Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
parent
a8b82d2b67
commit
0fed5c1193
@ -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,
|
||||
}
|
||||
|
||||
|
||||
|
@ -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
|
||||
|
Loading…
x
Reference in New Issue
Block a user