SQL: Move metrics tracking inside PlanExecutor (#38259)

Move metrics in one place, from the transport layer inside the
PlanExecutor
Remove unused class

Close #38258

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/PlanExecutor.java

@@ -27,9 +27,12 @@
 import org.elasticsearch.xpack.sql.session.SchemaRowSet;
 import org.elasticsearch.xpack.sql.session.SqlSession;
 import org.elasticsearch.xpack.sql.stats.Metrics;
+import org.elasticsearch.xpack.sql.stats.QueryMetric;
 
 import java.util.List;
 
+import static org.elasticsearch.action.ActionListener.wrap;
+
 public class PlanExecutor {
     private final Client client;
     private final NamedWriteableRegistry writableRegistry;
@@ -64,7 +67,9 @@ private SqlSession newSession(Configuration cfg) {
     }
 
     public void searchSource(Configuration cfg, String sql, List<SqlTypedParamValue> params, ActionListener<SearchSourceBuilder> listener) {
-        newSession(cfg).sqlExecutable(sql, params, ActionListener.wrap(exec -> {
+        metrics.translate();
+
+        newSession(cfg).sqlExecutable(sql, params, wrap(exec -> {
             if (exec instanceof EsQueryExec) {
                 EsQueryExec e = (EsQueryExec) exec;
                 listener.onResponse(SourceGenerator.sourceBuilder(e.queryContainer(), cfg.filter(), cfg.pageSize()));
@@ -87,11 +92,24 @@ public void searchSource(Configuration cfg, String sql, List<SqlTypedParamValue>
     }
 
     public void sql(Configuration cfg, String sql, List<SqlTypedParamValue> params, ActionListener<SchemaRowSet> listener) {
-        newSession(cfg).sql(sql, params, listener);
+        QueryMetric metric = QueryMetric.from(cfg.mode(), cfg.clientId());
+        metrics.total(metric);
+
+        newSession(cfg).sql(sql, params, wrap(listener::onResponse, ex -> {
+            metrics.failed(metric);
+            listener.onFailure(ex);
+        }));
     }
 
     public void nextPage(Configuration cfg, Cursor cursor, ActionListener<RowSet> listener) {
-        cursor.nextPage(cfg, client, writableRegistry, listener);
+        QueryMetric metric = QueryMetric.from(cfg.mode(), cfg.clientId());
+        metrics.total(metric);
+        metrics.paging(metric);
+
+        cursor.nextPage(cfg, client, writableRegistry, wrap(listener::onResponse, ex -> {
+            metrics.failed(metric);
+            listener.onFailure(ex);
+        }));
     }
 
     public void cleanCursor(Configuration cfg, Cursor cursor, ActionListener<Boolean> listener) {

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlClearCursorAction.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.xpack.sql.session.Cursor;
 import org.elasticsearch.xpack.sql.session.Cursors;
 import org.elasticsearch.xpack.sql.util.DateUtils;
+import org.elasticsearch.xpack.sql.util.StringUtils;
 
 import static org.elasticsearch.xpack.sql.action.SqlClearCursorAction.NAME;
 
@@ -46,7 +47,7 @@ public static void operation(PlanExecutor planExecutor, SqlClearCursorRequest re
         Cursor cursor = Cursors.decodeFromString(request.getCursor());
         planExecutor.cleanCursor(
                 new Configuration(DateUtils.UTC, Protocol.FETCH_SIZE, Protocol.REQUEST_TIMEOUT, Protocol.PAGE_TIMEOUT, null,
-                        request.mode(), "", ""),
+                        request.mode(), StringUtils.EMPTY, StringUtils.EMPTY, StringUtils.EMPTY),
                 cursor, ActionListener.wrap(
                 success -> listener.onResponse(new SqlClearCursorResponse(success)), listener::onFailure));
     }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlQueryAction.java

@@ -29,13 +29,13 @@
 import org.elasticsearch.xpack.sql.session.Cursors;
 import org.elasticsearch.xpack.sql.session.RowSet;
 import org.elasticsearch.xpack.sql.session.SchemaRowSet;
-import org.elasticsearch.xpack.sql.stats.QueryMetric;
 import org.elasticsearch.xpack.sql.type.Schema;
 
 import java.util.ArrayList;
 import java.util.List;
 
 import static java.util.Collections.unmodifiableList;
+import static org.elasticsearch.action.ActionListener.wrap;
 import static org.elasticsearch.xpack.sql.plugin.Transports.clusterName;
 import static org.elasticsearch.xpack.sql.plugin.Transports.username;
 
@@ -72,27 +72,14 @@ public static void operation(PlanExecutor planExecutor, SqlQueryRequest request,
         // The configuration is always created however when dealing with the next page, only the timeouts are relevant
         // the rest having default values (since the query is already created)
         Configuration cfg = new Configuration(request.zoneId(), request.fetchSize(), request.requestTimeout(), request.pageTimeout(),
-                request.filter(), request.mode(), username, clusterName);
-
-        // mode() shouldn't be null
-        QueryMetric metric = QueryMetric.from(request.mode(), request.clientId());
-        planExecutor.metrics().total(metric);
+                request.filter(), request.mode(), request.clientId(), username, clusterName);
 
         if (Strings.hasText(request.cursor()) == false) {
             planExecutor.sql(cfg, request.query(), request.params(),
-                    ActionListener.wrap(rowSet -> listener.onResponse(createResponse(request, rowSet)),
-                            e -> {
-                                planExecutor.metrics().failed(metric);
-                                listener.onFailure(e);
-                            }));
+                    wrap(rowSet -> listener.onResponse(createResponse(request, rowSet)), listener::onFailure));
         } else {
-            planExecutor.metrics().paging(metric);
             planExecutor.nextPage(cfg, Cursors.decodeFromString(request.cursor()),
-                    ActionListener.wrap(rowSet -> listener.onResponse(createResponse(request.mode(), rowSet, null)),
-                            e -> {
-                                planExecutor.metrics().failed(metric);
-                                listener.onFailure(e);
-                            }));
+                    wrap(rowSet -> listener.onResponse(createResponse(request.mode(), rowSet, null)), listener::onFailure));
         }
     }
 

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlTranslateAction.java

@@ -52,9 +52,9 @@ public TransportSqlTranslateAction(Settings settings, ClusterService clusterServ
     protected void doExecute(Task task, SqlTranslateRequest request, ActionListener<SqlTranslateResponse> listener) {
         sqlLicenseChecker.checkIfSqlAllowed(request.mode());
 
-        planExecutor.metrics().translate();
         Configuration cfg = new Configuration(request.zoneId(), request.fetchSize(),
-                request.requestTimeout(), request.pageTimeout(), request.filter(), request.mode(),
+                request.requestTimeout(), request.pageTimeout(), request.filter(), 
+                request.mode(), request.clientId(),
                 username(securityContext), clusterName(clusterService));
 
         planExecutor.searchSource(cfg, request.query(), request.params(), ActionListener.wrap(

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/session/Configuration.java

@@ -20,21 +20,24 @@ public class Configuration {
     private final TimeValue requestTimeout;
     private final TimeValue pageTimeout;
     private final Mode mode;
+    private final String clientId;
     private final String username;
     private final String clusterName;
     private final ZonedDateTime now;
 
     @Nullable
     private QueryBuilder filter;
 
-    public Configuration(ZoneId zi, int pageSize, TimeValue requestTimeout, TimeValue pageTimeout, QueryBuilder filter, Mode mode,
+    public Configuration(ZoneId zi, int pageSize, TimeValue requestTimeout, TimeValue pageTimeout, QueryBuilder filter,
+                         Mode mode, String clientId,
                          String username, String clusterName) {
         this.zoneId = zi.normalized();
         this.pageSize = pageSize;
         this.requestTimeout = requestTimeout;
         this.pageTimeout = pageTimeout;
         this.filter = filter;
         this.mode = mode == null ? Mode.PLAIN : mode;
+        this.clientId = clientId;
         this.username = username;
         this.clusterName = clusterName;
         this.now = ZonedDateTime.now(zoneId);
@@ -63,6 +66,10 @@ public Mode mode() {
         return mode;
     }
 
+    public String clientId() {
+        return clientId;
+    }
+
     public String username() {
         return username;
     }

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/TestUtils.java

@@ -20,7 +20,8 @@ public class TestUtils {
     private TestUtils() {}
 
     public static final Configuration TEST_CFG = new Configuration(DateUtils.UTC, Protocol.FETCH_SIZE,
-            Protocol.REQUEST_TIMEOUT, Protocol.PAGE_TIMEOUT, null, Mode.PLAIN, null, null);
+            Protocol.REQUEST_TIMEOUT, Protocol.PAGE_TIMEOUT, null, Mode.PLAIN,
+            null, null, null);
 
     /**
      * Returns the current UTC date-time with milliseconds precision.

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/FunctionRegistryTests.java

@@ -239,6 +239,7 @@ private Configuration randomConfiguration() {
                 null,
                 randomFrom(Mode.values()),
                 randomAlphaOfLength(10),
+                randomAlphaOfLength(10),
                 randomAlphaOfLength(10));
     }
 
@@ -250,6 +251,7 @@ private Configuration randomConfiguration(ZoneId providedZoneId) {
                 null,
                 randomFrom(Mode.values()),
                 randomAlphaOfLength(10),
+                randomAlphaOfLength(10),
                 randomAlphaOfLength(10));
     }
 

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/DatabaseFunctionTests.java

@@ -29,7 +29,9 @@ public void testDatabaseFunctionOutput() {
         EsIndex test = new EsIndex("test", TypesTests.loadMapping("mapping-basic.json", true));
         Analyzer analyzer = new Analyzer(
                 new Configuration(DateUtils.UTC, Protocol.FETCH_SIZE, Protocol.REQUEST_TIMEOUT,
-                                  Protocol.PAGE_TIMEOUT, null, randomFrom(Mode.values()), null, clusterName),
+                                  Protocol.PAGE_TIMEOUT, null, 
+                                  randomFrom(Mode.values()), randomAlphaOfLength(10),
+                                  null, clusterName),
                 new FunctionRegistry(),
                 IndexResolution.valid(test),
                 new Verifier(new Metrics())

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/UserFunctionTests.java

@@ -28,7 +28,9 @@ public void testNoUsernameFunctionOutput() {
         EsIndex test = new EsIndex("test", TypesTests.loadMapping("mapping-basic.json", true));
         Analyzer analyzer = new Analyzer(
                 new Configuration(DateUtils.UTC, Protocol.FETCH_SIZE, Protocol.REQUEST_TIMEOUT,
-                                  Protocol.PAGE_TIMEOUT, null, randomFrom(Mode.values()), null, randomAlphaOfLengthBetween(1, 15)),
+                                  Protocol.PAGE_TIMEOUT, null, 
+                                  randomFrom(Mode.values()), randomAlphaOfLength(10), 
+                                  null, randomAlphaOfLengthBetween(1, 15)),
                 new FunctionRegistry(),
                 IndexResolution.valid(test),
                 new Verifier(new Metrics())