From f79906ca750a22803ccb1ed1d5f659b063bb5fa7 Mon Sep 17 00:00:00 2001
From: Christian Haudum <christian.haudum@gmail.com>
Date: Tue, 25 Mar 2025 11:31:52 +0100
Subject: [PATCH] chore(engine): Add utility function `canExecuteWithNewEngine`
 (#16892)

The function determines whether a query can be executed by the new execution engine.

In the first step, the new execution engine will only be able to execute limited/filter queries without any transformation stages.


Signed-off-by: Christian Haudum <christian.haudum@gmail.com>
---
 pkg/engine/engine.go      | 25 +++++++++++
 pkg/engine/engine_test.go | 87 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 112 insertions(+)
 create mode 100644 pkg/engine/engine.go
 create mode 100644 pkg/engine/engine_test.go

diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go
new file mode 100644
index 0000000000..2bbeb8e948
--- /dev/null
+++ b/pkg/engine/engine.go
@@ -0,0 +1,25 @@
+package engine
+
+import "github.com/grafana/loki/v3/pkg/logql/syntax"
+
+// canExecuteWithNewEngine determines whether a query can be executed by the new execution engine.
+func canExecuteWithNewEngine(expr syntax.Expr) bool {
+	switch expr := expr.(type) {
+	case syntax.SampleExpr:
+		return false
+	case syntax.LogSelectorExpr:
+		ret := true
+		expr.Walk(func(e syntax.Expr) {
+			switch e.(type) {
+			case *syntax.LineParserExpr, *syntax.LogfmtParserExpr, *syntax.LogfmtExpressionParserExpr, *syntax.JSONExpressionParserExpr:
+				ret = false
+			case *syntax.LineFmtExpr, *syntax.LabelFmtExpr:
+				ret = false
+			case *syntax.KeepLabelsExpr, *syntax.DropLabelsExpr:
+				ret = false
+			}
+		})
+		return ret
+	}
+	return false
+}
diff --git a/pkg/engine/engine_test.go b/pkg/engine/engine_test.go
new file mode 100644
index 0000000000..cf85674835
--- /dev/null
+++ b/pkg/engine/engine_test.go
@@ -0,0 +1,87 @@
+package engine
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/loki/v3/pkg/logql/syntax"
+)
+
+func TestCanExecuteWithNewEngine(t *testing.T) {
+	for _, tt := range []struct {
+		statement string
+		expected  bool
+	}{
+		{
+			statement: `{env="prod"}`,
+			expected:  true,
+		},
+		{
+			statement: `{env="prod"} |= "metrics.go"`,
+			expected:  true,
+		},
+		{
+			statement: `{env="prod"} |= "metrics.go"`,
+			expected:  true,
+		},
+		{
+			statement: `{env="prod"} | tenant="loki"`,
+			expected:  true,
+		},
+		{
+			statement: `{env="prod"} | tenant="loki" != "foo"`,
+			expected:  true,
+		},
+		{
+			statement: `{env="prod"} | json`,
+			expected:  false,
+		},
+		{
+			statement: `{env="prod"} | json foo="bar"`,
+			expected:  false,
+		},
+		{
+			statement: `{env="prod"} | logfmt`,
+			expected:  false,
+		},
+		{
+			statement: `{env="prod"} | logfmt foo="bar"`,
+			expected:  false,
+		},
+		{
+			statement: `{env="prod"} | pattern "<_> foo=<foo> <_>"`,
+			expected:  false,
+		},
+		{
+			statement: `{env="prod"} | regexp ".* foo=(?P<foo>.+) .*"`,
+			expected:  false,
+		},
+		{
+			statement: `{env="prod"} | unpack`,
+			expected:  false,
+		},
+		{
+			statement: `{env="prod"} |= "metrics.go" | logfmt`,
+			expected:  false,
+		},
+		{
+			statement: `{env="prod"} | line_format "{.cluster}"`,
+			expected:  false,
+		},
+		{
+			statement: `{env="prod"} | label_format cluster="us"`,
+			expected:  false,
+		},
+		{
+			statement: `sum(rate({env="prod"}[1m]))`,
+			expected:  false,
+		},
+	} {
+		t.Run(tt.statement, func(t *testing.T) {
+			expr := syntax.MustParseExpr(tt.statement)
+			canExecute := canExecuteWithNewEngine(expr)
+			require.Equal(t, tt.expected, canExecute)
+		})
+	}
+}