diff --git a/api/configs/middleware/__init__.py b/api/configs/middleware/__init__.py index 9d30badb3d..1b015b3267 100644 --- a/api/configs/middleware/__init__.py +++ b/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, } diff --git a/api/tests/unit_tests/configs/test_dify_config.py b/api/tests/unit_tests/configs/test_dify_config.py index fafaf0419a..cac0a688cd 100644 --- a/api/tests/unit_tests/configs/test_dify_config.py +++ b/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" - -@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