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

Signed-off-by: -LAN- <laipz8200@outlook.com>

api/configs/middleware/__init__.py

@@ -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,
         }
 
 

api/tests/unit_tests/configs/test_dify_config.py

@@ -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"
 
-
+def test_dify_config(monkeypatch):
-@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):
     # 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(config["CODE_EXECUTION_ENDPOINT"]) == "http://127.0.0.1:8194/"
-    assert str(URL(str(config["CODE_EXECUTION_ENDPOINT"])) / "v1") == "http://sandbox:8194/v1"
+    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):
+def test_inner_api_config_exist(monkeypatch):
-    config = DifyConfig(_env_file=example_env_file)
+    # 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