diff --git a/pixeltable-mcp/base-sdk/tools.py b/pixeltable-mcp/base-sdk/tools.py index 028ca2bda..8724446be 100644 --- a/pixeltable-mcp/base-sdk/tools.py +++ b/pixeltable-mcp/base-sdk/tools.py @@ -1,9 +1,135 @@ +import ast +import operator +import re + import pixeltable as pxt from mcp.server.fastmcp import FastMCP mcp = FastMCP("Pixeltable") +# Safe expression evaluator to avoid eval()/exec() with untrusted input. +# Only allows attribute access on the table object, basic operators, and literals. + +_SAFE_OPERATORS = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.FloorDiv: operator.floordiv, + ast.Mod: operator.mod, + ast.Eq: operator.eq, + ast.NotEq: operator.ne, + ast.Lt: operator.lt, + ast.LtE: operator.le, + ast.Gt: operator.gt, + ast.GtE: operator.ge, + ast.USub: operator.neg, + ast.Not: operator.invert, # Use ~ (invert) instead of not for Pixeltable expressions +} + +# Validate identifiers to prevent injection through crafted attribute names +_IDENTIFIER_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$") + + +def _safe_eval(expression: str, table): + """Evaluate a simple expression safely against a Pixeltable table. + + Supports: table.column references, arithmetic/comparison operators, + string/number/boolean literals, and parenthesized sub-expressions. + + Raises ValueError for any unsupported or potentially dangerous construct. + """ + if len(expression) > 2000: + raise ValueError("Expression too long (max 2000 characters).") + try: + tree = ast.parse(expression, mode="eval") + except SyntaxError as e: + raise ValueError(f"Invalid expression syntax: {e}") from e + + return _eval_node(tree.body, {"table": table}) + + +def _eval_node(node, namespace): + """Recursively evaluate an AST node using only safe operations.""" + + # Literals: numbers, strings, booleans, None + if isinstance(node, ast.Constant): + if not isinstance(node.value, (int, float, str, bool, type(None))): + raise ValueError( + f"Unsupported literal type: {type(node.value).__name__}" + ) + return node.value + + # Variable lookup (only 'table' allowed; True/False/None are ast.Constant on 3.8+) + if isinstance(node, ast.Name): + if node.id == "table": + return namespace["table"] + raise ValueError( + f"Unsupported variable '{node.id}'. Only 'table' references are allowed." + ) + + # Attribute access: table.column_name + if isinstance(node, ast.Attribute): + value = _eval_node(node.value, namespace) + attr = node.attr + if not _IDENTIFIER_RE.match(attr): + raise ValueError(f"Invalid attribute name: '{attr}'") + if attr.startswith("_"): + raise ValueError(f"Access to private attribute '{attr}' is not allowed.") + if not hasattr(value, attr): + raise ValueError(f"Attribute '{attr}' not found.") + return getattr(value, attr) + + # Binary operations: +, -, *, /, //, %, ==, !=, <, <=, >, >= + if isinstance(node, ast.BinOp): + op_func = _SAFE_OPERATORS.get(type(node.op)) + if op_func is None: + raise ValueError(f"Unsupported operator: {type(node.op).__name__}") + left = _eval_node(node.left, namespace) + right = _eval_node(node.right, namespace) + return op_func(left, right) + + # Comparison operations (handles chained comparisons like a < b < c correctly) + if isinstance(node, ast.Compare): + left = _eval_node(node.left, namespace) + result = None + for op, comparator in zip(node.ops, node.comparators): + op_func = _SAFE_OPERATORS.get(type(op)) + if op_func is None: + raise ValueError(f"Unsupported comparison: {type(op).__name__}") + right = _eval_node(comparator, namespace) + cmp = op_func(left, right) + result = cmp if result is None else (result & cmp) + left = right + return result + + # Unary operations: -, not + if isinstance(node, ast.UnaryOp): + op_func = _SAFE_OPERATORS.get(type(node.op)) + if op_func is None: + raise ValueError(f"Unsupported unary operator: {type(node.op).__name__}") + return op_func(_eval_node(node.operand, namespace)) + + # Boolean operations: and, or + if isinstance(node, ast.BoolOp): + if isinstance(node.op, ast.And): + result = _eval_node(node.values[0], namespace) + for v in node.values[1:]: + result = result & _eval_node(v, namespace) + return result + if isinstance(node.op, ast.Or): + result = _eval_node(node.values[0], namespace) + for v in node.values[1:]: + result = result | _eval_node(v, namespace) + return result + + raise ValueError( + f"Unsupported expression type: {type(node).__name__}. " + "Only column references, operators, and literals are allowed." + ) + + @mcp.tool() def create_table(table_name: str, columns: dict[str, str]) -> str: """Create a table in Pixeltable. @@ -108,17 +234,15 @@ def add_computed_column(table_name: str, column_name: str, expression: str) -> s if table is None: return f"Error: Table {table_name} not found." - # Replace 'table.' with the actual table reference - modified_expr = expression.replace("table.", f"table.") - - # Construct the expression - column_expr = eval(modified_expr, {"table": table}) + column_expr = _safe_eval(expression, table) # Add the computed column with kwargs format kwargs = {column_name: column_expr} table.add_computed_column(**kwargs) return f"Computed column '{column_name}' added successfully to table '{table_name}'." + except ValueError as e: + return f"Error: Invalid expression: {e}" except Exception as e: return f"Error adding computed column: {str(e)}" @@ -143,11 +267,7 @@ def create_view(view_name: str, table_name: str, filter_expr: str = None) -> str return f"Error: Table {table_name} not found." if filter_expr: - # Replace 'table.' with the actual table reference - modified_expr = filter_expr.replace("table.", f"{table.name}.") - - # Use eval to convert the string expression to a Python expression - filter_condition = eval(modified_expr) + filter_condition = _safe_eval(filter_expr, table) # Create the view with the filter view = pxt.create_view(view_name, table.where(filter_condition)) @@ -156,6 +276,8 @@ def create_view(view_name: str, table_name: str, filter_expr: str = None) -> str view = pxt.create_view(view_name, table) return f"View '{view_name}' created successfully." + except ValueError as e: + return f"Error: Invalid filter expression: {e}" except Exception as e: return f"Error creating view: {str(e)}" @@ -196,12 +318,13 @@ def execute_query( # Apply where clause if provided if where_expr: - modified_expr = where_expr.replace("table.", f"{data_source.name}.") - where_condition = eval(modified_expr) + where_condition = _safe_eval(where_expr, data_source) query = query.where(where_condition) # Apply order by if provided if order_by_column: + if not _IDENTIFIER_RE.match(order_by_column): + return f"Error: Invalid column name: '{order_by_column}'" # Handle ordering on a specific column if hasattr(data_source, order_by_column): order_col = getattr(data_source, order_by_column) @@ -217,6 +340,8 @@ def execute_query( if select_columns: select_args = [] for col_name in select_columns: + if not _IDENTIFIER_RE.match(col_name): + return f"Error: Invalid column name: '{col_name}'" if hasattr(data_source, col_name): select_args.append(getattr(data_source, col_name)) else: @@ -232,25 +357,35 @@ def execute_query( result_str = result.to_pandas().to_string() return f"Query executed successfully:\n\n{result_str}" + except ValueError as e: + return f"Error: Invalid expression: {e}" except Exception as e: return f"Error executing query: {str(e)}" @mcp.tool() -def create_query(query_name: str, table_name: str, query_function: str) -> str: - """Create a named query in Pixeltable. +def create_query( + query_name: str, + table_name: str, + select_columns: list[str] = None, + where_expr: str = None, +) -> str: + """Create a named query as a persistent view in Pixeltable. Args: - query_name: The name of the query to create. + query_name: The name of the query to create (stored as a Pixeltable view). table_name: The name of the table the query will operate on. - query_function: A string representation of the query function body. - Should be a function that returns a query result. + select_columns: Optional list of column names to select. + If None, selects all columns. + where_expr: Optional filter expression as a string. + The expression should refer to columns using 'table.column_name'. Example: create_query( "get_active_users", "users", - "return users.where(users.is_active == True).select(users.name, users.email)" + select_columns=["name", "email"], + where_expr="table.is_active == True" ) """ try: @@ -258,16 +393,29 @@ def create_query(query_name: str, table_name: str, query_function: str) -> str: if table is None: return f"Error: Table {table_name} not found." - # Create a function definition - func_def = f""" - @{table.name}.query - def {query_name}(): - {query_function} - """ + # Build the query using Pixeltable's API directly instead of exec() + query = table + + if where_expr: + filter_condition = _safe_eval(where_expr, table) + query = query.where(filter_condition) + + if select_columns: + select_args = [] + for col_name in select_columns: + if not _IDENTIFIER_RE.match(col_name): + return f"Error: Invalid column name: '{col_name}'" + if hasattr(table, col_name): + select_args.append(getattr(table, col_name)) + else: + return f"Error: Column '{col_name}' not found in table '{table_name}'." + query = query.select(*select_args) - # Execute the function definition - exec(func_def) + # Persist the query as a Pixeltable view so it can be retrieved later + pxt.create_view(query_name, query) return f"Query '{query_name}' created successfully for table '{table_name}'." + except ValueError as e: + return f"Error: Invalid expression: {e}" except Exception as e: return f"Error creating query: {str(e)}"