Add Schema Retrieval to MySQL MCP server #73

src/mysql-mcp-server/.gitignore

@@ -0,0 +1 @@
+local_config.json

src/mysql-mcp-server/LICENSE.txt

@@ -0,0 +1,35 @@
+Copyright (c) 2025 Oracle and/or its affiliates.
+
+The Universal Permissive License (UPL), Version 1.0
+
+Subject to the condition set forth below, permission is hereby granted to any
+person obtaining a copy of this software, associated documentation and/or data
+(collectively the "Software"), free of charge and under any and all copyright
+rights in the Software, and any and all patent rights owned or freely
+licensable by each licensor hereunder covering either (i) the unmodified
+Software as contributed to or provided by such licensor, or (ii) the Larger
+Works (as defined below), to deal in both
+
+(a) the Software, and
+(b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
+one is included with the Software (each a "Larger Work" to which the Software
+is contributed by such licensors),
+
+without restriction, including without limitation the rights to copy, create
+derivative works of, display, perform, and distribute the Software and make,
+use, sell, offer for sale, import, export, have made, and have sold the
+Software and the Larger Work(s), and to sublicense the foregoing rights on
+either these or other terms.
+
+This license is subject to the following condition:
+The above copyright notice and either this complete permission notice or at
+a minimum a reference to the UPL must be included in all copies or
+substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

src/mysql-mcp-server/README.md

@@ -35,37 +35,18 @@ A Python-based MCP (Model Context Protocol) server that provides a suite of tool
   - List objects in a bucket
 
 ## Prerequisites
-
-- Python 3.x
-- `fastmcp`
-- `oci` SDK
-- `mysql-connector-python` SDK
 - Valid database connection file. Resolution order:
   1) Path specified by environment variable `MYSQL_MCP_CONFIG` (absolute or relative to this module)
-  2) `src/mysql-mcp-server/local_config.json` (default)
+  2) `local_config.json` (default)
 - Valid OCI configuration file (`~/.oci/config`) or environment variables
 
