feat: 基于PostgreSQL官方解析器重构SQL安全验证机制，增强查询防护能力

go.mod

@@ -22,6 +22,7 @@ require (
 	github.com/ollama/ollama v0.11.4
 	github.com/panjf2000/ants/v2 v2.11.2
 	github.com/parquet-go/parquet-go v0.25.0
+	github.com/pganalyze/pg_query_go/v6 v6.1.0
 	github.com/pgvector/pgvector-go v0.3.0
 	github.com/qdrant/go-client v1.16.1
 	github.com/redis/go-redis/v9 v9.14.0

go.sum

@@ -135,9 +135,11 @@ github.com/golang-jwt/jwt/v5 v5.3.0 h1:pv4AsKCKKZuqlgs5sUmn4x8UlGa0kEVt/puTpKx9v
 github.com/golang-jwt/jwt/v5 v5.3.0/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=
 github.com/golang-migrate/migrate/v4 v4.19.0 h1:RcjOnCGz3Or6HQYEJ/EEVLfWnmw9KnoigPSjzhCuaSE=
 github.com/golang-migrate/migrate/v4 v4.19.0/go.mod h1:9dyEcu+hO+G9hPSw8AIg50yg622pXJsoHItQnDGZkI0=
+github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
 github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
 github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
 github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
@@ -249,6 +251,8 @@ github.com/parquet-go/parquet-go v0.25.0 h1:GwKy11MuF+al/lV6nUsFw8w8HCiPOSAx1/y8
 github.com/parquet-go/parquet-go v0.25.0/go.mod h1:OqBBRGBl7+llplCvDMql8dEKaDqjaFA/VAPw+OJiNiw=
 github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4=
 github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY=
+github.com/pganalyze/pg_query_go/v6 v6.1.0 h1:jG5ZLhcVgL1FAw4C/0VNQaVmX1SUJx71wBGdtTtBvls=
+github.com/pganalyze/pg_query_go/v6 v6.1.0/go.mod h1:nvTHIuoud6e1SfrUaFwHqT0i4b5Nr+1rPWVds3B5+50=
 github.com/pgvector/pgvector-go v0.3.0 h1:Ij+Yt78R//uYqs3Zk35evZFvr+G0blW0OUN+Q2D1RWc=
 github.com/pgvector/pgvector-go v0.3.0/go.mod h1:duFy+PXWfW7QQd5ibqutBO4GxLsUZ9RVXhFZGIBsWSA=
 github.com/pierrec/lz4/v4 v4.1.21 h1:yOVMLb6qSIDP67pl/5F7RepeKYu/VmTyEXvuMI5d9mQ=
@@ -470,6 +474,8 @@ google.golang.org/genproto/googleapis/rpc v0.0.0-20251111163417-95abcf5c77ba h1:
 google.golang.org/genproto/googleapis/rpc v0.0.0-20251111163417-95abcf5c77ba/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
 google.golang.org/grpc v1.76.0 h1:UnVkv1+uMLYXoIz6o7chp59WfQUYA2ex/BXQ9rHZu7A=
 google.golang.org/grpc v1.76.0/go.mod h1:Ju12QI8M6iQJtbcsV+awF5a4hfJMLi4X0JLo94ULZ6c=
+google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
+google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
 google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE=
 google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

internal/agent/tools/database_query.go

@@ -9,9 +9,82 @@ import (
 
 	"github.com/Tencent/WeKnora/internal/logger"
 	"github.com/Tencent/WeKnora/internal/types"
+	pg_query "github.com/pganalyze/pg_query_go/v6"
 	"gorm.io/gorm"
 )
 
+// SQLSecurityValidator provides comprehensive SQL injection protection using PostgreSQL's official parser
+type SQLSecurityValidator struct {
+	allowedTables    map[string]bool
+	allowedFunctions map[string]bool
+	tenantID         uint64
+}
+
+// NewSQLSecurityValidator creates a new SQL security validator
+func NewSQLSecurityValidator(tenantID uint64) *SQLSecurityValidator {
+	return &SQLSecurityValidator{
+		allowedTables: map[string]bool{
+			"tenants":         true,
+			"knowledge_bases": true,
+			"knowledges":      true,
+			"sessions":        true,
+			"messages":        true,
+			"chunks":          true,
+			"embeddings":      true,
+			"models":          true,
+		},
+		// Whitelist of allowed SQL functions (aggregates and safe functions only)
+		allowedFunctions: map[string]bool{
+			// Aggregate functions
+			"count":            true,
+			"sum":              true,
+			"avg":              true,
+			"min":              true,
+			"max":              true,
+			"array_agg":        true,
+			"string_agg":       true,
+			"bool_and":         true,
+			"bool_or":          true,
+			"json_agg":         true,
+			"jsonb_agg":        true,
+			"json_object_agg":  true,
+			"jsonb_object_agg": true,
+			// Safe scalar functions
+			"coalesce":          true,
+			"nullif":            true,
+			"greatest":          true,
+			"least":             true,
+			"abs":               true,
+			"ceil":              true,
+			"floor":             true,
+			"round":             true,
+			"length":            true,
+			"lower":             true,
+			"upper":             true,
+			"trim":              true,
+			"ltrim":             true,
+			"rtrim":             true,
+			"substring":         true,
+			"concat":            true,
+			"concat_ws":         true,
+			"replace":           true,
+			"left":              true,
+			"right":             true,
+			"now":               true,
+			"current_date":      true,
+			"current_timestamp": true,
+			"date_trunc":        true,
+			"extract":           true,
+			"to_char":           true,
+			"to_date":           true,
+			"to_timestamp":      true,
+			"date_part":         true,
+			"age":               true,
+		},
+		tenantID: tenantID,
+	}
+}
+
 // DatabaseQueryTool allows AI to query the database with auto-injected tenant_id for security
 type DatabaseQueryTool struct {
 	BaseTool
@@ -247,62 +320,429 @@ func (t *DatabaseQueryTool) Execute(ctx context.Context, args map[string]interfa
 
 // validateAndSecureSQL validates the SQL query and injects tenant_id conditions
 func (t *DatabaseQueryTool) validateAndSecureSQL(sqlQuery string) (string, error) {
-	// Normalize SQL: trim and convert to lowercase for analysis
-	normalizedSQL := strings.TrimSpace(sqlQuery)
-	lowerSQL := strings.ToLower(normalizedSQL)
+	validator := NewSQLSecurityValidator(t.tenantID)
+	return validator.ValidateAndSecure(sqlQuery)
+}
 
-	// 1. Check if it's a SELECT statement
-	if !strings.HasPrefix(lowerSQL, "select") {
+// ValidateAndSecure performs comprehensive SQL validation using PostgreSQL's official parser
+func (v *SQLSecurityValidator) ValidateAndSecure(sqlQuery string) (string, error) {
+	// Phase 1: Basic input validation
+	if err := v.validateInput(sqlQuery); err != nil {
+		return "", err
+	}
+
+	// Phase 2: Parse SQL using PostgreSQL's official parser
+	result, err := pg_query.Parse(sqlQuery)
+	if err != nil {
+		return "", fmt.Errorf("SQL parse error: %v", err)
+	}
+
+	// Phase 3: Validate that we have exactly one statement
+	if len(result.Stmts) == 0 {
+		return "", fmt.Errorf("empty query")
+	}
+	if len(result.Stmts) > 1 {
+		return "", fmt.Errorf("multiple statements are not allowed")
+	}
+
+	stmt := result.Stmts[0].Stmt
+
+	// Phase 4: Ensure it's a SELECT statement
+	selectStmt := stmt.GetSelectStmt()
+	if selectStmt == nil {
 		return "", fmt.Errorf("only SELECT queries are allowed")
 	}
 
-	// 2. Check for dangerous keywords
-	dangerousKeywords := []string{
-		"drop", "delete", "insert", "update", "alter", "create",
-		"truncate", "replace", "execute", "exec", "grant", "revoke",
+	// Phase 5: Validate the SELECT statement recursively
+	tablesInQuery, err := v.validateSelectStmt(selectStmt)
+	if err != nil {
+		return "", err
 	}
-	for _, keyword := range dangerousKeywords {
-		// Use word boundary to avoid false positives (e.g., "description" contains "script")
-		pattern := fmt.Sprintf(`\b%s\b`, keyword)
-		matched, _ := regexp.MatchString(pattern, lowerSQL)
-		if matched {
-			return "", fmt.Errorf("dangerous keyword detected: %s", keyword)
+
+	// Phase 6: Normalize SQL (removes comments, standardizes format)
+	normalizedSQL, err := pg_query.Deparse(result)
+	if err != nil {
+		return "", fmt.Errorf("failed to normalize SQL: %v", err)
+	}
+
+	// Phase 7: Inject tenant_id conditions
+	securedSQL := v.injectTenantConditions(normalizedSQL, tablesInQuery)
+
+	return securedSQL, nil
+}
+
+// validateInput performs basic input validation
+func (v *SQLSecurityValidator) validateInput(sql string) error {
+	// Check for null bytes
+	if strings.Contains(sql, "\x00") {
+		return fmt.Errorf("invalid character in SQL query")
+	}
+
+	// Check length limits
+	if len(sql) < 6 {
+		return fmt.Errorf("SQL query too short")
+	}
+	if len(sql) > 4096 {
+		return fmt.Errorf("SQL query too long (max 4096 characters)")
+	}
+
+	return nil
+}
+
+// validateSelectStmt validates a SELECT statement and extracts table information
+func (v *SQLSecurityValidator) validateSelectStmt(stmt *pg_query.SelectStmt) (map[string]string, error) {
+	tablesInQuery := make(map[string]string) // table name -> alias
+
+	// Check for UNION/INTERSECT/EXCEPT (compound queries)
+	if stmt.Op != pg_query.SetOperation_SETOP_NONE {
+		return nil, fmt.Errorf("compound queries (UNION/INTERSECT/EXCEPT) are not allowed")
+	}
+
+	// Check for WITH clause (CTEs) - could be used for complex attacks
+	if stmt.WithClause != nil {
+		return nil, fmt.Errorf("WITH clause (CTEs) is not allowed")
+	}
+
+	// Check for INTO clause (SELECT INTO)
+	if stmt.IntoClause != nil {
+		return nil, fmt.Errorf("SELECT INTO is not allowed")
+	}
+
+	// Check for LOCKING clause (FOR UPDATE, etc.)
+	if len(stmt.LockingClause) > 0 {
+		return nil, fmt.Errorf("locking clauses (FOR UPDATE, etc.) are not allowed")
+	}
+
+	// Validate FROM clause
+	for _, fromItem := range stmt.FromClause {
+		if err := v.validateFromItem(fromItem, tablesInQuery); err != nil {
+			return nil, err
 		}
 	}
 
-	// 3. Check for allowed tables
-	allowedTables := []string{
-		"tenants", "knowledge_bases", "knowledges", "sessions",
-		"messages", "chunks", "embeddings", "models",
-	}
-
-	// Extract table names from FROM and JOIN clauses
-	tablePattern := regexp.MustCompile(`(?i)\b(?:from|join)\s+([a-z_]+)(?:\s+as\s+[a-z_]+|\s+[a-z_]+)?`)
-	matches := tablePattern.FindAllStringSubmatch(lowerSQL, -1)
-
-	tablesInQuery := make(map[string]bool)
-	for _, match := range matches {
-		if len(match) > 1 {
-			tableName := strings.ToLower(match[1])
-			tablesInQuery[tableName] = true
+	// Validate target list (SELECT columns)
+	for _, target := range stmt.TargetList {
+		if err := v.validateNode(target); err != nil {
+			return nil, err
 		}
 	}
 
-	// Verify all tables are allowed
-	for tableName := range tablesInQuery {
-		allowed := false
-		for _, allowedTable := range allowedTables {
-			if tableName == allowedTable {
-				allowed = true
-				break
+	// Validate WHERE clause
+	if stmt.WhereClause != nil {
+		if err := v.validateNode(stmt.WhereClause); err != nil {
+			return nil, err
+		}
+	}
+
+	// Validate GROUP BY clause
+	for _, groupBy := range stmt.GroupClause {
+		if err := v.validateNode(groupBy); err != nil {
+			return nil, err
+		}
+	}
+
+	// Validate HAVING clause
+	if stmt.HavingClause != nil {
+		if err := v.validateNode(stmt.HavingClause); err != nil {
+			return nil, err
+		}
+	}
+
+	// Validate ORDER BY clause
+	for _, sortBy := range stmt.SortClause {
+		if err := v.validateNode(sortBy); err != nil {
+			return nil, err
+		}
+	}
+
+	// Ensure at least one valid table is referenced
+	if len(tablesInQuery) == 0 {
+		return nil, fmt.Errorf("no valid table found in query")
+	}
+
+	return tablesInQuery, nil
+}
+
+// validateFromItem validates a FROM clause item
+func (v *SQLSecurityValidator) validateFromItem(node *pg_query.Node, tables map[string]string) error {
+	if node == nil {
+		return nil
+	}
+
+	// Handle RangeVar (simple table reference)
+	if rv := node.GetRangeVar(); rv != nil {
+		tableName := strings.ToLower(rv.Relname)
+
+		// Check for schema qualification (e.g., pg_catalog.pg_class)
+		if rv.Schemaname != "" {
+			schemaName := strings.ToLower(rv.Schemaname)
+			// Block all schema-qualified access except public
+			if schemaName != "public" {
+				return fmt.Errorf("access to schema '%s' is not allowed", rv.Schemaname)
 			}
 		}
-		if !allowed {
-			return "", fmt.Errorf("table not allowed: %s", tableName)
+
+		// Validate table name against whitelist
+		if !v.allowedTables[tableName] {
+			return fmt.Errorf("table not allowed: %s", rv.Relname)
+		}
+
+		// Get alias
+		alias := tableName
+		if rv.Alias != nil && rv.Alias.Aliasname != "" {
+			alias = strings.ToLower(rv.Alias.Aliasname)
+		}
+		tables[tableName] = alias
+		return nil
+	}
+
+	// Handle JoinExpr (JOIN)
+	if je := node.GetJoinExpr(); je != nil {
+		if err := v.validateFromItem(je.Larg, tables); err != nil {
+			return err
+		}
+		if err := v.validateFromItem(je.Rarg, tables); err != nil {
+			return err
+		}
+		if je.Quals != nil {
+			if err := v.validateNode(je.Quals); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+
+	// Handle RangeSubselect (subquery in FROM) - NOT ALLOWED
+	if node.GetRangeSubselect() != nil {
+		return fmt.Errorf("subqueries in FROM clause are not allowed")
+	}
+
+	// Handle RangeFunction (function in FROM) - NOT ALLOWED
+	if node.GetRangeFunction() != nil {
+		return fmt.Errorf("functions in FROM clause are not allowed")
+	}
+
+	return nil
+}
+
+// validateNode recursively validates AST nodes for security issues
+func (v *SQLSecurityValidator) validateNode(node *pg_query.Node) error {
+	if node == nil {
+		return nil
+	}
+
+	// Check for subqueries (SubLink)
+	if sl := node.GetSubLink(); sl != nil {
+		return fmt.Errorf("subqueries are not allowed")
+	}
+
+	// Check for function calls
+	if fc := node.GetFuncCall(); fc != nil {
+		return v.validateFuncCall(fc)
+	}
+
+	// Check for column references with schema
+	if cr := node.GetColumnRef(); cr != nil {
+		return v.validateColumnRef(cr)
+	}
+
+	// Check for type casts (could be used for attacks)
+	if tc := node.GetTypeCast(); tc != nil {
+		if err := v.validateNode(tc.Arg); err != nil {
+			return err
+		}
+		// Validate the target type
+		if tc.TypeName != nil {
+			typeName := v.getTypeName(tc.TypeName)
+			if strings.HasPrefix(strings.ToLower(typeName), "pg_") {
+				return fmt.Errorf("casting to system type '%s' is not allowed", typeName)
+			}
 		}
 	}
 
-	// 4. Inject tenant_id conditions for tables that have tenant_id
+	// Recursively check A_Expr (expressions)
+	if ae := node.GetAExpr(); ae != nil {
+		if err := v.validateNode(ae.Lexpr); err != nil {
+			return err
+		}
+		if err := v.validateNode(ae.Rexpr); err != nil {
+			return err
+		}
+	}
+
+	// Check BoolExpr (AND, OR, NOT)
+	if be := node.GetBoolExpr(); be != nil {
+		for _, arg := range be.Args {
+			if err := v.validateNode(arg); err != nil {
+				return err
+			}
+		}
+	}
+
+	// Check NullTest
+	if nt := node.GetNullTest(); nt != nil {
+		if err := v.validateNode(nt.Arg); err != nil {
+			return err
+		}
+	}
+
+	// Check CoalesceExpr
+	if ce := node.GetCoalesceExpr(); ce != nil {
+		for _, arg := range ce.Args {
+			if err := v.validateNode(arg); err != nil {
+				return err
+			}
+		}
+	}
+
+	// Check CaseExpr
+	if caseExpr := node.GetCaseExpr(); caseExpr != nil {
+		if err := v.validateNode(caseExpr.Arg); err != nil {
+			return err
+		}
+		for _, when := range caseExpr.Args {
+			if err := v.validateNode(when); err != nil {
+				return err
+			}
+		}
+		if err := v.validateNode(caseExpr.Defresult); err != nil {
+			return err
+		}
+	}
+
+	// Check CaseWhen
+	if cw := node.GetCaseWhen(); cw != nil {
+		if err := v.validateNode(cw.Expr); err != nil {
+			return err
+		}
+		if err := v.validateNode(cw.Result); err != nil {
+			return err
+		}
+	}
+
+	// Check ResTarget (SELECT list items)
+	if rt := node.GetResTarget(); rt != nil {
+		if err := v.validateNode(rt.Val); err != nil {
+			return err
+		}
+	}
+
+	// Check SortBy (ORDER BY items)
+	if sb := node.GetSortBy(); sb != nil {
+		if err := v.validateNode(sb.Node); err != nil {
+			return err
+		}
+	}
+
+	// Check List
+	if list := node.GetList(); list != nil {
+		for _, item := range list.Items {
+			if err := v.validateNode(item); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+// validateFuncCall validates a function call against the whitelist
+func (v *SQLSecurityValidator) validateFuncCall(fc *pg_query.FuncCall) error {
+	// Get function name
+	funcName := ""
+	for _, namePart := range fc.Funcname {
+		if s := namePart.GetString_(); s != nil {
+			funcName = strings.ToLower(s.Sval)
+		}
+	}
+
+	// Check for schema-qualified function calls
+	if len(fc.Funcname) > 1 {
+		// Get schema name
+		schemaName := ""
+		if s := fc.Funcname[0].GetString_(); s != nil {
+			schemaName = strings.ToLower(s.Sval)
+		}
+		// Block all schema-qualified function calls except pg_catalog for basic functions
+		if schemaName != "" && schemaName != "pg_catalog" {
+			return fmt.Errorf("schema-qualified function calls are not allowed: %s", schemaName)
+		}
+	}
+
+	// Block dangerous function prefixes
+	dangerousPrefixes := []string{
+		"pg_", "lo_", "dblink", "file_", "copy_",
+	}
+	for _, prefix := range dangerousPrefixes {
+		if strings.HasPrefix(funcName, prefix) {
+			return fmt.Errorf("function '%s' is not allowed (dangerous prefix)", funcName)
+		}
+	}
+
+	// Block specific dangerous functions
+	dangerousFunctions := map[string]bool{
+		"current_setting": true,
+		"set_config":      true,
+		"query_to_xml":    true,
+		"xpath":           true,
+		"xmlparse":        true,
+		"txid_current":    true,
+	}
+	if dangerousFunctions[funcName] {
+		return fmt.Errorf("function '%s' is not allowed", funcName)
+	}
+
+	// Check against whitelist
+	if !v.allowedFunctions[funcName] {
+		return fmt.Errorf("function not allowed: %s", funcName)
+	}
+
+	// Validate function arguments recursively
+	for _, arg := range fc.Args {
+		if err := v.validateNode(arg); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// validateColumnRef validates a column reference
+func (v *SQLSecurityValidator) validateColumnRef(cr *pg_query.ColumnRef) error {
+	// Check for system column access
+	for _, field := range cr.Fields {
+		if s := field.GetString_(); s != nil {
+			colName := strings.ToLower(s.Sval)
+			// Block access to system columns
+			systemColumns := []string{"xmin", "xmax", "cmin", "cmax", "ctid", "tableoid"}
+			for _, sysCol := range systemColumns {
+				if colName == sysCol {
+					return fmt.Errorf("access to system column '%s' is not allowed", colName)
+				}
+			}
+			// Block pg_ prefixed identifiers
+			if strings.HasPrefix(colName, "pg_") {
+				return fmt.Errorf("access to '%s' is not allowed", colName)
+			}
+		}
+	}
+	return nil
+}
+
+// getTypeName extracts the type name from a TypeName node
+func (v *SQLSecurityValidator) getTypeName(tn *pg_query.TypeName) string {
+	var parts []string
+	for _, name := range tn.Names {
+		if s := name.GetString_(); s != nil {
+			parts = append(parts, s.Sval)
+		}
+	}
+	return strings.Join(parts, ".")
+}
+
+// injectTenantConditions adds tenant_id filtering to the query
+func (v *SQLSecurityValidator) injectTenantConditions(sql string, tablesInQuery map[string]string) string {
+	// Tables that require tenant_id filtering
 	tablesWithTenantID := map[string]bool{
 		"tenants":         true,
 		"knowledge_bases": true,
@@ -311,65 +751,39 @@ func (t *DatabaseQueryTool) validateAndSecureSQL(sqlQuery string) (string, error
 		"chunks":          true,
 	}
 
-	// Build tenant_id injection
-	securedSQL := normalizedSQL
+	// Build tenant conditions
+	var conditions []string
+	for tableName, alias := range tablesInQuery {
+		if tablesWithTenantID[tableName] {
+			if tableName == "tenants" {
+				conditions = append(conditions, fmt.Sprintf("%s.id = %d", alias, v.tenantID))
+			} else {
+				conditions = append(conditions, fmt.Sprintf("%s.tenant_id = %d", alias, v.tenantID))
+			}
+		}
+	}
+
+	if len(conditions) == 0 {
+		return sql
+	}
+
+	tenantFilter := strings.Join(conditions, " AND ")
 
 	// Check if WHERE clause exists
-	wherePattern := regexp.MustCompile(`(?i)\bwhere\b`)
-	hasWhere := wherePattern.MatchString(securedSQL)
-
-	// Find table aliases
-	aliasPattern := regexp.MustCompile(`(?i)(?:from|join)\s+([a-z_]+)(?:\s+(?:as\s+)?([a-z_]+))?`)
-	aliasMatches := aliasPattern.FindAllStringSubmatch(lowerSQL, -1)
-
-	tableAliases := make(map[string]string) // table -> alias
-	for _, match := range aliasMatches {
-		if len(match) >= 3 && match[2] != "" {
-			// Has explicit alias
-			tableAliases[match[1]] = match[2]
-		} else if len(match) >= 2 {
-			// No alias, use table name itself
-			tableAliases[match[1]] = match[1]
-		}
+	wherePattern := regexp.MustCompile(`(?i)\bWHERE\b`)
+	if wherePattern.MatchString(sql) {
+		// Add to existing WHERE clause
+		return wherePattern.ReplaceAllString(sql, fmt.Sprintf("WHERE %s AND ", tenantFilter))
 	}
 
-	// Build tenant_id conditions
-	var tenantConditions []string
-	for tableName := range tablesInQuery {
-		if tablesWithTenantID[tableName] {
-			alias := tableAliases[tableName]
-			if alias == "" {
-				alias = tableName
-			}
-			// Special handling for tenants table - use id instead of tenant_id
-			if tableName == "tenants" {
-				tenantConditions = append(tenantConditions, fmt.Sprintf("%s.id = %d", alias, t.tenantID))
-			} else {
-				tenantConditions = append(tenantConditions, fmt.Sprintf("%s.tenant_id = %d", alias, t.tenantID))
-			}
-		}
+	// Add new WHERE clause before ORDER BY, GROUP BY, LIMIT, etc.
+	clausePattern := regexp.MustCompile(`(?i)\b(GROUP BY|ORDER BY|LIMIT|OFFSET|HAVING|FETCH)\b`)
+	if loc := clausePattern.FindStringIndex(sql); loc != nil {
+		return sql[:loc[0]] + fmt.Sprintf(" WHERE %s ", tenantFilter) + sql[loc[0]:]
 	}
 
-	if len(tenantConditions) > 0 {
-		tenantFilter := strings.Join(tenantConditions, " AND ")
-		if hasWhere {
-			// Add to existing WHERE clause
-			securedSQL = wherePattern.ReplaceAllString(securedSQL, fmt.Sprintf("WHERE %s AND ", tenantFilter))
-		} else {
-			// Add new WHERE clause before ORDER BY, GROUP BY, LIMIT, etc.
-			orderByPattern := regexp.MustCompile(`(?i)\b(group\s+by|order\s+by|limit|offset|having)\b`)
-			if orderByPattern.MatchString(securedSQL) {
-				securedSQL = orderByPattern.ReplaceAllStringFunc(securedSQL, func(match string) string {
-					return fmt.Sprintf(" WHERE %s %s", tenantFilter, match)
-				})
-			} else {
-				// Add WHERE clause at the end
-				securedSQL = fmt.Sprintf("%s WHERE %s", securedSQL, tenantFilter)
-			}
-		}
-	}
-
-	return securedSQL, nil
+	// Add WHERE clause at the end
+	return fmt.Sprintf("%s WHERE %s", sql, tenantFilter)
 }
 
 // formatQueryResults formats query results into readable text