-
Notifications
You must be signed in to change notification settings - Fork 17
Closed
Description
Hi
I have couple complex query which use TEMP MACRO, thus I need to keep DB connection between the MCP call.
Do you think it make sense to have the following option?
Thanks
diff --git a/src/mcp_server_duckdb/config.py b/src/mcp_server_duckdb/config.py
index 0f3e6e5..2f3e68f 100644
--- a/src/mcp_server_duckdb/config.py
+++ b/src/mcp_server_duckdb/config.py
@@ -19,6 +19,12 @@ class Config:
Run server in read-only mode.
"""
+ keep_connection: bool = False
+ """
+ Keep database connection open between queries.
+ Useful for TEMP objects like macros and tables, but may cause DB locking.
+ """
+
@staticmethod
def from_arguments() -> "Config":
"""
@@ -42,5 +48,16 @@ class Config:
"See: https://duckdb.org/docs/api/python/dbapi.html#read_only-connections",
)
+ parser.add_argument(
+ "--keep-connection",
+ action="store_true",
+ help="Keep database connection open between queries. "
+ "Useful for TEMP objects like macros and tables, but may cause DB locking.",
+ )
+
args = parser.parse_args()
- return Config(db_path=args.db_path, readonly=args.readonly)
+ return Config(
+ db_path=args.db_path,
+ readonly=args.readonly,
+ keep_connection=args.keep_connection
+ )
diff --git a/src/mcp_server_duckdb/server.py b/src/mcp_server_duckdb/server.py
index b711d92..5f5811e 100644
--- a/src/mcp_server_duckdb/server.py
+++ b/src/mcp_server_duckdb/server.py
@@ -34,13 +34,35 @@ class DuckDBDatabase:
duckdb.connect(config.db_path).close()
self.db_path = config.db_path
+ self._connection = None
def connect(self):
- return duckdb.connect(self.db_path, read_only=self.config.readonly)
+ if self.config.keep_connection:
+ if self._connection is None:
+ logger.info("Creating new persistent DuckDB connection")
+ self._connection = duckdb.connect(self.db_path, read_only=self.config.readonly)
+ return self._connection
+ else:
+ logger.debug("Creating new temporary DuckDB connection")
+ return duckdb.connect(self.db_path, read_only=self.config.readonly)
def execute_query(self, query: object, parameters: object = None) -> List[Any]:
- with closing(self.connect()) as connection:
- return connection.execute(query, parameters).fetchall()
+ if self.config.keep_connection:
+ connection = self.connect()
+ try:
+ return connection.execute(query, parameters).fetchall()
+ except Exception as e:
+ logger.error(f"Query execution error: {e}")
+ raise
+ else:
+ with closing(self.connect()) as connection:
+ return connection.execute(query, parameters).fetchall()
+
+ def __del__(self):
+ if self._connection is not None:
+ logger.info("Closing persistent DuckDB connection")
+ self._connection.close()
+ self._connection = None
async def main(config: Config):
~mcp-server-duckdb shang/mcp_server_keepconnection v1.0.1 ✚ ● ? git push origin
ERROR: Permission to ktanaka101/mcp-server-duckdb.git denied to risyou.
fatal: Could not read from remote repository.
Please make sure you have the correct access rights
and the repository exists.
ktanaka101
Metadata
Metadata
Assignees
Labels
No labels