refactor(querybuilder): derive the allowed table function message from the map

Values the map by the spelling to name back to the caller, so the message
comes off the same data the lookup uses and cannot drift from it. Drops
the test that existed only to catch that drift.

pkg/querybuilder/clickhouse_sql.go

@@ -3,6 +3,8 @@ package querybuilder
 import (
 	"context"
 	"log/slog"
+	"maps"
+	"slices"
 	"strings"
 
 	chparser "github.com/AfterShip/clickhouse-sql-parser/parser"
@@ -25,7 +27,23 @@ var internalDatabases = map[string]struct{}{
 	"information_schema": {},
 }
 
-// The parser's grammar has gaps against SQL that ClickHouse itself accepts. See TestErrIfStatementIsNotValid_ShouldPassButFails.
+// generatorTableFunctions compute their rows from their arguments alone. They open no file or socket, reach no other host, and name no table, database or dictionary, so none of them can read through anything the rules here exist to protect. Can be used to build a dense axis to join a sparse series against. Every other table function is refused.
+//
+// Keyed by the lowercased name so that matching is case-insensitive, valued by the spelling to name it back to the caller.
+//
+// TODO(@therealpandey): take a deployment level allow list on top of this, so an operator can permit more without a release.
+var generatorTableFunctions = map[string]string{
+	"numbers":         "numbers",
+	"numbers_mt":      "numbers_mt",
+	"zeros":           "zeros",
+	"zeros_mt":        "zeros_mt",
+	"generateseries":  "generateSeries",
+	"generate_series": "generate_series",
+}
+
+var generatorTableFunctionsMessage = "allowed table functions are " + strings.Join(slices.Sorted(maps.Values(generatorTableFunctions)), ", ")
+
+// The parser's grammar has gaps against SQL that ClickHouse itself accepts.
 func ErrIfStatementIsNotValid(query string) (err error) {
 	defer func() {
 		// The parser has a history of panicking on malformed input rather than returning an error.
@@ -52,8 +70,17 @@ func ErrIfStatementIsNotValid(query string) (err error) {
 	visitor := &chparser.DefaultASTVisitor{Visit: func(node chparser.Expr) error {
 		switch expr := node.(type) {
 		case *chparser.TableFunctionExpr:
-			// Source table functions remain usable in ClickHouse read-only mode.
-			return errors.NewInvalidInputf(CodeClickHouseSQLTableFunction, "ClickHouse table functions are not allowed in SQL queries: %s", chparser.Format(expr.Name))
+			// Source table functions remain usable in ClickHouse read-only mode. Arguments are
+			// visited before this, so a read smuggled into one is already refused by the time
+			// an allowed generator gets here.
+			name := chparser.Format(expr.Name)
+			if _, ok := generatorTableFunctions[strings.ToLower(name)]; ok {
+				return nil
+			}
+
+			return errors.
+				NewInvalidInputf(CodeClickHouseSQLTableFunction, "ClickHouse table functions are not allowed in SQL queries: %s", name).
+				WithAdditional(generatorTableFunctionsMessage)
 
 		case *chparser.TableIdentifier:
 			// Reading these is unaffected by ClickHouse read-only mode.

pkg/querybuilder/clickhouse_sql_test.go

@@ -54,6 +54,19 @@ func TestErrIfStatementIsNotValid_Pass(t *testing.T) {
 		{"StandardTrimSyntax", "SELECT trim(BOTH ' ' FROM body) FROM t"},
 		{"StandardSubstringSyntax", "SELECT substring(body FROM 2 FOR 3) FROM t"},
 		{"StandardOverlaySyntax", "SELECT overlay(body PLACING 'x' FROM 2) FROM t"},
+		// Row generators compute their rows from their arguments, so they read through nothing. This is the shape they get used for: a dense interval axis to CROSS JOIN a sparse series against.
+		{"NumbersTableFunction", "SELECT intervals.interval AS interval, active.cluster AS cluster, toFloat64(if(ts_data.has_data = 0, 0, 1)) AS value FROM ( SELECT DISTINCT JSONExtractString(labels, 'k8s.cluster.name') AS cluster FROM signoz_metrics.distributed_time_series_v4 WHERE metric_name = 'my_metric' AND unix_milli >= toUnixTimestamp(now() - INTERVAL 30 DAY) * 1000 HAVING cluster != '' ) AS active CROSS JOIN ( SELECT toStartOfInterval( toDateTime(toUnixTimestamp(now() - INTERVAL 30 MINUTE) + number * 60), INTERVAL 1 MINUTE ) AS interval FROM numbers(31) ) AS intervals LEFT JOIN ( SELECT toStartOfInterval( toDateTime(intDiv(s.unix_milli, 1000)), INTERVAL 1 MINUTE ) AS interval, JSONExtractString(ts.labels, 'k8s.cluster.name') AS cluster, 1 AS has_data FROM signoz_metrics.distributed_samples_v4 s INNER JOIN ( SELECT DISTINCT fingerprint, labels FROM signoz_metrics.distributed_time_series_v4 WHERE metric_name = 'my_metric' ) AS ts ON s.fingerprint = ts.fingerprint WHERE s.metric_name = 'my_metric' AND s.unix_milli >= toUnixTimestamp(now() - INTERVAL 30 MINUTE) * 1000 GROUP BY interval, cluster ) AS ts_data ON active.cluster = ts_data.cluster AND intervals.interval = ts_data.interval ORDER BY interval ASC"},
+		{"NumbersMtTableFunction", "SELECT * FROM numbers_mt(31)"},
+		{"ZerosTableFunction", "SELECT * FROM zeros(31)"},
+		{"ZerosMtTableFunction", "SELECT * FROM zeros_mt(31)"},
+		{"GenerateSeriesTableFunction", "SELECT * FROM generateSeries(1, 10)"},
+		{"GenerateSeriesSnakeCaseTableFunction", "SELECT * FROM generate_series(1, 10)"},
+		{"GeneratorTableFunctionUppercase", "SELECT * FROM NUMBERS(31)"},
+		{"GeneratorTableFunctionParenthesisedArgument", "SELECT * FROM NUMBERS((31))"},
+		{"GeneratorTableFunctionInJoin", "SELECT * FROM signoz_logs.distributed_logs_v2 AS l CROSS JOIN numbers(31) AS n"},
+		{"GeneratorTableFunctionInCommonTableExpression", "WITH axis AS (SELECT number FROM numbers(31)) SELECT * FROM axis"},
+		{"GeneratorTableFunctionInWhereSubquery", "SELECT * FROM t WHERE a IN (SELECT number FROM numbers(31))"},
+		{"GeneratorTableFunctionInUnion", "SELECT number FROM numbers(31) UNION ALL SELECT number FROM zeros(31)"},
 	}
 
 	for _, testCase := range testCases {
@@ -104,6 +117,20 @@ func TestErrIfStatementIsNotValid_Fail(t *testing.T) {
 		{"TableFunctionInCommonTableExpression", "WITH c AS (SELECT * FROM url('http://x', CSV, 'a String')) SELECT * FROM c", CodeClickHouseSQLTableFunction},
 		{"TableFunctionInWhereSubquery", "SELECT * FROM t WHERE a IN (SELECT * FROM file('/etc/passwd', CSV, 'a String'))", CodeClickHouseSQLTableFunction},
 		{"TableFunctionInUnion", "SELECT * FROM t UNION ALL SELECT * FROM url('http://x', CSV, 'a String')", CodeClickHouseSQLTableFunction},
+		// These reach the internal databases without ever naming one, so the table-function rule is the only thing that sees them.
+		{"MergeTableFunction", "SELECT * FROM merge('system', '.*')", CodeClickHouseSQLTableFunction},
+		{"RemoteTableFunction", "SELECT * FROM remote('other-host', 'system.users')", CodeClickHouseSQLTableFunction},
+		{"ClusterTableFunction", "SELECT * FROM cluster('c', 'system.users')", CodeClickHouseSQLTableFunction},
+		// Pure, but excluded: generateRandom streams rows the arguments do not bound, and values has no use here that an array literal does not already cover.
+		{"GenerateRandomTableFunction", "SELECT * FROM generateRandom('a UInt64')", CodeClickHouseSQLTableFunction},
+		{"ValuesTableFunction", "SELECT * FROM values('a UInt64', 1, 2)", CodeClickHouseSQLTableFunction},
+		// Arguments are visited before the table function itself, so allowing a generator does not give anyone a wrapper to smuggle a read through.
+		{"InternalDatabaseInsideAllowedTableFunction", "SELECT * FROM numbers((SELECT count() FROM system.users))", CodeClickHouseSQLInternalDatabase},
+		{"InternalDatabaseJoinedOntoAllowedTableFunction", "SELECT * FROM numbers(31) AS n JOIN system.users AS u ON 1 = 1", CodeClickHouseSQLInternalDatabase},
+		{"InternalDatabaseUnionedWithAllowedTableFunction", "SELECT number FROM numbers(31) UNION ALL SELECT name FROM system.users", CodeClickHouseSQLInternalDatabase},
+		{"RefusedTableFunctionJoinedOntoAllowedTableFunction", "SELECT * FROM numbers(31) AS n JOIN url('http://x', CSV, 'a String') AS u ON 1 = 1", CodeClickHouseSQLTableFunction},
+		{"RefusedTableFunctionInsideAllowedTableFunction", "SELECT * FROM numbers((SELECT count() FROM file('/etc/passwd', CSV, 'a String')))", CodeClickHouseSQLTableFunction},
+		{"InternalDatabaseInsideAllowedTableFunctionCommonTableExpression", "WITH axis AS (SELECT * FROM numbers((SELECT count() FROM system.users))) SELECT * FROM axis", CodeClickHouseSQLInternalDatabase},
 		// Internal databases, which hold grants and server metadata rather than telemetry.
 		{"SystemUsers", "SELECT * FROM system.users", CodeClickHouseSQLInternalDatabase},
 		{"SystemUppercase", "SELECT * FROM SYSTEM.USERS", CodeClickHouseSQLInternalDatabase},
@@ -126,29 +153,3 @@ func TestErrIfStatementIsNotValid_Fail(t *testing.T) {
 		})
 	}
 }
-
-// Queries ClickHouse runs that this rejects anyway. Each is a known false positive.
-func TestErrIfStatementIsNotValid_ShouldPassButFails(t *testing.T) {
-	testCases := []struct {
-		name         string
-		query        string
-		expectedCode errors.Code
-	}{
-		{
-			// numbers() generates rows rather than reading through anything, so the blanket
-			// table-function rule is stricter here than the threat it exists for.
-			name:         "NumbersTableFunction",
-			query:        "SELECT intervals.interval AS interval, active.cluster AS cluster, toFloat64(if(ts_data.has_data = 0, 0, 1)) AS value FROM ( SELECT DISTINCT JSONExtractString(labels, 'k8s.cluster.name') AS cluster FROM signoz_metrics.distributed_time_series_v4 WHERE metric_name = 'my_metric' AND unix_milli >= toUnixTimestamp(now() - INTERVAL 30 DAY) * 1000 HAVING cluster != '' ) AS active CROSS JOIN ( SELECT toStartOfInterval( toDateTime(toUnixTimestamp(now() - INTERVAL 30 MINUTE) + number * 60), INTERVAL 1 MINUTE ) AS interval FROM numbers(31) ) AS intervals LEFT JOIN ( SELECT toStartOfInterval( toDateTime(intDiv(s.unix_milli, 1000)), INTERVAL 1 MINUTE ) AS interval, JSONExtractString(ts.labels, 'k8s.cluster.name') AS cluster, 1 AS has_data FROM signoz_metrics.distributed_samples_v4 s INNER JOIN ( SELECT DISTINCT fingerprint, labels FROM signoz_metrics.distributed_time_series_v4 WHERE metric_name = 'my_metric' ) AS ts ON s.fingerprint = ts.fingerprint WHERE s.metric_name = 'my_metric' AND s.unix_milli >= toUnixTimestamp(now() - INTERVAL 30 MINUTE) * 1000 GROUP BY interval, cluster ) AS ts_data ON active.cluster = ts_data.cluster AND intervals.interval = ts_data.interval ORDER BY interval ASC",
-			expectedCode: CodeClickHouseSQLTableFunction,
-		},
-	}
-
-	for _, testCase := range testCases {
-		t.Run(testCase.name, func(t *testing.T) {
-			err := ErrIfStatementIsNotValid(testCase.query)
-
-			assert.Error(t, err)
-			assert.True(t, errors.Asc(err, testCase.expectedCode), "expected code %s, got %v", testCase.expectedCode, err)
-		})
-	}
-}