diff --git a/tests/memory/test_postgres_endpoint.py b/tests/memory/test_postgres_endpoint.py new file mode 100644 index 000000000..0965a4d88 --- /dev/null +++ b/tests/memory/test_postgres_endpoint.py @@ -0,0 +1,104 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from unittest.mock import Mock + +import pytest +from sqlalchemy.engine import make_url + +from veadk.configs.database_configs import PostgreSqlConfig +from veadk.memory.short_term_memory import ShortTermMemory +from veadk.memory.short_term_memory_backends.postgresql_backend import ( + PostgreSqlSTMBackend, +) + +BACKEND = "veadk.memory.short_term_memory_backends.postgresql_backend" + + +@pytest.fixture(autouse=True) +def isolated_database_environment(monkeypatch): + for key in tuple(os.environ): + if key.startswith("DATABASE_POSTGRESQL_"): + monkeypatch.delenv(key) + + +@pytest.mark.parametrize("async_driver", [False, True]) +@pytest.mark.parametrize( + "host,port,expected_host,expected_port", + [ + ("db.example.test", 5432, "db.example.test", 5432), + ("db.example.test", 6543, "db.example.test", 6543), + ("db.example.test:5432", 5432, "db.example.test", 5432), + ("db.example.test:6543", 5432, "db.example.test", 6543), + ("db.example.test:6543", 7654, "db.example.test", 6543), + ("127.0.0.1:5432", 5432, "127.0.0.1", 5432), + ("db.example.test:1", 5432, "db.example.test", 1), + ("db.example.test:65535", 5432, "db.example.test", 65535), + ], +) +def test_endpoint(host, port, expected_host, expected_port, async_driver, monkeypatch): + monkeypatch.setattr(f"{BACKEND}.should_use_async_db_drivers", lambda: async_driver) + config = PostgreSqlConfig( + host=host, port=port, user="fixture-user", password="fixture@:/", database="db" + ) + backend = PostgreSqlSTMBackend(postgresql_config=config) + url = make_url(backend._db_url) + assert (url.host, url.port) == (expected_host, expected_port) + assert url.drivername == ("postgresql+asyncpg" if async_driver else "postgresql") + assert (url.username, url.password, url.database) == ( + "fixture-user", + "fixture@:/", + "db", + ) + assert (config.host, config.port) == (host, port) + + +@pytest.mark.parametrize( + "host", + [ + "db.example.test:", + "db.example.test:abc", + "db.example.test:0", + "db.example.test:-1", + "db.example.test:65536", + "db.example.test:5432.0", + "db.example.test:123", + ":5432", + ], +) +def test_invalid_embedded_port(host): + with pytest.raises(ValueError, match="PostgreSQL"): + PostgreSqlSTMBackend(postgresql_config=PostgreSqlConfig(host=host)) + + +def test_environment_reaches_session_service(monkeypatch): + monkeypatch.setenv("DATABASE_POSTGRESQL_HOST", "db.example.test:6543") + monkeypatch.setenv("DATABASE_POSTGRESQL_PORT", "7654") + service = Mock() + monkeypatch.setattr(f"{BACKEND}.DatabaseSessionService", service) + ShortTermMemory(backend="postgresql") + url = make_url(service.call_args.kwargs["db_url"]) + assert (url.host, url.port) == ("db.example.test", 6543) + + +def test_explicit_db_url_bypasses_host_configuration(monkeypatch): + monkeypatch.setenv("DATABASE_POSTGRESQL_HOST", "db.example.test:invalid") + service = Mock() + monkeypatch.setattr( + "veadk.memory.short_term_memory.DatabaseSessionService", service + ) + url = "sqlite+aiosqlite:///:memory:" + ShortTermMemory(db_url=url) + service.assert_called_once_with(db_url=url) diff --git a/veadk/memory/short_term_memory_backends/postgresql_backend.py b/veadk/memory/short_term_memory_backends/postgresql_backend.py index afc385108..a16a62d7f 100644 --- a/veadk/memory/short_term_memory_backends/postgresql_backend.py +++ b/veadk/memory/short_term_memory_backends/postgresql_backend.py @@ -104,10 +104,22 @@ class PostgreSqlSTMBackend(BaseShortTermMemoryBackend): def model_post_init(self, context: Any) -> None: encoded_username = quote_plus(self.postgresql_config.user) encoded_password = quote_plus(self.postgresql_config.password) + host = self.postgresql_config.host + port = self.postgresql_config.port + if host.count(":") == 1: + host, embedded_port = host.rsplit(":", 1) + if not host or not embedded_port.isascii() or not embedded_port.isdecimal(): + raise ValueError( + "PostgreSQL host:port requires a hostname and a numeric port." + ) + # A port in the host takes precedence over the separate setting. + port = int(embedded_port) + if not 1 <= port <= 65535: + raise ValueError("PostgreSQL port must be between 1 and 65535.") if should_use_async_db_drivers(): - self._db_url = f"postgresql+asyncpg://{encoded_username}:{encoded_password}@{self.postgresql_config.host}:{self.postgresql_config.port}/{self.postgresql_config.database}" + self._db_url = f"postgresql+asyncpg://{encoded_username}:{encoded_password}@{host}:{port}/{self.postgresql_config.database}" else: - self._db_url = f"postgresql://{encoded_username}:{encoded_password}@{self.postgresql_config.host}:{self.postgresql_config.port}/{self.postgresql_config.database}" + self._db_url = f"postgresql://{encoded_username}:{encoded_password}@{host}:{port}/{self.postgresql_config.database}" @cached_property @override