Zaturn

Instructions

Implementation Reference

• zaturn/tools/core.py:69-95 (handler)

  The handler function for the 'run_query' tool. Takes source_id and query, executes the query using query_utils.execute_query, and formats results as a markdown table.
  def run_query(self, source_id: Annotated[ str, Field(description='The data source to run the query on') ], query: Annotated[ str, Field(description='SQL query to run on the data source') ] ) -> str: """ Run query against specified source For both csv and parquet sources, use DuckDB SQL syntax Use 'CSV' as the table name for csv sources. Use 'PARQUET' as the table name for parquet sources. This will return a dataframe with the results. """ try: source = self.data_sources.get(source_id) if not source: return f"Source {source_id} Not Found" df = query_utils.execute_query(source, query) return df.to_markdown(index=False) except Exception as e: return str(e)
• zaturn/mcp/__init__.py:92-93 (registration)

  Registers 'run_query' (included in zaturn_tools.tools) as an MCP tool using FastMCP.add_tool(Tool.from_function(tool_function)) in the ZaturnMCP server setup.
  for tool_function in zaturn_tools.tools: zaturn_mcp.add_tool(Tool.from_function(tool_function))
• zaturn/tools/query_utils.py:74-155 (helper)

  Helper function execute_query that implements the actual query execution logic for various data sources (SQLite, MySQL, PostgreSQL, etc.) in read-only mode, called directly by the run_query handler.
  def execute_query(source: dict, query: str): """Run the query using the appropriate engine and read only config""" url = source['url'] match source['source_type']: case "sqlite": with sqlalchemy.create_engine(url).connect() as conn: conn.execute(sqlalchemy.text('PRAGMA query_only = ON;')) result = conn.execute(sqlalchemy.text(query)) return pd.DataFrame(result) case "mysql": engine = sqlalchemy.create_engine(url) with Session(engine) as session: session.autoflush = False session.autocommit = False session.flush = lambda *args: None session.execute(sqlalchemy.text('SET SESSION TRANSACTION READ ONLY;')) result = session.execute(sqlalchemy.text(query)) return pd.DataFrame(result) case "mssql": engine = sqlalchemy.create_engine(url) with Session(engine) as session: # no known way to ensure read-only here # please use read-only credentials result = session.execute(sqlalchemy.text(query)) return pd.DataFrame(result) case "postgresql": engine = sqlalchemy.create_engine(url) with engine.connect() as conn: conn = conn.execution_options( isolation_level="SERIALIZABLE", postgresql_readonly=True, postgresql_deferrable=True, ) with conn.begin(): result = conn.execute(sqlalchemy.text(query)) return pd.DataFrame(result) case "clickhouse": client = clickhouse_connect.get_client(dsn=url) client.query('SET readonly=1;') return client.query_df(query, use_extended_dtypes=False) case "duckdb": conn = duckdb.connect(url, read_only=True) return conn.execute(query).df() case "csv": conn = duckdb.connect(database=':memory:') conn.execute(f"CREATE VIEW CSV AS SELECT * FROM read_csv('{url}')") return conn.execute(query).df() case "parquet": conn = duckdb.connect(database=':memory:') conn.execute(f"CREATE VIEW PARQUET AS SELECT * FROM read_parquet('{url}')") return conn.execute(query).df() case "bigquery": credentials = None if config.BIGQUERY_SERVICE_ACCOUNT_FILE: credentials = service_account.Credentials.from_service_account_file( config.BIGQUERY_SERVICE_ACCOUNT_FILE, ) chunks = source['url'].split('://')[1].split('/') bq_client = bigquery.Client( credentials = credentials, default_query_job_config = bigquery.QueryJobConfig( default_dataset = f'{chunks[0]}.{chunks[1]}' ) ) query_job = bq_client.query(query) return query_job.result().to_dataframe() case _: raise Exception("Unsupported Source")
• zaturn/tools/core.py:12-16 (registration)

  Initial registration of the run_query method in the Core class tools list, which is later aggregated into higher-level tool collections.
  self.tools = [ self.list_data_sources, self.describe_table, self.run_query, ]
• zaturn/tools/__init__.py:6-10 (registration)

  Aggregates tools from core.Core (including run_query) into ZaturnTools.tools list, used by MCP server.
  def __init__(self, data_sources): self.tools = [ *core.Core(data_sources).tools, *visualizations.Visualizations(data_sources).tools, ]