fix(memory): accept PostgreSQL host:port configuration

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:１２３",
+        ":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)

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