Fix database connection pool exhaustion from undisposed SQLAlchemy engines

backend/dataline/services/connection.py

@@ -62,6 +62,7 @@ async def delete_connection(self, session: AsyncSession, connection_id: UUID) ->
 
     async def get_db_from_dsn(self, dsn: str) -> SQLDatabase:
         # Check if connection can be established before saving it
+        db = None
         try:
             db = SQLDatabase.from_uri(dsn)
             database = db._engine.url.database
@@ -72,9 +73,14 @@ async def get_db_from_dsn(self, dsn: str) -> SQLDatabase:
             return db
 
         except OperationalError as exc:
+            # Dispose the first failed engine if it exists
+            if db is not None:
+                db.dispose()
+
             # Try again replacing localhost with host.docker.internal to connect with DBs running in docker
             if "localhost" in dsn:
                 dsn = dsn.replace("localhost", "host.docker.internal")
+                db = None  # Reset db reference
                 try:
                     db = SQLDatabase.from_uri(dsn)
                     database = db._engine.url.database
@@ -84,15 +90,22 @@ async def get_db_from_dsn(self, dsn: str) -> SQLDatabase:
 
                     return db
                 except OperationalError as e:
+                    if db is not None:
+                        db.dispose()
                     logger.error(e)
                     raise ValidationError("Failed to connect to database, please check your DSN.")
                 except Exception as e:
+                    if db is not None:
+                        db.dispose()
                     forward_connection_errors(e)
 
             logger.error(exc)
             raise ValidationError("Failed to connect to database, please check your DSN.")
 
         except Exception as e:
+            # Dispose on any other error
+            if db is not None:
+                db.dispose()
             forward_connection_errors(e)
             logger.error(e)
             raise ValidationError("Failed to connect to database, please check your DSN.")
