diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0aace65ff6b..7a00e6941be 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4342,38 +4342,70 @@ def is_valid_api_key(key: str) -> bool: def construct_database_url_from_env_vars() -> Optional[str]: """ Construct a DATABASE_URL from individual environment variables. + + Required env vars: + DATABASE_HOST - hostname (may include :port) + DATABASE_USERNAME - database user + DATABASE_NAME - database name + + Optional env vars: + DATABASE_PASSWORD - database password + DATABASE_PORT - port number (ignored if host already contains :port) + DATABASE_SCHEMA - schema name (appended as ?schema=...) + DATABASE_SSL_MODE - PostgreSQL sslmode (disable|allow|prefer|require|verify-ca|verify-full) + DATABASE_SSL_ROOT_CERT - path to CA certificate file for SSL verification + Returns: Optional[str]: The constructed DATABASE_URL or None if required variables are missing """ import urllib.parse - # Check if all required variables are provided database_host = os.getenv("DATABASE_HOST") database_username = os.getenv("DATABASE_USERNAME") database_password = os.getenv("DATABASE_PASSWORD") database_name = os.getenv("DATABASE_NAME") + database_port = os.getenv("DATABASE_PORT") database_schema = os.getenv("DATABASE_SCHEMA") + database_ssl_mode = os.getenv("DATABASE_SSL_MODE") + database_ssl_root_cert = os.getenv("DATABASE_SSL_ROOT_CERT") - if database_host and database_username and database_name: - # Handle the problem of special character escaping in the database URL - database_username_enc = urllib.parse.quote_plus(database_username) - database_password_enc = ( - urllib.parse.quote_plus(database_password) if database_password else "" - ) - database_name_enc = urllib.parse.quote_plus(database_name) + if not (database_host and database_username and database_name): + return None - # Construct DATABASE_URL from the provided variables - if database_password: - database_url = f"postgresql://{database_username_enc}:{database_password_enc}@{database_host}/{database_name_enc}" - else: - database_url = f"postgresql://{database_username_enc}@{database_host}/{database_name_enc}" + # RFC 3986 percent-encoding for URI userinfo and path components + _enc = lambda s: urllib.parse.quote(s, safe="") - if database_schema: - database_url += f"?schema={database_schema}" + username_enc = _enc(database_username) + password_enc = _enc(database_password) if database_password else "" + database_name_enc = _enc(database_name) - return database_url + # Build host[:port] — DATABASE_PORT is used only when the host + # does not already contain a port (e.g. "myhost:5432"). + host_part = database_host + if database_port and ":" not in database_host: + host_part = f"{database_host}:{database_port}" - return None + # Build authority section (user[:pass]@host[:port]) + if database_password: + authority = f"{username_enc}:{password_enc}@{host_part}" + else: + authority = f"{username_enc}@{host_part}" + + database_url = f"postgresql://{authority}/{database_name_enc}" + + # Collect query parameters + query_params: list[str] = [] + if database_schema: + query_params.append(f"schema={_enc(database_schema)}") + if database_ssl_mode: + query_params.append(f"sslmode={_enc(database_ssl_mode)}") + if database_ssl_root_cert: + query_params.append(f"sslrootcert={database_ssl_root_cert}") + + if query_params: + database_url += "?" + "&".join(query_params) + + return database_url async def get_available_models_for_user( diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 12065ad5b4d..b019041c553 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -415,6 +415,166 @@ def test_construct_database_url_from_env_vars(self): result = construct_database_url_from_env_vars() assert result is None + @patch.dict(os.environ, {}, clear=True) + def test_construct_database_url_with_separate_port(self): + """Test DATABASE_PORT as a separate env var""" + from litellm.proxy.utils import construct_database_url_from_env_vars + + # Port as separate env var + test_env = { + "DATABASE_HOST": "myhost.example.com", + "DATABASE_PORT": "5433", + "DATABASE_USERNAME": "testuser", + "DATABASE_PASSWORD": "testpass", + "DATABASE_NAME": "testdb", + } + with patch.dict(os.environ, test_env): + result = construct_database_url_from_env_vars() + assert result == "postgresql://testuser:testpass@myhost.example.com:5433/testdb" + + # Port embedded in host — DATABASE_PORT should be ignored + test_env_embedded = { + "DATABASE_HOST": "myhost.example.com:5432", + "DATABASE_PORT": "9999", + "DATABASE_USERNAME": "testuser", + "DATABASE_PASSWORD": "testpass", + "DATABASE_NAME": "testdb", + } + with patch.dict(os.environ, test_env_embedded): + result = construct_database_url_from_env_vars() + assert result == "postgresql://testuser:testpass@myhost.example.com:5432/testdb" + + # No port at all — default PostgreSQL behavior (no :port in URL) + test_env_no_port = { + "DATABASE_HOST": "myhost.example.com", + "DATABASE_USERNAME": "testuser", + "DATABASE_PASSWORD": "testpass", + "DATABASE_NAME": "testdb", + } + with patch.dict(os.environ, test_env_no_port): + result = construct_database_url_from_env_vars() + assert result == "postgresql://testuser:testpass@myhost.example.com/testdb" + + @patch.dict(os.environ, {}, clear=True) + def test_construct_database_url_with_ssl_params(self): + """Test SSL mode and root cert query parameters""" + from litellm.proxy.utils import construct_database_url_from_env_vars + + # sslmode only + test_env = { + "DATABASE_HOST": "myhost.example.com", + "DATABASE_PORT": "5432", + "DATABASE_USERNAME": "testuser", + "DATABASE_PASSWORD": "testpass", + "DATABASE_NAME": "testdb", + "DATABASE_SSL_MODE": "require", + } + with patch.dict(os.environ, test_env): + result = construct_database_url_from_env_vars() + assert result == "postgresql://testuser:testpass@myhost.example.com:5432/testdb?sslmode=require" + + # sslmode + sslrootcert + test_env_full_ssl = { + "DATABASE_HOST": "myhost.example.com", + "DATABASE_PORT": "5432", + "DATABASE_USERNAME": "testuser", + "DATABASE_PASSWORD": "testpass", + "DATABASE_NAME": "testdb", + "DATABASE_SSL_MODE": "verify-ca", + "DATABASE_SSL_ROOT_CERT": "/certs/ca.pem", + } + with patch.dict(os.environ, test_env_full_ssl): + result = construct_database_url_from_env_vars() + assert result == "postgresql://testuser:testpass@myhost.example.com:5432/testdb?sslmode=verify-ca&sslrootcert=/certs/ca.pem" + + @patch.dict(os.environ, {}, clear=True) + def test_construct_database_url_with_schema_encoded(self): + """Test that DATABASE_SCHEMA is properly URL-encoded""" + from litellm.proxy.utils import construct_database_url_from_env_vars + + # Schema with special characters + test_env = { + "DATABASE_HOST": "localhost:5432", + "DATABASE_USERNAME": "testuser", + "DATABASE_PASSWORD": "testpass", + "DATABASE_NAME": "testdb", + "DATABASE_SCHEMA": "my schema", + } + with patch.dict(os.environ, test_env): + result = construct_database_url_from_env_vars() + assert result == "postgresql://testuser:testpass@localhost:5432/testdb?schema=my%20schema" + + # Normal schema (no encoding needed) + test_env_normal = { + "DATABASE_HOST": "localhost:5432", + "DATABASE_USERNAME": "testuser", + "DATABASE_PASSWORD": "testpass", + "DATABASE_NAME": "testdb", + "DATABASE_SCHEMA": "public", + } + with patch.dict(os.environ, test_env_normal): + result = construct_database_url_from_env_vars() + assert result == "postgresql://testuser:testpass@localhost:5432/testdb?schema=public" + + @patch.dict(os.environ, {}, clear=True) + def test_construct_database_url_multiple_query_params(self): + """Test combining schema + SSL query parameters""" + from litellm.proxy.utils import construct_database_url_from_env_vars + + test_env = { + "DATABASE_HOST": "prod-db.example.com", + "DATABASE_PORT": "5432", + "DATABASE_USERNAME": "appuser", + "DATABASE_PASSWORD": "secret", + "DATABASE_NAME": "mydb", + "DATABASE_SCHEMA": "app", + "DATABASE_SSL_MODE": "verify-full", + "DATABASE_SSL_ROOT_CERT": "/etc/ssl/certs/rds-ca.pem", + } + with patch.dict(os.environ, test_env): + result = construct_database_url_from_env_vars() + assert result == ( + "postgresql://appuser:secret@prod-db.example.com:5432/mydb" + "?schema=app&sslmode=verify-full&sslrootcert=/etc/ssl/certs/rds-ca.pem" + ) + + @patch.dict(os.environ, {}, clear=True) + def test_construct_database_url_heavy_special_chars_password(self): + """Test password with URI-special characters that broke selfhosted deployments""" + from litellm.proxy.utils import construct_database_url_from_env_vars + + # Real-world case: password with @, #, $, ! + test_env = { + "DATABASE_HOST": "pgsql-prod.postgres.database.azure.com", + "DATABASE_PORT": "5432", + "DATABASE_USERNAME": "admin", + "DATABASE_PASSWORD": "pgHgsh@gh4#Mj$!", + "DATABASE_NAME": "litellm", + "DATABASE_SSL_MODE": "require", + } + with patch.dict(os.environ, test_env): + result = construct_database_url_from_env_vars() + assert result == ( + "postgresql://admin:pgHgsh%40gh4%23Mj%24%21" + "@pgsql-prod.postgres.database.azure.com:5432/litellm" + "?sslmode=require" + ) + + # Edge case: password with all problematic URI chars + test_env_extreme = { + "DATABASE_HOST": "localhost", + "DATABASE_PORT": "5432", + "DATABASE_USERNAME": "user", + "DATABASE_PASSWORD": "p@ss:w0rd/with?q=1&a=2#frag%20", + "DATABASE_NAME": "testdb", + } + with patch.dict(os.environ, test_env_extreme): + result = construct_database_url_from_env_vars() + assert result == ( + "postgresql://user:p%40ss%3Aw0rd%2Fwith%3Fq%3D1%26a%3D2%23frag%2520" + "@localhost:5432/testdb" + ) + @patch("uvicorn.run") @patch("builtins.print") def test_run_server_no_config_passed(self, mock_print, mock_uvicorn_run):