-## Installation
-
-1. Clone this repository.
-2. Install required dependencies using pip:
-   ```
-   pip install -r requirements.txt
-   ```
-   This will install `oci`, `fastmcp`, `mysql-connector-python`, and all other dependencies.
-3. Set up your OCI config file at ~/.oci/config
 
 ## OCI Configuration
 
 The server requires a valid OCI config file with proper credentials.
 The default location is ~/.oci/config. For instructions on setting up this file,
 see the [OCI SDK documentation](https://docs.oracle.com/en-us/iaas/Content/API/Concepts/sdkconfig.htm).
 
-## Required Python Packages
-
-- `oci`
-- `fastmcp`
-- `mysql-connector-python`
 
 ## Supported Database Modes
 
@@ -79,9 +60,12 @@ Installation is dependent on the MCP Client being used, it usually consists of a
 {
   "mcpServers": {
     "mysqltools": {
-      "command": "C:\\Python\\python.exe",
+      "command": "uv",
       "args": [
-        "C:\\Users\\user1\\mysql-mcp-server\\mysql_mcp_server.py"
+        "--directory",
+        "C:\\ABSOLUTE\\PATH\\TO\\PARENT\\FOLDER\\mysql-mcp-server",
+        "run",
+        "oracle.mysql_mcp_server"
       ]
     }
   }
@@ -95,9 +79,12 @@ Example with TENANCY_ID_OVERRIDE::
 {
   "mcpServers": {
     "mysqltools": {
-      "command": "C:\\Python\\python.exe",
+      "command": "uv",
       "args": [
-        "C:\\Users\\user1\\mysql-mcp-server\\mysql_mcp_server.py"
+        "--directory",
+        "C:\\ABSOLUTE\\PATH\\TO\\PARENT\\FOLDER\\mysql-mcp-server",
+        "run",
+        "oracle.mysql_mcp_server"
       ],
       "env": {
         "TENANCY_ID_OVERRIDE": "ocid1.tenancy.oc1..deadbeef"
@@ -196,7 +183,7 @@ Note:
 The server runs using stdio transport and can be started by running:
 
 ```bash
-python mysql_mcp_server.py
+uv run oracle.mysql_mcp_server
 ```
 
 ## API Tools

src/mysql-mcp-server/oracle/mysql_mcp_server/consts.py

@@ -0,0 +1,9 @@
+"""
+Copyright (c) 2025, Oracle and/or its affiliates.
+Licensed under the Universal Permissive License v1.0 as shown at
+https://oss.oracle.com/licenses/upl.
+"""
+
+MIN_CONTEXT_SIZE = 10
+DEFAULT_CONTEXT_SIZE = 20
+MAX_CONTEXT_SIZE = 100

src/mysql-mcp-server/mysql_mcp_server.py → src/mysql-mcp-server/oracle/mysql_mcp_server/server.py

@@ -12,17 +12,15 @@
 from fastmcp import FastMCP
 from mysql import connector
 from mysql.connector.abstracts import MySQLConnectionAbstract
-from utils import (
+from oracle.mysql_mcp_server.utils import (
     DatabaseConnectionError,
     Mode,
     OciInfo,
     get_ssh_command,
     load_mysql_config,
 )
 
-MIN_CONTEXT_SIZE = 10
-DEFAULT_CONTEXT_SIZE = 20
-MAX_CONTEXT_SIZE = 100
+from oracle.mysql_mcp_server.consts import MIN_CONTEXT_SIZE, DEFAULT_CONTEXT_SIZE, MAX_CONTEXT_SIZE
 
 ###############################################################
 # Start setup
@@ -828,18 +826,21 @@ def ask_nl_sql(connection_id: str, question: str) -> str:
         set_response = _execute_sql_tool(db_connection, "SET @response = NULL;")
         if check_error(set_response):
             return json.dumps({"error": f"Error with NL_SQL: {set_response}"})
-
+        if db_connection.database is not None:
+            call_nl_sql = f"CALL sys.NL_SQL(%s, @response, JSON_OBJECT('schemas', JSON_ARRAY(\"{db_connection.database}\"), 'execute', FALSE, 'model_id', 'meta.llama-4-maverick-17b-128e-instruct-fp8'))"
+        else:
+            call_nl_sql = "CALL sys.NL_SQL(%s, @response, NULL)"
         nl2sql_response = _execute_sql_tool(
             db_connection,
-            f"CALL sys.NL_SQL(%s, @response, NULL)",
+            call_nl_sql,
             params=[question],
         )
         if check_error(nl2sql_response):
             return json.dumps({"error": f"Error with NL_SQL: {nl2sql_response}"})
 
         fetch_response = _execute_sql_tool(db_connection, "SELECT @response;")
         if check_error(fetch_response):
-            return json.dumps({"error": f"Error with ML_RAG: {fetch_response}"})
+            return json.dumps({"error": f"Error with NL_SQL: {fetch_response}"})
 
         try:
             response = json.loads(fetch_one(fetch_response))
@@ -849,11 +850,90 @@ def ask_nl_sql(connection_id: str, question: str) -> str:
             return json.dumps({"error": "Unexpected response format from NL_SQL"})
 
 
+@mcp.tool()
+def retrieve_relevant_schema_information(connection_id: str, question: str) -> str:
+    """
+    [MCP Tool] Retrieve relevant schemas and tables for a given natural language question.
+
+    This tool analyzes the input question and, from the provided list of schema and/or table names, 
+    identifies only those that are relevant with respect to the question. 
+    It can optionally consider table and column comments for improved semantic matching. The results contain only the relevant schemas and tables in a JSON object.
+
+    Args:
+        connection_id (str): MySQL connection key.
+        input (str): Input question for use with ML_SCHEMA_RETRIEVAL.
+
+    Returns:
+        JSON object of the form:
+        {
+            "tables": ["schema1.tableA", "schema2.tableB", ...],
+            "schemas": ["schema1","schema",...],
+            "schema_metadata": [
+                {
+                "fkey_info": [
+                    {
+                    "fkey_column": "customer_id",
+                    "fkey_ref_col_name": "customer_id",
+                    "fkey_ref_tbl_name": "customers"
+                    }
+                ],
+                "table_name": "customer_addresses",
+                "column_info": [
+                    {
+                    "data_type": "int",
+                    "column_name": "address_id",
+                    "column_comment": ""
+                    },
+                    {
+                    "data_type": "int",
+                    "column_name": "customer_id",
+                    "column_comment": ""
+                    },...]
+                "table_schema": "db_anatoly",
+                "table_comment": ""
+            }
+        }
+        Where both arrays include only the schemas/tables relevant to the question.
+
+    MCP usage example:
+        - name: retrieve_relevant_schema_information
+          arguments: {
+            "input": "Which tables store customer address data?",
+          }
+        # Example output:
+        # {
+        #     "tables": ["schema1.customers", "schema2.tableB"],
+              "schemas": ["schema1","schema"]
+        # }
+    """
+    with _get_database_connection_cm(connection_id) as db_connection:
+        set_response = _execute_sql_tool(db_connection, "SET @response = NULL;")
+        if check_error(set_response):
+            return json.dumps({"error": f"Error with ML_RETRIEVE_SCHEMA: {set_response}"})
+
+        ml_retrieval_response = _execute_sql_tool(
+            db_connection,
+            f"CALL sys.ML_SQL_SCHEMA_METADATA(%s, NULL, TRUE, @response)",
+            params=[question],
+        )
+        if check_error(ml_retrieval_response):
+            return json.dumps({"error": f"Error with ML_RETRIEVE_SCHEMA: {ml_retrieval_response}"})
+
+        fetch_response = _execute_sql_tool(db_connection, "SELECT @response;")
+        if check_error(fetch_response):
+            return json.dumps({"error": f"Error with ML_RETRIEVE_SCHEMA: {fetch_response}"})
+
+        try:
+            response = json.loads(fetch_one(fetch_response))
+            return json.dumps(response)
+        except:
+            return json.dumps({"error": "Unexpected response format from ML_RETRIEVE_SCHEMA"})
+
+
 """
 Object store
 """
 
-
 def verify_compartment_access(compartments):
     access_report = {}
 
@@ -1034,7 +1114,9 @@ def object_storage_list_objects(namespace: str, bucket_name: str) -> str:
 
     return str(list_object_response.data.objects)
 
-
-if __name__ == "__main__":
-    # Initialize and run the server
+def main():
+    """Run the MCP server with CLI argument support."""
     mcp.run(transport="stdio")
+
+if __name__ == '__main__':
+    main()

src/mysql-mcp-server/test_mysql_mcp_server.py → src/mysql-mcp-server/oracle/mysql_mcp_server/tests/test_mysql_mcp_server.py

@@ -13,17 +13,16 @@
 import uuid
 from unittest import mock
 
-import mysql_mcp_server as m
-from utils import get_ssh_command, fill_config_defaults
-
+from oracle.mysql_mcp_server.utils import get_ssh_command, fill_config_defaults, Mode
+import oracle.mysql_mcp_server.server as m
 SKIP_ESTABLISHED = False
 
 
 def get_server_module():
     # Dynamically load the server module but shim fastmcp so @mcp.tool() returns the original function
     # This avoids FunctionTool non-callable wrappers when importing FastMCP from 'fastmcp'.
 
-    server_path = os.path.join(os.path.dirname(__file__), "mysql_mcp_server.py")
+    server_path = os.path.join(os.path.dirname(__file__), "../server.py")
     if not os.path.exists(server_path):
         raise FileNotFoundError(f"Server file not found at {server_path}")
 
@@ -93,13 +92,13 @@ def _is_valid_entry(e):
 
 
 try:
-    get_first_valid_connection_id(m.Mode.MYSQL_AI)
+    get_first_valid_connection_id(Mode.MYSQL_AI)
     SKIP_MYSQL_AI = False
 except:
     SKIP_MYSQL_AI = True
 
 try:
-    get_first_valid_connection_id(m.Mode.OCI)
+    get_first_valid_connection_id(Mode.OCI)
     SKIP_OCI = False
 except:
     SKIP_OCI = True
@@ -241,9 +240,11 @@ def test_load_uses_env_path_when_set(self):
         self.assertEqual(cfg, data)
 
     def test_load_uses_module_local_config_when_env_unset(self):
-        # Determine the utils module directory the function belongs to
-        utils_dir = os.path.dirname(os.path.abspath(m.load_mysql_config.__code__.co_filename))
-        local_path = os.path.join(utils_dir, "local_config.json")
+
+        module_dir = os.path.dirname(os.path.abspath(__file__))
+        local_path = os.path.abspath(
+        os.path.join(module_dir, "..", "..", "..", "local_config.json")
+    )
 
         def isfile_side_effect(path):
             return path == local_path

src/mysql-mcp-server/utils.py → src/mysql-mcp-server/oracle/mysql_mcp_server/utils.py

@@ -117,13 +117,13 @@ def fill_config_defaults(config : dict) -> dict:
 
     return config
 
-def load_mysql_config(config_path="config.json"):
+def load_mysql_config():
     """
     Load and validate the MySQL MCP server configuration file.
 
     Resolution:
       1) If MYSQL_MCP_CONFIG is set, load that absolute path.
-      2) Otherwise, load <module_dir>/local_config.json.
+      2) Otherwise, load <root__dir>/local_config.json.
 
     Notes:
       - No other fallbacks (e.g., config.json or CWD) are used.
@@ -139,13 +139,15 @@ def load_mysql_config(config_path="config.json"):
     env_path = os.environ.get('MYSQL_MCP_CONFIG')
     module_dir = os.path.dirname(os.path.abspath(__file__))
 
-    local_config_path = os.path.join(module_dir, "local_config.json")
+    local_config_path = os.path.abspath(
+        os.path.join(module_dir, "..", "..", "local_config.json")
+    )
     config_file = env_path if env_path else local_config_path
 
     if not os.path.isfile(config_file):
         raise Exception(
             f"Config file not found at {config_file}. "
-            "Set MYSQL_MCP_CONFIG to an absolute path or place local_config.json next to utils.py."
+            "Set MYSQL_MCP_CONFIG to an absolute path or create local_config.json in project root directory."
         )
 
     with open(config_file, "r") as f: