diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/PlanExecutor.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/PlanExecutor.java index c3fd9613bfcce..815c85b7fed8b 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/PlanExecutor.java +++ b/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 params, ActionListener 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 } public void sql(Configuration cfg, String sql, List params, ActionListener 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 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 listener) { diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlClearCursorAction.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlClearCursorAction.java index cce721e78fd1f..8bfe08e078476 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlClearCursorAction.java +++ b/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)); } diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlQueryAction.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlQueryAction.java index c0c1f67c05901..a4955b740b6c2 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlQueryAction.java +++ b/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)); } } diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlTranslateAction.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlTranslateAction.java index 840cc7f9cac4d..0bda719111589 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TransportSqlTranslateAction.java +++ b/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 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( diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/session/Configuration.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/session/Configuration.java index 6eb6ad19ad49c..e386cbb3b3205 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/session/Configuration.java +++ b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/session/Configuration.java @@ -20,6 +20,7 @@ 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; @@ -27,7 +28,8 @@ public class Configuration { @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; @@ -35,6 +37,7 @@ public Configuration(ZoneId zi, int pageSize, TimeValue requestTimeout, TimeValu 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; } diff --git a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/TestUtils.java b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/TestUtils.java index c704285f4eba0..be7f42d3f0c78 100644 --- a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/TestUtils.java +++ b/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. diff --git a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/FunctionRegistryTests.java b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/FunctionRegistryTests.java index 8c5446f5ee0d7..101f4dfe78c4e 100644 --- a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/FunctionRegistryTests.java +++ b/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)); } diff --git a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/DatabaseFunctionTests.java b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/DatabaseFunctionTests.java index de2fc69a263f1..86e4baf9fdc06 100644 --- a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/DatabaseFunctionTests.java +++ b/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()) diff --git a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/UserFunctionTests.java b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/UserFunctionTests.java index 7b1e86af5d513..f4f48cb735045 100644 --- a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/expression/function/scalar/UserFunctionTests.java +++ b/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())