Merge pull request #3 from VictoriaMetrics/subqueries

Implement subqueries support

Author: Amper

README.md

@@ -112,7 +112,7 @@ Supported highlights:
   - and date helpers (`CURRENT_DATE`, `CURRENT_TIMESTAMP`).
 - `WHERE` with comparison operators, `BETWEEN`, `IN`, `LIKE`, `IS (NOT) NULL`
 - `ORDER BY`, `LIMIT`, `OFFSET`, `DISTINCT`
-- Common Table Expressions (CTE) using `WITH` keyword
+- Common Table Expressions (CTE) using `WITH` keyword and subqueries
 - `GROUP BY`, `HAVING`, `COUNT/SUM/AVG/MIN/MAX`
 - Window functions (`OVER (PARTITION BY ... ORDER BY ...)`) for `SUM` and `COUNT`
 - `JOIN` (inner/left) on equality predicates, including subqueries.

cmd/sql-to-logsql/web/ui/src/components/docs/Docs.tsx

@@ -54,7 +54,7 @@ export function Docs() {
              <p>
                <ul className={"list-disc pl-4 pt-2"}>
                  <li><code>SELECT, DISTINCT, AS, OVER, PARTITION BY</code></li>
-                 <li><code>FROM, WITH</code></li>
+                 <li><code>FROM, WITH, subqueries</code></li>
                  <li><code>WHERE, AND, OR</code></li>
                  <li><code>LEFT JOIN / JOIN / INNER JOIN</code></li>
                  <li><code>LIKE, NOT LIKE, BETWEEN, IN, NOT IN, IS NULL, IS NOT NULL</code></li>

cmd/sql-to-logsql/web/ui/src/components/sql-editor/examples.ts

@@ -86,6 +86,19 @@ ORDER BY messages_count DESC`
 SELECT UPPER(container), total
 FROM container_stats
 WHERE container IS NOT NULL
+ORDER BY total DESC`,
+    },
+    {
+        id: "subqueries",
+        title: "Subqueries",
+        sql: `SELECT UPPER(container), total
+FROM (
+    SELECT kubernetes.container_name AS container, COUNT(*) AS total
+    FROM logs
+    GROUP BY kubernetes.container_name
+    LIMIT 20
+) container_stats
+WHERE container IS NOT NULL
 ORDER BY total DESC`,
     },
     {

lib/logsql/select.go

@@ -51,6 +51,7 @@ type selectTranslatorVisitor struct {
 	sp *store.Provider
 
 	bindings           map[string]*tableBinding
+	autoAliasCounter   int
 	baseAlias          string
 	pendingLeftFilter  []ast.Expr
 	aggResults         map[string]string
@@ -172,6 +173,7 @@ func (v *selectTranslatorVisitor) translateSimpleSelect(stmt *ast.SelectStatemen
 	}
 
 	v.bindings = make(map[string]*tableBinding)
+	v.autoAliasCounter = 0
 	v.pendingLeftFilter = nil
 	v.aggResults = nil
 	v.baseAlias = ""
@@ -399,6 +401,11 @@ func (v *selectTranslatorVisitor) processFrom(from ast.TableExpr) ([]string, err
 			return nil, err
 		}
 		return nil, nil
+	case *ast.SubqueryTable:
+		if err := v.registerBaseSubquery(t); err != nil {
+			return nil, err
+		}
+		return nil, nil
 	case *ast.JoinExpr:
 		return v.processJoin(t)
 	default:
@@ -487,6 +494,43 @@ func (v *selectTranslatorVisitor) registerBaseTable(table *ast.TableName) error
 	return nil
 }
 
+func (v *selectTranslatorVisitor) registerBaseSubquery(table *ast.SubqueryTable) error {
+	if table == nil || table.Select == nil {
+		return &TranslationError{
+			Code:    http.StatusBadRequest,
+			Message: "translator: invalid subquery reference",
+		}
+	}
+	alias := strings.TrimSpace(table.Alias)
+	if alias == "" {
+		alias = v.generateSubqueryAlias("base")
+	}
+	aliasLower := strings.ToLower(alias)
+	if v.baseAlias != "" && v.baseAlias != aliasLower {
+		return &TranslationError{
+			Code:    http.StatusBadRequest,
+			Message: "translator: multiple base tables are not supported",
+		}
+	}
+	subQuery, err := translateSelectStatementToLogsQLWithContext(table.Select, translationContext{
+		sp:   v.sp,
+		ctes: v.availableCTEs,
+	})
+	if err != nil {
+		return &TranslationError{
+			Code:    http.StatusBadRequest,
+			Message: fmt.Sprintf("translator: failed to translate subquery: %s", err),
+			Err:     err,
+		}
+	}
+	v.baseAlias = aliasLower
+	v.baseUsesPipeline = true
+	v.basePipeline = subQuery
+	v.baseFilter = ""
+	v.registerBinding(aliasLower, true)
+	return nil
+}
+
 func (v *selectTranslatorVisitor) registerBinding(alias string, isBase bool) {
 	key := strings.ToLower(alias)
 	if key == "" {
@@ -495,6 +539,21 @@ func (v *selectTranslatorVisitor) registerBinding(alias string, isBase bool) {
 	v.bindings[key] = &tableBinding{alias: key, isBase: isBase}
 }
 
+func (v *selectTranslatorVisitor) generateSubqueryAlias(prefix string) string {
+	base := strings.TrimSpace(prefix)
+	if base == "" {
+		base = "subquery"
+	}
+	base = strings.ToLower(base)
+	for {
+		v.autoAliasCounter++
+		candidate := fmt.Sprintf("__%s_%d", base, v.autoAliasCounter)
+		if _, exists := v.bindings[candidate]; !exists {
+			return candidate
+		}
+	}
+}
+
 func (v *selectTranslatorVisitor) processJoin(join *ast.JoinExpr) ([]string, error) {
 	if join == nil {
 		return nil, &TranslationError{
@@ -509,16 +568,21 @@ func (v *selectTranslatorVisitor) processJoin(join *ast.JoinExpr) ([]string, err
 		}
 	}
 
-	leftTable, ok := join.Left.(*ast.TableName)
-	if !ok {
+	switch left := join.Left.(type) {
+	case *ast.TableName:
+		if err := v.registerBaseTable(left); err != nil {
+			return nil, err
+		}
+	case *ast.SubqueryTable:
+		if err := v.registerBaseSubquery(left); err != nil {
+			return nil, err
+		}
+	default:
 		return nil, &TranslationError{
 			Code:    http.StatusBadRequest,
 			Message: "translator: JOIN left side must be table reference",
 		}
 	}
-	if err := v.registerBaseTable(leftTable); err != nil {
-		return nil, err
-	}
 
 	var rightAlias string
 	var rightQuery string
@@ -606,10 +670,7 @@ func (v *selectTranslatorVisitor) processJoin(join *ast.JoinExpr) ([]string, err
 	case *ast.SubqueryTable:
 		alias := strings.TrimSpace(rt.Alias)
 		if alias == "" {
-			return nil, &TranslationError{
-				Code:    http.StatusBadRequest,
-				Message: "translator: JOIN subquery requires alias",
-			}
+			alias = v.generateSubqueryAlias("join")
 		}
 		rightAlias = strings.ToLower(alias)
 		if _, exists := v.bindings[rightAlias]; exists {
@@ -716,8 +777,8 @@ func (v *selectTranslatorVisitor) extractJoinSpec(cond ast.JoinCondition, rightA
 
 			switch {
 			case leftIsIdent && rightIsIdent:
-				leftQual := v.qualifierForIdentifier(leftIdent)
-				rightQual := v.qualifierForIdentifier(rightIdent)
+				leftQual := v.qualifierForIdentifierWithDefault(leftIdent, v.baseAlias)
+				rightQual := v.qualifierForIdentifierWithDefault(rightIdent, rightAlias)
 				if leftQual == v.baseAlias && rightQual == rightAlias {
 					leftField, err := v.normalizeIdentifier(leftIdent)
 					if err != nil {
@@ -757,8 +818,8 @@ func (v *selectTranslatorVisitor) extractJoinSpec(cond ast.JoinCondition, rightA
 			}
 		}
 
-		leftAliases := v.aliasesForExpr(bin.Left)
-		rightAliases := v.aliasesForExpr(bin.Right)
+		leftAliases := v.aliasesForExprWithDefault(bin.Left, v.baseAlias)
+		rightAliases := v.aliasesForExprWithDefault(bin.Right, rightAlias)
 
 		if v.isAliasOnly(leftAliases, v.baseAlias) && len(rightAliases) == 0 {
 			leftFilters = append(leftFilters, expr)
@@ -806,22 +867,23 @@ func flattenAnd(expr ast.Expr) []ast.Expr {
 	return []ast.Expr{expr}
 }
 
-func (v *selectTranslatorVisitor) qualifierForIdentifier(ident *ast.Identifier) string {
+func (v *selectTranslatorVisitor) qualifierForIdentifierWithDefault(ident *ast.Identifier, fallback string) string {
 	if ident == nil || len(ident.Parts) == 0 {
-		return v.baseAlias
+		return fallback
 	}
 	first := strings.ToLower(ident.Parts[0])
 	if _, ok := v.bindings[first]; ok {
 		return first
 	}
-	return v.baseAlias
+	return fallback
 }
 
-func (v *selectTranslatorVisitor) aliasesForExpr(expr ast.Expr) map[string]struct{} {
+func (v *selectTranslatorVisitor) aliasesForExprWithDefault(expr ast.Expr, fallback string) map[string]struct{} {
 	aliases := make(map[string]struct{})
 	walkExpr(expr, func(e ast.Expr) {
 		if id, ok := e.(*ast.Identifier); ok {
-			aliases[v.qualifierForIdentifier(id)] = struct{}{}
+			alias := v.qualifierForIdentifierWithDefault(id, fallback)
+			aliases[alias] = struct{}{}
 		}
 	})
 	delete(aliases, "")
@@ -883,7 +945,13 @@ func (v *selectTranslatorVisitor) ensureBaseAliasesOnly(expr ast.Expr) error {
 }
 
 func (v *selectTranslatorVisitor) ensureAliases(expr ast.Expr, allowed map[string]struct{}) error {
-	aliases := v.aliasesForExpr(expr)
+	fallback := v.baseAlias
+	if len(allowed) == 1 {
+		for alias := range allowed {
+			fallback = alias
+		}
+	}
+	aliases := v.aliasesForExprWithDefault(expr, fallback)
 	for alias := range aliases {
 		if alias == "" {
 			continue

lib/logsql/select_test.go

@@ -324,6 +324,54 @@ SELECT user FROM recent_errors WHERE service = 'api'`,
 			sql:      "SELECT LOWER(user) AS user_lower, COUNT(*) AS total FROM logs GROUP BY user_lower",
 			expected: "* | format \"<lc:user>\" as group_1 | stats by (group_1) count() total | rename group_1 as user_lower",
 		},
+		{
+			name: "subquery as base table",
+			sql: `SELECT *
+FROM (
+    SELECT *
+    FROM logs
+    WHERE level = 'error'
+) AS recent_errors`,
+			expected: "level:error",
+		},
+		{
+			name: "subquery as base table without alias",
+			sql: `SELECT *
+FROM (
+    SELECT *
+    FROM logs
+    WHERE level = 'error'
+)`,
+			expected: "level:error",
+		},
+		{
+			name: "subquery as base with filter",
+			sql: `SELECT recent.user, recent.fail_count
+FROM (
+    SELECT user, COUNT(*) AS fail_count
+    FROM logs
+    WHERE level = 'error'
+    GROUP BY user
+) AS recent
+WHERE recent.fail_count > 10
+ORDER BY recent.fail_count DESC
+LIMIT 5`,
+			expected: "level:error | stats by (user) count() fail_count | filter fail_count:>10 | fields user, fail_count | sort by (fail_count desc) | limit 5",
+		},
+		{
+			name: "subquery as base with filter without alias",
+			sql: `SELECT user, fail_count
+FROM (
+    SELECT user, COUNT(*) AS fail_count
+    FROM logs
+    WHERE level = 'error'
+    GROUP BY user
+)
+WHERE fail_count > 10
+ORDER BY fail_count DESC
+LIMIT 5`,
+			expected: "level:error | stats by (user) count() fail_count | filter fail_count:>10 | fields user, fail_count | sort by (fail_count desc) | limit 5",
+		},
 		{
 			name: "join with subquery",
 			sql: `SELECT l.user, m.fail_count
@@ -336,6 +384,21 @@ INNER JOIN (
 ) AS m ON l.user = m.user
 WHERE l.level = 'error'
 ORDER BY m.fail_count DESC
+LIMIT 5`,
+			expected: "level:error | join by (user) (level:error | stats by (user) count() fail_count) inner | fields user, fail_count | sort by (fail_count desc) | limit 5",
+		},
+		{
+			name: "join with subquery without alias",
+			sql: `SELECT l.user, fail_count
+FROM logs AS l
+INNER JOIN (
+    SELECT user, COUNT(*) AS fail_count
+    FROM logs
+    WHERE level = 'error'
+    GROUP BY user
+) ON l.user = user
+WHERE l.level = 'error'
+ORDER BY fail_count DESC
 LIMIT 5`,
 			expected: "level:error | join by (user) (level:error | stats by (user) count() fail_count) inner | fields user, fail_count | sort by (fail_count desc) | limit 5",
 		},
@@ -379,6 +442,21 @@ func TestToLogsQLWithConfig(t *testing.T) {
 		}
 	})
 
+	t.Run("join with subquery base", func(t *testing.T) {
+		sql := `SELECT recent.user, a.level
+FROM (
+    SELECT user
+    FROM logs
+    WHERE level = 'error'
+) AS recent
+INNER JOIN api AS a ON recent.user = a.user`
+		got := mustTranslateWithTables(t, sql, tables)
+		expected := "level:error | fields user | join by (user) (service:api) inner | fields user, level"
+		if got != expected {
+			t.Fatalf("unexpected query:\nexpected: %s\n     got: %s", expected, got)
+		}
+	})
+
 	t.Run("unknown table", func(t *testing.T) {
 		_, err := translateWithTables(t, "SELECT * FROM missing", tables)
 		if err == nil {