DataPilot MCP Server

Instructions

Implementation Reference

• src/main.py:205-231 (handler)

  Primary MCP tool handler registered with @mcp.tool(). Handles context logging, error handling, and delegates to SnowflakeClient.get_table_sample.

  @mcp.tool()
  async def get_table_sample(table_name: str, limit: int = 10, ctx: Context = None) -> QueryResult:
      """Get a sample of data from a table"""
      await ctx.info(f"Getting sample data from table: {table_name} (limit: {limit})")
      
      try:
          client = await get_snowflake_client()
          result = await client.get_table_sample(table_name, limit)
          
          if result.success:
              await ctx.info(f"Retrieved {result.row_count} sample rows")
          else:
              await ctx.error(f"Failed to get sample: {result.error}")
          
          return result
          
      except Exception as e:
          logger.error(f"Error getting table sample: {str(e)}")
          await ctx.error(f"Failed to get table sample: {str(e)}")
          return QueryResult(
              success=False,
              data=[],
              columns=[],
              row_count=0,
              error=str(e)
          )

• src/snowflake_client.py:183-187 (helper)

  Core helper method in SnowflakeClient class that constructs the SQL query for sampling table data and executes it using the client's execute_query method.

  async def get_table_sample(self, table_name: str, limit: int = 10) -> QueryResult:
      """Get a sample of data from a table"""
      query = f"SELECT * FROM {table_name} LIMIT {limit}"
      return await self.execute_query(query)

• src/models.py:44-54 (schema)

  Pydantic model defining the output schema (QueryResult) returned by the get_table_sample tool.

  class QueryResult(BaseModel):
      """Result of a SQL query execution"""
      success: bool
      data: List[Dict[str, Any]]
      columns: List[str]
      row_count: int
      execution_time_ms: Optional[int] = None
      query_id: Optional[str] = None
      warehouse_used: Optional[str] = None
      error: Optional[str] = None