Skip to content

Commit 8a3bb5e

Browse files
committed
fix(orm): harden ConnectionFactory against bad port and unknown driver
build_url()/make() previously surfaced bad connection configs as raw, low-level errors: - mysql/postgres with a missing/empty 'port' produced a URL with a bare trailing colon that crashed at engine creation with ValueError: invalid literal for int() with base 10: ''. - an unsupported driver raised a raw KeyError from DRIVER_URLS[driver]. - make()'s 'Unsupported driver' fallback was unreachable dead code, since create_engine()/build_url() always raised the raw KeyError first. Hardening: - Add per-driver default ports (mysql 3306, postgres 5432); an explicit port still takes precedence, so the happy path is byte-for-byte unchanged. - build_url() and make() now raise the existing framework-level DriverNotFound with a message listing the supported drivers. The check in make() runs before any engine is built, so it is reachable and meaningful; the dead match fallback is removed in favour of a CONNECTIONS lookup. The documented 'url' passthrough still short-circuits field assembly, so supplying a full url bypasses both driver and port validation. Regression tests cover: mysql/postgres missing and empty-string ports, explicit-port precedence, unsupported driver via build_url() and make(), url passthrough, and unchanged valid sqlite/mysql/postgres URLs.
1 parent 31053d2 commit 8a3bb5e

3 files changed

Lines changed: 136 additions & 18 deletions

File tree

fastapi_startkit/src/fastapi_startkit/masoniteorm/connections/factory.py

Lines changed: 31 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from sqlalchemy.pool import NullPool
55
from sqlalchemy.ext.asyncio import create_async_engine, AsyncEngine
66

7+
from fastapi_startkit.exceptions.exceptions import DriverNotFound
78
from fastapi_startkit.masoniteorm.connections.connection import Connection
89
from fastapi_startkit.masoniteorm.connections.sqlite_connection import SQliteConnection
910
from fastapi_startkit.masoniteorm.connections.postgres_connection import (
@@ -19,13 +20,35 @@ class ConnectionFactory:
1920
"postgres": "postgresql+asyncpg",
2021
}
2122

23+
# Sensible per-driver defaults so a missing/empty "port" no longer surfaces as a
24+
# raw SQLAlchemy `ValueError: invalid literal for int()` at engine-creation time.
25+
DEFAULT_PORTS = {
26+
"mysql": 3306,
27+
"postgres": 5432,
28+
}
29+
30+
CONNECTIONS = {
31+
"sqlite": SQliteConnection,
32+
"mysql": MySQLConnection,
33+
"postgres": PostgresConnection,
34+
}
35+
36+
@classmethod
37+
def _unsupported_driver_error(cls, driver: Any) -> DriverNotFound:
38+
supported = ", ".join(sorted(cls.DRIVER_URLS))
39+
return DriverNotFound(
40+
f"Unsupported database driver {driver!r}. Supported drivers are: {supported}."
41+
)
42+
2243
@classmethod
2344
def build_url(cls, config: dict) -> str:
2445
if url := config.get("url"):
2546
return str(url)
2647

2748
driver = config["driver"]
28-
scheme = cls.DRIVER_URLS[driver]
49+
scheme = cls.DRIVER_URLS.get(driver)
50+
if scheme is None:
51+
raise cls._unsupported_driver_error(driver)
2952
db = config.get("database", "")
3053

3154
if driver == "sqlite":
@@ -34,7 +57,7 @@ def build_url(cls, config: dict) -> str:
3457
user = config.get("username", "")
3558
pwd = config.get("password", "")
3659
host = config.get("host", "localhost")
37-
port = config.get("port", "")
60+
port = config.get("port") or cls.DEFAULT_PORTS[driver]
3861
return f"{scheme}://{user}:{pwd}@{host}:{port}/{db}"
3962

4063
@classmethod
@@ -50,15 +73,11 @@ def create_engine(cls, cfg: dict) -> AsyncEngine:
5073
kwargs["poolclass"] = StaticPool
5174
return create_async_engine(url, **kwargs)
5275

53-
def make(self, config: dict, name: str) -> type[Connection]:
54-
engine = self.create_engine(config)
76+
def make(self, config: dict, name: str) -> Connection:
5577
driver = config["driver"]
56-
match driver:
57-
case "sqlite":
58-
return SQliteConnection(engine, config)
59-
case "postgres":
60-
return PostgresConnection(engine, config)
61-
case "mysql":
62-
return MySQLConnection(engine, config)
78+
connection_class = type(self).CONNECTIONS.get(driver)
79+
if connection_class is None:
80+
raise type(self)._unsupported_driver_error(driver)
6381

64-
raise ValueError(f"Unsupported driver: {driver}")
82+
engine = self.create_engine(config)
83+
return connection_class(engine, config)

fastapi_startkit/tests/masoniteorm/config/test_db_url.py

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import unittest
22

3+
from fastapi_startkit.exceptions.exceptions import DriverNotFound
34
from fastapi_startkit.masoniteorm.connections.factory import ConnectionFactory
45

56

