Merge pull request #3 from KxSystems/fix-date-formatting

fix date formatting

Author: cbeagan-kx

src/mcp_server/main.py

@@ -0,0 +1,4 @@
+from mcp_server.server import main
+
+if __name__ == "__main__":
+    main()

src/mcp_server/resources/kdbx_database_tables.py

@@ -19,7 +19,7 @@ async def kdbx_describe_table_impl(table: str) -> List[TextContent]:
         List[TextContent]: Table description with metadata and sample data
     """
     try:
-        conn = get_kdb_connection(config.kdbx_config)
+        conn = get_kdb_connection()
 
         total_records = conn('{count get x}', table).py()
 
@@ -62,7 +62,7 @@ async def kdbx_describe_table_impl(table: str) -> List[TextContent]:
 
 async def kdbx_describe_tables_impl() -> List[TextContent]:
     try:
-        conn = get_kdb_connection(config.kdbx_config)
+        conn = get_kdb_connection()
 
         available_tables = conn.tables(None).py()
 

src/mcp_server/server.py

@@ -215,9 +215,7 @@ def parse_args():
 
 
 def main():
-
     args = parse_args()
-
      # Determine transport mode (only set if CLI flag is provided)
     args.kdbx_mcp_transport = None
     if args.streamable_http:
@@ -238,6 +236,3 @@ def main():
 
     server = McpServer(config)
     server.run()
-
-if __name__ == "__main__":
-    main()

src/mcp_server/tools/kdbx_run_sql_query.py

@@ -1,11 +1,12 @@
 import logging
+import pykx as kx
+import json
 from typing import Dict, Any
-from mcp_server.utils.kdbx import run_kdbx_sql_query
+from mcp_server.utils.kdbx import get_kdb_connection
 
 logger = logging.getLogger(__name__)
 MAX_ROWS_RETURNED = 1000
 
-
 async def run_query_impl(sqlSelectQuery: str) -> Dict[str, Any]:
     try:
         dangerous_keywords = ['INSERT', 'DROP', 'DELETE', 'TRUNCATE', 'ALTER', 'CREATE']
@@ -15,15 +16,14 @@ async def run_query_impl(sqlSelectQuery: str) -> Dict[str, Any]:
             if keyword in query_upper and not query_upper.startswith('SELECT'):
                 raise ValueError(f"Query contains dangerous keyword: {keyword}")
 
-        result = run_kdbx_sql_query(sqlSelectQuery, MAX_ROWS_RETURNED)
+        conn = get_kdb_connection()
+        # below query gets kdbx table data back as json for correct conversion of different datatypes
+        result = conn('{r:.s.e x;`rowCount`data!(count r;.j.j y sublist r)}', kx.CharVector(sqlSelectQuery), MAX_ROWS_RETURNED)
         total = int(result['rowCount'])
-        table = result['data'].pd()
-
-        if table.empty:
+        if 0==total:
             return {"status": "success", "data": [], "message": "No rows returned"}
-
-        rows = table.to_dict(orient="records")
-
+        # parse json result
+        rows = json.loads(result['data'].py().decode('utf-8'))
         if total > MAX_ROWS_RETURNED:
             logger.info(f"Table has {total} rows. Query returned truncated data to {MAX_ROWS_RETURNED} rows.")
             return {

src/mcp_server/utils/kdbx.py

@@ -7,8 +7,24 @@
 
 logger = logging.getLogger(__name__)
 
+def get_kdb_connection() -> kx.QConnection:
+    try:
+        conn = kdb_sync_connection(config.kdbx_config)
+        conn('') # check if conn is live for existing connection from cache
+        return conn
+    except Exception as e:
+        if "Attempted to use a closed IPC connection" in str(e):
+            logger.warning("KDB-X connection was closed. Reinitializing...")
+            cleanup_kdb_connection()
+            conn = kdb_sync_connection(config.kdbx_config)
+            conn('')
+            return conn
+        else:
+            logger.error(f"Error in creating KDBX connection: {e}")
+            raise
+
 @lru_cache()
-def get_kdb_connection(config: Optional[KDBConfig] = None) -> kx.QConnection:
+def kdb_sync_connection(config: Optional[KDBConfig] = None) -> kx.QConnection:
     if config is None:
         config = default_kdb_config
 
@@ -34,21 +50,6 @@ def get_kdb_connection(config: Optional[KDBConfig] = None) -> kx.QConnection:
     logger.error(f"Failed to connect to KDB")
     raise
 
-def run_kdbx_sql_query(query:str, max_rows:int) -> any:
-    try:
-        conn = get_kdb_connection(config.kdbx_config)
-        result = conn('{r:.s.e x;`rowCount`data!(count r;y sublist r)}', kx.CharVector(query), max_rows)
-        return result
-    except Exception as e:
-        if "Attempted to use a closed IPC connection" in str(e):
-            logger.warning("KDB-X connection was closed. Reinitializing...")
-            cleanup_kdb_connection()
-            conn = get_kdb_connection(config.kdbx_config)
-            return conn('{r:.s.e x;`rowCount`data!(count r;y sublist r)}', kx.CharVector(query), max_rows)
-        else:
-            logger.error(f"Error running KDB query: {e}")
-            raise
-
 def cleanup_kdb_connection():
-    get_kdb_connection.cache_clear()
-    logger.info("KDB connection cache cleared")
+    kdb_sync_connection.cache_clear()
+    logger.info("KDBX connection cache cleared")