package clickhouse import ( "fmt" "regexp" "strings" "github.com/go-kratos/kratos/v2/log" ) type QueryBuilder struct { table string columns []string distinct bool conditions []string orderBy []string groupBy []string having []string joins []string with []string union []string limit int offset int params []interface{} // 用于存储参数 useIndex string // 索引提示 cacheResult bool // 是否缓存查询结果 debug bool // 是否启用调试 log *log.Helper } // NewQueryBuilder 创建一个新的 QueryBuilder 实例 func NewQueryBuilder(table string, log *log.Helper) *QueryBuilder { return &QueryBuilder{ log: log, table: table, params: []interface{}{}, } } // EnableDebug 启用调试模式 func (qb *QueryBuilder) EnableDebug() *QueryBuilder { qb.debug = true return qb } // logDebug 打印调试信息 func (qb *QueryBuilder) logDebug(message string) { if qb.debug { qb.log.Debug("[QueryBuilder Debug]:", message) } } // Select 设置查询的列 func (qb *QueryBuilder) Select(columns ...string) *QueryBuilder { for _, column := range columns { if !isValidIdentifier(column) { panic("Invalid column name") } } qb.columns = columns return qb } // Distinct 设置 DISTINCT 查询 func (qb *QueryBuilder) Distinct() *QueryBuilder { qb.distinct = true return qb } // Where 添加查询条件并支持参数化 func (qb *QueryBuilder) Where(condition string, args ...interface{}) *QueryBuilder { if !isValidCondition(condition) { panic("Invalid condition") } qb.conditions = append(qb.conditions, condition) qb.params = append(qb.params, args...) return qb } // OrderBy 设置排序条件 func (qb *QueryBuilder) OrderBy(order string) *QueryBuilder { qb.orderBy = append(qb.orderBy, order) return qb } // GroupBy 设置分组条件 func (qb *QueryBuilder) GroupBy(columns ...string) *QueryBuilder { qb.groupBy = append(qb.groupBy, columns...) return qb } // Having 添加分组后的过滤条件并支持参数化 func (qb *QueryBuilder) Having(condition string, args ...interface{}) *QueryBuilder { qb.having = append(qb.having, condition) qb.params = append(qb.params, args...) return qb } // Join 添加 JOIN 操作 func (qb *QueryBuilder) Join(joinType, table, onCondition string) *QueryBuilder { join := fmt.Sprintf("%s JOIN %s ON %s", joinType, table, onCondition) qb.joins = append(qb.joins, join) return qb } // With 添加 WITH 子句 func (qb *QueryBuilder) With(expression string) *QueryBuilder { qb.with = append(qb.with, expression) return qb } // Union 添加 UNION 操作 func (qb *QueryBuilder) Union(query string) *QueryBuilder { qb.union = append(qb.union, query) return qb } // Limit 设置查询结果的限制数量 func (qb *QueryBuilder) Limit(limit int) *QueryBuilder { qb.limit = limit return qb } // Offset 设置查询结果的偏移量 func (qb *QueryBuilder) Offset(offset int) *QueryBuilder { qb.offset = offset return qb } // UseIndex 设置索引提示 func (qb *QueryBuilder) UseIndex(index string) *QueryBuilder { qb.useIndex = index return qb } // CacheResult 启用查询结果缓存 func (qb *QueryBuilder) CacheResult() *QueryBuilder { qb.cacheResult = true return qb } // ArrayJoin 添加 ARRAY JOIN 子句 func (qb *QueryBuilder) ArrayJoin(expression string) *QueryBuilder { qb.joins = append(qb.joins, fmt.Sprintf("ARRAY JOIN %s", expression)) return qb } // Final 添加 FINAL 修饰符 func (qb *QueryBuilder) Final() *QueryBuilder { qb.table = fmt.Sprintf("%s FINAL", qb.table) return qb } // Sample 添加 SAMPLE 子句 func (qb *QueryBuilder) Sample(sampleRate float64) *QueryBuilder { qb.table = fmt.Sprintf("%s SAMPLE %f", qb.table, sampleRate) return qb } // LimitBy 添加 LIMIT BY 子句 func (qb *QueryBuilder) LimitBy(limit int, columns ...string) *QueryBuilder { qb.limit = limit qb.orderBy = append(qb.orderBy, fmt.Sprintf("LIMIT BY %d (%s)", limit, strings.Join(columns, ", "))) return qb } // PreWhere 添加 PREWHERE 子句 func (qb *QueryBuilder) PreWhere(condition string, args ...interface{}) *QueryBuilder { qb.conditions = append([]string{condition}, qb.conditions...) qb.params = append(args, qb.params...) return qb } // Format 添加 FORMAT 子句 func (qb *QueryBuilder) Format(format string) *QueryBuilder { qb.union = append(qb.union, fmt.Sprintf("FORMAT %s", format)) return qb } // Build 构建最终的 SQL 查询 func (qb *QueryBuilder) Build() (string, []interface{}) { query := "" if qb.cacheResult { query += "/* CACHE */ " } query += "SELECT " if qb.distinct { query += "DISTINCT " } query += qb.buildColumns() query += fmt.Sprintf(" FROM %s", qb.table) if qb.useIndex != "" { query += fmt.Sprintf(" USE INDEX (%s)", qb.useIndex) } if len(qb.conditions) > 0 { query += fmt.Sprintf(" WHERE %s", strings.Join(qb.conditions, " AND ")) } if len(qb.groupBy) > 0 { query += fmt.Sprintf(" GROUP BY %s", strings.Join(qb.groupBy, ", ")) } if len(qb.having) > 0 { query += fmt.Sprintf(" HAVING %s", strings.Join(qb.having, " AND ")) } if len(qb.orderBy) > 0 { query += fmt.Sprintf(" ORDER BY %s", strings.Join(qb.orderBy, ", ")) } if qb.limit > 0 { query += fmt.Sprintf(" LIMIT %d", qb.limit) } if qb.offset > 0 { query += fmt.Sprintf(" OFFSET %d", qb.offset) } return query, qb.params } func (qb *QueryBuilder) buildColumns() string { if len(qb.columns) == 0 { return "*" } return strings.Join(qb.columns, ", ") } // isValidIdentifier 验证表名或列名是否合法 func isValidIdentifier(identifier string) bool { // 仅允许字母、数字、下划线,且不能以数字开头 matched, _ := regexp.MatchString(`^[a-zA-Z_][a-zA-Z0-9_]*$`, identifier) return matched } // isValidCondition 验证条件语句是否合法 func isValidCondition(condition string) bool { // 简单验证条件中是否包含危险字符 return !strings.Contains(condition, ";") && !strings.Contains(condition, "--") }