Respect size parameter in aggregations (1/3 - refactor)

quesma/clickhouse/quesma_communicator.go

@@ -58,6 +58,7 @@ func (lm *LogManager) ProcessQuery(ctx context.Context, table *Table, query *mod
 		return nil, err
 	}
 
+	// will become: rows, err := executeQuery(ctx, lm, query.StringFromColumns(colNames, true), columns, rowToScan)
 	rows, err := executeQuery(ctx, lm, query.StringFromColumns(colNames), columns, rowToScan)
 	if err == nil {
 		for _, row := range rows {

quesma/eql/query_translator.go

@@ -30,7 +30,7 @@ func (cw *ClickhouseEQLQueryTranslator) applySizeLimit(size int) int {
 	return size
 }
 
-func (cw *ClickhouseEQLQueryTranslator) BuildNRowsQuery(fieldName string, simpleQuery queryparser.SimpleQuery, limit int) *model.Query {
+func (cw *ClickhouseEQLQueryTranslator) BuildNRowsQuery(fieldName string, simpleQuery model.SimpleQuery, limit int) *model.Query {
 	suffixClauses := make([]string, 0)
 	if len(simpleQuery.SortFields) > 0 {
 		suffixClauses = append(suffixClauses, "ORDER BY "+queryparser.AsQueryString(simpleQuery.SortFields))
@@ -88,13 +88,13 @@ func (cw *ClickhouseEQLQueryTranslator) MakeSearchResponse(ResultSet []model.Que
 	}, nil
 }
 
-func (cw *ClickhouseEQLQueryTranslator) ParseQuery(queryAsJson string) (query queryparser.SimpleQuery, searchQueryInfo model.SearchQueryInfo, highlighter model.Highlighter, err error) {
+func (cw *ClickhouseEQLQueryTranslator) ParseQuery(queryAsJson string) (query model.SimpleQuery, searchQueryInfo model.SearchQueryInfo, highlighter model.Highlighter, err error) {
 
 	// no highlighting here
 	highlighter = queryparser.NewEmptyHighlighter()
 
 	searchQueryInfo.Typ = model.ListAllFields
-	query.Sql = queryparser.Statement{}
+	query.Sql = model.Statement{}
 
 	queryAsMap := make(map[string]interface{})
 	err = json.Unmarshal([]byte(queryAsJson), &queryAsMap)
@@ -158,7 +158,7 @@ func (cw *ClickhouseEQLQueryTranslator) MakeResponseAggregation(aggregations []m
 	panic("EQL does not support aggregations")
 }
 
-func (cw *ClickhouseEQLQueryTranslator) BuildFacetsQuery(fieldName string, simpleQuery queryparser.SimpleQuery, limit int) *model.Query {
+func (cw *ClickhouseEQLQueryTranslator) BuildFacetsQuery(fieldName string, simpleQuery model.SimpleQuery, limit int) *model.Query {
 	panic("EQL does not support facets")
 }
 

quesma/logger/logger.go

@@ -112,6 +112,7 @@ func InitSimpleLoggerForTests() {
 			Out:        os.Stderr,
 			TimeFormat: time.StampMilli,
 		}).
+		Level(zerolog.DebugLevel).
 		With().
 		Timestamp().
 		Logger()

quesma/model/bucket_aggregations/filters.go

@@ -7,11 +7,25 @@ import (
 )
 
 type Filters struct {
-	ctx context.Context
+	ctx     context.Context
+	Filters []Filter
 }
 
-func NewFilters(ctx context.Context) Filters {
-	return Filters{ctx}
+func NewFiltersEmpty(ctx context.Context) Filters {
+	return Filters{ctx: ctx}
+}
+
+func NewFilters(ctx context.Context, filters []Filter) Filters {
+	return Filters{ctx: ctx, Filters: filters}
+}
+
+type Filter struct {
+	Name string
+	Sql  model.SimpleQuery
+}
+
+func NewFilter(name string, sql model.SimpleQuery) Filter {
+	return Filter{Name: name, Sql: sql}
 }
 
 func (query Filters) IsBucketAggregation() bool {

quesma/model/query.go

@@ -41,13 +41,27 @@ type Query struct {
 	Aggregators     []Aggregator // keeps names of aggregators, e.g. "0", "1", "2", "suggestions". Needed for JSON response.
 	Type            QueryType
 	SortFields      []SortField // fields to sort by
+	SubSelect       string
 	// dictionary to add as 'meta' field in the response.
 	// WARNING: it's probably not passed everywhere where it's needed, just in one place.
 	// But it works for the test + our dashboards, so let's fix it later if necessary.
 	// NoMetadataField (nil) is a valid option and means no meta field in the response.
 	Metadata JsonMap
 }
 
+/*
+type subQuery struct {
+	sql       string
+	innerJoin string
+	name      string
+}
+
+func newSubQuery(sql, innerJoin, name string) subQuery {
+	return subQuery{sql: sql, innerJoin: innerJoin, name: name}
+}
+ erase if still not used after 'respect size in aggregations' issue
+*/
+
 var NoMetadataField JsonMap = nil
 
 // returns string with SQL query

quesma/model/simple_query.go

@@ -0,0 +1,112 @@
+package model
+
+import (
+	"context"
+	"mitmproxy/quesma/logger"
+)
+
+type SimpleQuery struct {
+	Sql        Statement
+	CanParse   bool
+	FieldName  string
+	SortFields []SortField
+}
+
+func NewSimpleQuery(sql Statement, canParse bool) SimpleQuery {
+	return SimpleQuery{Sql: sql, CanParse: canParse}
+}
+
+func NewSimpleQueryWithFieldName(sql Statement, canParse bool, fieldName string) SimpleQuery {
+	return SimpleQuery{Sql: sql, CanParse: canParse, FieldName: fieldName}
+}
+
+type Statement struct {
+	Stmt       string
+	IsCompound bool // "a" -> not compound, "a AND b" -> compound. Used to not make unnecessary brackets (not always, but usually)
+	FieldName  string
+}
+
+func NewSimpleStatement(stmt string) Statement {
+	return Statement{Stmt: stmt, IsCompound: false}
+}
+
+func NewCompoundStatement(stmt, fieldName string) Statement {
+	return Statement{Stmt: stmt, IsCompound: true, FieldName: fieldName}
+}
+
+func NewCompoundStatementNoFieldName(stmt string) Statement {
+	return Statement{Stmt: stmt, IsCompound: true}
+}
+
+// Added to the generated SQL where the query is fine, but we're sure no rows will match it
+var AlwaysFalseStatement = NewSimpleStatement("false")
+
+func And(andStmts []Statement) Statement {
+	return combineStatements(andStmts, "AND")
+}
+
+func Or(orStmts []Statement) Statement {
+	return combineStatements(orStmts, "OR")
+}
+
+// sep = "AND" or "OR"
+func combineStatements(stmts []Statement, sep string) Statement {
+	stmts = FilterNonEmpty(stmts)
+	if len(stmts) > 1 {
+		stmts = quoteWithBracketsIfCompound(stmts)
+		var fieldName string
+		sql := ""
+		for i, stmt := range stmts {
+			sql += stmt.Stmt
+			if i < len(stmts)-1 {
+				sql += " " + sep + " "
+			}
+			if stmt.FieldName != "" {
+				fieldName = stmt.FieldName
+			}
+		}
+		return NewCompoundStatement(sql, fieldName)
+	}
+	if len(stmts) == 1 {
+		return stmts[0]
+	}
+	return NewSimpleStatement("")
+}
+
+func CombineWheres(ctx context.Context, where1, where2 SimpleQuery) SimpleQuery {
+	combined := SimpleQuery{
+		Sql:      And([]Statement{where1.Sql, where2.Sql}),
+		CanParse: where1.CanParse && where2.CanParse,
+	}
+	if len(where1.FieldName) > 0 && len(where2.FieldName) > 0 && where1.FieldName != where2.FieldName {
+		logger.WarnWithCtx(ctx).Msgf("combining 2 where clauses with different field names: %s, %s, where queries: %v %v", where1.FieldName, where2.FieldName, where1, where2)
+	}
+	if len(where1.FieldName) > 0 {
+		combined.FieldName = where1.FieldName
+	} else {
+		combined.FieldName = where2.FieldName
+	}
+	return combined
+}
+
+func FilterNonEmpty(slice []Statement) []Statement {
+	i := 0
+	for _, el := range slice {
+		if len(el.Stmt) > 0 {
+			slice[i] = el
+			i++
+		}
+	}
+	return slice[:i]
+}
+
+// used to combine statements with AND/OR
+// [a, b, a AND b] ==> ["a", "b", "(a AND b)"]
+func quoteWithBracketsIfCompound(slice []Statement) []Statement {
+	for i := range slice {
+		if slice[i].IsCompound {
+			slice[i].Stmt = "(" + slice[i].Stmt + ")"
+		}
+	}
+	return slice
+}

quesma/model/simple_query_test.go

@@ -0,0 +1,78 @@
+package model
+
+import (
+	"github.com/stretchr/testify/assert"
+	"strconv"
+	"strings"
+	"testing"
+)
+
+func TestFilterNonEmpty(t *testing.T) {
+	tests := []struct {
+		array    []Statement
+		filtered []Statement
+	}{
+		{
+			[]Statement{NewSimpleStatement(""), NewSimpleStatement("")},
+			[]Statement{},
+		},
+		{
+			[]Statement{NewSimpleStatement(""), NewSimpleStatement("a"), NewCompoundStatementNoFieldName("")},
+			[]Statement{NewSimpleStatement("a")},
+		},
+		{
+			[]Statement{NewCompoundStatementNoFieldName("a"), NewSimpleStatement("b"), NewCompoundStatement("c", "d")},
+			[]Statement{NewCompoundStatementNoFieldName("a"), NewSimpleStatement("b"), NewCompoundStatement("c", "d")},
+		},
+	}
+	for i, tt := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			assert.Equal(t, tt.filtered, FilterNonEmpty(tt.array))
+		})
+	}
+}
+
+func TestOrAndAnd(t *testing.T) {
+	tests := []struct {
+		stmts []Statement
+		want  Statement
+	}{
+		{
+			[]Statement{NewSimpleStatement("a"), NewSimpleStatement("b"), NewSimpleStatement("c")},
+			NewCompoundStatementNoFieldName("a AND b AND c"),
+		},
+		{
+			[]Statement{NewSimpleStatement("a"), NewSimpleStatement(""), NewCompoundStatementNoFieldName(""), NewCompoundStatementNoFieldName("b")},
+			NewCompoundStatementNoFieldName("a AND (b)"),
+		},
+		{
+			[]Statement{NewSimpleStatement(""), NewSimpleStatement(""), NewSimpleStatement("a"), NewCompoundStatementNoFieldName(""), NewSimpleStatement(""), NewCompoundStatementNoFieldName("")},
+			NewSimpleStatement("a"),
+		},
+		{
+			[]Statement{NewSimpleStatement(""), NewSimpleStatement(""), NewSimpleStatement(""), NewSimpleStatement("")},
+			NewSimpleStatement(""),
+		},
+		{
+			[]Statement{NewCompoundStatementNoFieldName("a AND b"), NewCompoundStatementNoFieldName("c AND d"), NewCompoundStatement("e AND f", "field")},
+			NewCompoundStatement("(a AND b) AND (c AND d) AND (e AND f)", "field"),
+		},
+	}
+	// copy, because and() and or() modify the slice
+	for i, tt := range tests {
+		t.Run("AND "+strconv.Itoa(i), func(t *testing.T) {
+			b := make([]Statement, len(tt.stmts))
+			copy(b, tt.stmts)
+			assert.Equal(t, tt.want, And(b))
+		})
+	}
+	for i, tt := range tests {
+		t.Run("OR "+strconv.Itoa(i), func(t *testing.T) {
+			tt.want.Stmt = strings.ReplaceAll(tt.want.Stmt, "AND", "OR")
+			for i := range tt.stmts {
+				tt.stmts[i].Stmt = strings.ReplaceAll(tt.stmts[i].Stmt, "AND", "OR")
+			}
+			assert.Equal(t, tt.want, Or(tt.stmts))
+		})
+	}
+}