@@ -76,3 +77,101 @@ def test_direct_url_passthrough_takes_precedence(self):
7677
}
7778
url = ConnectionFactory.build_url(config)
7879
self.assertEqual(url, "mysql+aiomysql://admin:pw@prod-host:3306/live")
80+
81+
82+
class TestConnectionFactoryMissingPort(unittest.TestCase):
83+
"""A missing/empty port must fall back to a sane per-driver default instead of
84+
producing a URL with a bare trailing colon that crashes at engine creation with
85+
`ValueError: invalid literal for int() with base 10: ''`."""
86+
87+
def test_mysql_missing_port_uses_default_3306(self):
88+
config = {
89+
"driver": "mysql",
90+
"host": "localhost",
91+
"database": "mydb",
92+
"username": "root",
93+
"password": "secret",
94+
}
95+
url = ConnectionFactory.build_url(config)
96+
self.assertEqual(url, "mysql+aiomysql://root:secret@localhost:3306/mydb")
97+
98+
def test_postgres_missing_port_uses_default_5432(self):
99+
config = {
100+
"driver": "postgres",
101+
"host": "db.example.com",
102+
"database": "mydb",
103+
"username": "user",
104+
"password": "pass",
105+
}
106+
url = ConnectionFactory.build_url(config)
107+
self.assertEqual(url, "postgresql+asyncpg://user:pass@db.example.com:5432/mydb")
108+
109+
def test_mysql_empty_string_port_uses_default(self):
110+
config = {
111+
"driver": "mysql",
112+
"host": "localhost",
113+
"port": "",
114+
"database": "mydb",
115+
"username": "root",
116+
"password": "secret",
117+
}
118+
url = ConnectionFactory.build_url(config)
119+
self.assertEqual(url, "mysql+aiomysql://root:secret@localhost:3306/mydb")
120+
121+
def test_postgres_empty_string_port_uses_default(self):
122+
config = {
123+
"driver": "postgres",
124+
"host": "localhost",
125+
"port": "",
126+
"database": "mydb",
127+
"username": "user",
128+
"password": "pass",
129+
}
130+
url = ConnectionFactory.build_url(config)
131+
self.assertEqual(url, "postgresql+asyncpg://user:pass@localhost:5432/mydb")
132+
133+
def test_explicit_port_is_preserved(self):
134+
"""An explicitly configured port must win over the default (happy path)."""
135+
config = {
136+
"driver": "postgres",
137+
"host": "localhost",
138+
"port": 6543,
139+
"database": "mydb",
140+
"username": "user",
141+
"password": "pass",
142+
}
143+
url = ConnectionFactory.build_url(config)
144+
self.assertEqual(url, "postgresql+asyncpg://user:pass@localhost:6543/mydb")
145+
146+
147+
class TestConnectionFactoryUnsupportedDriver(unittest.TestCase):
148+
"""An unknown driver must raise a friendly framework error, not a raw KeyError."""
149+
150+
def test_build_url_unknown_driver_raises_driver_not_found(self):
151+
with self.assertRaises(DriverNotFound) as ctx:
152+
ConnectionFactory.build_url({"driver": "oracle", "database": "mydb"})
153+
message = str(ctx.exception)
154+
self.assertIn("oracle", message)
155+
# The message should guide the user toward the supported drivers.
156+
self.assertIn("sqlite", message)
157+
self.assertIn("mysql", message)
158+
self.assertIn("postgres", message)
159+
160+
def test_build_url_does_not_raise_raw_keyerror(self):
161+
with self.assertRaises(DriverNotFound):
162+
ConnectionFactory.build_url({"driver": "cassandra"})
163+
164+
def test_make_unknown_driver_raises_driver_not_found(self):
165+
"""make()'s driver check is reachable and fires before any engine is built."""
166+
with self.assertRaises(DriverNotFound):
167+
ConnectionFactory().make({"driver": "oracle", "database": "mydb"}, "oracle")
168+
169+
def test_url_passthrough_bypasses_driver_validation(self):
170+
"""The documented `url` stopgap short-circuits field assembly entirely, so an
171+
unknown driver name is irrelevant when a full url is supplied."""
172+
config = {
173+
"driver": "oracle",
174+
"url": "sqlite+aiosqlite:///db.sqlite3",
175+
}
176+
url = ConnectionFactory.build_url(config)
177+
self.assertEqual(url, "sqlite+aiosqlite:///db.sqlite3")

fastapi_startkit/tests/masoniteorm/models/test_model_attributes.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
import pendulum
22
import pytest
3-
from unittest.mock import MagicMock, patch
43

54
from fastapi_startkit.carbon import Carbon
5+
from fastapi_startkit.exceptions.exceptions import DriverNotFound
66
from fastapi_startkit.masoniteorm.models.fields import DateTimeField
77
from fastapi_startkit.masoniteorm.connections.factory import ConnectionFactory
88
from fastapi_startkit.masoniteorm.connections.manager import DatabaseManager
@@ -86,8 +86,9 @@ def test_connection_is_cached(self, db):
8686
assert conn1 is conn2
8787

8888
def test_connection_raises_for_missing_driver(self):
89-
# create_engine is called before the driver switch, so mock it to isolate
90-
# the "Unsupported driver" branch in ConnectionFactory.make().
89+
# make() now validates the driver before any engine is built, so an
90+
# unsupported driver fails fast with a friendly framework error rather
91+
# than a raw KeyError/ValueError from deep inside SQLAlchemy.
9192
factory = ConnectionFactory()
9293
bad_config = {
9394
"default": "mssql",
@@ -96,9 +97,8 @@ def test_connection_raises_for_missing_driver(self):
9697
},
9798
}
9899
dm = DatabaseManager(factory, bad_config)
99-
with patch.object(ConnectionFactory, "create_engine", return_value=MagicMock()):
100-
with pytest.raises(ValueError, match="Unsupported driver"):
101-
dm.connection("mssql")
100+
with pytest.raises(DriverNotFound, match="mssql"):
101+
dm.connection("mssql")
102102

103103

104104
# ---------------------------------------------------------------------------

0 commit comments

Comments
 (0)