@@ -123,14 +136,18 @@ async def update_connection(
 
             # Check if connection can be established before saving it
             db = await self.get_db_from_dsn(data.dsn)
-            update.dsn = str(db._engine.url.render_as_string(hide_password=False))
-            update.database = db._engine.url.database
-            update.dialect = db.dialect
-            current_connection = await self.get_connection(session, connection_uuid)
-            old_options = (
-                ConnectionOptions.model_validate(current_connection.options) if current_connection.options else None
-            )
-            update.options = self.merge_options(old_options, db)
+            try:
+                update.dsn = str(db._engine.url.render_as_string(hide_password=False))
+                update.database = db._engine.url.database
+                update.dialect = db.dialect
+                current_connection = await self.get_connection(session, connection_uuid)
+                old_options = (
+                    ConnectionOptions.model_validate(current_connection.options) if current_connection.options else None
+                )
+                update.options = self.merge_options(old_options, db)
+            finally:
+                # Dispose engine to release database connections
+                db.dispose()
         elif data.options:
             # only modify options if dsn hasn't changed
             update.options = data.options
@@ -151,34 +168,38 @@ async def create_connection(
     ) -> ConnectionOut:
         # Check if connection can be established before saving it
         db = await self.get_db_from_dsn(dsn)
-        # get potentially modified dsn (eg. if localhost was replaced with host.docker.internal)
-        dsn = str(db._engine.url.render_as_string(hide_password=False))
-        if not connection_type:
-            connection_type = db.dialect
-
-        # Check if connection already exists
-        await self.check_dsn_already_exists(session, dsn)
-        connection_schemas: list[ConnectionSchema] = [
-            ConnectionSchema(
-                name=schema,
-                tables=[ConnecitonSchemaTable(name=table, enabled=True) for table in tables],
-                enabled=True,
+        try:
+            # get potentially modified dsn (eg. if localhost was replaced with host.docker.internal)
+            dsn = str(db._engine.url.render_as_string(hide_password=False))
+            if not connection_type:
+                connection_type = db.dialect
+
+            # Check if connection already exists
+            await self.check_dsn_already_exists(session, dsn)
+            connection_schemas: list[ConnectionSchema] = [
+                ConnectionSchema(
+                    name=schema,
+                    tables=[ConnecitonSchemaTable(name=table, enabled=True) for table in tables],
+                    enabled=True,
+                )
+                for schema, tables in db._all_tables_per_schema.items()
+            ]
+            connection = await self.connection_repo.create(
+                session,
+                ConnectionCreate(
+                    dsn=dsn,
+                    database=db._engine.url.database,
+                    name=name,
+                    dialect=db.dialect,
+                    type=connection_type,
+                    is_sample=is_sample,
+                    options=ConnectionOptions(schemas=connection_schemas),
+                ),
             )
-            for schema, tables in db._all_tables_per_schema.items()
-        ]
-        connection = await self.connection_repo.create(
-            session,
-            ConnectionCreate(
-                dsn=dsn,
-                database=db._engine.url.database,
-                name=name,
-                dialect=db.dialect,
-                type=connection_type,
-                is_sample=is_sample,
-                options=ConnectionOptions(schemas=connection_schemas),
-            ),
-        )
-        return ConnectionOut.model_validate(connection)
+            return ConnectionOut.model_validate(connection)
+        finally:
+            # Dispose engine to release database connections
+            db.dispose()
 
     async def create_sqlite_connection(
         self, session: AsyncSession, file: BinaryIO, name: str, is_sample: bool = False
@@ -328,13 +349,16 @@ async def refresh_connection_schema(self, session: AsyncSession, connection_id:
 
         # Get the latest schema information
         db = await self.get_db_from_dsn(connection.dsn)
+        try:
+            old_options = ConnectionOptions.model_validate(connection.options) if connection.options else None
+            new_options = self.merge_options(old_options, db)
 
-        old_options = ConnectionOptions.model_validate(connection.options) if connection.options else None
-        new_options = self.merge_options(old_options, db)
-
-        # Update the connection with new options
-        updated_connection = await self.connection_repo.update_by_uuid(
-            session, connection_id, ConnectionUpdate(options=new_options)
-        )
+            # Update the connection with new options
+            updated_connection = await self.connection_repo.update_by_uuid(
+                session, connection_id, ConnectionUpdate(options=new_options)
+            )
 
-        return ConnectionOut.model_validate(updated_connection)
+            return ConnectionOut.model_validate(updated_connection)
+        finally:
+            # Dispose engine to release database connections
+            db.dispose()

backend/dataline/services/llm_flow/graph.py

@@ -55,6 +55,11 @@ def __init__(self, connection: ConnectionProtocol) -> None:
         self.tool_executor = ToolExecutor(tools=all_tools)
         self.tracer = None  # no tracing by default
 
+    def dispose(self) -> None:
+        """Dispose of the database engine and close all connections."""
+        if hasattr(self, "db") and self.db is not None:
+            self.db.dispose()
+
     async def query(
         self, query: str, options: QueryOptions, history: Sequence[BaseMessage] | None = None
     ) -> AsyncGenerator[tuple[Sequence[BaseMessage] | None, Sequence[ResultType] | None], None]:
@@ -84,11 +89,15 @@ async def query(
         config: RunnableConfig | None = {"callbacks": [self.tracer]} if self.tracer is not None else None
         current_results: Sequence[ResultType] | None
         current_messages: Sequence[BaseMessage] | None
-        async for chunk in app.astream(initial_state, config=config):
-            for tool, tool_chunk in chunk.items():
-                current_results = tool_chunk.get("results")
-                current_messages = tool_chunk.get("messages")
-                yield (current_messages, current_results)
+        try:
+            async for chunk in app.astream(initial_state, config=config):
+                for tool, tool_chunk in chunk.items():
+                    current_results = tool_chunk.get("results")
+                    current_messages = tool_chunk.get("messages")
+                    yield (current_messages, current_results)
+        finally:
+            # Always dispose of the engine after query completion to prevent connection leaks
+            self.dispose()
 
     def build_graph(self) -> StateGraph:
         # Create the graph

backend/dataline/services/llm_flow/utils.py

@@ -106,6 +106,21 @@ def from_uri(
     ) -> Self:
         """Construct a SQLAlchemy engine from URI."""
         _engine_args = engine_args or {}
+        # Set conservative connection pool limits to prevent exhausting database connections
+        # pool_size: number of connections to maintain in the pool
+        # max_overflow: number of connections that can be created beyond pool_size (only for QueuePool)
+        # pool_pre_ping: verify connections before using them (recover from stale connections)
+
+        # Check if this is SQLite (which uses SingletonThreadPool and doesn't support max_overflow)
+        is_sqlite = database_uri.startswith("sqlite")
+
+        if "pool_size" not in _engine_args:
+            _engine_args["pool_size"] = 2
+        # max_overflow only works with QueuePool (not with SQLite's SingletonThreadPool)
+        if "max_overflow" not in _engine_args and not is_sqlite:
+            _engine_args["max_overflow"] = 3
+        if "pool_pre_ping" not in _engine_args:
+            _engine_args["pool_pre_ping"] = True
         engine = create_engine(database_uri, **_engine_args)
         return cls(engine, schemas=schemas, **kwargs)
 
@@ -168,6 +183,11 @@ def from_dataline_connection(
             **kwargs,
         )
 
+    def dispose(self) -> None:
+        """Dispose of the database engine and close all connections in the pool."""
+        if self._engine is not None:
+            self._engine.dispose()
+
     def get_table_info(self, table_names: list[str] | None = None) -> str:
         """Get information about specified tables.