Skip to content

Keep connection for TEMP MACRO #23

@risyou

Description

@risyou

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions