Clickzetta MCP Server

import argparse import asyncio import os import dotenv from clickzetta import connect from . import server def parse_args(): parser = argparse.ArgumentParser() # Add arguments parser.add_argument( "--allow_write", required=False, default=False, action="store_true", help="Allow write operations on the database" ) parser.add_argument("--log_dir", required=False, default=None, help="Directory to log to") parser.add_argument("--log_level", required=False, default="INFO", help="Logging level") parser.add_argument( "--prefetch", action="store_true", dest="prefetch", default=True, help="Prefetch table descriptions (when enabled, list_tables and describe_table are disabled)", ) parser.add_argument( "--no-prefetch", action="store_false", dest="prefetch", help="Don't prefetch table descriptions", ) parser.add_argument( "--exclude_tools", required=False, default=[], nargs="+", help="List of tools to exclude", ) # First, get all the arguments we don't know about args, unknown = parser.parse_known_args() # Create a dictionary to store our key-value pairs connection_args = {} # Iterate through unknown args in pairs for i in range(0, len(unknown), 2): if i + 1 >= len(unknown): break key = unknown[i] value = unknown[i + 1] # Make sure it's a keyword argument (starts with --) if key.startswith("--"): key = key[2:] # Remove the '--' connection_args[key] = value # Now we can add the known args to kwargs server_args = { "allow_write": args.allow_write, "log_dir": args.log_dir, "log_level": args.log_level, "prefetch": args.prefetch, "exclude_tools": args.exclude_tools, } return server_args, connection_args def main(): """Main entry point for the package.""" dotenv.load_dotenv() default_connection_args = { "service": "", # Service Endpoint "instance": "", # Instance ID "workspace": "", # Workspace/Database Name "schema": "", # Schema Name "username": "", # Login username "password": "", # Login password "vcluster": "" # Virtual Cluster Name } connection_args_from_env = { k: os.getenv("CLICKZETTA_" + k.upper()) for k in default_connection_args if os.getenv("CLICKZETTA_" + k.upper()) is not None } server_args, connection_args = parse_args() connection_args = {**connection_args_from_env, **connection_args} assert ( "workspace" in connection_args ), 'You must provide the account identifier as "--workspace" argument or "CLICKZETTA_WORKSPACE" environment variable. This MCP server can only operate on a single database.' assert ( "schema" in connection_args ), 'You must provide the username as "--schema" argument or "CLICKZETTA_SCHEMA" environment variable. This MCP server can only operate on a single schema.' asyncio.run( server.main( connection_args=connection_args, allow_write=server_args["allow_write"], log_dir=server_args["log_dir"], prefetch=server_args["prefetch"], log_level=server_args["log_level"], exclude_tools=server_args["exclude_tools"], ) ) # def main(): # """Main entry point for the package.""" # dotenv.load_dotenv() # # 显式列出所有可能的连接参数键 # possible_connection_keys = [ # '--workspace', # '--schema', # '--username', # '--password', # '--service', # '--instance', # '--vcluster' # ] # # 从环境变量获取连接参数 # connection_args_from_env = { # k: os.getenv(f"CLICKZETTA_{k.upper()}") # for k in possible_connection_keys # if os.getenv(f"CLICKZETTA_{k.upper()}") is not None # } # server_args, connection_args = parse_args() # # 合并环境变量参数和命令行参数 # connection_args = {**connection_args_from_env, **connection_args} # # 必须参数检查 # # assert 'workspace' in connection_args, ( # # 'You must provide the database identifier as "--workspace" argument ' # # 'or "CLICKZETTA_WORKSPACE" environment variable.' # # ) # # assert 'schema' in connection_args, ( # # 'You must provide the schema as "--schema" argument ' # # 'or "CLICKZETTA_SCHEMA" environment variable.' # # ) # asyncio.run( # server.main( # connection_args=connection_args, # allow_write=server_args["allow_write"], # log_dir=server_args["log_dir"], # prefetch=server_args["prefetch"], # log_level=server_args["log_level"], # exclude_tools=server_args["exclude_tools"], # ) # ) # Optionally expose other important items at package level __all__ = ["main", "server", "write_detector"] if __name__ == "__main__": main()