DataPilot MCP Server

Instructions

Implementation Reference

• src/main.py:337-352 (handler)

  The MCP tool handler for 'suggest_query_optimizations'. It handles context logging, retrieves the OpenAI client, calls the optimization method, and returns the suggestions.

  @mcp.tool()
  async def suggest_query_optimizations(query: str, ctx: Context) -> str:
      """Get AI-powered suggestions for optimizing a SQL query"""
      await ctx.info(f"Analyzing query for optimization: {query[:100]}...")
      
      try:
          openai = await get_openai_client()
          optimizations = await openai.suggest_optimizations(query)
          await ctx.info("Generated optimization suggestions")
          return optimizations
          
      except Exception as e:
          logger.error(f"Error suggesting optimizations: {str(e)}")
          await ctx.error(f"Failed to suggest optimizations: {str(e)}")
          raise

• src/openai_client.py:128-170 (helper)

  The supporting method in OpenAIClient class that implements the core AI logic for query optimization using a tailored prompt and OpenAI chat completions.

  async def suggest_optimizations(self, query: str) -> str:
      """Suggest optimizations for a SQL query"""
      
      system_prompt = """
      You are a Snowflake SQL performance expert. Analyze the provided query and suggest optimizations.
      
      Consider:
      - Index usage and clustering keys
      - JOIN optimization
      - WHERE clause efficiency
      - Warehouse sizing recommendations
      - Query structure improvements
      - Snowflake-specific optimizations (clustering, materialized views, etc.)
      
      Provide specific, actionable recommendations.
      """
      
      user_prompt = f"""
      SQL Query to optimize:
      {query}
      
      Please provide optimization suggestions.
      """
      
      try:
          response = await self.client.chat.completions.create(
              model=self.model,
              messages=[
                  {"role": "system", "content": system_prompt},
                  {"role": "user", "content": user_prompt}
              ],
              temperature=0.2,
              max_tokens=800
          )
          
          optimization_suggestions = response.choices[0].message.content.strip()
          logger.info("Generated optimization suggestions")
          return optimization_suggestions
          
      except Exception as e:
          logger.error(f"Error generating optimizations: {str(e)}")
          raise Exception(f"Failed to generate optimization suggestions: {str(e)}")