Skip to content

Commit 011791c

Browse files
author
Sebastian
committed
Better code documentation
1 parent b3e9283 commit 011791c

File tree

1 file changed

+41
-15
lines changed

1 file changed

+41
-15
lines changed

mcpgateway/db.py

Lines changed: 41 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
create_engine,
3535
func,
3636
select,
37+
make_url,
3738
)
3839
from sqlalchemy.event import listen
3940
from sqlalchemy.exc import SQLAlchemyError
@@ -49,27 +50,52 @@
4950
from mcpgateway.config import settings
5051
from mcpgateway.types import ResourceContent
5152

52-
base_connect_args = {}
53-
# Add the SQLite-only flag when appropriate
54-
if settings.database_url.startswith("sqlite"):
55-
base_connect_args["check_same_thread"] = False
56-
elif settings.database_url.startswith("postgres"):
57-
base_connect_args = {
58-
"keepalives": 1,
59-
"keepalives_idle": 30,
60-
"keepalives_interval": 5,
61-
"keepalives_count": 5,
62-
}
63-
64-
# Create SQLAlchemy engine with connection pooling
53+
54+
# ---------------------------------------------------------------------------
55+
# 1. Parse the URL so we can inspect backend ("postgresql", "sqlite", …)
56+
# and the specific driver ("psycopg2", "asyncpg", empty string = default).
57+
# ---------------------------------------------------------------------------
58+
url = make_url(settings.database_url)
59+
backend = url.get_backend_name() # e.g. 'postgresql', 'sqlite'
60+
driver = url.get_driver_name() or "default"
61+
62+
# Start with an empty dict and add options only when the driver can accept
63+
# them; this prevents unexpected TypeError at connect time.
64+
connect_args: dict[str, object] = {}
65+
66+
# ---------------------------------------------------------------------------
67+
# 2. PostgreSQL (synchronous psycopg2 only)
68+
# The keep-alive parameters below are recognised exclusively by libpq /
69+
# psycopg2 and let the kernel detect broken network links quickly.
70+
# ---------------------------------------------------------------------------
71+
if backend == "postgresql" and driver in ("psycopg2", "default", ""):
72+
connect_args.update(
73+
keepalives=1, # enable TCP keep-alive probes
74+
keepalives_idle=30, # seconds of idleness before first probe
75+
keepalives_interval=5, # seconds between probes
76+
keepalives_count=5, # drop the link after N failed probes
77+
)
78+
79+
# ---------------------------------------------------------------------------
80+
# 3. SQLite (optional) – only one extra flag and it is *SQLite-specific*.
81+
# ---------------------------------------------------------------------------
82+
elif backend == "sqlite":
83+
# Allow pooled connections to hop across threads.
84+
connect_args["check_same_thread"] = False
85+
86+
# 4. Other backends (MySQL, MSSQL, etc.) leave `connect_args` empty.
87+
88+
# ---------------------------------------------------------------------------
89+
# 5. Build the Engine with a single, clean connect_args mapping.
90+
# ---------------------------------------------------------------------------
6591
engine = create_engine(
6692
settings.database_url,
67-
pool_pre_ping=True,
93+
pool_pre_ping=True, # quick liveness check per checkout
6894
pool_size=settings.db_pool_size,
6995
max_overflow=settings.db_max_overflow,
7096
pool_timeout=settings.db_pool_timeout,
7197
pool_recycle=settings.db_pool_recycle,
72-
connect_args=base_connect_args
98+
connect_args=connect_args,
7399
)
74100

75101
# Session factory

0 commit comments

Comments
 (0)