StarTree MCP Server for Apache Pinot

import base64 from fnmatch import fnmatch import json import re from threading import Lock from typing import Any, Dict, Tuple import pandas as pd from pinotdb import connect import requests from .config import PinotConfig, get_logger, reload_table_filters_from_file logger = get_logger() def get_auth_credentials(config: PinotConfig) -> Tuple[str | None, str | None]: """Extract authentication credentials for PinotDB connection""" if config.token: if config.token.startswith("Bearer "): return "", config.token # Empty username, full Bearer token as password else: return "", config.token elif config.username and config.password: return config.username, config.password return None, None def test_connection_query(connection) -> None: """Test connection with a simple query""" test_cursor = connection.cursor() test_cursor.execute("SELECT 1") test_result = test_cursor.fetchall() logger.debug(f"Connection test successful: {test_result}") # URL pattern constants class PinotEndpoints: QUERY_SQL = "query/sql" TABLES = "tables" SCHEMAS = "schemas" TABLE_SIZE = "tables/{}/size" SEGMENTS = "segments/{}" SEGMENT_METADATA = "segments/{}/metadata" SEGMENT_DETAIL = "segments/{}_{}/{}/metadata?columns=*" TABLE_CONFIG = "tableConfigs/{}" def create_connection(config: PinotConfig) -> connect: """Create Pinot connection with proper authentication handling""" try: auth_username, auth_password = get_auth_credentials(config) logger.debug( f"Creating connection to {config.broker_host}:{config.broker_port} " f"with MSQE={config.use_msqe}" ) auth_method = "token" if config.token else "username/password" logger.debug(f"Database: {config.database}, Auth method: {auth_method}") connection = connect( host=config.broker_host, port=config.broker_port, path="/query/sql", scheme=config.broker_scheme, username=auth_username, password=auth_password, use_multistage_engine=config.use_msqe, database=config.database, extra_conn_args={ "timeout": config.query_timeout, "verify": True, "retries": 3, "backoff_factor": 1.0, }, ) test_connection_query(connection) return connection except Exception as e: logger.error(f"Failed to create Pinot connection: {e}") logger.error( f"Connection details - Host: {config.broker_host}, " f"Port: {config.broker_port}, Scheme: {config.broker_scheme}" ) raise class PinotClient: def __init__(self, config: PinotConfig): self.config = config self.insights: list[str] = [] self._conn = None # Store filters separately to avoid mutating config self._included_tables = config.included_tables self._config_lock = Lock() # For thread-safe filter updates def reload_table_filters(self) -> dict[str, Any]: """Reload table filters from the configured filter file without restarting. This allows dynamic updates to the table access list by: 1. Editing the YAML filter file 2. Calling this method to reload the configuration Returns: dict: Status information with previous and new filter counts Raises: ValueError: If no table filter file is configured FileNotFoundError: If the filter file doesn't exist yaml.YAMLError: If the file contains invalid YAML """ if not self.config.table_filter_file: raise ValueError( "No table filter file configured. " "Set PINOT_TABLE_FILTER_FILE to enable hot-reload." ) logger.info(f"Reloading table filters from {self.config.table_filter_file}") # Load new filters (validates file exists and parses YAML) new_filters = reload_table_filters_from_file(self.config.table_filter_file) # Atomically update the filters with lock with self._config_lock: old_filters = self._included_tables old_count = len(old_filters) if old_filters else 0 self._included_tables = new_filters new_count = len(new_filters) if new_filters else 0 logger.info(f"Table filters reloaded: {old_count} -> {new_count} tables") return { "status": "success", "message": "Table filters reloaded successfully", "previous_filter_count": old_count, "new_filter_count": new_count, "previous_filters": old_filters, "new_filters": new_filters, } def _create_auth_headers(self) -> Dict[str, str]: """Create HTTP headers with authentication based on configuration""" headers = {"accept": "application/json", "Content-Type": "application/json"} if self.config.token: headers["Authorization"] = self.config.token elif self.config.username and self.config.password: creds_str = f"{self.config.username}:{self.config.password}" credentials = base64.b64encode(creds_str.encode()).decode() headers["Authorization"] = f"Basic {credentials}" if self.config.database: headers["database"] = self.config.database return headers def http_request( self, url: str, method: str = "GET", json_data: Dict = None, ) -> requests.Response: """Make HTTP request with authentication headers and timeout handling""" headers = self._create_auth_headers() try: if method.upper() == "POST": response = requests.post( url, headers=headers, json=json_data, timeout=( self.config.connection_timeout, self.config.request_timeout, ), verify=True, ) else: response = requests.get( url, headers=headers, timeout=( self.config.connection_timeout, self.config.request_timeout, ), verify=True, ) response.raise_for_status() return response except requests.exceptions.Timeout: logger.error(f"HTTP request timeout for {url}") raise except Exception as e: logger.error(f"HTTP request failed for {url}: {e}") raise def get_connection(self): """Get or create a reusable connection""" try: if self._conn is None: self._conn = create_connection(self.config) else: # Test if connection is still alive test_connection_query(self._conn) return self._conn except Exception as e: logger.warning(f"Connection test failed, creating new connection: {e}") self._conn = create_connection(self.config) return self._conn def test_connection(self) -> dict[str, Any]: """Test the connection and return diagnostic information""" result = { "connection_test": False, "query_test": False, "tables_test": False, "error": None, "config": { "broker_host": self.config.broker_host, "broker_port": self.config.broker_port, "broker_scheme": self.config.broker_scheme, "controller_url": self.config.controller_url, "database": self.config.database, "use_msqe": self.config.use_msqe, "has_token": bool(self.config.token), "has_username": bool(self.config.username), "timeout_config": { "connection": self.config.connection_timeout, "request": self.config.request_timeout, "query": self.config.query_timeout, }, }, } try: # Test basic connection conn = self.get_connection() result["connection_test"] = True # Test simple query curs = conn.cursor() curs.execute("SELECT 1 as test_column") test_result = curs.fetchall() result["query_test"] = True result["query_result"] = test_result # Test tables listing tables = self.get_tables() result["tables_test"] = True result["tables_count"] = len(tables) result["sample_tables"] = tables[:5] if tables else [] except Exception as e: result["error"] = str(e) logger.error(f"Connection test failed: {e}") return result def execute_query_http(self, query: str) -> list[dict[str, Any]]: """Alternative query execution using HTTP requests directly to broker""" broker_url = f"{self.config.broker_scheme}://{self.config.broker_host}:{self.config.broker_port}/{PinotEndpoints.QUERY_SQL}" logger.debug(f"Executing query via HTTP: {query[:100]}...") payload = { "sql": query, "queryOptions": f"timeoutMs={self.config.query_timeout * 1000}", } response = self.http_request(broker_url, "POST", payload) result_data = response.json() # Check for query errors in response if "exceptions" in result_data and result_data["exceptions"]: raise Exception(f"Query error: {result_data['exceptions']}") # Parse the result into pandas-like format if "resultTable" in result_data: columns = result_data["resultTable"]["dataSchema"]["columnNames"] rows = result_data["resultTable"]["rows"] # Convert to list of dictionaries result = [dict(zip(columns, row)) for row in rows] logger.debug(f"HTTP query returned {len(result)} rows") return result else: logger.warning("No resultTable in response, returning empty result") return [] def execute_query( self, query: str, params: dict[str, Any] | None = None, ) -> list[dict[str, Any]]: logger.debug(f"Executing query: {query[:100]}...") # Log first 100 chars # Validate table access authorization self._validate_table_access(query) # Use HTTP as primary method since it works reliably with authenticated clusters try: return self.execute_query_http(query) except Exception as e: logger.warning(f"HTTP query failed: {e}, trying PinotDB fallback") try: return self.execute_query_pinotdb(query, params) except Exception as pinotdb_error: error_msg = ( f"Both HTTP and PinotDB queries failed. " f"HTTP: {e}, PinotDB: {pinotdb_error}" ) logger.error(error_msg) raise def preprocess_query(self, query: str) -> str: """Preprocess query by removing database prefix and adding timeout options""" # Remove database prefix if present if self.config.database and f"{self.config.database}." in query: query = query.replace(f"{self.config.database}.", "") logger.debug(f"Removed database prefix, query now: {query[:100]}...") # Add query timeout hint if not present if "SET timeoutMs" not in query.upper() and "OPTION" not in query.upper(): timeout_ms = self.config.query_timeout * 1000 # Convert to milliseconds if query.strip().endswith(";"): query = query.rstrip(";") query = f"{query} OPTION(timeoutMs={timeout_ms})" logger.debug(f"Added timeout option: {timeout_ms}ms") return query def execute_query_pinotdb( self, query: str, params: dict[str, Any] | None = None, ) -> list[dict[str, Any]]: """Original pinotdb-based query execution""" logger.debug(f"Executing query via PinotDB: {query[:100]}...") try: current_conn = self.get_connection() curs = current_conn.cursor() query = self.preprocess_query(query) logger.debug(f"Final query: {query}") curs.execute(query) # Get column names and fetch results columns = [item[0] for item in curs.description] if curs.description else [] df = pd.DataFrame(curs.fetchall(), columns=columns) result = df.to_dict(orient="records") logger.debug(f"Query executed successfully, returned {len(result)} rows") return result except Exception as e: logger.error(f"Query execution failed: {e}") logger.error(f"Query was: {query}") # Reset connection on error self._conn = None raise def _matches_patterns(self, table: str, patterns: list[str]) -> bool: """Check if table matches any pattern.""" return any(fnmatch(table, pattern) for pattern in patterns) def _is_table_filtering_enabled(self) -> bool: """Check if table filtering is configured and enabled. Returns: bool: True if filtering is enabled (included_tables is configured), False otherwise (None, empty list, or any falsy value) """ return bool(self._included_tables) def _extract_sql_table_names(self, query: str) -> list[str]: """Extract table names from a SQL query. Handles table references in FROM, JOIN, and subquery clauses. Supports quoted identifiers (double quotes, backticks). Args: query: SQL query string Returns: list[str]: Unique list of table names found in the query """ # Remove comments and normalize whitespace query = re.sub(r"--.*?$", "", query, flags=re.MULTILINE) query = re.sub(r"/\*.*?\*/", "", query, flags=re.DOTALL) query = " ".join(query.split()) matches = [] # Pattern 1: Unquoted tables (after FROM/JOIN or comma-separated) # Matches: FROM table, JOIN table, table1, table2 # Uses negative lookahead to exclude SQL keywords (LEFT, RIGHT, INNER, etc.) unquoted_pattern = ( r"(?:\b(?:FROM|JOIN)\s+|,\s*)" r"(?:[\w.]+\.)?" r"(?!(?:LEFT|RIGHT|INNER|OUTER|FULL|CROSS|ON|WHERE|GROUP|ORDER|" r"HAVING|LIMIT)\b)" r"(\w+)" ) matches.extend(re.findall(unquoted_pattern, query, re.IGNORECASE)) # Pattern 2: Double-quoted tables (after FROM/JOIN or comma-separated) # Matches: FROM "table name", "quoted_table", "another table" double_quoted_pattern = r'(?:\b(?:FROM|JOIN)\s+|,\s*)(?:[\w.]+\.)?"([^"]+)"' matches.extend(re.findall(double_quoted_pattern, query, re.IGNORECASE)) # Pattern 3: Backtick-quoted tables (after FROM/JOIN or comma-separated) # Matches: FROM `table_name`, `quoted table`, `another table` backtick_pattern = r"(?:\b(?:FROM|JOIN)\s+|,\s*)(?:[\w.]+\.)?`([^`]+)`" matches.extend(re.findall(backtick_pattern, query, re.IGNORECASE)) return list(set(matches)) def _validate_table_name_access(self, table_name: str) -> None: """Validate that a table name is allowed by filtering rules. Args: table_name: Table name to validate Raises: ValueError: If table is not in included_tables filter """ if not self._is_table_filtering_enabled(): return if not self._matches_patterns(table_name, self._included_tables): allowed = ", ".join(self._included_tables) raise ValueError( f"Access denied to table '{table_name}'. Allowed tables: {allowed}" ) def _extract_and_validate_name_from_json(self, json_str: str, key: str) -> None: """Extract and validate table/schema name from JSON. Args: json_str: JSON string containing table or schema name key: JSON key to extract ("tableName" or "schemaName") Raises: ValueError: If name extraction fails or access is denied """ try: data = json.loads(json_str) name = data.get(key) if not name: raise ValueError(f"Missing required field '{key}' in JSON") self._validate_table_name_access(name) except json.JSONDecodeError as e: raise ValueError(f"Invalid JSON: {e}") def _validate_table_access(self, query: str) -> None: """Validate that query only accesses allowed tables. Args: query: SQL query string to validate Raises: ValueError: If query references tables not in included_tables filter """ if not self._is_table_filtering_enabled(): return table_names = self._extract_sql_table_names(query) if not table_names: return unauthorized_tables = [ table for table in table_names if not self._matches_patterns(table, self._included_tables) ] if unauthorized_tables: allowed = ", ".join(self._included_tables) unauthorized = ", ".join(unauthorized_tables) raise ValueError( f"Query references unauthorized tables: {unauthorized}. " f"Allowed tables: {allowed}" ) def _filter_tables(self, tables: list[str]) -> list[str]: """Filter tables based on included_tables configuration.""" if not tables or not self._is_table_filtering_enabled(): return tables return [t for t in tables if self._matches_patterns(t, self._included_tables)] def get_tables(self, params: dict[str, Any] | None = None) -> list[str]: url = f"{self.config.controller_url}/{PinotEndpoints.TABLES}" logger.debug(f"Fetching tables from: {url}") response = self.http_request(url) tables = response.json()["tables"] logger.debug(f"Successfully fetched {len(tables)} tables") return self._filter_tables(tables) def get_table_detail( self, tableName: str, params: dict[str, Any] | None = None, ) -> dict[str, Any]: self._validate_table_name_access(tableName) endpoint = PinotEndpoints.TABLE_SIZE.format(tableName) url = f"{self.config.controller_url}/{endpoint}" logger.debug(f"Fetching table details for {tableName} from: {url}") response = self.http_request(url) return response.json() def get_segment_metadata_detail( self, tableName: str, params: dict[str, Any] | None = None, ) -> dict[str, Any]: self._validate_table_name_access(tableName) endpoint = PinotEndpoints.SEGMENT_METADATA.format(tableName) url = f"{self.config.controller_url}/{endpoint}" logger.debug(f"Fetching segment metadata for {tableName} from: {url}") response = self.http_request(url) return response.json() def get_segments( self, tableName: str, params: dict[str, Any] | None = None, ) -> dict[str, Any]: self._validate_table_name_access(tableName) endpoint = PinotEndpoints.SEGMENTS.format(tableName) url = f"{self.config.controller_url}/{endpoint}" logger.debug(f"Fetching segments for {tableName} from: {url}") response = self.http_request(url) return response.json() def get_index_column_detail( self, tableName: str, segmentName: str, params: dict[str, Any] | None = None, ) -> dict[str, Any]: self._validate_table_name_access(tableName) for type_suffix in ["REALTIME", "OFFLINE"]: endpoint = PinotEndpoints.SEGMENT_DETAIL.format( tableName, type_suffix, segmentName ) url = f"{self.config.controller_url}/{endpoint}" logger.debug(f"Trying to fetch index column details from: {url}") try: response = self.http_request(url) return response.json() except Exception as e: error_msg = ( f"Failed to fetch index column details for " f"{tableName}_{type_suffix}/{segmentName}: {e}" ) logger.error(error_msg) continue raise ValueError("Index column detail not found") def get_tableconfig_schema_detail( self, tableName: str, params: dict[str, Any] | None = None, ) -> dict[str, Any]: self._validate_table_name_access(tableName) endpoint = PinotEndpoints.TABLE_CONFIG.format(tableName) url = f"{self.config.controller_url}/{endpoint}" logger.debug(f"Fetching table config for {tableName} from: {url}") response = self.http_request(url) return response.json() def create_schema( self, schemaJson: str, override: bool = True, force: bool = False, ) -> dict[str, Any]: self._extract_and_validate_name_from_json(schemaJson, "schemaName") url = f"{self.config.controller_url}/{PinotEndpoints.SCHEMAS}" params = {"override": str(override).lower(), "force": str(force).lower()} headers = self._create_auth_headers() headers["Content-Type"] = "application/json" response = requests.post( url, headers=headers, params=params, data=schemaJson, timeout=(self.config.connection_timeout, self.config.request_timeout), verify=True, ) response.raise_for_status() try: return response.json() except requests.exceptions.JSONDecodeError: return { "status": "success", "message": "Schema creation request processed.", "response_body": response.text, } def update_schema( self, schemaName: str, schemaJson: str, reload: bool = False, force: bool = False, ) -> dict[str, Any]: self._validate_table_name_access(schemaName) url = f"{self.config.controller_url}/{PinotEndpoints.SCHEMAS}/{schemaName}" params = {"reload": str(reload).lower(), "force": str(force).lower()} headers = self._create_auth_headers() headers["Content-Type"] = "application/json" response = requests.put( url, headers=headers, params=params, data=schemaJson, timeout=(self.config.connection_timeout, self.config.request_timeout), verify=True, ) response.raise_for_status() try: return response.json() except requests.exceptions.JSONDecodeError: return { "status": "success", "message": "Schema update request processed.", "response_body": response.text, } def get_schema(self, schemaName: str) -> dict[str, Any]: self._validate_table_name_access(schemaName) url = f"{self.config.controller_url}/{PinotEndpoints.SCHEMAS}/{schemaName}" headers = self._create_auth_headers() response = requests.get( url, headers=headers, timeout=(self.config.connection_timeout, self.config.request_timeout), verify=True, ) response.raise_for_status() return response.json() def create_table_config( self, tableConfigJson: str, validationTypesToSkip: str | None = None, ) -> dict[str, Any]: self._extract_and_validate_name_from_json(tableConfigJson, "tableName") url = f"{self.config.controller_url}/{PinotEndpoints.TABLES}" params: dict[str, str] = {} if validationTypesToSkip: params["validationTypesToSkip"] = validationTypesToSkip headers = self._create_auth_headers() headers["Content-Type"] = "application/json" response = requests.post( url, headers=headers, params=params, data=tableConfigJson, timeout=(self.config.connection_timeout, self.config.request_timeout), verify=True, ) response.raise_for_status() try: return response.json() except requests.exceptions.JSONDecodeError: return { "status": "success", "message": "Table config creation request processed.", "response_body": response.text, } def update_table_config( self, tableName: str, tableConfigJson: str, validationTypesToSkip: str | None = None, ) -> dict[str, Any]: self._validate_table_name_access(tableName) url = f"{self.config.controller_url}/{PinotEndpoints.TABLES}/{tableName}" params: dict[str, str] = {} if validationTypesToSkip: params["validationTypesToSkip"] = validationTypesToSkip headers = self._create_auth_headers() headers["Content-Type"] = "application/json" response = requests.put( url, headers=headers, params=params, data=tableConfigJson, timeout=(self.config.connection_timeout, self.config.request_timeout), verify=True, ) response.raise_for_status() try: return response.json() except requests.exceptions.JSONDecodeError: return { "status": "success", "message": "Table config update request processed.", "response_body": response.text, } def get_table_config( self, tableName: str, tableType: str | None = None, ) -> dict[str, Any]: self._validate_table_name_access(tableName) url = f"{self.config.controller_url}/{PinotEndpoints.TABLES}/{tableName}" params: dict[str, str] = {} if tableType: params["type"] = tableType headers = self._create_auth_headers() response = requests.get( url, headers=headers, params=params, timeout=(self.config.connection_timeout, self.config.request_timeout), verify=True, ) response.raise_for_status() raw_response = response.json() if tableType and tableType.upper() in raw_response: return raw_response[tableType.upper()] if not tableType and ("OFFLINE" in raw_response or "REALTIME" in raw_response): return raw_response return raw_response