Go로 PostgreSQL 프록시 만들기 (14) - AST 파서와 쿼리 방화벽

-- 문자열 파서로는 SELECT로 분류되지만, 실제로는 write 쿼리
WITH updated AS (
    UPDATE users SET active = false WHERE last_login < '2025-01-01'
    RETURNING id
)
SELECT count(*) FROM updated;

-- 서브쿼리가 포함된 복잡한 케이스
SELECT * FROM orders WHERE user_id IN (
    SELECT id FROM users WHERE deleted_at IS NOT NULL
);

-- DDL인지 DML인지 문자열로 구분하기 어려운 케이스
CREATE TABLE IF NOT EXISTS backup AS SELECT * FROM users;

func ParseSQL(query string) (*pg_query.ParseResult, error) {
    result, err := pg_query.Parse(query)
    if err != nil {
        return nil, fmt.Errorf("ast parse: %w", err)
    }
    return result, nil
}

func WalkNodes(node *pg_query.Node, visitor func(*pg_query.Node) bool) {
    if node == nil {
        return
    }
    if !visitor(node) {
        return
    }
    // 노드 타입별로 자식 노드를 재귀 순회
    switch n := node.Node.(type) {
    case *pg_query.Node_SelectStmt:
        for _, target := range n.SelectStmt.TargetList {
            WalkNodes(target, visitor)
        }
        for _, from := range n.SelectStmt.FromClause {
            WalkNodes(from, visitor)
        }
        WalkNodes(n.SelectStmt.WhereClause, visitor)
        // ...
    case *pg_query.Node_JoinExpr:
        WalkNodes(n.JoinExpr.Larg, visitor)
        WalkNodes(n.JoinExpr.Rarg, visitor)
        WalkNodes(n.JoinExpr.Quals, visitor)
    case *pg_query.Node_SubLink:
        WalkNodes(n.SubLink.Subselect, visitor)
    case *pg_query.Node_CommonTableExpr:
        WalkNodes(n.CommonTableExpr.Ctequery, visitor)
    // InsertStmt, UpdateStmt, DeleteStmt, ...
    }
}

func ClassifyAST(query string) QueryType {
    // 1. hint 주석 체크 (/*read*/, /*write*/)
    if hint := parseHint(query); hint != QueryUnknown {
        return hint
    }

    // 2. AST 파싱
    result, err := ParseSQL(query)
    if err != nil {
        // 파싱 실패 시 기존 문자열 파서로 fallback
        return Classify(query)
    }

    // 3. AST 순회하며 write 노드 탐지
    for _, stmt := range result.Stmts {
        if isWriteNode(stmt.Stmt) {
            return QueryWrite
        }
    }

    return QueryRead
}

func isWriteNode(node *pg_query.Node) bool {
    switch n := node.Node.(type) {
    case *pg_query.Node_InsertStmt:
        return true
    case *pg_query.Node_UpdateStmt:
        return true
    case *pg_query.Node_DeleteStmt:
        return true
    case *pg_query.Node_CreateStmt:
        return true
    case *pg_query.Node_AlterTableStmt:
        return true
    case *pg_query.Node_DropStmt:
        return true
    case *pg_query.Node_TruncateStmt:
        return true
    case *pg_query.Node_GrantStmt:
        return true
    case *pg_query.Node_IndexStmt:
        return true
    case *pg_query.Node_ViewStmt:
        return true
    // DDL 20+ 노드 타입 추가 커버
    case *pg_query.Node_SelectStmt:
        // CTE 내부에 write가 있는지 재귀 검사
        for _, cte := range n.SelectStmt.WithClause.GetCtes() {
            if cteExpr, ok := cte.Node.(*pg_query.Node_CommonTableExpr); ok {
                if isWriteNode(cteExpr.CommonTableExpr.Ctequery) {
                    return true
                }
            }
        }
        return false
    default:
        return false
    }
}

WITH x AS (UPDATE accounts SET balance = 0 RETURNING id)
SELECT * FROM x;

func ExtractTablesAST(query string) []string {
    result, err := ParseSQL(query)
    if err != nil {
        return ExtractTables(query) // fallback
    }

    tables := make(map[string]struct{})
    for _, stmt := range result.Stmts {
        WalkNodes(stmt.Stmt, func(node *pg_query.Node) bool {
            if rangeVar := extractRangeVar(node); rangeVar != "" {
                tables[rangeVar] = struct{}{}
            }
            return true
        })
        // CTE 내부의 write 테이블도 추출
        extractCTEWriteTables(stmt.Stmt, tables)
    }

    result := make([]string, 0, len(tables))
    for t := range tables {
        result = append(result, t)
    }
    return result
}

func CheckFirewall(query string, cfg FirewallConfig) FirewallResult {
    result, err := ParseSQL(query)
    if err != nil {
        // fail-open: 파싱 불가 시 허용
        return FirewallResult{Allowed: true}
    }

    for _, stmt := range result.Stmts {
        if reason := checkStatement(stmt.Stmt, cfg); reason != "" {
            return FirewallResult{
                Allowed: false,
                Reason:  reason,
            }
        }
    }

    return FirewallResult{Allowed: true}
}

func checkStatement(node *pg_query.Node, cfg FirewallConfig) string {
    switch n := node.Node.(type) {
    case *pg_query.Node_DeleteStmt:
        if cfg.BlockDeleteWithoutWhere && n.DeleteStmt.WhereClause == nil {
            return "DELETE without WHERE clause is blocked"
        }
    case *pg_query.Node_UpdateStmt:
        if cfg.BlockUpdateWithoutWhere && n.UpdateStmt.WhereClause == nil {
            return "UPDATE without WHERE clause is blocked"
        }
    case *pg_query.Node_DropStmt:
        if cfg.BlockDropTable {
            return "DROP TABLE is blocked"
        }
    case *pg_query.Node_TruncateStmt:
        if cfg.BlockTruncate {
            return "TRUNCATE is blocked"
        }
    }
    return ""
}

func SemanticCacheKey(query string) uint64 {
    tree, err := pg_query.Parse(query)
    if err != nil {
        return CacheKey(query) // fallback
    }
    deparsed, err := pg_query.Deparse(tree)
    if err != nil {
        return CacheKey(query)
    }
    h := fnv.New64a()
    h.Write([]byte(deparsed))
    return h.Sum64()
}

SELECT id, name FROM users WHERE age > 25;
select  id, name  from  users  where  age > 25;

-- 입력
SELECT * FROM users WHERE id = 42 AND name = 'alice';
-- 정규화 결과
SELECT * FROM users WHERE id = $1 AND name = $2;

routing:
  ast_parser: true      # AST 기반 쿼리 분류 활성화

firewall:
  enabled: true
  block_delete_without_where: true
  block_update_without_where: true
  block_drop_table: false       # 필요에 따라 활성화
  block_truncate: false         # 필요에 따라 활성화