From 30f52ee441ecf91d18155d10f9943c2bd823fdee Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 9 Sep 2026 11:43:30 +0800 Subject: [PATCH 1/5] Revert "refactor(sql): split native tasks support into follow-up" This reverts commit 097f6474bd866e097ddaba7218940eccd9b60fc2. --- docs/querying/sql-metadata-tables.md | 15 +- .../query/NativeSysTasksQueryTest.java | 344 ++++++++++++++++++ ...SQLServerMetadataStorageActionHandler.java | 6 + .../MySQLMetadataStorageActionHandler.java | 6 + ...ostgreSQLMetadataStorageActionHandler.java | 6 + .../overlord/MetadataTaskStorage.java | 25 ++ .../indexing/overlord/TaskQueryTool.java | 135 +++++-- .../druid/indexing/overlord/TaskStorage.java | 23 ++ .../overlord/task/TasksTableDataProvider.java | 108 ++++++ .../DerbyMetadataStorageActionHandler.java | 6 + .../MetadataStorageActionHandler.java | 12 + .../SQLMetadataStorageActionHandler.java | 267 ++++++-------- .../org/apache/druid/metadata/SqlDialect.java | 79 ++++ .../druid/metadata/SqlPredicateBuilder.java | 133 +++++++ .../metadata/TaskStorageQueryFilter.java | 129 +++++++ .../task/TasksTableDataProviderTest.java | 311 ++++++++++++++++ .../SQLMetadataStorageActionHandlerTest.java | 161 ++++++++ .../apache/druid/metadata/SqlDialectTest.java | 97 +++++ .../system/module/SystemTableModule.java | 7 +- .../system/table/TaskTableDescriptor.java | 105 ++++++ .../handler/SystemTableQueryClientTest.java | 19 +- .../org/apache/druid/cli/CliOverlord.java | 12 + .../sql/calcite/schema/NativeTasksTable.java | 70 ++++ .../sql/calcite/schema/SystemSchema.java | 32 +- .../calcite/schema/SystemSchemaProvider.java | 3 +- .../sql/calcite/schema/SystemSchemaTest.java | 5 +- .../schema/SystemTableDataProviderTest.java | 12 +- 27 files changed, 1895 insertions(+), 233 deletions(-) create mode 100644 embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/NativeSysTasksQueryTest.java create mode 100644 indexing-service/src/main/java/org/apache/druid/indexing/overlord/task/TasksTableDataProvider.java create mode 100644 indexing-service/src/main/java/org/apache/druid/metadata/SqlDialect.java create mode 100644 indexing-service/src/main/java/org/apache/druid/metadata/SqlPredicateBuilder.java create mode 100644 indexing-service/src/main/java/org/apache/druid/metadata/TaskStorageQueryFilter.java create mode 100644 indexing-service/src/test/java/org/apache/druid/indexing/overlord/task/TasksTableDataProviderTest.java create mode 100644 indexing-service/src/test/java/org/apache/druid/metadata/SqlDialectTest.java create mode 100644 server/src/main/java/org/apache/druid/server/system/table/TaskTableDescriptor.java create mode 100644 sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeTasksTable.java diff --git a/docs/querying/sql-metadata-tables.md b/docs/querying/sql-metadata-tables.md index 8b47335617fe..dc9c4cf03bde 100644 --- a/docs/querying/sql-metadata-tables.md +++ b/docs/querying/sql-metadata-tables.md @@ -147,7 +147,7 @@ The native SQL engine can plan supported system tables as native datasources. To ```json { - "query": "SELECT COUNT(DISTINCT server) FROM sys.server_properties", + "query": "SELECT COUNT(DISTINCT task_id) FROM sys.tasks", "context": { "useNativeQueryForSystemTables": true } @@ -159,7 +159,7 @@ execution with `SET`: ```sql SET useNativeQueryForSystemTables = 'true'; -SELECT COUNT(DISTINCT server) FROM sys.server_properties; +SELECT COUNT(DISTINCT datasource) FROM sys.tasks; ``` Native system-table execution is available when the resolved SQL engine is `native`. You don't need to explicitly set @@ -168,17 +168,13 @@ execution: |Table|Source of rows| |-----|--------------| +|[`sys.tasks`](#tasks-table)|The Overlord that owns task state. Supported filters are pushed into task storage when the configured task storage implementation supports filter pushdown.| |[`sys.server_properties`](#server_properties-table)|The Druid server processes discovered in the cluster. Filters on `server` and `service_name` can avoid reading properties from nodes that don't match.| After Druid retrieves the system-table rows, the native engine applies the remaining filters, expressions, aggregations, sorting, and result processing. A system table that doesn't advertise native query support continues to use its existing SQL-layer execution path, even when `useNativeQueryForSystemTables` is `true`. -The parameter defaults to `false`. During a rolling upgrade, leave it disabled until the Broker and all nodes -that serve the native system tables have been upgraded. After the upgrade, you can enable it by query or set -`druid.query.default.context.useNativeQueryForSystemTables=true` on Brokers as the cluster-wide default. For more -information, see [SQL query context](sql-query-context.md). - Native system-table queries sent to the Router use distributed Broker execution by default. To execute a native system-table Scan against only the contacted node, set the HTTP header `X-Druid-Native-Query-Route: local`. Local execution uses the authenticated request identity and applies the table's @@ -186,6 +182,11 @@ authorization rules. The Broker uses the same header for remote node fan-out req one of the selected nodes, it executes that node Scan in-process without an HTTP request. The header controls routing and doesn't grant additional permissions. +The parameter defaults to `false`. During a rolling upgrade, leave it disabled until the Broker and all nodes +that serve the native system tables have been upgraded. After the upgrade, you can enable it by query or set +`druid.query.default.context.useNativeQueryForSystemTables=true` on Brokers as the cluster-wide default. For more +information, see [SQL query context](sql-query-context.md). + ### SEGMENTS table Segments table provides details on all Druid segments, whether they are published yet or not. diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/NativeSysTasksQueryTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/NativeSysTasksQueryTest.java new file mode 100644 index 000000000000..89679c7776c9 --- /dev/null +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/NativeSysTasksQueryTest.java @@ -0,0 +1,344 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.testing.embedded.query; + +import org.apache.druid.indexing.common.task.NoopTask; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.query.QueryContexts; +import org.apache.druid.sql.calcite.planner.PlannerContext; +import org.apache.druid.sql.calcite.run.NativeSqlEngine; +import org.apache.druid.testing.embedded.EmbeddedBroker; +import org.apache.druid.testing.embedded.EmbeddedCoordinator; +import org.apache.druid.testing.embedded.EmbeddedDruidCluster; +import org.apache.druid.testing.embedded.EmbeddedIndexer; +import org.apache.druid.testing.embedded.EmbeddedOverlord; +import org.apache.druid.testing.embedded.junit5.EmbeddedClusterTestBase; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import java.util.Arrays; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +public class NativeSysTasksQueryTest extends EmbeddedClusterTestBase +{ + private static final String TASK_PREFIX = "native_sys_mvp_"; + + private final EmbeddedOverlord overlord = new EmbeddedOverlord(); + private final EmbeddedBroker broker = new EmbeddedBroker(); + + @Override + protected EmbeddedDruidCluster createCluster() + { + return EmbeddedDruidCluster.withEmbeddedDerbyAndZookeeper() + .useLatchableEmitter() + .addServer(new EmbeddedCoordinator()) + .addServer(new EmbeddedIndexer() + .addProperty("druid.worker.capacity", "5")) + .addServer(overlord) + .addServer(broker); + } + + @BeforeAll + public void createTasks() + { + createTasks("a", "native_sys_a", 2); + createTasks("b", "native_sys_b", 3); + } + + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testGroupByUsesOverlordProvider(final String plannerStrategy) + { + final String result = cluster.runSql( + "SELECT datasource, COUNT(*) " + + "FROM sys.tasks " + + "WHERE task_id = 'native_sys_mvp_a_0' AND datasource = 'native_sys_a' " + + "GROUP BY datasource", + nativeQueryContext(plannerStrategy) + ); + + final Set rows = Arrays.stream(result.split("\\n")).collect(Collectors.toSet()); + Assertions.assertEquals(Set.of("native_sys_a,1"), rows); + } + + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testNativeAggregationsSupportDistinctCount(final String plannerStrategy) + { + final String result = cluster.runSql( + "SELECT COUNT(*), COUNT(DISTINCT task_id), COUNT(DISTINCT datasource), SUM(1) " + + "FROM sys.tasks " + + "WHERE datasource IN ('native_sys_a', 'native_sys_b')", + nativeQueryContext(plannerStrategy) + ); + + Assertions.assertEquals("5,5,2,5", result); + } + + /** + * Verifies execution when native planning represents the inner aggregation as a query datasource. + * + * 
{@code
+   * SELECT COUNT(*)
+   * FROM
+   * (
+   *   SELECT
+   *     task_id,
+   *     COUNT(*) AS task_count
+   *   FROM sys.tasks
+   *   WHERE datasource IN ('native_sys_a', 'native_sys_b')
+   *   GROUP BY task_id
+   * )
+   * WHERE task_count > 0
+   * }
+ */ + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testNestedTaskAggregation(final String plannerStrategy) + { + final String result = cluster.runSql( + "SELECT COUNT(*) " + + "FROM (" + + " SELECT task_id, COUNT(*) AS task_count " + + " FROM sys.tasks " + + " WHERE datasource IN ('native_sys_a', 'native_sys_b') " + + " GROUP BY task_id" + + ") " + + "WHERE task_count > 0", + nativeQueryContext(plannerStrategy) + ); + + Assertions.assertEquals("5", result); + } + + /** + * Verifies that an outer aggregation can consume grouped rows from a native system-table subquery. + * + * 
{@code
+   * SELECT SUM(task_count)
+   * FROM
+   * (
+   *   SELECT
+   *     datasource,
+   *     COUNT(*) AS task_count
+   *   FROM sys.tasks
+   *   WHERE datasource IN ('native_sys_a', 'native_sys_b')
+   *   GROUP BY datasource
+   * )
+   * }
+ */ + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testOuterAggregationOverTaskSubquery(final String plannerStrategy) + { + final String result = cluster.runSql( + "SELECT SUM(task_count) " + + "FROM (" + + " SELECT datasource, COUNT(*) AS task_count " + + " FROM sys.tasks " + + " WHERE datasource IN ('native_sys_a', 'native_sys_b') " + + " GROUP BY datasource" + + ")", + nativeQueryContext(plannerStrategy) + ); + + Assertions.assertEquals("5", result); + } + + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testInAndOrFilters(final String plannerStrategy) + { + assertTaskFilterQuery( + "task_id IN ('native_sys_mvp_a_0', 'native_sys_mvp_b_2') " + + "OR task_id = 'native_sys_mvp_missing'", + Set.of("native_sys_mvp_a_0", "native_sys_mvp_b_2"), + plannerStrategy + ); + } + + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testRangeAndLikeFilters(final String plannerStrategy) + { + assertTaskFilterQuery( + "task_id >= 'native_sys_mvp_b_0' " + + "AND task_id < 'native_sys_mvp_b_2' " + + "AND task_id LIKE 'native_sys_mvp_b_%'", + Set.of("native_sys_mvp_b_0", "native_sys_mvp_b_1"), + plannerStrategy + ); + } + + /** + * Verifies that expression virtual columns required by a pushed node filter are preserved. + * + * 
{@code
+   * SELECT task_id
+   * FROM sys.tasks
+   * WHERE UPPER(task_id) = 'NATIVE_SYS_MVP_A_0'
+   * }
+ */ + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testExpressionFilterPreservesVirtualColumn(final String plannerStrategy) + { + assertTaskFilterQuery( + "UPPER(task_id) = 'NATIVE_SYS_MVP_A_0'", + Set.of("native_sys_mvp_a_0"), + plannerStrategy + ); + } + + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testNegatedLikeFilter(final String plannerStrategy) + { + assertTaskFilterQuery( + "datasource = 'native_sys_b' AND task_id NOT LIKE '%b_2'", + Set.of("native_sys_mvp_b_0", "native_sys_mvp_b_1"), + plannerStrategy + ); + } + + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testStatusLikeFilterRemainsForResidualEvaluation(final String plannerStrategy) + { + assertTaskFilterQuery( + "task_id LIKE 'native_sys_mvp_%' AND status LIKE '%'", + Set.of( + "native_sys_mvp_a_0", + "native_sys_mvp_a_1", + "native_sys_mvp_b_0", + "native_sys_mvp_b_1", + "native_sys_mvp_b_2" + ), + plannerStrategy + ); + } + + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testWebConsoleTasksQueryUsesOverlordProvider(final String plannerStrategy) + { + final String result = cluster.runSql( + "WITH tasks AS (SELECT\n" + + " \"task_id\", \"group_id\", \"type\", \"datasource\", \"created_time\", \"location\", " + + "\"duration\", \"error_msg\",\n" + + " CASE WHEN \"error_msg\" IN ('Shutdown request from user', " + + "'Canceled: Query canceled by user or by task shutdown.') THEN 'CANCELED' " + + "WHEN \"status\" = 'RUNNING' THEN \"runner_status\" ELSE \"status\" END AS \"status\"\n" + + " FROM sys.tasks\n" + + ")\n" + + "SELECT \"task_id\", \"group_id\", \"type\", \"datasource\", \"created_time\", \"location\", " + + "\"duration\", \"error_msg\", \"status\"\n" + + "FROM tasks\n" + + "ORDER BY\n" + + " (CASE \"status\" WHEN 'RUNNING' THEN 4 WHEN 'PENDING' THEN 3 WHEN 'WAITING' THEN 2 ELSE 1 END) DESC,\n" + + " \"created_time\" DESC", + nativeQueryContext(plannerStrategy) + ); + + final Set taskIds = Arrays.stream(result.split("\\n")) + .map(row -> row.substring(0, row.indexOf(','))) + .collect(Collectors.toSet()); + Assertions.assertEquals( + Set.of( + "native_sys_mvp_a_0", + "native_sys_mvp_a_1", + "native_sys_mvp_b_0", + "native_sys_mvp_b_1", + "native_sys_mvp_b_2" + ), + taskIds + ); + } + + private void createTasks(final String suffix, final String datasource, final int count) + { + for (int i = 0; i < count; i++) { + final String taskId = TASK_PREFIX + suffix + "_" + i; + cluster.callApi().runTask(new NoopTask(taskId, null, datasource, 1L, 0L, null), overlord); + } + } + + private void assertTaskFilterQuery( + final String whereClause, + final Set expectedTaskIds, + final String plannerStrategy + ) + { + // EmbeddedServiceClient formats SQL with the selected broker address, so escape LIKE wildcards for that step. + final String escapedWhereClause = StringUtils.replace(whereClause, "%", "%%"); + final String result = cluster.runSql( + "SELECT task_id FROM sys.tasks WHERE " + escapedWhereClause, + nativeQueryContext(plannerStrategy) + ); + + Assertions.assertEquals(expectedTaskIds, Arrays.stream(result.split("\\n")).collect(Collectors.toSet())); + } + + private static Map nativeQueryContext(final String plannerStrategy) + { + return Map.of( + QueryContexts.ENGINE, + NativeSqlEngine.NAME, + PlannerContext.CTX_USE_NATIVE_QUERY_FOR_SYSTEM_TABLES, + true, + QueryContexts.CTX_NATIVE_QUERY_SQL_PLANNING_MODE, + plannerStrategy + ); + } +} diff --git a/extensions-contrib/sqlserver-metadata-storage/src/main/java/org/apache/druid/metadata/SQLServerMetadataStorageActionHandler.java b/extensions-contrib/sqlserver-metadata-storage/src/main/java/org/apache/druid/metadata/SQLServerMetadataStorageActionHandler.java index cf491c041731..130d2e839680 100644 --- a/extensions-contrib/sqlserver-metadata-storage/src/main/java/org/apache/druid/metadata/SQLServerMetadataStorageActionHandler.java +++ b/extensions-contrib/sqlserver-metadata-storage/src/main/java/org/apache/druid/metadata/SQLServerMetadataStorageActionHandler.java @@ -38,4 +38,10 @@ protected String decorateSqlWithLimit(String sql) { return "SELECT TOP (:n)" + sql.substring("SELECT".length()); } + + @Override + protected SqlDialect getSqlDialect() + { + return SqlDialect.SQL_SERVER; + } } diff --git a/extensions-core/mysql-metadata-storage/src/main/java/org/apache/druid/metadata/MySQLMetadataStorageActionHandler.java b/extensions-core/mysql-metadata-storage/src/main/java/org/apache/druid/metadata/MySQLMetadataStorageActionHandler.java index 54a0938dec3f..a443ad79fac2 100644 --- a/extensions-core/mysql-metadata-storage/src/main/java/org/apache/druid/metadata/MySQLMetadataStorageActionHandler.java +++ b/extensions-core/mysql-metadata-storage/src/main/java/org/apache/druid/metadata/MySQLMetadataStorageActionHandler.java @@ -38,4 +38,10 @@ protected String decorateSqlWithLimit(String sql) { return sql + " LIMIT :n"; } + + @Override + protected SqlDialect getSqlDialect() + { + return SqlDialect.MYSQL; + } } diff --git a/extensions-core/postgresql-metadata-storage/src/main/java/org/apache/druid/metadata/PostgreSQLMetadataStorageActionHandler.java b/extensions-core/postgresql-metadata-storage/src/main/java/org/apache/druid/metadata/PostgreSQLMetadataStorageActionHandler.java index 4181633463ea..629be66c3cd4 100644 --- a/extensions-core/postgresql-metadata-storage/src/main/java/org/apache/druid/metadata/PostgreSQLMetadataStorageActionHandler.java +++ b/extensions-core/postgresql-metadata-storage/src/main/java/org/apache/druid/metadata/PostgreSQLMetadataStorageActionHandler.java @@ -39,4 +39,10 @@ protected String decorateSqlWithLimit(String sql) return sql + " LIMIT :n"; } + @Override + protected SqlDialect getSqlDialect() + { + return SqlDialect.POSTGRESQL; + } + } diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/MetadataTaskStorage.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/MetadataTaskStorage.java index 35858fdc7347..f09f4477e1e8 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/MetadataTaskStorage.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/MetadataTaskStorage.java @@ -42,6 +42,7 @@ import org.apache.druid.metadata.TaskLookup; import org.apache.druid.metadata.TaskLookup.ActiveTaskLookup; import org.apache.druid.metadata.TaskLookup.TaskLookupType; +import org.apache.druid.metadata.TaskStorageQueryFilter; import org.joda.time.DateTime; import javax.annotation.Nullable; @@ -232,6 +233,30 @@ public List getTaskStatusPlusList( .collect(Collectors.toList()); } + @Override + public boolean supportsTaskStatusQueryFilter() + { + return true; + } + + @Override + public List getTaskStatusPlusListWithFilter( + final Map taskLookups, + final TaskStorageQueryFilter filter + ) + { + Map processedTaskLookups = + TaskStorageUtils.processTaskLookups( + taskLookups, + DateTimes.nowUtc().minus(config.getRecentlyFinishedThreshold()) + ); + + return handler.getTaskStatusListWithFilter(processedTaskLookups, filter) + .stream() + .map(TaskStatusPlus::fromTaskIdentifierInfo) + .collect(Collectors.toList()); + } + @Override public void addLock(final String taskid, final TaskLock taskLock) { diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueryTool.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueryTool.java index 9be83c597ff9..822f879557c7 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueryTool.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueryTool.java @@ -21,7 +21,6 @@ import com.google.common.base.Optional; import com.google.common.base.Supplier; -import com.google.common.collect.ImmutableMap; import com.google.common.collect.Maps; import com.google.inject.Inject; import org.apache.druid.indexer.TaskInfo; @@ -33,17 +32,20 @@ import org.apache.druid.indexing.overlord.http.TotalWorkerCapacityResponse; import org.apache.druid.indexing.overlord.setup.WorkerBehaviorConfig; import org.apache.druid.java.util.common.IAE; +import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.Intervals; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.metadata.LockFilterPolicy; import org.apache.druid.metadata.TaskLookup; import org.apache.druid.metadata.TaskLookup.TaskLookupType; +import org.apache.druid.metadata.TaskStorageQueryFilter; import org.joda.time.Duration; import org.joda.time.Interval; import javax.annotation.Nullable; import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -195,12 +197,84 @@ public List getTaskStatusPlusList( @Nullable Integer maxCompletedTasks, @Nullable String type ) + { + return getTaskStatusPlusList( + state, + createdTimeInterval, + maxCompletedTasks, + dataSource, + type, + null + ); + } + + public boolean supportsTaskStatusQueryFilter() + { + return storage.supportsTaskStatusQueryFilter(); + } + + public List getTaskStatusPlusList( + final TaskStateLookup state, + final TaskStorageQueryFilter filter + ) + { + final Optional taskRunner = taskMaster.getTaskRunner(); + if (!taskRunner.isPresent()) { + throw new ISE("Task status filter queries must run on the current Overlord leader"); + } + final TaskRunner selectedTaskRunner = taskRunner.get(); + final List taskStatuses = getTaskStatusPlusListWithRunner( + state, + null, + null, + null, + null, + filter, + selectedTaskRunner + ); + final Optional currentTaskRunner = taskMaster.getTaskRunner(); + if (!taskMaster.isHalfOrFullLeader() + || !currentTaskRunner.isPresent() + || currentTaskRunner.get() != selectedTaskRunner) { + throw new ISE("Overlord leadership changed during the task status filter query"); + } + return taskStatuses; + } + + private List getTaskStatusPlusList( + final TaskStateLookup state, + @Nullable final String createdTimeInterval, + @Nullable final Integer maxCompletedTasks, + @Nullable final String dataSource, + @Nullable final String type, + @Nullable final TaskStorageQueryFilter filter + ) { Optional taskRunnerOptional = taskMaster.getTaskRunner(); if (!taskRunnerOptional.isPresent()) { return Collections.emptyList(); } - final TaskRunner taskRunner = taskRunnerOptional.get(); + return getTaskStatusPlusListWithRunner( + state, + createdTimeInterval, + maxCompletedTasks, + dataSource, + type, + filter, + taskRunnerOptional.get() + ); + } + + private List getTaskStatusPlusListWithRunner( + final TaskStateLookup state, + @Nullable final String createdTimeInterval, + @Nullable final Integer maxCompletedTasks, + @Nullable final String dataSource, + @Nullable final String type, + @Nullable final TaskStorageQueryFilter filter, + final TaskRunner taskRunner + ) + { final Duration createdTimeDuration; if (createdTimeInterval != null) { @@ -216,13 +290,20 @@ public List getTaskStatusPlusList( // This way, we can use the snapshot from taskStorage as the source of truth for the set of tasks to process // and use the snapshot from taskRunner as a reference for potential task state updates happened // after the first snapshotting. - Stream taskStatusPlusStream = retrieveTaskStatusesFromMetadataStore( + final Map taskLookups = createTaskLookups( state, - dataSource, createdTimeDuration, maxCompletedTasks, - type + filter ); + Stream taskStatusPlusStream = filter == null + ? storage.getTaskStatusPlusList(taskLookups, dataSource).stream() + : storage.getTaskStatusPlusListWithFilter(taskLookups, filter).stream(); + if (type != null) { + taskStatusPlusStream = taskStatusPlusStream.filter( + statusPlus -> type.equals(statusPlus == null ? null : statusPlus.getType()) + ); + } final Map runnerWorkItems = getTaskRunnerWorkItems( taskRunner, state, @@ -235,7 +316,7 @@ public List getTaskStatusPlusList( taskStatusPlusStream = taskStatusPlusStream .filter(statusPlus -> runnerWorkItems.containsKey(statusPlus.getId())); } - final List taskStatusPlusList = taskStatusPlusStream.collect(Collectors.toList()); + final List taskStatusPlusList = taskStatusPlusStream.toList(); // Separate complete and active tasks from taskStorage. // Note that taskStorage can return only either complete tasks or active tasks per TaskLookupType. @@ -282,26 +363,24 @@ public List getTaskStatusPlusList( return taskStatuses; } - private Stream retrieveTaskStatusesFromMetadataStore( - TaskStateLookup state, - @Nullable String dataSource, - Duration createdTimeDuration, - @Nullable Integer maxCompletedTasks, - @Nullable String type + private static Map createTaskLookups( + final TaskStateLookup state, + final Duration createdTimeDuration, + @Nullable final Integer maxCompletedTasks, + @Nullable final TaskStorageQueryFilter filter ) { - final Map taskLookups; + final Map taskLookups = new HashMap<>(); switch (state) { case ALL: - taskLookups = ImmutableMap.of( - TaskLookupType.ACTIVE, - TaskLookup.ActiveTaskLookup.getInstance(), + taskLookups.put(TaskLookupType.ACTIVE, TaskLookup.ActiveTaskLookup.getInstance()); + taskLookups.put( TaskLookupType.COMPLETE, TaskLookup.CompleteTaskLookup.of(maxCompletedTasks, createdTimeDuration) ); break; case COMPLETE: - taskLookups = ImmutableMap.of( + taskLookups.put( TaskLookupType.COMPLETE, TaskLookup.CompleteTaskLookup.of(maxCompletedTasks, createdTimeDuration) ); @@ -309,24 +388,22 @@ private Stream retrieveTaskStatusesFromMetadataStore( case WAITING: case PENDING: case RUNNING: - taskLookups = ImmutableMap.of( - TaskLookupType.ACTIVE, - TaskLookup.ActiveTaskLookup.getInstance() - ); + taskLookups.put(TaskLookupType.ACTIVE, TaskLookup.ActiveTaskLookup.getInstance()); break; default: throw new IAE("Unknown state: [%s]", state); } - final Stream taskStatusPlusStream - = storage.getTaskStatusPlusList(taskLookups, dataSource).stream(); - if (type != null) { - return taskStatusPlusStream.filter( - statusPlus -> type.equals(statusPlus == null ? null : statusPlus.getType()) - ); - } else { - return taskStatusPlusStream; + if (filter != null) { + if (!filter.includesActiveTasks()) { + taskLookups.remove(TaskLookupType.ACTIVE); + } + if (!filter.includesCompleteTasks()) { + taskLookups.remove(TaskLookupType.COMPLETE); + } } + + return taskLookups; } private Map getTaskRunnerWorkItems( diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskStorage.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskStorage.java index 3aac918b5ef3..3e3c894a62dc 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskStorage.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskStorage.java @@ -27,6 +27,7 @@ import org.apache.druid.indexing.common.task.Task; import org.apache.druid.metadata.TaskLookup; import org.apache.druid.metadata.TaskLookup.TaskLookupType; +import org.apache.druid.metadata.TaskStorageQueryFilter; import org.joda.time.DateTime; import javax.annotation.Nullable; @@ -162,6 +163,28 @@ List getTaskStatusPlusList( @Nullable String datasource ); + /** + * Returns whether this storage applies {@link TaskStorageQueryFilter} while retrieving task statuses. + * Implementations returning {@code false} may return a superset because the native query retains the original + * filter for final evaluation. + */ + default boolean supportsTaskStatusQueryFilter() + { + return false; + } + + /** + * Returns task statuses with validated native-query prefilters. Implementations that do not support these filters + * may return a superset because the native query retains the original filter for final evaluation. + */ + default List getTaskStatusPlusListWithFilter( + final Map taskLookups, + final TaskStorageQueryFilter filter + ) + { + return getTaskStatusPlusList(taskLookups, null); + } + /** * Returns a list of tasks stored in the storage facility as {@link TaskInfo}. No * particular order is guaranteed, but implementations are encouraged to return tasks in ascending order of creation. diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/task/TasksTableDataProvider.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/task/TasksTableDataProvider.java new file mode 100644 index 000000000000..512558a113a9 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/task/TasksTableDataProvider.java @@ -0,0 +1,108 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.overlord.task; + +import com.google.common.collect.Iterables; +import com.google.inject.Inject; +import org.apache.druid.indexer.TaskStatusPlus; +import org.apache.druid.indexing.overlord.TaskMaster; +import org.apache.druid.indexing.overlord.TaskQueryTool; +import org.apache.druid.indexing.overlord.http.TaskStateLookup; +import org.apache.druid.metadata.TaskStorageQueryFilter; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.system.SystemTableNotLeaderException; +import org.apache.druid.server.system.table.SystemTableDataProvider; +import org.apache.druid.server.system.table.SystemTablePushdownFilter; + +import java.util.List; + +/** Native row supplier for {@code sys.tasks}. */ +public class TasksTableDataProvider implements SystemTableDataProvider +{ + private static final List METADATA_STORAGE_PUSHDOWN_FILTERS = List.of( + new SystemTablePushdownFilter("task_id", "id"), + new SystemTablePushdownFilter("group_id", null), + new SystemTablePushdownFilter("type", null), + new SystemTablePushdownFilter("datasource", null), + new SystemTablePushdownFilter("created_time", "created_date"), + new SystemTablePushdownFilter("status", null) + ); + private static final List STATUS_PUSHDOWN_FILTERS = List.of( + new SystemTablePushdownFilter("status", null) + ); + + private final TaskQueryTool taskQueryTool; + private final TaskMaster taskMaster; + + @Inject + public TasksTableDataProvider( + final TaskQueryTool taskQueryTool, + final TaskMaster taskMaster + ) + { + this.taskQueryTool = taskQueryTool; + this.taskMaster = taskMaster; + } + + @Override + public List getPushdownFilters() + { + return taskQueryTool.supportsTaskStatusQueryFilter() + ? METADATA_STORAGE_PUSHDOWN_FILTERS + : STATUS_PUSHDOWN_FILTERS; + } + + @Override + public Iterable getRows( + final List filters, + final AuthenticationResult internalAuthenticationResult + ) + { + if (!taskMaster.isHalfOrFullLeader()) { + throw new SystemTableNotLeaderException("overlord"); + } + final List tasks = taskQueryTool.getTaskStatusPlusList( + TaskStateLookup.ALL, + new TaskStorageQueryFilter(filters) + ); + return Iterables.transform(tasks, TasksTableDataProvider::taskToRow); + } + + private static Object[] taskToRow(final TaskStatusPlus task) + { + return new Object[]{ + task.getId(), + task.getGroupId(), + task.getType(), + task.getDataSource(), + task.getCreatedTime() == null ? null : task.getCreatedTime().toString(), + task.getQueueInsertionTime() == null ? null : task.getQueueInsertionTime().toString(), + task.getStatusCode() == null ? null : task.getStatusCode().toString(), + task.getRunnerStatusCode() == null ? null : task.getRunnerStatusCode().toString(), + task.getDuration() == null ? 0L : task.getDuration(), + task.getLocation().getLocation(), + task.getLocation().getHost(), + (long) task.getLocation().getPort(), + (long) task.getLocation().getTlsPort(), + task.getErrorMsg() + }; + } +} diff --git a/indexing-service/src/main/java/org/apache/druid/metadata/DerbyMetadataStorageActionHandler.java b/indexing-service/src/main/java/org/apache/druid/metadata/DerbyMetadataStorageActionHandler.java index f60341e293e5..c3a1cb1597db 100644 --- a/indexing-service/src/main/java/org/apache/druid/metadata/DerbyMetadataStorageActionHandler.java +++ b/indexing-service/src/main/java/org/apache/druid/metadata/DerbyMetadataStorageActionHandler.java @@ -41,4 +41,10 @@ protected String decorateSqlWithLimit(String sql) return sql + " FETCH FIRST :n ROWS ONLY"; } + @Override + protected SqlDialect getSqlDialect() + { + return SqlDialect.DERBY; + } + } diff --git a/indexing-service/src/main/java/org/apache/druid/metadata/MetadataStorageActionHandler.java b/indexing-service/src/main/java/org/apache/druid/metadata/MetadataStorageActionHandler.java index 4f06027c8e9d..e2ddb390db72 100644 --- a/indexing-service/src/main/java/org/apache/druid/metadata/MetadataStorageActionHandler.java +++ b/indexing-service/src/main/java/org/apache/druid/metadata/MetadataStorageActionHandler.java @@ -129,6 +129,18 @@ List getTaskStatusList( @Nullable String datasource ); + /** + * Returns task statuses with validated native-query prefilters. Implementations that do not support these filters + * may return a superset because the native query retains the original filter for final evaluation. + */ + default List getTaskStatusListWithFilter( + final Map taskLookups, + final TaskStorageQueryFilter filter + ) + { + return getTaskStatusList(taskLookups, null); + } + default List getTaskInfos( TaskLookup taskLookup, @Nullable String datasource diff --git a/indexing-service/src/main/java/org/apache/druid/metadata/SQLMetadataStorageActionHandler.java b/indexing-service/src/main/java/org/apache/druid/metadata/SQLMetadataStorageActionHandler.java index 59fbe86a0ad6..13add69bda7e 100644 --- a/indexing-service/src/main/java/org/apache/druid/metadata/SQLMetadataStorageActionHandler.java +++ b/indexing-service/src/main/java/org/apache/druid/metadata/SQLMetadataStorageActionHandler.java @@ -73,6 +73,7 @@ public abstract class SQLMetadataStorageActionHandler { private static final EmittingLogger log = new EmittingLogger(SQLMetadataStorageActionHandler.class); private static final String CONTEXT_KEY_IS_TRANSIENT = "isTransient"; + private static final TaskStorageQueryFilter NO_FILTERS = new TaskStorageQueryFilter(List.of()); private final SQLMetadataConnector connector; private final ObjectMapper jsonMapper; @@ -305,7 +306,12 @@ public List getTaskInfos( final Query> query; switch (entry.getKey()) { case ACTIVE: - query = createActiveTaskStreamingQuery(handle, dataSource); + query = createActiveTaskStreamingQuery( + handle, + dataSource, + NO_FILTERS, + true + ); tasks.addAll(query.map(taskInfoMapper).list()); break; case COMPLETE: @@ -314,7 +320,9 @@ public List getTaskInfos( handle, completeTaskLookup.getTasksCreatedPriorTo(), completeTaskLookup.getMaxTaskStatuses(), - dataSource + dataSource, + NO_FILTERS, + true ); tasks.addAll(query.map(taskInfoMapper).list()); break; @@ -344,7 +352,25 @@ public List getTaskStatusList( log.info(e, "Exception getting task migration future"); } } - return getTaskStatusList(taskLookups, dataSource, fetchPayload); + return getTaskStatusList(taskLookups, dataSource, NO_FILTERS, fetchPayload); + } + + @Override + public List getTaskStatusListWithFilter( + final Map taskLookups, + final TaskStorageQueryFilter filter + ) + { + boolean fetchPayload = true; + if (taskMigrationCompleteFuture != null && taskMigrationCompleteFuture.isDone()) { + try { + fetchPayload = !taskMigrationCompleteFuture.get(); + } + catch (Exception e) { + log.info(e, "Exception getting task migration future"); + } + } + return getTaskStatusList(taskLookups, null, filter, fetchPayload); } @VisibleForTesting @@ -354,8 +380,30 @@ List getTaskStatusList( boolean fetchPayload ) { - ResultSetMapper resultSetMapper = + return getTaskStatusList(taskLookups, dataSource, NO_FILTERS, fetchPayload); + } + + @VisibleForTesting + List getTaskStatusListWithFilter( + final Map taskLookups, + final TaskStorageQueryFilter pushdownFilters, + final boolean fetchPayload + ) + { + return getTaskStatusList(taskLookups, null, pushdownFilters, fetchPayload); + } + + private List getTaskStatusList( + final Map taskLookups, + @Nullable final String dataSource, + final TaskStorageQueryFilter filter, + final boolean fetchPayload + ) + { + final ResultSetMapper resultSetMapper = fetchPayload ? taskStatusMapperFromPayload : taskStatusMapper; + // Status is represented by the ACTIVE and COMPLETE lookups, not by a physical metadata column. + final TaskStorageQueryFilter metadataStorageFilter = filter.withoutStatusFilters(); return getConnector().retryTransaction( (handle, status) -> { final List taskMetadataInfos = new ArrayList<>(); @@ -363,18 +411,27 @@ List getTaskStatusList( final Query> query; switch (entry.getKey()) { case ACTIVE: - query = fetchPayload - ? createActiveTaskStreamingQuery(handle, dataSource) - : createActiveTaskSummaryStreamingQuery(handle, dataSource); + if (!filter.includesActiveTasks()) { + break; + } + query = createActiveTaskStreamingQuery(handle, dataSource, metadataStorageFilter, fetchPayload); taskMetadataInfos.addAll(query.map(resultSetMapper).list()); break; case COMPLETE: + if (!filter.includesCompleteTasks()) { + break; + } CompleteTaskLookup completeTaskLookup = (CompleteTaskLookup) entry.getValue(); DateTime priorTo = completeTaskLookup.getTasksCreatedPriorTo(); Integer limit = completeTaskLookup.getMaxTaskStatuses(); - query = fetchPayload - ? createCompletedTaskStreamingQuery(handle, priorTo, limit, dataSource) - : createCompletedTaskSummaryStreamingQuery(handle, priorTo, limit, dataSource); + query = createCompletedTaskStreamingQuery( + handle, + priorTo, + limit, + dataSource, + metadataStorageFilter, + fetchPayload + ); taskMetadataInfos.addAll(query.map(resultSetMapper).list()); break; default: @@ -412,81 +469,36 @@ private DruidException wrapInDruidException(String taskId, Throwable t) } /** - * Fetches the columns needed to build TaskStatusPlus for completed tasks - * Please note that this requires completion of data migration to avoid empty values for task type and groupId - * Recommended for GET /tasks API + * Fetches completed tasks, selecting the task payload only while it is still needed to derive task type and group ID. * Uses streaming SQL query to avoid fetching too many rows at once into memory * @param handle db handle - * @param dataSource datasource to which the tasks belong. null if we don't want to filter - * @return Query object for TaskStatusPlus for completed tasks of interest - */ - private Query> createCompletedTaskSummaryStreamingQuery( - Handle handle, - DateTime timestamp, - @Nullable Integer maxNumStatuses, - @Nullable String dataSource - ) - { - String sql = StringUtils.format( - "SELECT " - + " id, " - + " created_date, " - + " datasource, " - + " group_id, " - + " type, " - + " status_payload " - + "FROM " - + " %s " - + "WHERE " - + getWhereClauseForInactiveStatusesSinceQuery(dataSource) - + "ORDER BY created_date DESC", - getEntryTable() - ); - - if (maxNumStatuses != null) { - sql = decorateSqlWithLimit(sql); - } - Query> query = handle.createQuery(sql) - .bind("start", timestamp.toString()) - .setFetchSize(connector.getStreamingFetchSize()); - - if (maxNumStatuses != null) { - query = query.bind("n", maxNumStatuses); - } - if (dataSource != null) { - query = query.bind("ds", dataSource); - } - return query; - } - - /** - * Fetches the columns needed to build a Task object with payload for completed tasks - * This requires the task payload which can be large. Please use only when necessary. - * For example for ingestion tasks view before migration of the new columns - * Uses streaming SQL query to avoid fetching too many rows at once into memory - * @param handle db handle - * @param dataSource datasource to which the tasks belong. null if we don't want to filter - * @return Query object for completed TaskInfos of interest + * @param filter validated native-query prefilters + * @param fetchPayload whether to select the task payload instead of the migrated task columns + * @return query for completed tasks of interest */ private Query> createCompletedTaskStreamingQuery( - Handle handle, - DateTime timestamp, - @Nullable Integer maxNumStatuses, - @Nullable String dataSource + final Handle handle, + final DateTime timestamp, + @Nullable final Integer maxNumStatuses, + @Nullable final String dataSource, + final TaskStorageQueryFilter filter, + final boolean fetchPayload ) { + final SqlPredicateBuilder filterSqlBuilder = new SqlPredicateBuilder( + dataSource, + filter.getFilters(), + getSqlDialect() + ); + final String columns = fetchPayload + ? "id, status_payload, created_date, datasource, payload" + : "id, created_date, datasource, group_id, type, status_payload"; String sql = StringUtils.format( - "SELECT " - + " id, " - + " status_payload, " - + " created_date, " - + " datasource, " - + " payload " - + "FROM " - + " %s " + "SELECT %s FROM %s " + "WHERE " - + getWhereClauseForInactiveStatusesSinceQuery(dataSource) + + "active = FALSE AND created_date >= :start " + filterSqlBuilder.getSql() + "ORDER BY created_date DESC", + columns, getEntryTable() ); @@ -500,99 +512,52 @@ private Query> createCompletedTaskStreamingQuery( if (maxNumStatuses != null) { query = query.bind("n", maxNumStatuses); } - if (dataSource != null) { - query = query.bind("ds", dataSource); - } - return query; + return filterSqlBuilder.bind(query); } protected abstract String decorateSqlWithLimit(String sql); - private String getWhereClauseForInactiveStatusesSinceQuery(@Nullable String datasource) + /** Returns the metadata-database capabilities used for task-filter SQL generation. */ + protected SqlDialect getSqlDialect() { - String sql = StringUtils.format("active = FALSE AND created_date >= :start "); - if (datasource != null) { - sql += " AND datasource = :ds "; - } - return sql; + return SqlDialect.NONE; } /** - * Fetches the columns needed to build TaskStatusPlus for active tasks - * Please note that this requires completion of data migration to avoid empty values for task type and groupId - * Recommended for GET /tasks API + * Fetches active tasks, selecting the task payload only while it is still needed to derive task type and group ID. * Uses streaming SQL query to avoid fetching too many rows at once into memory * @param handle db handle - * @param dataSource datasource to which the tasks belong. null if we don't want to filter - * @return Query object for TaskStatusPlus for active tasks of interest + * @param filter validated native-query prefilters + * @param fetchPayload whether to select the task payload instead of the migrated task columns + * @return query for active tasks of interest */ - private Query> createActiveTaskSummaryStreamingQuery(Handle handle, @Nullable String dataSource) + private Query> createActiveTaskStreamingQuery( + final Handle handle, + @Nullable final String dataSource, + final TaskStorageQueryFilter filter, + final boolean fetchPayload + ) { - String sql = StringUtils.format( - "SELECT " - + " id, " - + " status_payload, " - + " group_id, " - + " type, " - + " datasource, " - + " created_date " - + "FROM " - + " %s " - + "WHERE " - + getWhereClauseForActiveStatusesQuery(dataSource) - + "ORDER BY created_date", - entryTable + final SqlPredicateBuilder filterSqlBuilder = new SqlPredicateBuilder( + dataSource, + filter.getFilters(), + getSqlDialect() ); - - Query> query = handle.createQuery(sql) - .setFetchSize(connector.getStreamingFetchSize()); - if (dataSource != null) { - query = query.bind("ds", dataSource); - } - return query; - } - - /** - * Fetches the columns needed to build Task objects with payload for active tasks - * This requires the task payload which can be large. Please use only when necessary. - * For example for ingestion tasks view before migration of the new columns - * Uses streaming SQL query to avoid fetching too many rows at once into memory - * @param handle db handle - * @param dataSource datasource to which the tasks belong. null if we don't want to filter - * @return Query object for active TaskInfos of interest - */ - private Query> createActiveTaskStreamingQuery(Handle handle, @Nullable String dataSource) - { - String sql = StringUtils.format( - "SELECT " - + " id, " - + " status_payload, " - + " payload, " - + " datasource, " - + " created_date " - + "FROM " - + " %s " + final String columns = fetchPayload + ? "id, status_payload, payload, datasource, created_date" + : "id, status_payload, group_id, type, datasource, created_date"; + final String sql = StringUtils.format( + "SELECT %s FROM %s " + "WHERE " - + getWhereClauseForActiveStatusesQuery(dataSource) + + "active = TRUE " + filterSqlBuilder.getSql() + "ORDER BY created_date", + columns, entryTable ); - Query> query = handle.createQuery(sql) - .setFetchSize(connector.getStreamingFetchSize()); - if (dataSource != null) { - query = query.bind("ds", dataSource); - } - return query; - } - - private String getWhereClauseForActiveStatusesQuery(String dataSource) - { - String sql = StringUtils.format("active = TRUE "); - if (dataSource != null) { - sql += " AND datasource = :ds "; - } - return sql; + final Query> query = handle.createQuery(sql) + .setFetchSize(connector.getStreamingFetchSize()); + return filterSqlBuilder.bind(query); } private class TaskStatusMapperFromPayload implements ResultSetMapper diff --git a/indexing-service/src/main/java/org/apache/druid/metadata/SqlDialect.java b/indexing-service/src/main/java/org/apache/druid/metadata/SqlDialect.java new file mode 100644 index 000000000000..f3a2cd0c3426 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/metadata/SqlDialect.java @@ -0,0 +1,79 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.metadata; + +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.filter.EqualityFilter; +import org.apache.druid.query.filter.InDimFilter; +import org.apache.druid.query.filter.OrDimFilter; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.query.filter.TypedInFilter; +import org.apache.druid.segment.column.ColumnType; + +import java.util.Objects; + +/** + * Metadata-database capabilities for translating validated task-storage filters into SQL predicates. + * + *

Every supported translation must return a superset of the rows matched by the corresponding Druid filter. + * The original Druid filter remains on the Broker query and removes any false positives. A dialect must decline a + * translation if its database comparison, collation, or null semantics could introduce false negatives.

+ */ +public enum SqlDialect +{ + NONE(false), + DERBY(true), + MYSQL(true), + POSTGRESQL(true), + SQL_SERVER(true); + + private final boolean supportsEquality; + + SqlDialect(final boolean supportsEquality) + { + this.supportsEquality = supportsEquality; + } + + /** + * Equality and {@code IN} predicates use bound string parameters and are safe candidate filters for the configured + * metadata stores. + */ + public boolean supports(final DimFilter filter) + { + return switch (filter) { + case SelectorDimFilter selector -> supportsEquality + && selector.getValue() != null + && selector.getExtractionFn() == null; + case EqualityFilter equality -> supportsEquality + && ColumnType.STRING.equals(equality.getMatchValueType()) + && equality.getMatchValue() instanceof String; + case InDimFilter in -> supportsEquality + && in.getExtractionFn() == null + && !in.getValues().isEmpty() + && in.getValues().stream().allMatch(Objects::nonNull); + case TypedInFilter in -> supportsEquality + && ColumnType.STRING.equals(in.getMatchValueType()) + && !in.getSortedValues().isEmpty() + && in.getSortedValues().stream().allMatch(String.class::isInstance); + case OrDimFilter or -> !or.getFields().isEmpty() && or.getFields().stream().allMatch(this::supports); + default -> false; + }; + } +} diff --git a/indexing-service/src/main/java/org/apache/druid/metadata/SqlPredicateBuilder.java b/indexing-service/src/main/java/org/apache/druid/metadata/SqlPredicateBuilder.java new file mode 100644 index 000000000000..5181d2f1cf19 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/metadata/SqlPredicateBuilder.java @@ -0,0 +1,133 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.metadata; + +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.server.system.table.SystemTablePushdownFilter; +import org.skife.jdbi.v2.Query; + +import javax.annotation.Nullable; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * Converts validated task storage filters to a SQL predicate and its corresponding bindings. + * + *

SQL generation and parameter binding intentionally happen together for each filter type so that a generated + * placeholder cannot get out of sync with the bindings applied to the query.

+ */ +final class SqlPredicateBuilder +{ + private final StringBuilder sql = new StringBuilder(); + private final Map bindings = new LinkedHashMap<>(); + + SqlPredicateBuilder( + @Nullable final String dataSource, + final List filters, + final SqlDialect dialect + ) + { + if (dataSource != null) { + appendStringValues("datasource", "legacy_datasource_filter", List.of(dataSource), false); + } + for (int i = 0; i < filters.size(); i++) { + final DimFilter filter = filters.get(i); + // Planning extracts possible storage prefilters without removing the original Druid filter from either the + // component Scan or the Broker query. Skipping a dialect-unsupported predicate therefore returns more candidate + // rows but cannot change the final result. + if (dialect.supports(filter)) { + appendStringValuesFilter(filter, "task_filter_" + i); + } + } + } + + String getSql() + { + return sql.toString(); + } + + Query bind(Query query) + { + for (final Map.Entry binding : bindings.entrySet()) { + query = query.bind(binding.getKey(), binding.getValue()); + } + return query; + } + + private void appendStringValuesFilter(final DimFilter filter, final String parameterPrefix) + { + final String column = SystemTablePushdownFilter.getStringValuesColumn(filter); + appendStringValues( + column, + parameterPrefix, + SystemTablePushdownFilter.getStringValues(filter).stream().sorted().toList(), + isTaskMigrationNullableColumn(column) + ); + } + + private void appendStringValues( + final String column, + final String parameterPrefix, + final List values, + final boolean includeNull + ) + { + sql.append(" AND "); + if (includeNull) { + sql.append('('); + } + sql.append(column); + if (values.size() == 1) { + sql.append(" = :").append(parameterPrefix).append("_0"); + } else { + sql.append(" IN ("); + appendParameters(parameterPrefix, values); + sql.append(')'); + } + if (includeNull) { + sql.append(" OR ").append(column).append(" IS NULL)"); + } + sql.append(' '); + addBindings(parameterPrefix, values); + } + + private void appendParameters(final String parameterPrefix, final List values) + { + for (int i = 0; i < values.size(); i++) { + if (i > 0) { + sql.append(", "); + } + sql.append(':').append(parameterPrefix).append('_').append(i); + } + } + + private void addBindings(final String parameterPrefix, final List values) + { + for (int i = 0; i < values.size(); i++) { + bindings.put(parameterPrefix + '_' + i, values.get(i)); + } + } + + private static boolean isTaskMigrationNullableColumn(final String column) + { + return "group_id".equals(column) || "type".equals(column); + } +} diff --git a/indexing-service/src/main/java/org/apache/druid/metadata/TaskStorageQueryFilter.java b/indexing-service/src/main/java/org/apache/druid/metadata/TaskStorageQueryFilter.java new file mode 100644 index 000000000000..6a3ccc304751 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/metadata/TaskStorageQueryFilter.java @@ -0,0 +1,129 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.metadata; + +import org.apache.druid.indexer.TaskState; +import org.apache.druid.query.filter.BoundDimFilter; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.filter.EqualityFilter; +import org.apache.druid.query.filter.InDimFilter; +import org.apache.druid.query.filter.LikeDimFilter; +import org.apache.druid.query.filter.NotDimFilter; +import org.apache.druid.query.filter.OrDimFilter; +import org.apache.druid.query.filter.RangeFilter; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.query.filter.TypedInFilter; +import org.apache.druid.server.system.table.SystemTablePushdownFilter; + +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +/** Immutable wrapper for the validated native filters pushed into task metadata queries. */ +public final class TaskStorageQueryFilter +{ + private final List filters; + private final boolean includesActiveTasks; + private final boolean includesCompleteTasks; + + public TaskStorageQueryFilter(final List filters) + { + this.filters = List.copyOf(filters); + final Set possibleTaskActiveness = computePossibleTaskActiveness(); + this.includesActiveTasks = possibleTaskActiveness.contains(true); + this.includesCompleteTasks = possibleTaskActiveness.contains(false); + } + + public List getFilters() + { + return filters; + } + + public String getStringValuesColumn(final DimFilter filter) + { + return SystemTablePushdownFilter.getStringValuesColumn(filter); + } + + public Set getStringValues(final DimFilter filter) + { + return SystemTablePushdownFilter.getStringValues(filter); + } + + TaskStorageQueryFilter withoutStatusFilters() + { + return new TaskStorageQueryFilter( + filters.stream().filter(filter -> !"status".equals(getColumn(filter))).toList() + ); + } + + public boolean includesActiveTasks() + { + return includesActiveTasks; + } + + public boolean includesCompleteTasks() + { + return includesCompleteTasks; + } + + private Set getPossibleTaskActiveness(final DimFilter statusFilter) + { + final Set possibleActiveness = new HashSet<>(); + for (final String status : getStringValues(statusFilter)) { + if (TaskState.RUNNING.name().equals(status)) { + possibleActiveness.add(true); + } else if (TaskState.SUCCESS.name().equals(status) || TaskState.FAILED.name().equals(status)) { + possibleActiveness.add(false); + } + } + return possibleActiveness; + } + + private boolean isStringValuesFilter(final DimFilter filter) + { + return filter instanceof SelectorDimFilter + || filter instanceof EqualityFilter + || filter instanceof InDimFilter + || filter instanceof TypedInFilter + || filter instanceof OrDimFilter; + } + + private Set computePossibleTaskActiveness() + { + final Set possibleActiveness = new HashSet<>(Set.of(true, false)); + for (final DimFilter filter : filters) { + if (isStringValuesFilter(filter) && "status".equals(getStringValuesColumn(filter))) { + possibleActiveness.retainAll(getPossibleTaskActiveness(filter)); + } + } + return possibleActiveness; + } + + private String getColumn(final DimFilter filter) + { + return switch (filter) { + case BoundDimFilter bound -> bound.getDimension(); + case RangeFilter range -> range.getColumn(); + case LikeDimFilter like -> like.getDimension(); + case NotDimFilter not -> getColumn(not.getField()); + default -> getStringValuesColumn(filter); + }; + } +} diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/task/TasksTableDataProviderTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/task/TasksTableDataProviderTest.java new file mode 100644 index 000000000000..c46e55ef78e6 --- /dev/null +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/task/TasksTableDataProviderTest.java @@ -0,0 +1,311 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.indexing.overlord.task; + +import com.google.common.base.Optional; +import org.apache.druid.indexing.overlord.GlobalTaskLockbox; +import org.apache.druid.indexing.overlord.TaskMaster; +import org.apache.druid.indexing.overlord.TaskQueryTool; +import org.apache.druid.indexing.overlord.TaskRunner; +import org.apache.druid.indexing.overlord.TaskStorage; +import org.apache.druid.indexing.overlord.http.TaskStateLookup; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.metadata.TaskStorageQueryFilter; +import org.apache.druid.query.Druids; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryDataSource; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.query.filter.AndDimFilter; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.filter.EqualityFilter; +import org.apache.druid.query.filter.LikeDimFilter; +import org.apache.druid.query.filter.OrDimFilter; +import org.apache.druid.query.filter.RangeFilter; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.query.filter.TypedInFilter; +import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.spec.LegacySegmentSpec; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.system.SystemTableNotLeaderException; +import org.apache.druid.server.system.table.SystemTablePushdownFilter; +import org.apache.druid.server.system.table.TaskTableDescriptor; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; + +import java.util.Collections; +import java.util.List; +import java.util.Set; + +public class TasksTableDataProviderTest +{ + @Test + public void testAdvertisesAllPushdownFiltersForSupportingStorage() + { + final TaskQueryTool taskQueryTool = Mockito.mock(TaskQueryTool.class); + Mockito.when(taskQueryTool.supportsTaskStatusQueryFilter()).thenReturn(true); + + final List filters = dataProvider(taskQueryTool).getPushdownFilters(); + + Assertions.assertEquals(6, filters.size()); + } + + @Test + public void testAdvertisesOnlyStatusPushdownForUnsupportedStorage() + { + final TaskQueryTool taskQueryTool = Mockito.mock(TaskQueryTool.class); + + Assertions.assertEquals( + List.of(new SystemTablePushdownFilter("status", null)), + dataProvider(taskQueryTool).getPushdownFilters() + ); + } + + @Test + public void testPassesExtractedFiltersDirectlyToTaskQueryTool() + { + final TaskQueryTool taskQueryTool = Mockito.mock(TaskQueryTool.class); + final TaskMaster taskMaster = Mockito.mock(TaskMaster.class); + Mockito.when(taskMaster.isHalfOrFullLeader()).thenReturn(true); + final List filters = List.of(new SelectorDimFilter("datasource", "native_sys_a", null)); + Mockito.when( + taskQueryTool.getTaskStatusPlusList( + Mockito.eq(TaskStateLookup.ALL), + Mockito.any(TaskStorageQueryFilter.class) + ) + ) + .thenReturn(Collections.emptyList()); + final TasksTableDataProvider supplier = new TasksTableDataProvider( + taskQueryTool, + taskMaster + ); + + supplier.getRows(filters, Mockito.mock(AuthenticationResult.class)); + + final ArgumentCaptor filtersCaptor = ArgumentCaptor.forClass(TaskStorageQueryFilter.class); + Mockito.verify(taskQueryTool).getTaskStatusPlusList(Mockito.eq(TaskStateLookup.ALL), filtersCaptor.capture()); + Assertions.assertEquals(filters, filtersCaptor.getValue().getFilters()); + } + + /** A standby Overlord rejects {@code sys.tasks} instead of returning an incorrectly empty table. */ + @Test + public void testRejectsRowsOnStandbyOverlord() + { + final TaskMaster taskMaster = Mockito.mock(TaskMaster.class); + final TasksTableDataProvider provider = new TasksTableDataProvider( + Mockito.mock(TaskQueryTool.class), + taskMaster + ); + + Assertions.assertThrows( + SystemTableNotLeaderException.class, + () -> provider.getRows(Collections.emptyList(), Mockito.mock(AuthenticationResult.class)) + ); + } + + /** A leadership loss before the task-runner snapshot fails instead of returning an incorrectly empty table. */ + @Test + public void testRejectsRowsWhenTaskRunnerDisappearsDuringLeaderTransition() + { + final TaskMaster taskMaster = Mockito.mock(TaskMaster.class); + Mockito.when(taskMaster.isHalfOrFullLeader()).thenReturn(true); + Mockito.when(taskMaster.getTaskRunner()).thenReturn(Optional.absent()); + final TaskQueryTool taskQueryTool = new TaskQueryTool( + Mockito.mock(TaskStorage.class), + Mockito.mock(GlobalTaskLockbox.class), + taskMaster, + null + ); + final TasksTableDataProvider provider = new TasksTableDataProvider( + taskQueryTool, + taskMaster + ); + + Assertions.assertThrows( + IllegalStateException.class, + () -> provider.getRows(Collections.emptyList(), Mockito.mock(AuthenticationResult.class)) + ); + } + + /** Losing the selected task runner while reading also fails instead of returning a stale or empty table. */ + @Test + public void testRejectsRowsWhenTaskRunnerDisappearsAfterSnapshot() + { + final TaskMaster taskMaster = Mockito.mock(TaskMaster.class); + final TaskRunner taskRunner = Mockito.mock(TaskRunner.class); + final TaskStorage taskStorage = Mockito.mock(TaskStorage.class); + Mockito.when(taskMaster.isHalfOrFullLeader()).thenReturn(true); + Mockito.when(taskMaster.getTaskRunner()).thenReturn(Optional.of(taskRunner), Optional.absent()); + Mockito.when(taskRunner.getKnownTasks()).thenReturn(Collections.emptyList()); + Mockito.when( + taskStorage.getTaskStatusPlusListWithFilter( + Mockito.anyMap(), + Mockito.any(TaskStorageQueryFilter.class) + ) + ).thenReturn(Collections.emptyList()); + final TaskQueryTool taskQueryTool = new TaskQueryTool( + taskStorage, + Mockito.mock(GlobalTaskLockbox.class), + taskMaster, + null + ); + final TasksTableDataProvider provider = new TasksTableDataProvider( + taskQueryTool, + taskMaster + ); + + Assertions.assertThrows( + IllegalStateException.class, + () -> provider.getRows(Collections.emptyList(), Mockito.mock(AuthenticationResult.class)) + ); + } + + @Test + public void testExtractsExactDataSourceFromAndFilter() + { + final ScanQuery query = scanQuery( + new AndDimFilter( + new SelectorDimFilter("datasource", "native_sys_a", null), + new LikeDimFilter("task_id", "native_sys_mvp_%", null, null) + ) + ); + + final List extracted = extract(query); + + Assertions.assertEquals(2, extracted.size()); + Assertions.assertInstanceOf(SelectorDimFilter.class, extracted.get(0)); + Assertions.assertEquals(Set.of("native_sys_a"), new TaskStorageQueryFilter(extracted).getStringValues(extracted.get(0))); + Assertions.assertEquals("id", ((LikeDimFilter) extracted.get(1)).getDimension()); + } + + @Test + public void testPushesMultipleDataSourcesFromOrFilter() + { + final ScanQuery query = scanQuery( + new OrDimFilter( + new SelectorDimFilter("datasource", "native_sys_a", null), + new SelectorDimFilter("datasource", "native_sys_b", null) + ) + ); + + final List extracted = extract(query); + + Assertions.assertEquals(1, extracted.size()); + Assertions.assertEquals( + Set.of("native_sys_a", "native_sys_b"), + new TaskStorageQueryFilter(extracted).getStringValues(extracted.get(0)) + ); + } + + @Test + public void testExtractsNativeEqualityFilter() + { + final EqualityFilter filter = new EqualityFilter("datasource", ColumnType.STRING, "native_sys_a", null); + + final List extracted = extract(scanQuery(filter)); + + Assertions.assertEquals(List.of(filter), extracted); + Assertions.assertEquals(Set.of("native_sys_a"), new TaskStorageQueryFilter(extracted).getStringValues(extracted.get(0))); + } + + @Test + public void testExtractsDataSourceFromWindowLeafScan() + { + final ScanQuery scanQuery = scanQuery(new SelectorDimFilter("datasource", "native_sys_a", null)); + final WindowOperatorQuery windowQuery = new WindowOperatorQuery( + new QueryDataSource(scanQuery), + new LegacySegmentSpec(Intervals.ETERNITY), + Collections.emptyMap(), + RowSignature.empty(), + Collections.emptyList(), + null + ); + + final List extracted = extract(windowQuery); + + Assertions.assertEquals(1, extracted.size()); + Assertions.assertEquals(Set.of("native_sys_a"), new TaskStorageQueryFilter(extracted).getStringValues(extracted.get(0))); + } + + @Test + public void testExtractsAllStorageBackedTaskFilters() + { + final ScanQuery query = scanQuery( + new AndDimFilter( + new TypedInFilter("task_id", ColumnType.STRING, List.of("task-a", "task-b"), null, null), + new SelectorDimFilter("group_id", "group-a", null), + new SelectorDimFilter("type", "noop", null), + new SelectorDimFilter("datasource", "native_sys_a", null), + new EqualityFilter("created_time", ColumnType.STRING, "2026-01-02T00:00:00.000Z", null), + new RangeFilter( + "created_time", + ColumnType.STRING, + "2026-01-01T00:00:00.000Z", + "2026-02-01T00:00:00.000Z", + false, + true, + null + ), + new SelectorDimFilter("status", "SUCCESS", null), + new LikeDimFilter("error_msg", "%ignored%", null, null) + ) + ); + + final List extracted = extract(query); + final TaskStorageQueryFilter pushdownFilters = new TaskStorageQueryFilter(extracted); + + Assertions.assertEquals(7, extracted.size()); + Assertions.assertEquals("id", pushdownFilters.getStringValuesColumn(extracted.get(0))); + Assertions.assertEquals(Set.of("task-a", "task-b"), pushdownFilters.getStringValues(extracted.get(0))); + Assertions.assertEquals(Set.of("group-a"), pushdownFilters.getStringValues(extracted.get(1))); + Assertions.assertEquals(Set.of("noop"), pushdownFilters.getStringValues(extracted.get(2))); + Assertions.assertEquals(Set.of("native_sys_a"), pushdownFilters.getStringValues(extracted.get(3))); + Assertions.assertEquals("created_date", pushdownFilters.getStringValuesColumn(extracted.get(4))); + Assertions.assertEquals("created_date", ((RangeFilter) extracted.get(5)).getColumn()); + Assertions.assertFalse(pushdownFilters.includesActiveTasks()); + Assertions.assertTrue(pushdownFilters.includesCompleteTasks()); + } + + private static List extract(final Query query) + { + final TaskQueryTool taskQueryTool = Mockito.mock(TaskQueryTool.class); + Mockito.when(taskQueryTool.supportsTaskStatusQueryFilter()).thenReturn(true); + final TasksTableDataProvider dataSupplier = dataProvider(taskQueryTool); + return SystemTablePushdownFilter.extract(query, dataSupplier.getPushdownFilters()); + } + + private static TasksTableDataProvider dataProvider(final TaskQueryTool taskQueryTool) + { + return new TasksTableDataProvider(taskQueryTool, null); + } + + private static ScanQuery scanQuery(final DimFilter filter) + { + return Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource(TaskTableDescriptor.TABLE_NAME)) + .intervals(new LegacySegmentSpec(Intervals.ETERNITY)) + .filters(filter) + .build(); + } +} diff --git a/indexing-service/src/test/java/org/apache/druid/metadata/SQLMetadataStorageActionHandlerTest.java b/indexing-service/src/test/java/org/apache/druid/metadata/SQLMetadataStorageActionHandlerTest.java index 303a915e430a..c71e21bfa41b 100644 --- a/indexing-service/src/test/java/org/apache/druid/metadata/SQLMetadataStorageActionHandlerTest.java +++ b/indexing-service/src/test/java/org/apache/druid/metadata/SQLMetadataStorageActionHandlerTest.java @@ -42,6 +42,13 @@ import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.metadata.TaskLookup.ActiveTaskLookup; import org.apache.druid.metadata.TaskLookup.CompleteTaskLookup; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.filter.EqualityFilter; +import org.apache.druid.query.filter.InDimFilter; +import org.apache.druid.query.filter.OrDimFilter; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.query.filter.TypedInFilter; +import org.apache.druid.segment.column.ColumnType; import org.joda.time.DateTime; import org.joda.time.DateTimeZone; import org.joda.time.Duration; @@ -451,6 +458,160 @@ public void testGetTaskStatusPlusListInternal() verifyTaskInfoToMetadataInfo(activeAltered, taskMetadataInfos, false); } + @Test + public void testGetTaskStatusListWithStorageFilters() + { + insertStorageFilterTestTasks(); + + assertFilteredTaskIds( + Set.of("task-1", "task-3"), + List.of(new TypedInFilter("id", ColumnType.STRING, List.of("task-1", "task-3"), null, null)) + ); + // Null group_id and type values are retained as a migration-safe superset. The native filter remains residual. + assertFilteredTaskIds( + Set.of("task-1", "task-3", "task-legacy"), + List.of(new SelectorDimFilter("group_id", "group-a", null)) + ); + assertFilteredTaskIds( + Set.of("task-1", "task-legacy"), + List.of(new SelectorDimFilter("type", "noop", null)) + ); + assertFilteredTaskIds( + Set.of("task-2"), + List.of(new SelectorDimFilter("datasource", "datasource-b", null)) + ); + assertFilteredTaskIds( + Set.of("task-3"), + List.of(new SelectorDimFilter("created_date", "2014-01-30T00:00:00.000Z", null)) + ); + assertFilteredTaskIds( + Set.of("task-1", "task-legacy"), + List.of(new SelectorDimFilter("status", TaskState.RUNNING.name(), null)) + ); + // SUCCESS and FAILED share the completed-task lookup; the native residual filter distinguishes them. + assertFilteredTaskIds( + Set.of("task-2", "task-3"), + List.of(new SelectorDimFilter("status", TaskState.SUCCESS.name(), null)) + ); + } + + /** + * Verifies SQL generation and binding for string-value filters: + * + * - `id = 'task-2'` + * - `id IN ('task-1', 'task-3')` + * - `id = 'task-1' OR id = 'task-3'` + */ + @Test + public void testGetTaskStatusListWithStringValueFilterVariants() + { + insertStorageFilterTestTasks(); + + assertFilteredTaskIds( + Set.of("task-2"), + List.of(new EqualityFilter("id", ColumnType.STRING, "task-2", null)) + ); + assertFilteredTaskIds( + Set.of("task-1", "task-3"), + List.of(new InDimFilter("id", List.of("task-1", "task-3"), null)) + ); + assertFilteredTaskIds( + Set.of("task-1", "task-3"), + List.of( + new OrDimFilter( + new SelectorDimFilter("id", "task-1", null), + new SelectorDimFilter("id", "task-3", null) + ) + ) + ); + } + + /** + * Verifies multiple predicates and the legacy datasource argument in the same metadata query: + * + * - `datasource = 'datasource-a' AND id IN ('task-1', 'task-3')` + * - legacy datasource filter `datasource = 'datasource-a'` + */ + @Test + public void testGetTaskStatusListWithCombinedFiltersAndLegacyDatasource() + { + insertStorageFilterTestTasks(); + + assertFilteredTaskIds( + Set.of("task-1", "task-3", "task-legacy"), + List.of( + new SelectorDimFilter("datasource", "datasource-a", null), + new InDimFilter("id", List.of("task-1", "task-3", "task-legacy"), null) + ) + ); + assertFilteredTaskIds(Set.of("task-1", "task-3", "task-legacy"), "datasource-a"); + } + + private void assertFilteredTaskIds( + final Set expectedTaskIds, + final List pushdownFilters + ) + { + final Set actualTaskIds = handler + .getTaskStatusListWithFilter(taskLookups(), new TaskStorageQueryFilter(pushdownFilters), false) + .stream() + .map(taskStatus -> taskStatus.getTaskIdentifier().getId()) + .collect(Collectors.toSet()); + Assertions.assertEquals(expectedTaskIds, actualTaskIds); + } + + private void assertFilteredTaskIds(final Set expectedTaskIds, final String dataSource) + { + final Set actualTaskIds = handler.getTaskStatusList(taskLookups(), dataSource, false) + .stream() + .map(taskStatus -> taskStatus.getTaskIdentifier().getId()) + .collect(Collectors.toSet()); + Assertions.assertEquals(expectedTaskIds, actualTaskIds); + } + + private static Map taskLookups() + { + final Map taskLookups = new HashMap<>(); + taskLookups.put(TaskLookup.TaskLookupType.ACTIVE, ActiveTaskLookup.getInstance()); + taskLookups.put( + TaskLookup.TaskLookupType.COMPLETE, + CompleteTaskLookup.withTasksCreatedPriorTo(null, DateTimes.of("2014-01-01")) + ); + return taskLookups; + } + + private void insertStorageFilterTestTasks() + { + insertTask("task-1", "group-a", "noop", "datasource-a", "2014-01-10", TaskState.RUNNING, true); + insertTask("task-2", "group-b", "other", "datasource-b", "2014-01-20", TaskState.SUCCESS, true); + insertTask("task-3", "group-a", "other", "datasource-a", "2014-01-30", TaskState.FAILED, true); + insertTask("task-legacy", "group-a", "noop", "datasource-a", "2014-01-15", TaskState.RUNNING, false); + } + + private void insertTask( + final String id, + final String groupId, + final String storedType, + final String dataSource, + final String createdTime, + final TaskState taskState, + final boolean migrated + ) + { + final Task task = new NoopTask(id, groupId, dataSource, 1L, 0L, null); + final TaskStatus status = new TaskStatus(id, taskState, 0L, null, TaskLocation.unknown()); + handler.insert( + id, + DateTimes.of(createdTime), + dataSource, + task, + TaskState.RUNNING.equals(taskState), + status, + migrated ? storedType : null, + migrated ? groupId : null + ); + } + private Integer getUnmigratedTaskCount() { return handler.getConnector().retryWithHandle( diff --git a/indexing-service/src/test/java/org/apache/druid/metadata/SqlDialectTest.java b/indexing-service/src/test/java/org/apache/druid/metadata/SqlDialectTest.java new file mode 100644 index 000000000000..f59202e76182 --- /dev/null +++ b/indexing-service/src/test/java/org/apache/druid/metadata/SqlDialectTest.java @@ -0,0 +1,97 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.metadata; + +import org.apache.druid.query.filter.EqualityFilter; +import org.apache.druid.query.filter.InDimFilter; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.segment.column.ColumnType; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.List; +import java.util.stream.Stream; + +public class SqlDialectTest +{ + private static final List EQUALITY_VALUES = List.of( + "CaseSensitive", + "casesensitive", + "café", + "cafe\u0301", + "trailing-space ", + "", + "%_wildcards", + "backslash\\escape" + ); + + /** Equality and IN remain eligible for every configured metadata database across sensitive string values. */ + @ParameterizedTest(name = "{0}: equality/IN value [{1}]") + @MethodSource("dialectAndEqualityValues") + public void testEqualityAndInCapabilities( + final SqlDialect dialect, + final String value + ) + { + Assertions.assertTrue(dialect.supports(new SelectorDimFilter("id", value, null))); + Assertions.assertTrue(dialect.supports(new EqualityFilter("id", ColumnType.STRING, value, null))); + Assertions.assertTrue(dialect.supports(new InDimFilter("id", List.of(value, "other"), null))); + Assertions.assertEquals( + " AND id = :task_filter_0_0 ", + new SqlPredicateBuilder( + null, + List.of(new SelectorDimFilter("id", value, null)), + dialect + ).getSql() + ); + } + + /** Null equality is not eligible for metadata SQL and remains a Broker residual filter. */ + @ParameterizedTest + @EnumSource(SqlDialect.class) + public void testNullEqualityIsDeclined(final SqlDialect dialect) + { + Assertions.assertFalse(dialect.supports(new SelectorDimFilter("id", null, null))); + } + + /** An unknown metadata database declines even equality rather than risking an unsafe predicate. */ + @Test + public void testUnknownDialectDeclinesEquality() + { + Assertions.assertFalse( + SqlDialect.NONE.supports(new SelectorDimFilter("id", "task-1", null)) + ); + } + + private static Stream dialectAndEqualityValues() + { + return Stream.of( + SqlDialect.DERBY, + SqlDialect.MYSQL, + SqlDialect.POSTGRESQL, + SqlDialect.SQL_SERVER + ) + .flatMap(dialect -> EQUALITY_VALUES.stream().map(value -> Arguments.of(dialect, value))); + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/module/SystemTableModule.java b/server/src/main/java/org/apache/druid/server/system/module/SystemTableModule.java index 6e25c527f3ee..2f83d2347ea9 100644 --- a/server/src/main/java/org/apache/druid/server/system/module/SystemTableModule.java +++ b/server/src/main/java/org/apache/druid/server/system/module/SystemTableModule.java @@ -29,11 +29,13 @@ import org.apache.druid.server.system.table.ServerPropertiesTableDescriptor; import org.apache.druid.server.system.table.SystemTableDataProvider; import org.apache.druid.server.system.table.SystemTableDescriptor; +import org.apache.druid.server.system.table.TaskTableDescriptor; /** * Registers native system-table routing and the node-local server-properties supplier. * - *

Table-specific integrations contribute their own entries to the native system-table multibinders.

+ *

Table-specific integrations, such as the task supplier in indexing-service, contribute their own entries to the + * native system-table multibinders.

*/ public class SystemTableModule implements Module { @@ -48,6 +50,9 @@ public void configure(final Binder binder) final MapBinder descriptorBinder = MapBinder.newMapBinder(binder, String.class, SystemTableDescriptor.class); descriptorBinder.addBinding(ServerPropertiesTableDescriptor.TABLE_NAME) .toInstance(new ServerPropertiesTableDescriptor()); + descriptorBinder.addBinding(TaskTableDescriptor.TABLE_NAME) + .toInstance(new TaskTableDescriptor()); + final MapBinder dataProviderBinder = MapBinder.newMapBinder(binder, String.class, SystemTableDataProvider.class); dataProviderBinder.addBinding(ServerPropertiesTableDescriptor.TABLE_NAME) .to(ServerPropertiesTableDataProvider.class) diff --git a/server/src/main/java/org/apache/druid/server/system/table/TaskTableDescriptor.java b/server/src/main/java/org/apache/druid/server/system/table/TaskTableDescriptor.java new file mode 100644 index 000000000000..8015b596ce41 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/table/TaskTableDescriptor.java @@ -0,0 +1,105 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.server.system.table; + +import org.apache.druid.discovery.NodeRole; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.AuthorizationUtils; +import org.apache.druid.server.security.AuthorizerMapper; + +import java.util.Collections; +import java.util.Set; + +/** Descriptor for the native {@code sys.tasks} table. */ +public class TaskTableDescriptor implements SystemTableDescriptor +{ + public static final String TABLE_NAME = "tasks"; + public static final RowSignature ROW_SIGNATURE = RowSignature + .builder() + .add("task_id", ColumnType.STRING) + .add("group_id", ColumnType.STRING) + .add("type", ColumnType.STRING) + .add("datasource", ColumnType.STRING) + .add("created_time", ColumnType.STRING) + .add("queue_insertion_time", ColumnType.STRING) + .add("status", ColumnType.STRING) + .add("runner_status", ColumnType.STRING) + .add("duration", ColumnType.LONG) + .add("location", ColumnType.STRING) + .add("host", ColumnType.STRING) + .add("plaintext_port", ColumnType.LONG) + .add("tls_port", ColumnType.LONG) + .add("error_msg", ColumnType.STRING) + .build(); + + private static final Set NODE_ROLES = Set.of(NodeRole.OVERLORD); + private static final int DATASOURCE_COLUMN = ROW_SIGNATURE.indexOf("datasource"); + private static final SystemTableRowAuthorizer ROW_AUTHORIZER = new SystemTableRowAuthorizer() + { + @Override + public Iterable filterAuthorizedRows( + final Iterable rows, + final AuthenticationResult authenticationResult, + final AuthorizerMapper authorizerMapper + ) + { + return AuthorizationUtils.filterAuthorizedResources( + authenticationResult, + rows, + row -> Collections.singletonList( + AuthorizationUtils.DATASOURCE_READ_RA_GENERATOR.apply((String) row[DATASOURCE_COLUMN]) + ), + authorizerMapper + ); + } + }; + + @Override + public String getTableName() + { + return TABLE_NAME; + } + + @Override + public Set getNodeRoles() + { + return NODE_ROLES; + } + + @Override + public SystemTableRoutingMode getRoutingMode() + { + return SystemTableRoutingMode.LEADER_ONLY; + } + + @Override + public RowSignature getRowSignature() + { + return ROW_SIGNATURE; + } + + @Override + public SystemTableRowAuthorizer getRowAuthorizer() + { + return ROW_AUTHORIZER; + } +} diff --git a/server/src/test/java/org/apache/druid/server/system/handler/SystemTableQueryClientTest.java b/server/src/test/java/org/apache/druid/server/system/handler/SystemTableQueryClientTest.java index 3e7273efc209..d5931fcbb5ef 100644 --- a/server/src/test/java/org/apache/druid/server/system/handler/SystemTableQueryClientTest.java +++ b/server/src/test/java/org/apache/druid/server/system/handler/SystemTableQueryClientTest.java @@ -83,7 +83,7 @@ import org.apache.druid.server.system.SystemTableNotLeaderException; import org.apache.druid.server.system.table.ServerPropertiesTableDescriptor; import org.apache.druid.server.system.table.SystemTableDescriptor; -import org.apache.druid.server.system.table.SystemTableRoutingMode; +import org.apache.druid.server.system.table.TaskTableDescriptor; import org.joda.time.Duration; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -933,7 +933,7 @@ public void testServerPropertiesMidStreamTimeoutIsPropagated() @Test public void testLeaderOnlyQueryRetriesAfterLeadershipChange() { - final SystemTableDescriptor descriptor = new LeaderOnlySystemTableDescriptor(); + final SystemTableDescriptor descriptor = new TaskTableDescriptor(); final SystemTableNode oldLeader = systemTableNode(testNode(8081)); final SystemTableNode newLeader = systemTableNode(testNode(8082)); final SystemTableNodeLocator nodeLocator = Mockito.mock(SystemTableNodeLocator.class); @@ -1407,19 +1407,4 @@ public org.apache.druid.server.system.table.SystemTableRowAuthorizer getRowAutho return (rows, authenticationResult, authorizerMapper) -> rows; } } - - private static class LeaderOnlySystemTableDescriptor extends TestSystemTableDescriptor - { - @Override - public String getTableName() - { - return "tasks"; - } - - @Override - public SystemTableRoutingMode getRoutingMode() - { - return SystemTableRoutingMode.LEADER_ONLY; - } - } } diff --git a/services/src/main/java/org/apache/druid/cli/CliOverlord.java b/services/src/main/java/org/apache/druid/cli/CliOverlord.java index b6af7205cdd3..dce73573c3b9 100644 --- a/services/src/main/java/org/apache/druid/cli/CliOverlord.java +++ b/services/src/main/java/org/apache/druid/cli/CliOverlord.java @@ -110,6 +110,7 @@ import org.apache.druid.indexing.overlord.setup.WorkerBehaviorConfig; import org.apache.druid.indexing.overlord.supervisor.SupervisorManager; import org.apache.druid.indexing.overlord.supervisor.SupervisorResource; +import org.apache.druid.indexing.overlord.task.TasksTableDataProvider; import org.apache.druid.indexing.scheduledbatch.ScheduledBatchTaskManager; import org.apache.druid.indexing.worker.config.WorkerConfig; import org.apache.druid.indexing.worker.shuffle.DeepStorageIntermediaryDataManager; @@ -143,6 +144,8 @@ import org.apache.druid.server.security.AuthenticationUtils; import org.apache.druid.server.security.Authenticator; import org.apache.druid.server.security.AuthenticatorMapper; +import org.apache.druid.server.system.table.SystemTableDataProvider; +import org.apache.druid.server.system.table.TaskTableDescriptor; import org.apache.druid.storage.local.LocalTmpStorageConfig; import org.apache.druid.tasklogs.TaskLogStreamer; import org.apache.druid.tasklogs.TaskLogs; @@ -216,6 +219,15 @@ protected List getModules(final boolean standalone) @Override public void configure(Binder binder) { + final MapBinder dataProviderBinder = MapBinder.newMapBinder( + binder, + String.class, + SystemTableDataProvider.class + ); + dataProviderBinder.addBinding(TaskTableDescriptor.TABLE_NAME) + .to(TasksTableDataProvider.class) + .in(LazySingleton.class); + validateCentralizedDatasourceSchemaConfig(properties); if (standalone) { diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeTasksTable.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeTasksTable.java new file mode 100644 index 000000000000..fd5fbae972ee --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeTasksTable.java @@ -0,0 +1,70 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.sql.calcite.schema; + +import org.apache.calcite.plan.RelOptTable; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.logical.LogicalTableScan; +import org.apache.calcite.schema.Schema; +import org.apache.druid.query.DataSource; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.server.system.table.TaskTableDescriptor; +import org.apache.druid.sql.calcite.table.DruidTable; + +/** Native-query representation of {@code sys.tasks}. */ +class NativeTasksTable extends DruidTable +{ + private static final DataSource DATA_SOURCE = new SystemTableDataSource("tasks"); + + NativeTasksTable() + { + super(TaskTableDescriptor.ROW_SIGNATURE); + } + + @Override + public DataSource getDataSource() + { + return DATA_SOURCE; + } + + @Override + public boolean isJoinable() + { + return false; + } + + @Override + public boolean isBroadcast() + { + return false; + } + + @Override + public Schema.TableType getJdbcTableType() + { + return Schema.TableType.SYSTEM_TABLE; + } + + @Override + public RelNode toRel(RelOptTable.ToRelContext context, RelOptTable table) + { + return LogicalTableScan.create(context.getCluster(), table, context.getTableHints()); + } +} diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java index 94ef050892c3..050c66a88174 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java @@ -74,6 +74,7 @@ import org.apache.druid.server.security.Resource; import org.apache.druid.server.security.ResourceAction; import org.apache.druid.server.security.ResourceType; +import org.apache.druid.server.system.table.TaskTableDescriptor; import org.apache.druid.sql.calcite.planner.PlannerConfig; import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.calcite.run.NativeSqlEngine; @@ -106,7 +107,6 @@ public class SystemSchema extends AbstractTableSchema public static final String SEGMENTS_TABLE = "segments"; public static final String SERVERS_TABLE = "servers"; public static final String SERVER_SEGMENTS_TABLE = "server_segments"; - public static final String TASKS_TABLE = "tasks"; public static final String SUPERVISOR_TABLE = "supervisors"; public static final String QUERIES_TABLE = "queries"; @@ -225,24 +225,6 @@ public static DruidTable getNativeSystemTable(final RelOptTable table) .add("segment_id", ColumnType.STRING) .build(); - static final RowSignature TASKS_SIGNATURE = RowSignature - .builder() - .add("task_id", ColumnType.STRING) - .add("group_id", ColumnType.STRING) - .add("type", ColumnType.STRING) - .add("datasource", ColumnType.STRING) - .add("created_time", ColumnType.STRING) - .add("queue_insertion_time", ColumnType.STRING) - .add("status", ColumnType.STRING) - .add("runner_status", ColumnType.STRING) - .add("duration", ColumnType.LONG) - .add("location", ColumnType.STRING) - .add("host", ColumnType.STRING) - .add("plaintext_port", ColumnType.LONG) - .add("tls_port", ColumnType.LONG) - .add("error_msg", ColumnType.STRING) - .build(); - static final RowSignature SUPERVISOR_SIGNATURE = RowSignature .builder() .add("supervisor_id", ColumnType.STRING) @@ -348,7 +330,7 @@ public Table getTable(String name) authenticationResult ); case SERVER_SEGMENTS_TABLE -> new ServerSegmentsTable(serverView, authorizerMapper, authenticationResult); - case TASKS_TABLE -> new TasksTable(overlordClient, authorizerMapper, authenticationResult); + case TaskTableDescriptor.TABLE_NAME -> new TasksTable(overlordClient, authorizerMapper, authenticationResult); case SUPERVISOR_TABLE -> new SupervisorsTable(overlordClient, authorizerMapper, authenticationResult); case SystemServerPropertiesTable.TABLE_NAME -> new SystemServerPropertiesTable( druidNodeDiscoveryProvider, @@ -974,7 +956,7 @@ public Enumerable scan(DataContext root) /** * This table contains row per task. */ - static class TasksTable extends AbstractTable implements ScannableTable + static class TasksTable extends AbstractTable implements ScannableTable, NativeSystemTable { private final OverlordClient overlordClient; private final AuthorizerMapper authorizerMapper; @@ -994,7 +976,7 @@ public TasksTable( @Override public RelDataType getRowType(RelDataTypeFactory typeFactory) { - return RowSignatures.toRelDataType(TASKS_SIGNATURE, typeFactory); + return RowSignatures.toRelDataType(TaskTableDescriptor.ROW_SIGNATURE, typeFactory); } @Override @@ -1003,6 +985,12 @@ public TableType getJdbcTableType() return TableType.SYSTEM_TABLE; } + @Override + public DruidTable asNativeTable() + { + return new NativeTasksTable(); + } + @Override public Enumerable scan(DataContext root) { diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchemaProvider.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchemaProvider.java index fe03f36203fe..86727c4c842c 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchemaProvider.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchemaProvider.java @@ -32,6 +32,7 @@ import org.apache.druid.rpc.indexing.OverlordClient; import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.server.system.table.TaskTableDescriptor; import org.apache.druid.sql.calcite.planner.PlannerConfig; import org.apache.druid.sql.http.SqlEngineRegistry; @@ -94,7 +95,7 @@ public static Set computeAllTableNames(final PlannerConfig plannerConfig allTableNames.add(SystemSchema.SEGMENTS_TABLE); allTableNames.add(SystemSchema.SERVERS_TABLE); allTableNames.add(SystemSchema.SERVER_SEGMENTS_TABLE); - allTableNames.add(SystemSchema.TASKS_TABLE); + allTableNames.add(TaskTableDescriptor.TABLE_NAME); allTableNames.add(SystemSchema.SUPERVISOR_TABLE); allTableNames.add(SystemServerPropertiesTable.TABLE_NAME); diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemSchemaTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemSchemaTest.java index 0a83956c1528..fbb027fdda54 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemSchemaTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemSchemaTest.java @@ -107,6 +107,7 @@ import org.apache.druid.server.security.AuthorizerMapper; import org.apache.druid.server.security.NoopEscalator; import org.apache.druid.server.security.ResourceType; +import org.apache.druid.server.system.table.TaskTableDescriptor; import org.apache.druid.sql.calcite.planner.PlannerConfig; import org.apache.druid.sql.calcite.run.SqlEngine; import org.apache.druid.sql.calcite.schema.SystemSchema.QueriesTable; @@ -1479,10 +1480,10 @@ public void testTasksTable() throws Exception Assertions.assertEquals("192.168.1.6", row1[10]); Assertions.assertEquals(8100L, row1[11]); Assertions.assertEquals(-1L, row1[12]); - Assertions.assertEquals(null, row1[13]); + Assertions.assertNull(row1[13]); // Verify value types. - verifyTypes(rows, SystemSchema.TASKS_SIGNATURE); + verifyTypes(rows, TaskTableDescriptor.ROW_SIGNATURE); } @Test diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemTableDataProviderTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemTableDataProviderTest.java index 4fb7561d84ce..c9334532940d 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemTableDataProviderTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemTableDataProviderTest.java @@ -36,6 +36,12 @@ public class SystemTableDataProviderTest @Test public void testNativeTablesExposeSystemMetadata() { + final NativeTasksTable tasks = new NativeTasksTable(); + Assertions.assertEquals("tasks", ((SystemTableDataSource) tasks.getDataSource()).getTable()); + Assertions.assertFalse(tasks.isJoinable()); + Assertions.assertFalse(tasks.isBroadcast()); + Assertions.assertEquals(Schema.TableType.SYSTEM_TABLE, tasks.getJdbcTableType()); + final NativeServerPropertiesTable serverProperties = new NativeServerPropertiesTable(); Assertions.assertEquals( "server_properties", @@ -62,7 +68,7 @@ public void testSystemSchemaSelectsNativeTablesOnlyWhenEnabledForNativeEngine() EasyMock.expect(bindableContext.getEngine()).andReturn(bindableEngine).once(); EasyMock.expect(enabledContext.useNativeQueryForSystemTables()).andReturn(true).once(); EasyMock.expect(enabledContext.getEngine()).andReturn(nativeEngine).once(); - EasyMock.expect(table.unwrap(NativeSystemTable.class)).andReturn(NativeServerPropertiesTable::new).once(); + EasyMock.expect(table.unwrap(NativeSystemTable.class)).andReturn(NativeTasksTable::new).once(); EasyMock.replay(table, nativeEngine, bindableEngine, disabledContext, bindableContext, enabledContext); Assertions.assertFalse(SystemSchema.canUseNativeSystemTable(table, disabledContext)); @@ -76,10 +82,10 @@ public void testSystemSchemaSelectsNativeTablesOnlyWhenEnabledForNativeEngine() public void testSystemSchemaGetsNativeRepresentation() { final RelOptTable table = EasyMock.createMock(RelOptTable.class); - EasyMock.expect(table.unwrap(NativeSystemTable.class)).andReturn(NativeServerPropertiesTable::new).once(); + EasyMock.expect(table.unwrap(NativeSystemTable.class)).andReturn(NativeTasksTable::new).once(); EasyMock.replay(table); - Assertions.assertInstanceOf(NativeServerPropertiesTable.class, SystemSchema.getNativeSystemTable(table)); + Assertions.assertInstanceOf(NativeTasksTable.class, SystemSchema.getNativeSystemTable(table)); EasyMock.verify(table); } From dac11b40fbd09e790ec6865086d2d2ecc96c25a2 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 9 Sep 2026 11:51:38 +0800 Subject: [PATCH 2/5] perf(sql): benchmark native sys.tasks filtering --- .../metadata/SysTasksTableBenchmark.java | 282 ++++++++++++++++++ docs/querying/sql-metadata-tables.md | 10 +- .../query/NativeSysTasksQueryTest.java | 42 --- 3 files changed, 287 insertions(+), 47 deletions(-) create mode 100644 benchmarks/src/test/java/org/apache/druid/metadata/SysTasksTableBenchmark.java diff --git a/benchmarks/src/test/java/org/apache/druid/metadata/SysTasksTableBenchmark.java b/benchmarks/src/test/java/org/apache/druid/metadata/SysTasksTableBenchmark.java new file mode 100644 index 000000000000..5150c91040da --- /dev/null +++ b/benchmarks/src/test/java/org/apache/druid/metadata/SysTasksTableBenchmark.java @@ -0,0 +1,282 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.metadata; + +import org.apache.druid.indexer.TaskIdStatus; +import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.jackson.JacksonUtils; +import org.apache.druid.metadata.TaskLookup.ActiveTaskLookup; +import org.apache.druid.metadata.TaskLookup.CompleteTaskLookup; +import org.apache.druid.metadata.TaskLookup.TaskLookupType; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.filter.EqualityFilter; +import org.apache.druid.segment.column.ColumnType; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; +import org.skife.jdbi.v2.PreparedBatch; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.function.Predicate; + +/** + * Compares the metadata-read portion of the bindable and native {@code sys.tasks} execution paths. + * + *

The bindable path requests every task from the Overlord and evaluates the SQL predicate afterward. The native + * path converts the same predicate to a validated {@link TaskStorageQueryFilter}, applies it in metadata storage, and + * retains the predicate as a residual filter. This benchmark intentionally excludes SQL planning and HTTP transport + * so it isolates the benefit of metadata filter pushdown.

+ * + *

The synthetic table contains 100,000 migrated task records split evenly between active and completed tasks. + * Datasource, type, and group values have controlled cardinalities so each query returns a stable number of rows.

+ * + *

Query cases:

+ *
    + *
  • The unfiltered task-list query issued by the web console: + * 
    {@code
+ *     WITH tasks AS (
+ *       SELECT "task_id", "group_id", "type", "datasource", "created_time", "location", "duration", "error_msg",
+ *         CASE
+ *           WHEN "error_msg" IN (
+ *             'Shutdown request from user',
+ *             'Canceled: Query canceled by user or by task shutdown.'
+ *           ) THEN 'CANCELED'
+ *           WHEN "status" = 'RUNNING' THEN "runner_status"
+ *           ELSE "status"
+ *         END AS "status"
+ *       FROM sys.tasks
+ *     )
+ *     SELECT "task_id", "group_id", "type", "datasource", "created_time", "location", "duration", "error_msg", "status"
+ *     FROM tasks
+ *     ORDER BY
+ *       CASE "status"
+ *         WHEN 'RUNNING' THEN 4
+ *         WHEN 'PENDING' THEN 3
+ *         WHEN 'WAITING' THEN 2
+ *         ELSE 1
+ *       END DESC,
+ *       "created_time" DESC
+ *     }
    + *
    • + *
  • {@code SELECT * FROM sys.tasks WHERE task_id = 'task-050000'}
    • + *
  • {@code SELECT * FROM sys.tasks WHERE datasource = 'datasource-000'}
    • + *
  • {@code SELECT * FROM sys.tasks WHERE datasource = 'datasource-000' AND type = 'type-0'}
    • + *
  • {@code SELECT * FROM sys.tasks WHERE group_id = 'group-0000'}
    • + *
+ * + *

The native planner maps {@code task_id} to metadata column {@code id}. The other benchmarked columns retain + * their names. Equality filters on {@code type} and {@code group_id} are both pushed down, although the current task + * table schema does not index either column.

+ */ +@State(Scope.Benchmark) +@Fork(value = 1, jvmArgsAppend = {"-Xmx4g"}) +@Warmup(iterations = 3, time = 2) +@Measurement(iterations = 5, time = 2) +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +public class SysTasksTableBenchmark +{ + private static final int NUM_TASKS = 100_000; + private static final int INSERT_BATCH_SIZE = 5_000; + private static final String CREATED_TIME = "2025-01-01T00:00:00.000Z"; + private static final String TARGET_TASK_ID = "task-050000"; + private static final String TARGET_DATASOURCE = "datasource-000"; + private static final String TARGET_TYPE = "type-0"; + private static final String TARGET_GROUP_ID = "group-0000"; + + public enum QueryCase + { + FULL_LIST, + TASK_ID, + DATASOURCE, + DATASOURCE_AND_TYPE, + GROUP_ID + } + + @Param + private QueryCase queryCase; + + private TestDerbyConnector connector; + private DerbyMetadataStorageActionHandler handler; + private Map taskLookups; + private TaskStorageQueryFilter pushdownFilter; + private Predicate residualFilter; + + @Setup(Level.Trial) + public void setup() throws ClassNotFoundException + { + // The benchmark uber-jar has a consolidated JDBC service file that does not list Derby. + Class.forName("org.apache.derby.jdbc.EmbeddedDriver"); + connector = new TestDerbyConnector(); + connector.createDatabase(); + + final MetadataStorageTablesConfig tablesConfig = connector.getMetadataTablesConfig(); + connector.prepareTaskEntryTable(tablesConfig.getTasksTable()); + handler = new DerbyMetadataStorageActionHandler( + connector, + new DefaultObjectMapper(), + tablesConfig.getTasksTable(), + tablesConfig.getTaskLockTable() + ); + insertTasks(tablesConfig.getTasksTable()); + + taskLookups = Map.of( + TaskLookupType.ACTIVE, + ActiveTaskLookup.getInstance(), + TaskLookupType.COMPLETE, + CompleteTaskLookup.withTasksCreatedPriorTo(null, DateTimes.of("2020-01-01")) + ); + pushdownFilter = new TaskStorageQueryFilter(filtersFor(queryCase)); + residualFilter = residualFilterFor(queryCase); + } + + @TearDown(Level.Trial) + public void tearDown() + { + connector.tearDown(); + } + + /** Simulates the bindable path: retrieve all tasks, then apply the SQL predicate. */ + @Benchmark + public void bindable(Blackhole blackhole) + { + consumeMatchingRows(handler.getTaskStatusList(taskLookups, null, false), blackhole); + } + + /** Simulates the native path: push eligible predicates into metadata storage, then apply the residual predicate. */ + @Benchmark + public void nativeQuery(Blackhole blackhole) + { + consumeMatchingRows(handler.getTaskStatusListWithFilter(taskLookups, pushdownFilter, false), blackhole); + } + + private void consumeMatchingRows(final List candidates, final Blackhole blackhole) + { + long matchedRows = 0; + for (final TaskIdStatus candidate : candidates) { + if (residualFilter.test(candidate)) { + blackhole.consume(candidate); + matchedRows++; + } + } + blackhole.consume(matchedRows); + } + + private void insertTasks(final String tasksTable) + { + final byte[] payload = new byte[]{'{', '}'}; + final byte[] activeStatus = JacksonUtils.toBytes(new DefaultObjectMapper(), TaskStatus.running("benchmark")); + final byte[] completeStatus = JacksonUtils.toBytes(new DefaultObjectMapper(), TaskStatus.success("benchmark")); + final String sql = StringUtils.format( + "INSERT INTO %s (id, created_date, datasource, payload, status_payload, active, type, group_id) " + + "VALUES (:id, :created_date, :datasource, :payload, :status_payload, :active, :type, :group_id)", + tasksTable + ); + + for (int batchStart = 0; batchStart < NUM_TASKS; batchStart += INSERT_BATCH_SIZE) { + final int start = batchStart; + final int end = Math.min(batchStart + INSERT_BATCH_SIZE, NUM_TASKS); + connector.retryWithHandle(handle -> { + final PreparedBatch batch = handle.prepareBatch(sql); + for (int i = start; i < end; i++) { + final boolean active = i % 2 == 0; + batch.add() + .bind("id", taskId(i)) + .bind("created_date", CREATED_TIME) + .bind("datasource", dataSource(i)) + .bind("payload", payload) + .bind("status_payload", active ? activeStatus : completeStatus) + .bind("active", active) + .bind("type", type(i)) + .bind("group_id", groupId(i)); + } + batch.execute(); + return null; + }); + } + } + + private static List filtersFor(final QueryCase query) + { + return switch (query) { + case FULL_LIST -> List.of(); + case TASK_ID -> List.of(equality("id", TARGET_TASK_ID)); + case DATASOURCE -> List.of(equality("datasource", TARGET_DATASOURCE)); + case DATASOURCE_AND_TYPE -> List.of( + equality("datasource", TARGET_DATASOURCE), + equality("type", TARGET_TYPE) + ); + case GROUP_ID -> List.of(equality("group_id", TARGET_GROUP_ID)); + }; + } + + private static Predicate residualFilterFor(final QueryCase query) + { + return switch (query) { + case FULL_LIST -> ignored -> true; + case TASK_ID -> task -> TARGET_TASK_ID.equals(task.getTaskIdentifier().getId()); + case DATASOURCE -> task -> TARGET_DATASOURCE.equals(task.getDataSource()); + case DATASOURCE_AND_TYPE -> task -> TARGET_DATASOURCE.equals(task.getDataSource()) + && TARGET_TYPE.equals(task.getTaskIdentifier().getType()); + case GROUP_ID -> task -> TARGET_GROUP_ID.equals(task.getTaskIdentifier().getGroupId()); + }; + } + + private static EqualityFilter equality(final String column, final String value) + { + return new EqualityFilter(column, ColumnType.STRING, value, null); + } + + private static String taskId(final int taskNumber) + { + return StringUtils.format("task-%06d", taskNumber); + } + + private static String dataSource(final int taskNumber) + { + return StringUtils.format("datasource-%03d", taskNumber % 100); + } + + private static String type(final int taskNumber) + { + return StringUtils.format("type-%d", (taskNumber / 100) % 10); + } + + private static String groupId(final int taskNumber) + { + return StringUtils.format("group-%04d", taskNumber % 1_000); + } +} diff --git a/docs/querying/sql-metadata-tables.md b/docs/querying/sql-metadata-tables.md index dc9c4cf03bde..4b25cb432828 100644 --- a/docs/querying/sql-metadata-tables.md +++ b/docs/querying/sql-metadata-tables.md @@ -175,6 +175,11 @@ After Druid retrieves the system-table rows, the native engine applies the remai aggregations, sorting, and result processing. A system table that doesn't advertise native query support continues to use its existing SQL-layer execution path, even when `useNativeQueryForSystemTables` is `true`. +The parameter defaults to `false`. During a rolling upgrade, leave it disabled until the Broker and all nodes +that serve the native system tables have been upgraded. After the upgrade, you can enable it by query or set +`druid.query.default.context.useNativeQueryForSystemTables=true` on Brokers as the cluster-wide default. For more +information, see [SQL query context](sql-query-context.md). + Native system-table queries sent to the Router use distributed Broker execution by default. To execute a native system-table Scan against only the contacted node, set the HTTP header `X-Druid-Native-Query-Route: local`. Local execution uses the authenticated request identity and applies the table's @@ -182,11 +187,6 @@ authorization rules. The Broker uses the same header for remote node fan-out req one of the selected nodes, it executes that node Scan in-process without an HTTP request. The header controls routing and doesn't grant additional permissions. -The parameter defaults to `false`. During a rolling upgrade, leave it disabled until the Broker and all nodes -that serve the native system tables have been upgraded. After the upgrade, you can enable it by query or set -`druid.query.default.context.useNativeQueryForSystemTables=true` on Brokers as the cluster-wide default. For more -information, see [SQL query context](sql-query-context.md). - ### SEGMENTS table Segments table provides details on all Druid segments, whether they are published yet or not. diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/NativeSysTasksQueryTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/NativeSysTasksQueryTest.java index 89679c7776c9..05a28dc902a1 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/NativeSysTasksQueryTest.java +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/NativeSysTasksQueryTest.java @@ -102,23 +102,6 @@ public void testNativeAggregationsSupportDistinctCount(final String plannerStrat Assertions.assertEquals("5,5,2,5", result); } - /** - * Verifies execution when native planning represents the inner aggregation as a query datasource. - * - * 
{@code
-   * SELECT COUNT(*)
-   * FROM
-   * (
-   *   SELECT
-   *     task_id,
-   *     COUNT(*) AS task_count
-   *   FROM sys.tasks
-   *   WHERE datasource IN ('native_sys_a', 'native_sys_b')
-   *   GROUP BY task_id
-   * )
-   * WHERE task_count > 0
-   * }
- */ @ParameterizedTest(name = "plannerStrategy = {0}") @ValueSource(strings = { QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, @@ -141,22 +124,6 @@ public void testNestedTaskAggregation(final String plannerStrategy) Assertions.assertEquals("5", result); } - /** - * Verifies that an outer aggregation can consume grouped rows from a native system-table subquery. - * - * 
{@code
-   * SELECT SUM(task_count)
-   * FROM
-   * (
-   *   SELECT
-   *     datasource,
-   *     COUNT(*) AS task_count
-   *   FROM sys.tasks
-   *   WHERE datasource IN ('native_sys_a', 'native_sys_b')
-   *   GROUP BY datasource
-   * )
-   * }
- */ @ParameterizedTest(name = "plannerStrategy = {0}") @ValueSource(strings = { QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, @@ -209,15 +176,6 @@ public void testRangeAndLikeFilters(final String plannerStrategy) ); } - /** - * Verifies that expression virtual columns required by a pushed node filter are preserved. - * - * 
{@code
-   * SELECT task_id
-   * FROM sys.tasks
-   * WHERE UPPER(task_id) = 'NATIVE_SYS_MVP_A_0'
-   * }
- */ @ParameterizedTest(name = "plannerStrategy = {0}") @ValueSource(strings = { QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, From 7d314c36e11d2b055dc5c053c0889627fa8fa3a9 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 9 Sep 2026 12:42:23 +0800 Subject: [PATCH 3/5] perf(sql): benchmark sys.tasks through SQL execution --- .../metadata/SysTasksTableBenchmark.java | 515 +++++++++++++----- .../druid/sql/calcite/util/CalciteTests.java | 55 +- .../sql/calcite/util/QueryFrameworkUtils.java | 13 + .../sql/calcite/util/SqlTestFramework.java | 11 +- 4 files changed, 452 insertions(+), 142 deletions(-) diff --git a/benchmarks/src/test/java/org/apache/druid/metadata/SysTasksTableBenchmark.java b/benchmarks/src/test/java/org/apache/druid/metadata/SysTasksTableBenchmark.java index 5150c91040da..85ae22abdede 100644 --- a/benchmarks/src/test/java/org/apache/druid/metadata/SysTasksTableBenchmark.java +++ b/benchmarks/src/test/java/org/apache/druid/metadata/SysTasksTableBenchmark.java @@ -19,18 +19,68 @@ package org.apache.druid.metadata; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import org.apache.druid.error.DruidException; +import org.apache.druid.guice.DruidBinders; +import org.apache.druid.indexer.RunnerTaskState; import org.apache.druid.indexer.TaskIdStatus; +import org.apache.druid.indexer.TaskLocation; import org.apache.druid.indexer.TaskStatus; +import org.apache.druid.indexer.TaskStatusPlus; +import org.apache.druid.initialization.DruidModule; import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.java.util.common.CloseableIterators; import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.jackson.JacksonUtils; +import org.apache.druid.java.util.common.parsers.CloseableIterator; import org.apache.druid.metadata.TaskLookup.ActiveTaskLookup; import org.apache.druid.metadata.TaskLookup.CompleteTaskLookup; import org.apache.druid.metadata.TaskLookup.TaskLookupType; -import org.apache.druid.query.filter.DimFilter; -import org.apache.druid.query.filter.EqualityFilter; -import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.query.DataSource; +import org.apache.druid.query.InlineDataSource; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryDataSource; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.query.QuerySegmentWalker; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.query.Druids; +import org.apache.druid.query.policy.NoopPolicyEnforcer; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.rpc.indexing.NoopOverlordClient; +import org.apache.druid.rpc.indexing.OverlordClient; +import org.apache.druid.segment.join.JoinableFactoryWrapper; +import org.apache.druid.server.DataSourceQueryHandler; +import org.apache.druid.server.SpecificSegmentsQuerySegmentWalker; +import org.apache.druid.server.security.AuthConfig; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.system.table.SystemTablePushdownFilter; +import org.apache.druid.server.system.table.TaskTableDescriptor; +import org.apache.druid.sql.DirectStatement; +import org.apache.druid.sql.SqlQueryPlus; +import org.apache.druid.sql.SqlStatementFactory; +import org.apache.druid.sql.calcite.BaseCalciteQueryTest; +import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; +import org.apache.druid.sql.calcite.TempDirProducer; +import org.apache.druid.sql.calcite.planner.CalciteRulesManager; +import org.apache.druid.sql.calcite.planner.CatalogResolver; +import org.apache.druid.sql.calcite.planner.PlannerConfig; +import org.apache.druid.sql.calcite.planner.PlannerContext; +import org.apache.druid.sql.calcite.planner.PlannerFactory; +import org.apache.druid.sql.calcite.schema.DruidSchemaCatalogProvider; +import org.apache.druid.sql.calcite.schema.DruidSchemaProvider; +import org.apache.druid.sql.calcite.schema.LookupSchema; +import org.apache.druid.sql.calcite.schema.SystemSchemaProvider; +import org.apache.druid.sql.calcite.util.CalciteTests; +import org.apache.druid.sql.calcite.util.DruidModuleCollection; +import org.apache.druid.sql.calcite.util.QueryFrameworkUtils; +import org.apache.druid.sql.calcite.util.SqlTestFramework; +import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier; +import org.apache.druid.sql.calcite.view.ViewManager; +import org.apache.druid.sql.hook.DruidHookDispatcher; import org.openjdk.jmh.annotations.Benchmark; import org.openjdk.jmh.annotations.BenchmarkMode; import org.openjdk.jmh.annotations.Fork; @@ -47,59 +97,19 @@ import org.openjdk.jmh.infra.Blackhole; import org.skife.jdbi.v2.PreparedBatch; +import javax.annotation.Nullable; +import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.concurrent.TimeUnit; -import java.util.function.Predicate; /** - * Compares the metadata-read portion of the bindable and native {@code sys.tasks} execution paths. + * Executes SQL against 100,000 {@code sys.tasks} rows and compares the Bindable and native planning paths. * - *

The bindable path requests every task from the Overlord and evaluates the SQL predicate afterward. The native - * path converts the same predicate to a validated {@link TaskStorageQueryFilter}, applies it in metadata storage, and - * retains the predicate as a residual filter. This benchmark intentionally excludes SQL planning and HTTP transport - * so it isolates the benefit of metadata filter pushdown.

- * - *

The synthetic table contains 100,000 migrated task records split evenly between active and completed tasks. - * Datasource, type, and group values have controlled cardinalities so each query returns a stable number of rows.

- * - *

Query cases:

- *
    - *
  • The unfiltered task-list query issued by the web console: - * 
    {@code
- *     WITH tasks AS (
- *       SELECT "task_id", "group_id", "type", "datasource", "created_time", "location", "duration", "error_msg",
- *         CASE
- *           WHEN "error_msg" IN (
- *             'Shutdown request from user',
- *             'Canceled: Query canceled by user or by task shutdown.'
- *           ) THEN 'CANCELED'
- *           WHEN "status" = 'RUNNING' THEN "runner_status"
- *           ELSE "status"
- *         END AS "status"
- *       FROM sys.tasks
- *     )
- *     SELECT "task_id", "group_id", "type", "datasource", "created_time", "location", "duration", "error_msg", "status"
- *     FROM tasks
- *     ORDER BY
- *       CASE "status"
- *         WHEN 'RUNNING' THEN 4
- *         WHEN 'PENDING' THEN 3
- *         WHEN 'WAITING' THEN 2
- *         ELSE 1
- *       END DESC,
- *       "created_time" DESC
- *     }
    - *
    • - *
  • {@code SELECT * FROM sys.tasks WHERE task_id = 'task-050000'}
    • - *
  • {@code SELECT * FROM sys.tasks WHERE datasource = 'datasource-000'}
    • - *
  • {@code SELECT * FROM sys.tasks WHERE datasource = 'datasource-000' AND type = 'type-0'}
    • - *
  • {@code SELECT * FROM sys.tasks WHERE group_id = 'group-0000'}
    • - *
- * - *

The native planner maps {@code task_id} to metadata column {@code id}. The other benchmarked columns retain - * their names. Equality filters on {@code type} and {@code group_id} are both pushed down, although the current task - * table schema does not index either column.

+ *

Both paths execute the complete SQL statement, including planning, expressions, aggregation, sorting, and result + * materialization. The native path runs the Broker query over an in-process system-table handler so this benchmark + * excludes only node discovery and HTTP transport. Its task provider reads the same Derby metadata table as the + * Bindable Overlord client and applies the production metadata predicate builder.

*/ @State(Scope.Benchmark) @Fork(value = 1, jvmArgsAppend = {"-Xmx4g"}) @@ -107,7 +117,8 @@ @Measurement(iterations = 5, time = 2) @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MILLISECONDS) -public class SysTasksTableBenchmark +@SqlTestFrameworkConfig.ComponentSupplier(SysTasksTableBenchmark.BenchmarkComponentSupplier.class) +public class SysTasksTableBenchmark extends BaseCalciteQueryTest { private static final int NUM_TASKS = 100_000; private static final int INSERT_BATCH_SIZE = 5_000; @@ -117,89 +128,205 @@ public class SysTasksTableBenchmark private static final String TARGET_TYPE = "type-0"; private static final String TARGET_GROUP_ID = "group-0000"; - public enum QueryCase - { - FULL_LIST, - TASK_ID, - DATASOURCE, - DATASOURCE_AND_TYPE, - GROUP_ID - } + private static final List PUSHDOWN_FILTERS = List.of( + new SystemTablePushdownFilter("task_id", "id"), + new SystemTablePushdownFilter("group_id", null), + new SystemTablePushdownFilter("type", null), + new SystemTablePushdownFilter("datasource", null), + new SystemTablePushdownFilter("created_time", "created_date"), + new SystemTablePushdownFilter("status", null) + ); + + private static final String WEB_CONSOLE_TASKS_SQL = + "WITH tasks AS (SELECT\n" + + " \"task_id\", \"group_id\", \"type\", \"datasource\", \"created_time\", \"location\", " + + "\"duration\", \"error_msg\",\n" + + " CASE WHEN \"error_msg\" IN ('Shutdown request from user', " + + "'Canceled: Query canceled by user or by task shutdown.') THEN 'CANCELED' " + + "WHEN \"status\" = 'RUNNING' THEN \"runner_status\" ELSE \"status\" END AS \"status\"\n" + + " FROM sys.tasks\n" + + ")\n" + + "SELECT \"task_id\", \"group_id\", \"type\", \"datasource\", \"created_time\", \"location\", " + + "\"duration\", \"error_msg\", \"status\"\n" + + "FROM tasks\n" + + "ORDER BY\n" + + " (CASE \"status\" WHEN 'RUNNING' THEN 4 WHEN 'PENDING' THEN 3 WHEN 'WAITING' THEN 2 ELSE 1 END) DESC,\n" + + " \"created_time\" DESC"; + + private static BenchmarkTaskData taskData; - @Param - private QueryCase queryCase; + @Param({"false", "true"}) + private boolean useNativeQueryForSystemTables; private TestDerbyConnector connector; - private DerbyMetadataStorageActionHandler handler; - private Map taskLookups; - private TaskStorageQueryFilter pushdownFilter; - private Predicate residualFilter; + private SqlStatementFactory statementFactory; + private Map queryContext; @Setup(Level.Trial) - public void setup() throws ClassNotFoundException + public void setup() throws Exception { - // The benchmark uber-jar has a consolidated JDBC service file that does not list Derby. Class.forName("org.apache.derby.jdbc.EmbeddedDriver"); connector = new TestDerbyConnector(); connector.createDatabase(); final MetadataStorageTablesConfig tablesConfig = connector.getMetadataTablesConfig(); connector.prepareTaskEntryTable(tablesConfig.getTasksTable()); - handler = new DerbyMetadataStorageActionHandler( + final DerbyMetadataStorageActionHandler handler = new DerbyMetadataStorageActionHandler( connector, new DefaultObjectMapper(), tablesConfig.getTasksTable(), tablesConfig.getTaskLockTable() ); insertTasks(tablesConfig.getTasksTable()); + taskData = new BenchmarkTaskData(handler); - taskLookups = Map.of( - TaskLookupType.ACTIVE, - ActiveTaskLookup.getInstance(), - TaskLookupType.COMPLETE, - CompleteTaskLookup.withTasksCreatedPriorTo(null, DateTimes.of("2020-01-01")) + queryFrameworkRule.setConfig( + new SqlTestFrameworkConfig(Arrays.asList(SysTasksTableBenchmark.class.getAnnotations())) + ); + final SqlTestFramework framework = queryFramework(); + statementFactory = createStatementFactory(framework, taskData.createOverlordClient()); + queryContext = ImmutableMap.of( + PlannerContext.CTX_USE_NATIVE_QUERY_FOR_SYSTEM_TABLES, + useNativeQueryForSystemTables ); - pushdownFilter = new TaskStorageQueryFilter(filtersFor(queryCase)); - residualFilter = residualFilterFor(queryCase); } @TearDown(Level.Trial) - public void tearDown() + public void tearDown() throws Exception { + queryFramework().close(); + taskData = null; connector.tearDown(); } - /** Simulates the bindable path: retrieve all tasks, then apply the SQL predicate. */ + /** Executes the exact task-list query currently issued by the Web Console. */ @Benchmark - public void bindable(Blackhole blackhole) + public void webConsoleFullList(final Blackhole blackhole) { - consumeMatchingRows(handler.getTaskStatusList(taskLookups, null, false), blackhole); + executeSql(WEB_CONSOLE_TASKS_SQL, blackhole); } - /** Simulates the native path: push eligible predicates into metadata storage, then apply the residual predicate. */ + /** {@code SELECT * FROM sys.tasks WHERE task_id = 'task-050000'} */ @Benchmark - public void nativeQuery(Blackhole blackhole) + public void taskIdFilter(final Blackhole blackhole) { - consumeMatchingRows(handler.getTaskStatusListWithFilter(taskLookups, pushdownFilter, false), blackhole); + executeSql("SELECT * FROM sys.tasks WHERE task_id = '" + TARGET_TASK_ID + "'", blackhole); } - private void consumeMatchingRows(final List candidates, final Blackhole blackhole) + /** {@code SELECT * FROM sys.tasks WHERE datasource = 'datasource-000'} */ + @Benchmark + public void dataSourceFilter(final Blackhole blackhole) { - long matchedRows = 0; - for (final TaskIdStatus candidate : candidates) { - if (residualFilter.test(candidate)) { - blackhole.consume(candidate); - matchedRows++; - } + executeSql("SELECT * FROM sys.tasks WHERE datasource = '" + TARGET_DATASOURCE + "'", blackhole); + } + + /** {@code SELECT * FROM sys.tasks WHERE datasource = 'datasource-000' AND type = 'type-0'} */ + @Benchmark + public void dataSourceAndTypeFilter(final Blackhole blackhole) + { + executeSql( + "SELECT * FROM sys.tasks WHERE datasource = '" + TARGET_DATASOURCE + "' AND type = '" + TARGET_TYPE + "'", + blackhole + ); + } + + /** {@code SELECT * FROM sys.tasks WHERE group_id = 'group-0000'} */ + @Benchmark + public void groupIdFilter(final Blackhole blackhole) + { + executeSql("SELECT * FROM sys.tasks WHERE group_id = '" + TARGET_GROUP_ID + "'", blackhole); + } + + /** {@code SELECT CASE WHEN status = 'RUNNING' THEN runner_status ELSE status END, COUNT(*) FROM sys.tasks GROUP BY 1} */ + @Benchmark + public void groupByEffectiveStatus(final Blackhole blackhole) + { + executeSql( + "SELECT CASE WHEN status = 'RUNNING' THEN runner_status ELSE status END AS effective_status, COUNT(*) " + + "FROM sys.tasks GROUP BY 1", + blackhole + ); + } + + /** {@code SELECT datasource, COUNT(*) FROM sys.tasks GROUP BY datasource} */ + @Benchmark + public void groupByDataSource(final Blackhole blackhole) + { + executeSql("SELECT datasource, COUNT(*) FROM sys.tasks GROUP BY datasource", blackhole); + } + + /** {@code SELECT type, COUNT(*) FROM sys.tasks WHERE datasource = 'datasource-000' GROUP BY type} */ + @Benchmark + public void groupByTypeForDataSource(final Blackhole blackhole) + { + executeSql( + "SELECT type, COUNT(*) FROM sys.tasks WHERE datasource = '" + TARGET_DATASOURCE + "' GROUP BY type", + blackhole + ); + } + + private void executeSql(final String sql, final Blackhole blackhole) + { + final SqlQueryPlus sqlQuery = SqlQueryPlus.builder(sql) + .auth(CalciteTests.REGULAR_USER_AUTH_RESULT) + .systemDefaultContext(Map.of()) + .queryContext(queryContext) + .build(); + try (DirectStatement statement = statementFactory.directStatement(sqlQuery)) { + blackhole.consume(statement.execute().getResults().toList()); } - blackhole.consume(matchedRows); + } + + private SqlStatementFactory createStatementFactory( + final SqlTestFramework framework, + final OverlordClient overlordClient + ) + { + final PlannerConfig plannerConfig = PLANNER_CONFIG_DEFAULT; + final DruidSchemaProvider druidSchemaProvider = framework.injector().getInstance(DruidSchemaProvider.class); + final SystemSchemaProvider systemSchemaProvider = CalciteTests.createMockSystemSchemaProvider( + druidSchemaProvider.getSegmentMetadataCache(), + framework.injector().getInstance(org.apache.druid.client.TimelineServerView.class), + CalciteTests.TEST_AUTHORIZER_MAPPER, + plannerConfig, + overlordClient + ); + final DruidSchemaCatalogProvider schemaProvider = QueryFrameworkUtils.createMockRootSchemaProvider( + framework.injector().getInstance(ViewManager.class), + CalciteTests.TEST_AUTHORIZER_MAPPER, + druidSchemaProvider, + systemSchemaProvider, + framework.injector().getInstance(LookupSchema.class), + framework.operatorTable(), + plannerConfig + ); + final SqlTestFramework.QueryComponentSupplier componentSupplier = framework.injector().getInstance( + SqlTestFramework.QueryComponentSupplier.class + ); + final PlannerFactory plannerFactory = new PlannerFactory( + schemaProvider, + framework.operatorTable(), + framework.macroTable(), + plannerConfig, + CalciteTests.TEST_AUTHORIZER_MAPPER, + framework.queryJsonMapper(), + CalciteTests.DRUID_SCHEMA_NAME, + new CalciteRulesManager(componentSupplier.getPlannerComponentSupplier().extensionCalciteRules()), + framework.injector().getInstance(JoinableFactoryWrapper.class), + CatalogResolver.NULL_RESOLVER, + new AuthConfig(), + NoopPolicyEnforcer.instance(), + new DruidHookDispatcher() + ); + return QueryFrameworkUtils.createSqlMultiStatementFactory(framework.engine(), plannerFactory); } private void insertTasks(final String tasksTable) { final byte[] payload = new byte[]{'{', '}'}; - final byte[] activeStatus = JacksonUtils.toBytes(new DefaultObjectMapper(), TaskStatus.running("benchmark")); - final byte[] completeStatus = JacksonUtils.toBytes(new DefaultObjectMapper(), TaskStatus.success("benchmark")); + final ObjectMapper mapper = new DefaultObjectMapper(); + final byte[] activeStatus = JacksonUtils.toBytes(mapper, TaskStatus.running("benchmark")); + final byte[] completeStatus = JacksonUtils.toBytes(mapper, TaskStatus.success("benchmark")); final String sql = StringUtils.format( "INSERT INTO %s (id, created_date, datasource, payload, status_payload, active, type, group_id) " + "VALUES (:id, :created_date, :datasource, :payload, :status_payload, :active, :type, :group_id)", @@ -229,54 +356,198 @@ private void insertTasks(final String tasksTable) } } - private static List filtersFor(final QueryCase query) + private static String taskId(final int taskNumber) { - return switch (query) { - case FULL_LIST -> List.of(); - case TASK_ID -> List.of(equality("id", TARGET_TASK_ID)); - case DATASOURCE -> List.of(equality("datasource", TARGET_DATASOURCE)); - case DATASOURCE_AND_TYPE -> List.of( - equality("datasource", TARGET_DATASOURCE), - equality("type", TARGET_TYPE) - ); - case GROUP_ID -> List.of(equality("group_id", TARGET_GROUP_ID)); - }; + return StringUtils.format("task-%06d", taskNumber); } - private static Predicate residualFilterFor(final QueryCase query) + private static String dataSource(final int taskNumber) { - return switch (query) { - case FULL_LIST -> ignored -> true; - case TASK_ID -> task -> TARGET_TASK_ID.equals(task.getTaskIdentifier().getId()); - case DATASOURCE -> task -> TARGET_DATASOURCE.equals(task.getDataSource()); - case DATASOURCE_AND_TYPE -> task -> TARGET_DATASOURCE.equals(task.getDataSource()) - && TARGET_TYPE.equals(task.getTaskIdentifier().getType()); - case GROUP_ID -> task -> TARGET_GROUP_ID.equals(task.getTaskIdentifier().getGroupId()); - }; + return StringUtils.format("datasource-%03d", taskNumber % 100); } - private static EqualityFilter equality(final String column, final String value) + private static String type(final int taskNumber) { - return new EqualityFilter(column, ColumnType.STRING, value, null); + return StringUtils.format("type-%d", (taskNumber / 100) % 10); } - private static String taskId(final int taskNumber) + private static String groupId(final int taskNumber) { - return StringUtils.format("task-%06d", taskNumber); + return StringUtils.format("group-%04d", taskNumber % 1_000); } - private static String dataSource(final int taskNumber) + public static class BenchmarkComponentSupplier extends StandardComponentSupplier { - return StringUtils.format("datasource-%03d", taskNumber % 100); + public BenchmarkComponentSupplier(final TempDirProducer tempDirProducer) + { + super(tempDirProducer); + } + + @Override + public DruidModule getCoreModule() + { + return DruidModuleCollection.of( + super.getCoreModule(), + binder -> DruidBinders.dataSourceQueryHandlerBinder(binder) + .addBinding(SystemTableDataSource.class) + .to(BenchmarkSystemTableQueryHandler.class) + ); + } + + @Override + public SpecificSegmentsQuerySegmentWalker addSegmentsToWalker( + final SpecificSegmentsQuerySegmentWalker walker, + final ObjectMapper jsonMapper + ) + { + return walker; + } } - private static String type(final int taskNumber) + public static class BenchmarkSystemTableQueryHandler implements DataSourceQueryHandler { - return StringUtils.format("type-%d", (taskNumber / 100) % 10); + private final QuerySegmentWalker walker; + + @com.google.inject.Inject + public BenchmarkSystemTableQueryHandler(final QuerySegmentWalker walker) + { + this.walker = walker; + } + + @Override + public QueryRunner createRunner( + final Query query, + final AuthenticationResult authenticationResult, + final boolean executeLocally + ) + { + final BenchmarkTaskData currentTaskData = taskData; + if (currentTaskData == null) { + throw DruidException.defensive("Benchmark task data has not been initialized"); + } + final InlineDataSource inlineDataSource = InlineDataSource.fromIterable( + currentTaskData.getNativeRows(query), + TaskTableDescriptor.ROW_SIGNATURE + ); + final ScanQuery scanQuery = Druids.newScanQueryBuilder() + .dataSource(inlineDataSource) + .eternityInterval() + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .limit(Long.MAX_VALUE) + .columns(TaskTableDescriptor.ROW_SIGNATURE) + .build(); + final Query resolvedQuery = query.withDataSource( + replaceSystemTable(query.getDataSource(), new QueryDataSource(scanQuery)) + ); + final QueryRunner runner = resolvedQuery.getRunner(walker); + return (queryPlus, responseContext) -> runner.run(queryPlus.withQuery(resolvedQuery), responseContext); + } + + private static DataSource replaceSystemTable( + final DataSource dataSource, + final DataSource replacementDataSource + ) + { + if (dataSource instanceof SystemTableDataSource) { + return replacementDataSource; + } + return dataSource.withChildren( + dataSource.getChildren() + .stream() + .map(child -> replaceSystemTable(child, replacementDataSource)) + .toList() + ); + } } - private static String groupId(final int taskNumber) + private static class BenchmarkTaskData { - return StringUtils.format("group-%04d", taskNumber % 1_000); + private final DerbyMetadataStorageActionHandler handler; + private final Map taskLookups = Map.of( + TaskLookupType.ACTIVE, + ActiveTaskLookup.getInstance(), + TaskLookupType.COMPLETE, + CompleteTaskLookup.withTasksCreatedPriorTo(null, DateTimes.of("2020-01-01")) + ); + + private BenchmarkTaskData(final DerbyMetadataStorageActionHandler handler) + { + this.handler = handler; + } + + private OverlordClient createOverlordClient() + { + return new NoopOverlordClient() + { + @Override + public ListenableFuture> taskStatuses( + @Nullable final String state, + @Nullable final String dataSource, + @Nullable final Integer maxCompletedTasks + ) + { + return Futures.immediateFuture( + CloseableIterators.withEmptyBaggage( + handler.getTaskStatusList(taskLookups, null, false) + .stream() + .map(BenchmarkTaskData::toTaskStatusPlus) + .iterator() + ) + ); + } + }; + } + + private Iterable getNativeRows(final Query query) + { + final TaskStorageQueryFilter filter = new TaskStorageQueryFilter( + SystemTablePushdownFilter.extract(query, PUSHDOWN_FILTERS) + ); + return () -> handler.getTaskStatusListWithFilter(taskLookups, filter, false) + .stream() + .map(BenchmarkTaskData::toRow) + .iterator(); + } + + private static TaskStatusPlus toTaskStatusPlus(final TaskIdStatus task) + { + final TaskStatus status = task.getStatus(); + final RunnerTaskState runnerState = status.isRunnable() ? RunnerTaskState.RUNNING : RunnerTaskState.NONE; + return new TaskStatusPlus( + task.getTaskIdentifier().getId(), + task.getTaskIdentifier().getGroupId(), + task.getTaskIdentifier().getType(), + task.getCreatedTime(), + task.getCreatedTime(), + status.getStatusCode(), + runnerState, + status.getDuration(), + status.getLocation(), + task.getDataSource(), + status.getErrorMsg() + ); + } + + private static Object[] toRow(final TaskIdStatus task) + { + final TaskStatusPlus status = toTaskStatusPlus(task); + final TaskLocation location = status.getLocation(); + return new Object[]{ + status.getId(), + status.getGroupId(), + status.getType(), + status.getDataSource(), + status.getCreatedTime().toString(), + status.getQueueInsertionTime().toString(), + status.getStatusCode().toString(), + status.getRunnerStatusCode().toString(), + status.getDuration(), + location.getLocation(), + location.getHost(), + (long) location.getPort(), + (long) location.getTlsPort(), + status.getErrorMsg() + }; + } } } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java index d9731ac4d014..6541e0a80438 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java @@ -336,25 +336,7 @@ public static SystemSchemaProvider createMockSystemSchemaProvider( final PlannerConfig plannerConfig ) { - final DruidNode coordinatorNode = mockCoordinatorNode(); - FakeDruidNodeDiscoveryProvider provider = mockDruidNodeDiscoveryProvider(coordinatorNode); - final DruidNode overlordNode = new DruidNode("test-overlord", "dummy", false, 8090, null, true, false); - - final CoordinatorClient coordinatorClient = new NoopCoordinatorClient() - { - @Override - public ListenableFuture findCurrentLeader() - { - try { - return Futures.immediateFuture(new URI(coordinatorNode.getHostAndPortToUse())); - } - catch (URISyntaxException e) { - throw new RuntimeException(e); - } - } - }; - final OverlordClient overlordClient = new NoopOverlordClient() { @Override @@ -401,6 +383,43 @@ private TaskStatusPlus createTaskStatus(String id, String datasource, Long durat } }; + return createMockSystemSchemaProvider( + segmentMetadataCache, + timelineServerView, + authorizerMapper, + plannerConfig, + overlordClient + ); + } + + /** + * Creates the standard mock system schema with a caller-provided Overlord client. + */ + public static SystemSchemaProvider createMockSystemSchemaProvider( + final BrokerSegmentMetadataCache segmentMetadataCache, + final TimelineServerView timelineServerView, + final AuthorizerMapper authorizerMapper, + final PlannerConfig plannerConfig, + final OverlordClient overlordClient + ) + { + final DruidNode coordinatorNode = mockCoordinatorNode(); + final FakeDruidNodeDiscoveryProvider provider = mockDruidNodeDiscoveryProvider(coordinatorNode); + + final CoordinatorClient coordinatorClient = new NoopCoordinatorClient() + { + @Override + public ListenableFuture findCurrentLeader() + { + try { + return Futures.immediateFuture(new URI(coordinatorNode.getHostAndPortToUse())); + } + catch (URISyntaxException e) { + throw new RuntimeException(e); + } + } + }; + return new SystemSchemaProvider( segmentMetadataCache, new MetadataSegmentView( diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java index c15923bc4564..d25b928daaed 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/QueryFrameworkUtils.java @@ -25,6 +25,7 @@ import org.apache.calcite.avatica.remote.TypedValue; import org.apache.druid.client.InternalQueryConfig; import org.apache.druid.client.TimelineServerView; +import org.apache.druid.query.DataSource; import org.apache.druid.query.DefaultGenericQueryMetricsFactory; import org.apache.druid.query.DefaultQueryConfig; import org.apache.druid.query.GlobalTableDataSource; @@ -35,6 +36,7 @@ import org.apache.druid.segment.join.JoinableFactory; import org.apache.druid.segment.loading.SegmentCacheManager; import org.apache.druid.segment.metadata.CentralizedDatasourceSchemaConfig; +import org.apache.druid.server.DataSourceQueryHandler; import org.apache.druid.server.QueryLifecycleFactory; import org.apache.druid.server.QueryStackTests; import org.apache.druid.server.SegmentManager; @@ -86,6 +88,16 @@ public static QueryLifecycleFactory createMockQueryLifecycleFactory( final QueryRunnerFactoryConglomerate conglomerate, final AuthorizerMapper authorizerMapper ) + { + return createMockQueryLifecycleFactory(walker, conglomerate, authorizerMapper, Map.of()); + } + + public static QueryLifecycleFactory createMockQueryLifecycleFactory( + final QuerySegmentWalker walker, + final QueryRunnerFactoryConglomerate conglomerate, + final AuthorizerMapper authorizerMapper, + final Map, DataSourceQueryHandler> dataSourceQueryHandlers + ) { return new QueryLifecycleFactory( conglomerate, @@ -97,6 +109,7 @@ public static QueryLifecycleFactory createMockQueryLifecycleFactory( NoopPolicyEnforcer.instance(), authorizerMapper, new DefaultQueryConfig(Map.of()), + dataSourceQueryHandlers, null // BrokerConfigManager - null for tests ); } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index 6346e0691462..433afd91d7a6 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -58,6 +58,7 @@ import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.java.util.http.client.HttpClient; import org.apache.druid.math.expr.ExprMacroTable; +import org.apache.druid.query.DataSource; import org.apache.druid.query.DefaultGenericQueryMetricsFactory; import org.apache.druid.query.DruidProcessingConfig; import org.apache.druid.query.GlobalTableDataSource; @@ -82,6 +83,7 @@ import org.apache.druid.segment.join.JoinableFactoryWrapper; import org.apache.druid.segment.realtime.ChatHandlerProvider; import org.apache.druid.server.ClientQuerySegmentWalker; +import org.apache.druid.server.DataSourceQueryHandler; import org.apache.druid.server.DruidNode; import org.apache.druid.server.LocalQuerySegmentWalker; import org.apache.druid.server.QueryLifecycle; @@ -135,6 +137,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.Properties; import java.util.Set; @@ -1025,12 +1028,16 @@ public JoinableFactoryWrapper joinableFactoryWrapper(final Injector injector, Bu @Provides @LazySingleton - public QueryLifecycleFactory queryLifecycleFactory(final Injector injector) + public QueryLifecycleFactory queryLifecycleFactory( + final Injector injector, + final Map, DataSourceQueryHandler> dataSourceQueryHandlers + ) { return QueryFrameworkUtils.createMockQueryLifecycleFactory( injector.getInstance(QuerySegmentWalker.class), injector.getInstance(QueryRunnerFactoryConglomerate.class), - injector.getInstance(AuthorizerMapper.class) + injector.getInstance(AuthorizerMapper.class), + dataSourceQueryHandlers ); } From 0956d4c6bb2eacbca8c9dbf852f9bc3046249a7c Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 9 Sep 2026 12:49:59 +0800 Subject: [PATCH 4/5] fix(sql): handle tasks without datasource safely --- .../system/table/TaskTableDescriptor.java | 10 ++- .../system/table/TaskTableDescriptorTest.java | 64 +++++++++++++++++++ .../sql/calcite/schema/NativeTasksTable.java | 2 +- 3 files changed, 72 insertions(+), 4 deletions(-) create mode 100644 server/src/test/java/org/apache/druid/server/system/table/TaskTableDescriptorTest.java diff --git a/server/src/main/java/org/apache/druid/server/system/table/TaskTableDescriptor.java b/server/src/main/java/org/apache/druid/server/system/table/TaskTableDescriptor.java index 8015b596ce41..acd2c212bda7 100644 --- a/server/src/main/java/org/apache/druid/server/system/table/TaskTableDescriptor.java +++ b/server/src/main/java/org/apache/druid/server/system/table/TaskTableDescriptor.java @@ -65,9 +65,13 @@ public Iterable filterAuthorizedRows( return AuthorizationUtils.filterAuthorizedResources( authenticationResult, rows, - row -> Collections.singletonList( - AuthorizationUtils.DATASOURCE_READ_RA_GENERATOR.apply((String) row[DATASOURCE_COLUMN]) - ), + row -> { + final String dataSource = (String) row[DATASOURCE_COLUMN]; + // A null resource-action generator result makes AuthorizationUtils omit the row, including for allow-all. + return dataSource == null + ? null + : Collections.singletonList(AuthorizationUtils.DATASOURCE_READ_RA_GENERATOR.apply(dataSource)); + }, authorizerMapper ); } diff --git a/server/src/test/java/org/apache/druid/server/system/table/TaskTableDescriptorTest.java b/server/src/test/java/org/apache/druid/server/system/table/TaskTableDescriptorTest.java new file mode 100644 index 000000000000..e63b2bc0052b --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/system/table/TaskTableDescriptorTest.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.server.system.table; + +import com.google.common.collect.Lists; +import org.apache.druid.server.security.AllowAllAuthorizer; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.AuthorizerMapper; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; + +public class TaskTableDescriptorTest +{ + private static final int DATASOURCE_COLUMN = TaskTableDescriptor.ROW_SIGNATURE.indexOf("datasource"); + private static final AuthenticationResult AUTHENTICATION_RESULT = + new AuthenticationResult("test", "allow", "test", null); + private static final AuthorizerMapper AUTHORIZER_MAPPER = + new AuthorizerMapper(Map.of("allow", new AllowAllAuthorizer(null))); + + @Test + public void testRowAuthorizerDropsRowsWithoutDatasource() + { + final Object[] rowWithoutDatasource = taskRow(null); + final Object[] authorizedRow = taskRow("wikipedia"); + + final List result = Lists.newArrayList( + new TaskTableDescriptor().getRowAuthorizer().filterAuthorizedRows( + List.of(rowWithoutDatasource, authorizedRow), + AUTHENTICATION_RESULT, + AUTHORIZER_MAPPER + ) + ); + + Assertions.assertEquals(1, result.size()); + Assertions.assertSame(authorizedRow, result.getFirst()); + } + + private static Object[] taskRow(final String dataSource) + { + final Object[] row = new Object[TaskTableDescriptor.ROW_SIGNATURE.size()]; + row[DATASOURCE_COLUMN] = dataSource; + return row; + } +} diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeTasksTable.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeTasksTable.java index fd5fbae972ee..f5fdf0b2af93 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeTasksTable.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeTasksTable.java @@ -31,7 +31,7 @@ /** Native-query representation of {@code sys.tasks}. */ class NativeTasksTable extends DruidTable { - private static final DataSource DATA_SOURCE = new SystemTableDataSource("tasks"); + private static final DataSource DATA_SOURCE = new SystemTableDataSource(TaskTableDescriptor.TABLE_NAME); NativeTasksTable() { From d0cc7cd8ab85e85368120d99725aa3639cf1b897 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 11 Sep 2026 22:11:49 +0800 Subject: [PATCH 5/5] fix(indexing): avoid unsafe task time pushdown --- .../overlord/task/TasksTableDataProvider.java | 1 - .../task/TasksTableDataProviderTest.java | 16 ++++++++++++---- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/task/TasksTableDataProvider.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/task/TasksTableDataProvider.java index 512558a113a9..5f6ecb307e3d 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/task/TasksTableDataProvider.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/task/TasksTableDataProvider.java @@ -42,7 +42,6 @@ public class TasksTableDataProvider implements SystemTableDataProvider new SystemTablePushdownFilter("group_id", null), new SystemTablePushdownFilter("type", null), new SystemTablePushdownFilter("datasource", null), - new SystemTablePushdownFilter("created_time", "created_date"), new SystemTablePushdownFilter("status", null) ); private static final List STATUS_PUSHDOWN_FILTERS = List.of( diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/task/TasksTableDataProviderTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/task/TasksTableDataProviderTest.java index c46e55ef78e6..bdc4d4ad9296 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/task/TasksTableDataProviderTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/task/TasksTableDataProviderTest.java @@ -68,7 +68,16 @@ public void testAdvertisesAllPushdownFiltersForSupportingStorage() final List filters = dataProvider(taskQueryTool).getPushdownFilters(); - Assertions.assertEquals(6, filters.size()); + Assertions.assertEquals( + List.of( + new SystemTablePushdownFilter("task_id", "id"), + new SystemTablePushdownFilter("group_id", null), + new SystemTablePushdownFilter("type", null), + new SystemTablePushdownFilter("datasource", null), + new SystemTablePushdownFilter("status", null) + ), + filters + ); } @Test @@ -275,14 +284,13 @@ public void testExtractsAllStorageBackedTaskFilters() final List extracted = extract(query); final TaskStorageQueryFilter pushdownFilters = new TaskStorageQueryFilter(extracted); - Assertions.assertEquals(7, extracted.size()); + Assertions.assertEquals(5, extracted.size()); Assertions.assertEquals("id", pushdownFilters.getStringValuesColumn(extracted.get(0))); Assertions.assertEquals(Set.of("task-a", "task-b"), pushdownFilters.getStringValues(extracted.get(0))); Assertions.assertEquals(Set.of("group-a"), pushdownFilters.getStringValues(extracted.get(1))); Assertions.assertEquals(Set.of("noop"), pushdownFilters.getStringValues(extracted.get(2))); Assertions.assertEquals(Set.of("native_sys_a"), pushdownFilters.getStringValues(extracted.get(3))); - Assertions.assertEquals("created_date", pushdownFilters.getStringValuesColumn(extracted.get(4))); - Assertions.assertEquals("created_date", ((RangeFilter) extracted.get(5)).getColumn()); + Assertions.assertEquals("status", ((SelectorDimFilter) extracted.get(4)).getDimension()); Assertions.assertFalse(pushdownFilters.includesActiveTasks()); Assertions.assertTrue(pushdownFilters.includesCompleteTasks()); }