diff --git a/docs/querying/sql-metadata-tables.md b/docs/querying/sql-metadata-tables.md index 055305d81e50..8b47335617fe 100644 --- a/docs/querying/sql-metadata-tables.md +++ b/docs/querying/sql-metadata-tables.md @@ -135,10 +135,57 @@ WHERE "IS_AGGREGATOR" = 'YES' The "sys" schema provides visibility into Druid segments, servers and tasks. :::info - Note: "sys" tables do not currently support Druid-specific functions like `TIME_PARSE` and - `APPROX_QUANTILE_DS`. Only standard SQL functions can be used. + By default, "sys" tables use the SQL-layer execution path and support only standard SQL functions. You can enable + [native query execution](#native-query-execution) for supported system tables, which enables expressions and + aggregations that the native SQL engine can translate. ::: +### Native query execution + +The native SQL engine can plan supported system tables as native datasources. To enable this behavior for a query, set +`useNativeQueryForSystemTables` to `true` in the SQL query context: + +```json +{ + "query": "SELECT COUNT(DISTINCT server) FROM sys.server_properties", + "context": { + "useNativeQueryForSystemTables": true + } +} +``` + +For clients that support setting query context parameters with SQL statements, you can enable native system-table +execution with `SET`: + +```sql +SET useNativeQueryForSystemTables = 'true'; +SELECT COUNT(DISTINCT server) FROM sys.server_properties; +``` + +Native system-table execution is available when the resolved SQL engine is `native`. You don't need to explicitly set +the `engine` context parameter when `native` is already the default engine. The following tables support native query +execution: + +|Table|Source of rows| +|-----|--------------| +|[`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 +authorization rules. The Broker uses the same header for remote node fan-out requests. If the Broker itself is +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. + ### SEGMENTS table Segments table provides details on all Druid segments, whether they are published yet or not. diff --git a/docs/querying/sql-query-context.md b/docs/querying/sql-query-context.md index 88cafe1ced71..b5f813d17716 100644 --- a/docs/querying/sql-query-context.md +++ b/docs/querying/sql-query-context.md @@ -45,6 +45,7 @@ The table below lists the query context parameters you can use with Druid SQL. |`useLexicographicTopN`|If `true`, Druid can use [TopN queries](topnquery.md) with lexicographic dimension ordering. If `false`, Druid uses [GroupBy queries](groupbyquery.md) instead for lexicographic ordering. When both `useLexicographicTopN` and `useApproximateTopN` are `false`, TopN queries are never used.|`false`| |`enableTimeBoundaryPlanning`|If `true`, Druid converts SQL queries to [time boundary queries](timeboundaryquery.md) wherever possible. Time boundary queries are very efficient for min-max calculation on the `__time` column in a datasource. |`false`| |`useNativeQueryExplain`|If `true`, `EXPLAIN PLAN FOR` returns the explain plan as a JSON representation of equivalent native query, else it returns the original version of explain plan generated by Calcite.

This property is provided for backwards compatibility. We don't recommend setting this parameter unless your application depends on the older behavior.|`true`| +|`useNativeQueryForSystemTables`|If `true` and the resolved SQL engine is `native`, Druid uses native query planning for [system tables](sql-metadata-tables.md#native-query-execution) that advertise native query support. System tables without native query support continue to use their existing execution path.

During a rolling upgrade, leave this parameter disabled until the Broker and the components that serve the native system tables have been upgraded. You can set a cluster-wide default on Brokers with `druid.query.default.context.useNativeQueryForSystemTables`.|`false`| |`sqlFinalizeOuterSketches`|If `false` (default behavior in Druid 25.0.0 and later), `DS_HLL`, `DS_THETA`, and `DS_QUANTILES_SKETCH` return sketches in query results. If `true` (default behavior in Druid 24.0.1 and earlier), Druid finalizes sketches from these functions when they appear in query results.

This property is provided for backwards compatibility with behavior in Druid 24.0.1 and earlier. We don't recommend setting this parameter unless your application uses Druid 24.0.1 or earlier. Instead, use a function that doesn't return a sketch, such as `APPROX_COUNT_DISTINCT_DS_HLL`, `APPROX_COUNT_DISTINCT_DS_THETA`, `APPROX_QUANTILE_DS`, `DS_THETA_ESTIMATE`, or `DS_GET_QUANTILE`.|`false`| |`sqlUseBoundAndSelectors`|If `false` (default behavior in Druid 27.0.0 and later), the SQL planner uses [equality](./filters.md#equality-filter), [null](./filters.md#null-filter), and [range](./filters.md#range-filter) filters instead of [selector](./filters.md#selector-filter) and [bounds](./filters.md#bound-filter). For filtering `ARRAY` typed values, `sqlUseBoundAndSelectors` must be `false`. | `false`.| |`sqlUseExtractionFns`|If false, the SQL planner avoids using [`extractionFn`](dimensionspecs.md#extraction-functions) in favor of using other constructs such as [virtual columns](virtual-columns.md). This parameter is provided for compatibility with prior behavior, and may be removed in a future release.|false| diff --git a/docs/querying/sql-translation.md b/docs/querying/sql-translation.md index 8c2e323b6b48..3c3384148332 100644 --- a/docs/querying/sql-translation.md +++ b/docs/querying/sql-translation.md @@ -817,14 +817,20 @@ This query context parameter is a temporary solution to avoid the known issue. Druid does not support all SQL features. In particular, the following features are not supported. -- JOIN between native datasources (table, lookup, subquery) and [system tables](sql-metadata-tables.md). +- JOIN between native datasources (table, lookup, subquery) and [system tables](sql-metadata-tables.md) that use the + traditional SQL-layer execution path. System tables that support native execution can participate in these joins + when `useNativeQueryForSystemTables` is enabled. - JOIN conditions that are not an equality between expressions from the left- and right-hand sides. - JOIN conditions containing a constant value inside the condition. - JOIN conditions on a column which contains a multi-value dimension. -- ORDER BY for a non-aggregating query, except for `ORDER BY __time` or `ORDER BY __time DESC`, which are supported. - This restriction only applies to non-aggregating queries; you can ORDER BY any column in an aggregating query. +- ORDER BY for a non-aggregating query, except for `ORDER BY __time` or `ORDER BY __time DESC`, and native-enabled + system-table queries that use the native window-query path. This restriction only applies to non-aggregating queries; + you can ORDER BY any column in an aggregating query. - DDL and DML. -- Using Druid-specific functions like `TIME_PARSE` and `APPROX_QUANTILE_DS` on [system tables](sql-metadata-tables.md). +- Using Druid-specific functions like `TIME_PARSE` and `APPROX_QUANTILE_DS` on system tables that use the traditional + SQL-layer execution path. [System tables with native query support](sql-metadata-tables.md#native-query-execution) can + use expressions and aggregations that the native SQL engine can translate when `useNativeQueryForSystemTables` is + enabled. Additionally, some Druid native query features are not supported by the SQL language. Some unsupported Druid features include: @@ -834,5 +840,3 @@ include: - [Multi-value dimensions](sql-data-types.md#multi-value-strings) are only partially implemented in Druid SQL. There are known inconsistencies between their behavior in SQL queries and in native queries due to how they are currently treated by the SQL planner. - - diff --git a/docs/querying/sql.md b/docs/querying/sql.md index b74c36365025..f69b7affb109 100644 --- a/docs/querying/sql.md +++ b/docs/querying/sql.md @@ -79,11 +79,15 @@ datasources can be referenced as either `druid.dataSourceName` or simply `dataSo - [Lookups](datasource.md#lookup) from the `lookup` schema, for example `lookup.countries`. Note that lookups can also be queried using the [`LOOKUP` function](sql-scalar.md#string-functions). - [Subqueries](datasource.md#query). -- [Joins](datasource.md#join) between anything in this list, except between native datasources (table, lookup, -query) and system tables. The join condition must be an equality between expressions from the left- and right-hand side -of the join. -- [Metadata tables](sql-metadata-tables.md) from the `INFORMATION_SCHEMA` or `sys` schemas. Unlike the other options for the -FROM clause, metadata tables are not considered datasources. They exist only in the SQL layer. +- [Joins](datasource.md#join) between anything in this list. The traditional SQL-layer system-table path does not + support joins between native datasources (table, lookup, query) and system tables. When + [`useNativeQueryForSystemTables`](sql-metadata-tables.md#native-query-execution) is enabled, system tables that support + native execution can participate in native datasource joins. The join condition must be an equality between + expressions from the left- and right-hand side of the join. +- [Metadata tables](sql-metadata-tables.md) from the `INFORMATION_SCHEMA` or `sys` schemas. By default, metadata tables + exist only in the SQL layer and are not considered datasources. When + [`useNativeQueryForSystemTables`](sql-metadata-tables.md#native-query-execution) is enabled, supported `sys` tables are + represented as native system-table datasources. For more information about table, lookup, query, and join datasources, refer to the [Datasources](datasource.md) documentation. @@ -295,7 +299,8 @@ grouping expressions or aggregated values. It can only be used together with GRO The ORDER BY clause refers to columns that are present after execution of GROUP BY. It can be used to order the results based on either grouping expressions or aggregated values. ORDER BY can refer to an expression or a select clause ordinal position (like `ORDER BY 2` to order by the second selected column). For non-aggregation queries, ORDER BY -can only order by the `__time` column. For aggregation queries, ORDER BY can order by any column. +can normally only order by the `__time` column. Native-enabled system tables can also use the native window-query path +to order by other columns. For aggregation queries, ORDER BY can order by any column. ## LIMIT diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/NativeSysServerPropertiesQueryTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/NativeSysServerPropertiesQueryTest.java new file mode 100644 index 000000000000..18334a0547fa --- /dev/null +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/NativeSysServerPropertiesQueryTest.java @@ -0,0 +1,182 @@ +/* + * 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.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.EmbeddedHistorical; +import org.apache.druid.testing.embedded.EmbeddedIndexer; +import org.apache.druid.testing.embedded.EmbeddedOverlord; +import org.apache.druid.testing.embedded.EmbeddedRouter; +import org.apache.druid.testing.embedded.junit5.EmbeddedClusterTestBase; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import java.util.Map; + +public class NativeSysServerPropertiesQueryTest extends EmbeddedClusterTestBase +{ + private static final String SERVICE_NAME = "native/mvp/broker"; + private static final String COORDINATOR_PROPERTY = "native.sys.server.properties.coordinator"; + private static final String OVERLORD_PROPERTY = "native.sys.server.properties.overlord"; + private static final String BROKER_PROPERTY = "native.sys.server.properties.broker"; + private static final String HISTORICAL_PROPERTY = "native.sys.server.properties.historical"; + private static final String INDEXER_PROPERTY = "native.sys.server.properties.indexer"; + private static final String ROUTER_PROPERTY = "native.sys.server.properties.router"; + + private final EmbeddedCoordinator coordinator = new EmbeddedCoordinator() + .addProperty(COORDINATOR_PROPERTY, "enabled"); + + private final EmbeddedOverlord overlord = new EmbeddedOverlord() + .addProperty(OVERLORD_PROPERTY, "enabled"); + + private final EmbeddedBroker broker = new EmbeddedBroker() + .addProperty("druid.service", SERVICE_NAME) + .addProperty(BROKER_PROPERTY, "enabled"); + + private final EmbeddedHistorical historical = new EmbeddedHistorical() + .addProperty(HISTORICAL_PROPERTY, "enabled"); + + private final EmbeddedIndexer indexer = new EmbeddedIndexer() + .addProperty(INDEXER_PROPERTY, "enabled"); + + private final EmbeddedRouter router = new EmbeddedRouter() + .addProperty(ROUTER_PROPERTY, "enabled"); + + @Override + protected EmbeddedDruidCluster createCluster() + { + return EmbeddedDruidCluster.withEmbeddedDerbyAndZookeeper() + .useLatchableEmitter() + .addCommonProperty("druid.centralizedDatasourceSchema.enabled", "true") + .addServer(coordinator) + .addServer(overlord) + .addServer(broker) + .addServer(historical) + .addServer(indexer) + .addServer(router); + } + + /** + * Verifies that a native query for {@code sys.server_properties} reaches every persistent node type in the + * embedded cluster. In particular, the Router row proves that the Broker's native request included + * {@code X-Druid-Native-Query-Route: local}: without that header the Router would forward the request back to the + * Broker instead of reading its local system-table provider. + */ + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testServerPropertiesFansOutToAllNodes(final String plannerStrategy) + { + final String result = cluster.runSql( + "SELECT service_name, COUNT(*) " + + "FROM sys.server_properties " + + "WHERE property IN ('" + COORDINATOR_PROPERTY + "', '" + OVERLORD_PROPERTY + "', '" + BROKER_PROPERTY + + "', '" + HISTORICAL_PROPERTY + "', '" + INDEXER_PROPERTY + "', '" + ROUTER_PROPERTY + "') " + + "GROUP BY service_name ORDER BY service_name", + nativeQueryContext(plannerStrategy) + ); + + Assertions.assertEquals( + String.join( + "\n", + "druid/coordinator,1", + "druid/historical,1", + "druid/indexer,1", + "druid/overlord,1", + "druid/router,1", + SERVICE_NAME + ",1" + ), + result + ); + } + + @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 service_name), COUNT(DISTINCT server), " + + "COUNT(DISTINCT property), SUM(1) " + + "FROM sys.server_properties " + + "WHERE property IN ('" + COORDINATOR_PROPERTY + "', '" + OVERLORD_PROPERTY + "', '" + BROKER_PROPERTY + + "', '" + HISTORICAL_PROPERTY + "', '" + INDEXER_PROPERTY + "', '" + ROUTER_PROPERTY + "')", + nativeQueryContext(plannerStrategy) + ); + + Assertions.assertEquals("6,6,6,6,6", result); + } + + /** + * A filtered window-function query is planned as a {@code WindowOperatorQuery} over a system-table Scan, whose + * filter is pushed into the node scans. + */ + @ParameterizedTest(name = "plannerStrategy = {0}") + @ValueSource(strings = { + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_COUPLED, + QueryContexts.NATIVE_QUERY_SQL_PLANNING_MODE_DECOUPLED + }) + public void testNativeWindowFunction(final String plannerStrategy) + { + final String result = cluster.runSql( + "SELECT service_name, ROW_NUMBER() OVER (ORDER BY service_name) " + + "FROM sys.server_properties " + + "WHERE property IN ('" + COORDINATOR_PROPERTY + "', '" + OVERLORD_PROPERTY + "', '" + BROKER_PROPERTY + + "', '" + HISTORICAL_PROPERTY + "', '" + INDEXER_PROPERTY + "', '" + ROUTER_PROPERTY + "') " + + "ORDER BY service_name", + nativeQueryContext(plannerStrategy) + ); + + Assertions.assertEquals( + String.join( + "\n", + "druid/coordinator,1", + "druid/historical,2", + "druid/indexer,3", + "druid/overlord,4", + "druid/router,5", + SERVICE_NAME + ",6" + ), + result + ); + } + + 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/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordRedirectInfo.java b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordRedirectInfo.java index 41e1ef8b7c9a..de7b67b0d5d1 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordRedirectInfo.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/overlord/http/OverlordRedirectInfo.java @@ -24,6 +24,7 @@ import com.google.inject.Inject; import org.apache.druid.indexing.overlord.DruidOverlord; import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.query.SystemTableDataSource; import org.apache.druid.server.http.RedirectInfo; import java.net.URL; @@ -35,7 +36,9 @@ public class OverlordRedirectInfo implements RedirectInfo { private static final Set LOCAL_PATHS = ImmutableSet.of( "/druid/indexer/v1/leader", - "/druid/indexer/v1/isLeader" + "/druid/indexer/v1/isLeader", + "/druid/v2", + "/druid/v2/" ); private final DruidOverlord overlord; @@ -49,7 +52,10 @@ public OverlordRedirectInfo(DruidOverlord overlord) @Override public boolean doLocal(String requestURI) { - return (requestURI != null && LOCAL_PATHS.contains(requestURI)) || overlord.isLeader(); + return (requestURI != null + && (LOCAL_PATHS.contains(requestURI) + || requestURI.startsWith("/druid/v2/" + SystemTableDataSource.NODE_QUERY_ID_PREFIX))) + || overlord.isLeader(); } @Override diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordRedirectInfoTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordRedirectInfoTest.java index 7dde7dab0232..dbcf81bb7e61 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordRedirectInfoTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordRedirectInfoTest.java @@ -50,6 +50,7 @@ public void testDoLocalWhenLeading() Assertions.assertTrue(redirectInfo.doLocal(null)); Assertions.assertTrue(redirectInfo.doLocal("/druid/indexer/v1/leader")); Assertions.assertTrue(redirectInfo.doLocal("/druid/indexer/v1/isLeader")); + Assertions.assertTrue(redirectInfo.doLocal("/druid/v2")); Assertions.assertTrue(redirectInfo.doLocal("/druid/indexer/v1/other/path")); EasyMock.verify(overlord); } @@ -62,6 +63,7 @@ public void testDoLocalWhenNotLeading() Assertions.assertFalse(redirectInfo.doLocal(null)); Assertions.assertTrue(redirectInfo.doLocal("/druid/indexer/v1/leader")); Assertions.assertTrue(redirectInfo.doLocal("/druid/indexer/v1/isLeader")); + Assertions.assertTrue(redirectInfo.doLocal("/druid/v2")); Assertions.assertFalse(redirectInfo.doLocal("/druid/indexer/v1/other/path")); EasyMock.verify(overlord); } diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index 58b4265430b3..a42cc4534cd8 100644 --- a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -562,6 +562,33 @@ public void test_doPost_informationSchema() ); } + @Test + public void test_doPost_nativeCapableSystemTableUsesBindablePlan() + { + final MockAsyncContext asyncContext = new MockAsyncContext(); + final MockHttpServletResponse asyncResponse = new MockHttpServletResponse(); + asyncContext.response = asyncResponse; + + Mockito.when(httpServletRequest.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)) + .thenReturn(makeAuthenticationResult(CalciteTests.TEST_SUPERUSER_NAME)); + Mockito.when(httpServletRequest.startAsync()) + .thenReturn(asyncContext); + + final SqlQuery sqlQuery = new SqlQuery( + "SELECT COUNT(*) FROM sys.tasks", + ResultFormat.ARRAY, + false, + false, + false, + Map.of(QueryContexts.ENGINE, DartSqlEngine.NAME), + Collections.emptyList() + ); + + Assertions.assertNull(sqlResource.doPost(sqlQuery, httpServletRequest)); + Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); + Assertions.assertEquals("[[4]]\n", StringUtils.fromUtf8(asyncResponse.baos.toByteArray())); + } + @Test public void test_doPost_sysTableJoinedToDatasource() { diff --git a/processing/src/main/java/org/apache/druid/query/DataSource.java b/processing/src/main/java/org/apache/druid/query/DataSource.java index 9b0d65f2d83a..a89ec5a43a07 100644 --- a/processing/src/main/java/org/apache/druid/query/DataSource.java +++ b/processing/src/main/java/org/apache/druid/query/DataSource.java @@ -49,7 +49,8 @@ @JsonSubTypes.Type(value = GlobalTableDataSource.class, name = "globalTable"), @JsonSubTypes.Type(value = UnnestDataSource.class, name = "unnest"), @JsonSubTypes.Type(value = FilteredDataSource.class, name = "filter"), - @JsonSubTypes.Type(value = RestrictedDataSource.class, name = "restrict") + @JsonSubTypes.Type(value = RestrictedDataSource.class, name = "restrict"), + @JsonSubTypes.Type(value = SystemTableDataSource.class, name = "systemTable") }) public interface DataSource extends Cacheable { diff --git a/processing/src/main/java/org/apache/druid/query/SystemTableDataSource.java b/processing/src/main/java/org/apache/druid/query/SystemTableDataSource.java new file mode 100644 index 000000000000..2d2886619669 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/query/SystemTableDataSource.java @@ -0,0 +1,110 @@ +/* + * 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.query; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.Preconditions; + +import java.util.Collections; +import java.util.Objects; +import java.util.Set; + +/** + * Identifies a Druid system table whose rows are supplied by a service rather than by segments. + */ +public class SystemTableDataSource extends LeafDataSource +{ + public static final String NODE_QUERY_ID_PREFIX = "native-system-node-"; + + private final String table; + + @JsonCreator + public SystemTableDataSource(@JsonProperty("table") final String table) + { + this.table = Preconditions.checkNotNull(table, "table"); + } + + @JsonProperty + public String getTable() + { + return table; + } + + @Override + public Set getTableNames() + { + // QueryScheduler uses table names as cancellation authorization resources. Namespace the synthetic datasource + // name so a system table cannot collide with a regular Druid datasource that has the same unqualified name. + return Collections.singleton("sys." + table); + } + + @Override + public boolean isCacheable(boolean isBroker) + { + return false; + } + + @Override + public boolean isGlobal() + { + return false; + } + + @Override + public boolean isProcessable() + { + // The owning service must first resolve this datasource to an InlineDataSource. + return false; + } + + @Override + public byte[] getCacheKey() + { + return null; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + final SystemTableDataSource that = (SystemTableDataSource) o; + return Objects.equals(table, that.table); + } + + @Override + public int hashCode() + { + return Objects.hash(table); + } + + @Override + public String toString() + { + return "SystemTableDataSource{" + + "table='" + table + '\'' + + '}'; + } +} diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java index 270ab566814b..8fb513643cde 100644 --- a/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java +++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java @@ -32,6 +32,7 @@ import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.Pair; import org.apache.druid.java.util.common.UOE; +import org.apache.druid.query.ResourceLimitExceededException; import org.apache.druid.query.filter.Filter; import org.apache.druid.query.filter.ValueMatcher; import org.apache.druid.query.operator.ColumnWithDirection; @@ -281,7 +282,9 @@ private Pair materializeCursorFactory(CursorFactory cursor cursor.advance(); } for (; !cursor.isDoneOrInterrupted() && remainingRowsToFetch > 0; remainingRowsToFetch--) { - writer.addSelection(); + if (!writer.addSelection()) { + throw materializationLimitExceeded(); + } cursor.advance(); } @@ -395,10 +398,20 @@ private Pair naiveMaterialize(RowsAndColumns rac) continue; } remainingRowsToFetch--; - frameWriter.addSelection(); + if (!frameWriter.addSelection()) { + throw materializationLimitExceeded(); + } } return Pair.of(frameWriter.toByteArray(), sigBob.build()); } } + + private ResourceLimitExceededException materializationLimitExceeded() + { + return ResourceLimitExceededException.withMessage( + "RowsAndColumns materialization exceeded the configured frame capacity of [%,d] bytes", + allocatorCapacity + ); + } } diff --git a/processing/src/main/java/org/apache/druid/segment/FilteredSegment.java b/processing/src/main/java/org/apache/druid/segment/FilteredSegment.java index 909d80b38aaa..1a56bc48a3c7 100644 --- a/processing/src/main/java/org/apache/druid/segment/FilteredSegment.java +++ b/processing/src/main/java/org/apache/druid/segment/FilteredSegment.java @@ -20,6 +20,7 @@ package org.apache.druid.segment; import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.rowsandcols.CursorFactoryRowsAndColumns; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -44,6 +45,8 @@ public T as(@Nonnull Class clazz) { if (CursorFactory.class.equals(clazz)) { return (T) new FilteredCursorFactory(delegate.as(CursorFactory.class), filter); + } else if (CloseableShapeshifter.class.equals(clazz)) { + return (T) new CursorFactoryRowsAndColumns(as(CursorFactory.class)); } else if (TopNOptimizationInspector.class.equals(clazz)) { return (T) new SimpleTopNOptimizationInspector(filter == null); } diff --git a/processing/src/test/java/org/apache/druid/query/rowsandcols/semantic/EvaluateRowsAndColumnsTest.java b/processing/src/test/java/org/apache/druid/query/rowsandcols/semantic/EvaluateRowsAndColumnsTest.java index 810b1b83da21..eef947598c95 100644 --- a/processing/src/test/java/org/apache/druid/query/rowsandcols/semantic/EvaluateRowsAndColumnsTest.java +++ b/processing/src/test/java/org/apache/druid/query/rowsandcols/semantic/EvaluateRowsAndColumnsTest.java @@ -20,6 +20,7 @@ package org.apache.druid.query.rowsandcols.semantic; import com.google.common.collect.Lists; +import org.apache.druid.query.ResourceLimitExceededException; import org.apache.druid.query.expression.TestExprMacroTable; import org.apache.druid.query.operator.OffsetLimit; import org.apache.druid.query.operator.window.RowsAndColumnsHelper; @@ -118,4 +119,34 @@ public void testMaterializeColumns() .expectColumn("array", expectedArr, ColumnType.STRING_ARRAY) .validate(ras); } + + @Test + public void testMaterializationFailsInsteadOfReturningPartialRowsWhenFrameIsFull() + { + final Object[][] vals = new Object[][] {{0L, "x".repeat(64 * 1024)}}; + final RowSignature signature = RowSignature.builder() + .add("__time", ColumnType.LONG) + .add("value", ColumnType.STRING) + .build(); + final RowsAndColumns base = make(MapOfColumnsRowsAndColumns.fromRowObjects(vals, signature)); + + Assumptions.assumeTrue(base.as(CursorFactory.class) != null, "skipping: CursorFactory not supported"); + + final LazilyDecoratedRowsAndColumns rowsAndColumns = new LazilyDecoratedRowsAndColumns( + base, + null, + null, + null, + OffsetLimit.limit(1), + null, + null, + 1024L + ); + + final ResourceLimitExceededException exception = Assertions.assertThrows( + ResourceLimitExceededException.class, + rowsAndColumns::numRows + ); + Assertions.assertTrue(exception.getMessage().contains("configured frame capacity")); + } } diff --git a/processing/src/test/java/org/apache/druid/segment/FilteredSegmentTest.java b/processing/src/test/java/org/apache/druid/segment/FilteredSegmentTest.java new file mode 100644 index 000000000000..39cbda153464 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/segment/FilteredSegmentTest.java @@ -0,0 +1,56 @@ +/* + * 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.segment; + +import org.apache.druid.query.InlineDataSource; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.query.rowsandcols.RowsAndColumns; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; + +public class FilteredSegmentTest +{ + /** Window queries read segments as {@link RowsAndColumns}; a filtered segment provides them with its filter applied. */ + @Test + public void testShapeshiftsToFilteredRowsAndColumns() throws Exception + { + final RowSignature signature = RowSignature.builder().add("dim", ColumnType.STRING).build(); + final InlineDataSource inline = InlineDataSource.fromIterable( + List.of(new Object[]{"a"}, new Object[]{"b"}, new Object[]{"a"}), + signature + ); + final Segment segment = new FilteredSegment( + new ArrayListSegment<>(new ArrayList<>(inline.getRowsAsList()), inline.rowAdapter(), signature), + new SelectorDimFilter("dim", "b", null) + ); + + try (final CloseableShapeshifter shapeshifter = segment.as(CloseableShapeshifter.class)) { + Assertions.assertNotNull(shapeshifter); + final RowsAndColumns rac = Assertions.assertInstanceOf(RowsAndColumns.class, shapeshifter); + Assertions.assertEquals(1, rac.numRows()); + Assertions.assertEquals("b", rac.findColumn("dim").toAccessor().getObject(0)); + } + } +} diff --git a/server/src/main/java/org/apache/druid/client/DirectDruidClient.java b/server/src/main/java/org/apache/druid/client/DirectDruidClient.java index 2f963e0b270c..21addd635241 100644 --- a/server/src/main/java/org/apache/druid/client/DirectDruidClient.java +++ b/server/src/main/java/org/apache/druid/client/DirectDruidClient.java @@ -797,7 +797,8 @@ private void checkTotalBytesLimit(long bytes) .setHeader( HttpHeaders.Names.CONTENT_TYPE, isSmile ? SmileMediaTypes.APPLICATION_JACKSON_SMILE : MediaType.APPLICATION_JSON - ), + ) + .setHeader(QueryResource.HEADER_NATIVE_QUERY_ROUTE, QueryResource.NATIVE_QUERY_ROUTE_LOCAL), responseHandler, Duration.millis(timeLeft) ); @@ -880,8 +881,12 @@ private void cancelQuery(Query query, String cancelUrl) try { Future responseFuture = httpClient.go( new Request(HttpMethod.DELETE, new URL(cancelUrl)) - .setContent(objectMapper.writeValueAsBytes(query)) - .setHeader(HttpHeaders.Names.CONTENT_TYPE, isSmile ? SmileMediaTypes.APPLICATION_JACKSON_SMILE : MediaType.APPLICATION_JSON), + .setContent(objectMapper.writeValueAsBytes(query)) + .setHeader( + HttpHeaders.Names.CONTENT_TYPE, + isSmile ? SmileMediaTypes.APPLICATION_JACKSON_SMILE : MediaType.APPLICATION_JSON + ) + .setHeader(QueryResource.HEADER_NATIVE_QUERY_ROUTE, QueryResource.NATIVE_QUERY_ROUTE_LOCAL), StatusResponseHandler.getInstance(), Duration.standardSeconds(1)); diff --git a/server/src/main/java/org/apache/druid/guice/BrokerQueryResourceModule.java b/server/src/main/java/org/apache/druid/guice/BrokerQueryResourceModule.java new file mode 100644 index 000000000000..999e7547a631 --- /dev/null +++ b/server/src/main/java/org/apache/druid/guice/BrokerQueryResourceModule.java @@ -0,0 +1,40 @@ +/* + * 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.guice; + +import com.google.inject.Binder; +import org.apache.druid.initialization.DruidModule; +import org.apache.druid.server.BrokerQueryResource; +import org.apache.druid.server.ResponseContextConfig; +import org.apache.druid.server.metrics.QueryCountStatsProvider; + +/** Registers the Broker's specialized native query HTTP resource. */ +public class BrokerQueryResourceModule implements DruidModule +{ + @Override + public void configure(final Binder binder) + { + binder.bind(ResponseContextConfig.class).toInstance(ResponseContextConfig.newConfig(false)); + binder.bind(BrokerQueryResource.class).in(LazySingleton.class); + binder.bind(QueryCountStatsProvider.class).to(BrokerQueryResource.class).in(LazySingleton.class); + Jerseys.addResource(binder, BrokerQueryResource.class); + LifecycleModule.register(binder, BrokerQueryResource.class); + } +} diff --git a/server/src/main/java/org/apache/druid/guice/DruidBinders.java b/server/src/main/java/org/apache/druid/guice/DruidBinders.java index 5c969d4ddf6b..7699a2b613ef 100644 --- a/server/src/main/java/org/apache/druid/guice/DruidBinders.java +++ b/server/src/main/java/org/apache/druid/guice/DruidBinders.java @@ -32,6 +32,7 @@ import org.apache.druid.query.QueryToolChest; import org.apache.druid.segment.SegmentWrangler; import org.apache.druid.segment.join.JoinableFactory; +import org.apache.druid.server.DataSourceQueryHandler; import java.lang.annotation.Annotation; import java.util.Set; @@ -147,6 +148,17 @@ public static MapBinder, SegmentWrangler> segmentWra ); } + public static MapBinder, DataSourceQueryHandler> dataSourceQueryHandlerBinder( + Binder binder + ) + { + return MapBinder.newMapBinder( + binder, + new TypeLiteral<>() {}, + new TypeLiteral<>() {} + ); + } + public static Multibinder joinableFactoryMultiBinder(Binder binder) { return Multibinder.newSetBinder( diff --git a/server/src/main/java/org/apache/druid/guice/NativeQueryEngineModule.java b/server/src/main/java/org/apache/druid/guice/NativeQueryEngineModule.java new file mode 100644 index 000000000000..cb161c344cd2 --- /dev/null +++ b/server/src/main/java/org/apache/druid/guice/NativeQueryEngineModule.java @@ -0,0 +1,141 @@ +/* + * 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.guice; + +import com.google.common.collect.ImmutableSet; +import com.google.inject.Binder; +import com.google.inject.Module; +import com.google.inject.multibindings.OptionalBinder; +import com.google.inject.util.Modules; +import org.apache.druid.initialization.DruidModule; +import org.apache.druid.query.QuerySegmentWalker; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.server.NoopQuerySegmentWalker; +import org.apache.druid.server.system.module.SystemTableModule; +import org.apache.druid.server.system.module.SystemTableQueryResourceModule; + +import javax.annotation.Nullable; +import java.util.List; + +/** + * Facade for the modules that provide native query execution infrastructure. + */ +public class NativeQueryEngineModule implements DruidModule +{ + private final Module executionModule; + private final Module queryResourceModule; + + private NativeQueryEngineModule( + final Module executionModule, + final Module queryResourceModule + ) + { + this.executionModule = executionModule; + this.queryResourceModule = queryResourceModule; + } + + public static Builder builder() + { + return new Builder(); + } + + public static final class Builder + { + private Module overrideModule = binder -> {}; + @Nullable + private Module queryResourceModule; + private boolean scanOnly; + + private Builder() + { + } + + /** + * Uses the minimum infrastructure required to serve Scan queries. This profile does not install processing or + * merge-buffer dependencies. + */ + public Builder scanOnly() + { + scanOnly = true; + return this; + } + + /** Applies role-specific overrides to the query execution bindings. */ + public Builder withOverrideModule(final Module module) + { + overrideModule = module; + return this; + } + + /** Replaces the standard native query HTTP resource with a role-specific resource module. */ + public Builder withQueryResourceModule(final Module module) + { + queryResourceModule = module; + return this; + } + + public NativeQueryEngineModule build() + { + final Module queryRunnerFactoryModule; + final Module querySegmentWalkerModule; + if (scanOnly) { + queryRunnerFactoryModule = new QueryRunnerFactoryModule(ImmutableSet.of(ScanQuery.class)); + // Scan-only servers resolve node-local system tables without walking segments. Modules that need a real + // walker, such as SegmentSchemaCacheModule, replace this optional default with an explicit binding. + querySegmentWalkerModule = binder -> OptionalBinder.newOptionalBinder(binder, QuerySegmentWalker.class) + .setDefault() + .to(NoopQuerySegmentWalker.class) + .in(LazySingleton.class); + } else { + queryRunnerFactoryModule = new QueryRunnerFactoryModule(); + querySegmentWalkerModule = binder -> {}; + } + + return new NativeQueryEngineModule( + Modules.override( + Modules.combine( + new QueryableModule(), + queryRunnerFactoryModule, + querySegmentWalkerModule + ) + ).with(overrideModule), + queryResourceModule == null + ? scanOnly ? new SystemTableQueryResourceModule() : new QueryResourceModule() + : queryResourceModule + ); + } + } + + @Override + public void configure(final Binder binder) + { + binder.install(executionModule); + binder.install(new SegmentWranglerModule()); + binder.install(new JoinableFactoryModule()); + binder.install(new SystemTableModule()); + binder.install(queryResourceModule); + } + + @Override + public List getJacksonModules() + { + return new QueryableModule().getJacksonModules(); + } +} diff --git a/services/src/main/java/org/apache/druid/guice/QueryablePeonModule.java b/server/src/main/java/org/apache/druid/guice/QueryResourceModule.java similarity index 80% rename from services/src/main/java/org/apache/druid/guice/QueryablePeonModule.java rename to server/src/main/java/org/apache/druid/guice/QueryResourceModule.java index 0d9c316b3ecf..3902946e0771 100644 --- a/services/src/main/java/org/apache/druid/guice/QueryablePeonModule.java +++ b/server/src/main/java/org/apache/druid/guice/QueryResourceModule.java @@ -22,13 +22,16 @@ import com.google.inject.Binder; import org.apache.druid.initialization.DruidModule; import org.apache.druid.server.QueryResource; +import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.metrics.QueryCountStatsProvider; -public class QueryablePeonModule implements DruidModule +/** Registers the standard native query HTTP resource. */ +public class QueryResourceModule implements DruidModule { @Override - public void configure(Binder binder) + public void configure(final Binder binder) { + binder.bind(ResponseContextConfig.class).toInstance(ResponseContextConfig.newConfig(true)); binder.bind(QueryCountStatsProvider.class).to(QueryResource.class); Jerseys.addResource(binder, QueryResource.class); LifecycleModule.register(binder, QueryResource.class); diff --git a/server/src/main/java/org/apache/druid/guice/QueryRunnerFactoryModule.java b/server/src/main/java/org/apache/druid/guice/QueryRunnerFactoryModule.java index e991fac51b5d..0d7f2d2e6592 100644 --- a/server/src/main/java/org/apache/druid/guice/QueryRunnerFactoryModule.java +++ b/server/src/main/java/org/apache/druid/guice/QueryRunnerFactoryModule.java @@ -19,7 +19,9 @@ package org.apache.druid.guice; +import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.inject.Binder; import com.google.inject.Key; import com.google.inject.Provides; @@ -52,6 +54,7 @@ import org.apache.druid.server.QuerySchedulerProvider; import java.util.Map; +import java.util.Set; /** */ @@ -70,6 +73,32 @@ public class QueryRunnerFactoryModule extends QueryToolChestModule .put(WindowOperatorQuery.class, WindowOperatorQueryQueryRunnerFactory.class) .build(); + private static final Set> SUPPORTED_QUERY_TYPES = + ImmutableSet.>builder() + .addAll(MAPPINGS.keySet()) + .add(UnionQuery.class) + .build(); + + private final Set> queryTypes; + + public QueryRunnerFactoryModule() + { + this(SUPPORTED_QUERY_TYPES); + } + + public QueryRunnerFactoryModule(final Set> queryTypes) + { + super(queryTypes); + Preconditions.checkArgument( + SUPPORTED_QUERY_TYPES.containsAll(queryTypes), + "Unsupported query types[%s]", + ImmutableSet.copyOf(queryTypes).stream() + .filter(queryType -> !SUPPORTED_QUERY_TYPES.contains(queryType)) + .toList() + ); + this.queryTypes = ImmutableSet.copyOf(queryTypes); + } + @Override public void configure(Binder binder) { @@ -86,12 +115,18 @@ public void configure(Binder binder) ); for (Map.Entry>, Class>> entry : MAPPINGS.entrySet()) { - queryFactoryBinder.addBinding(entry.getKey()).to(entry.getValue()); - binder.bind(entry.getValue()).in(LazySingleton.class); + if (queryTypes.contains(entry.getKey())) { + queryFactoryBinder.addBinding(entry.getKey()).to(entry.getValue()); + binder.bind(entry.getValue()).in(LazySingleton.class); + } } - DruidBinders.queryBinder(binder) - .bindQueryLogic(UnionQuery.class, UnionQueryLogic.class); + if (queryTypes.contains(UnionQuery.class)) { + DruidBinders.queryBinder(binder) + .bindQueryLogic(UnionQuery.class, UnionQueryLogic.class); + } else { + DruidBinders.queryBinder(binder); + } } @LazySingleton diff --git a/server/src/main/java/org/apache/druid/guice/QueryToolChestModule.java b/server/src/main/java/org/apache/druid/guice/QueryToolChestModule.java index 72f0e142b297..d1ecdb6009fb 100644 --- a/server/src/main/java/org/apache/druid/guice/QueryToolChestModule.java +++ b/server/src/main/java/org/apache/druid/guice/QueryToolChestModule.java @@ -20,6 +20,7 @@ package org.apache.druid.guice; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.inject.Binder; import com.google.inject.Key; import com.google.inject.Module; @@ -64,6 +65,7 @@ import org.apache.druid.query.topn.TopNQueryQueryToolChest; import java.util.Map; +import java.util.Set; /** */ @@ -75,7 +77,7 @@ public class QueryToolChestModule implements Module public static final String TOPN_QUERY_METRICS_FACTORY_PROPERTY = "druid.query.topN.queryMetricsFactory"; public static final String SEARCH_QUERY_METRICS_FACTORY_PROPERTY = "druid.query.search.queryMetricsFactory"; - public final Map, Class> mappings = + private static final Map, Class> MAPPINGS = ImmutableMap., Class>builder() .put(DataSourceMetadataQuery.class, DataSourceQueryQueryToolChest.class) .put(GroupByQuery.class, GroupByQueryQueryToolChest.class) @@ -88,14 +90,30 @@ public class QueryToolChestModule implements Module .put(WindowOperatorQuery.class, WindowOperatorQueryQueryToolChest.class) .build(); + public final Map, Class> mappings = MAPPINGS; + + private final Set> queryTypes; + + public QueryToolChestModule() + { + this(MAPPINGS.keySet()); + } + + public QueryToolChestModule(final Set> queryTypes) + { + this.queryTypes = ImmutableSet.copyOf(queryTypes); + } + @Override public void configure(Binder binder) { MapBinder, QueryToolChest> toolChests = DruidBinders.queryToolChestBinder(binder); - for (Map.Entry, Class> entry : mappings.entrySet()) { - toolChests.addBinding(entry.getKey()).to(entry.getValue()); - binder.bind(entry.getValue()).in(LazySingleton.class); + for (Map.Entry, Class> entry : MAPPINGS.entrySet()) { + if (queryTypes.contains(entry.getKey())) { + toolChests.addBinding(entry.getKey()).to(entry.getValue()); + binder.bind(entry.getValue()).in(LazySingleton.class); + } } binder.bind(QueryToolChestWarehouse.class).to(ConglomerateBackedToolChestWarehouse.class); @@ -103,11 +121,21 @@ public void configure(Binder binder) JsonConfigProvider.bind(binder, "druid.query.default", DefaultQueryConfig.class); // DefaultQueryContext defaults to the static DefaultQueryConfig; brokers override this binding. binder.bind(QueryConfigProvider.class).to(DefaultQueryConfig.class); - JsonConfigProvider.bind(binder, "druid.query.groupBy", GroupByQueryConfig.class); - JsonConfigProvider.bind(binder, "druid.query.search", SearchQueryConfig.class); - JsonConfigProvider.bind(binder, "druid.query.topN", TopNQueryConfig.class); - JsonConfigProvider.bind(binder, "druid.query.segmentMetadata", SegmentMetadataQueryConfig.class); - JsonConfigProvider.bind(binder, "druid.query.scan", ScanQueryConfig.class); + if (queryTypes.contains(GroupByQuery.class)) { + JsonConfigProvider.bind(binder, "druid.query.groupBy", GroupByQueryConfig.class); + } + if (queryTypes.contains(SearchQuery.class)) { + JsonConfigProvider.bind(binder, "druid.query.search", SearchQueryConfig.class); + } + if (queryTypes.contains(TopNQuery.class)) { + JsonConfigProvider.bind(binder, "druid.query.topN", TopNQueryConfig.class); + } + if (queryTypes.contains(SegmentMetadataQuery.class)) { + JsonConfigProvider.bind(binder, "druid.query.segmentMetadata", SegmentMetadataQueryConfig.class); + } + if (queryTypes.contains(ScanQuery.class)) { + JsonConfigProvider.bind(binder, "druid.query.scan", ScanQueryConfig.class); + } PolyBind.createChoice( binder, @@ -120,48 +148,56 @@ public void configure(Binder binder) .addBinding("default") .to(DefaultGenericQueryMetricsFactory.class); - PolyBind.createChoice( - binder, - GROUPBY_QUERY_METRICS_FACTORY_PROPERTY, - Key.get(GroupByQueryMetricsFactory.class), - Key.get(DefaultGroupByQueryMetricsFactory.class) - ); - PolyBind - .optionBinder(binder, Key.get(GroupByQueryMetricsFactory.class)) - .addBinding("default") - .to(DefaultGroupByQueryMetricsFactory.class); + if (queryTypes.contains(GroupByQuery.class)) { + PolyBind.createChoice( + binder, + GROUPBY_QUERY_METRICS_FACTORY_PROPERTY, + Key.get(GroupByQueryMetricsFactory.class), + Key.get(DefaultGroupByQueryMetricsFactory.class) + ); + PolyBind + .optionBinder(binder, Key.get(GroupByQueryMetricsFactory.class)) + .addBinding("default") + .to(DefaultGroupByQueryMetricsFactory.class); + } - PolyBind.createChoice( - binder, - TIMESERIES_QUERY_METRICS_FACTORY_PROPERTY, - Key.get(TimeseriesQueryMetricsFactory.class), - Key.get(DefaultTimeseriesQueryMetricsFactory.class) - ); - PolyBind - .optionBinder(binder, Key.get(TimeseriesQueryMetricsFactory.class)) - .addBinding("default") - .to(DefaultTimeseriesQueryMetricsFactory.class); + if (queryTypes.contains(TimeseriesQuery.class)) { + PolyBind.createChoice( + binder, + TIMESERIES_QUERY_METRICS_FACTORY_PROPERTY, + Key.get(TimeseriesQueryMetricsFactory.class), + Key.get(DefaultTimeseriesQueryMetricsFactory.class) + ); + PolyBind + .optionBinder(binder, Key.get(TimeseriesQueryMetricsFactory.class)) + .addBinding("default") + .to(DefaultTimeseriesQueryMetricsFactory.class); + } - PolyBind.createChoice( - binder, - TOPN_QUERY_METRICS_FACTORY_PROPERTY, - Key.get(TopNQueryMetricsFactory.class), - Key.get(DefaultTopNQueryMetricsFactory.class) - ); - PolyBind - .optionBinder(binder, Key.get(TopNQueryMetricsFactory.class)) - .addBinding("default") - .to(DefaultTopNQueryMetricsFactory.class); + if (queryTypes.contains(TopNQuery.class)) { + PolyBind.createChoice( + binder, + TOPN_QUERY_METRICS_FACTORY_PROPERTY, + Key.get(TopNQueryMetricsFactory.class), + Key.get(DefaultTopNQueryMetricsFactory.class) + ); + PolyBind + .optionBinder(binder, Key.get(TopNQueryMetricsFactory.class)) + .addBinding("default") + .to(DefaultTopNQueryMetricsFactory.class); + } - PolyBind.createChoice( - binder, - SEARCH_QUERY_METRICS_FACTORY_PROPERTY, - Key.get(SearchQueryMetricsFactory.class), - Key.get(DefaultSearchQueryMetricsFactory.class) - ); - PolyBind - .optionBinder(binder, Key.get(SearchQueryMetricsFactory.class)) - .addBinding("default") - .to(DefaultSearchQueryMetricsFactory.class); + if (queryTypes.contains(SearchQuery.class)) { + PolyBind.createChoice( + binder, + SEARCH_QUERY_METRICS_FACTORY_PROPERTY, + Key.get(SearchQueryMetricsFactory.class), + Key.get(DefaultSearchQueryMetricsFactory.class) + ); + PolyBind + .optionBinder(binder, Key.get(SearchQueryMetricsFactory.class)) + .addBinding("default") + .to(DefaultSearchQueryMetricsFactory.class); + } } } diff --git a/server/src/main/java/org/apache/druid/guice/QueryableModule.java b/server/src/main/java/org/apache/druid/guice/QueryableModule.java index 7bd194132144..2a462734f7af 100644 --- a/server/src/main/java/org/apache/druid/guice/QueryableModule.java +++ b/server/src/main/java/org/apache/druid/guice/QueryableModule.java @@ -45,6 +45,7 @@ public class QueryableModule implements DruidModule @Override public void configure(Binder binder) { + DruidBinders.dataSourceQueryHandlerBinder(binder); binder.bind(RequestLogger.class).toProvider(RequestLoggerProvider.class).in(ManageLifecycle.class); JsonConfigProvider.bindWithDefault( binder, diff --git a/server/src/main/java/org/apache/druid/server/DataSourceQueryHandler.java b/server/src/main/java/org/apache/druid/server/DataSourceQueryHandler.java new file mode 100644 index 000000000000..6ad2cef59e3c --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/DataSourceQueryHandler.java @@ -0,0 +1,44 @@ +/* + * 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; + +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.server.security.AuthenticationResult; + +/** + * Creates a runner for a datasource whose execution is not provided by the normal segment walker. The query may + * contain the registered datasource below its root; the handler is responsible for resolving every matching vertex + * before delegating the remaining query to normal native execution. + */ +public interface DataSourceQueryHandler +{ + /** + * Creates the datasource runner for an authenticated request. + * + * @param executeLocally whether the request selected local execution with + * {@link QueryResource#HEADER_NATIVE_QUERY_ROUTE} + */ + QueryRunner createRunner( + Query query, + AuthenticationResult authenticationResult, + boolean executeLocally + ); +} diff --git a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java index 0f970a7fe0ce..ad7f8aac8500 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java @@ -34,6 +34,7 @@ import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.query.BaseQuery; +import org.apache.druid.query.DataSource; import org.apache.druid.query.DruidMetrics; import org.apache.druid.query.GenericQueryMetricsFactory; import org.apache.druid.query.Query; @@ -41,6 +42,7 @@ import org.apache.druid.query.QueryInterruptedException; import org.apache.druid.query.QueryMetrics; import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; import org.apache.druid.query.QueryRunnerFactoryConglomerate; import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.QueryTimeoutException; @@ -62,6 +64,7 @@ import javax.annotation.Nullable; import javax.servlet.http.HttpServletRequest; +import java.util.Collections; import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; @@ -99,12 +102,14 @@ public class QueryLifecycle private final AuthConfig authConfig; private final PolicyEnforcer policyEnforcer; private final QueryConfigSnapshot configSnapshot; + private final Map, DataSourceQueryHandler> dataSourceQueryHandlers; private final long startMs; private final long startNs; private State state = State.NEW; private AuthenticationResult authenticationResult; private QueryToolChest toolChest; + private boolean executeNativeQueryLocally; @MonotonicNonNull private Query baseQuery; @@ -124,6 +129,7 @@ public QueryLifecycle( final AuthConfig authConfig, final PolicyEnforcer policyEnforcer, final QueryConfigSnapshot configSnapshot, + final Map, DataSourceQueryHandler> dataSourceQueryHandlers, final long startMs, final long startNs ) @@ -137,10 +143,41 @@ public QueryLifecycle( this.authConfig = authConfig; this.policyEnforcer = policyEnforcer; this.configSnapshot = configSnapshot; + this.dataSourceQueryHandlers = dataSourceQueryHandlers; this.startMs = startMs; this.startNs = startNs; } + public QueryLifecycle( + final QueryRunnerFactoryConglomerate conglomerate, + final QuerySegmentWalker texasRanger, + final GenericQueryMetricsFactory queryMetricsFactory, + final ServiceEmitter emitter, + final RequestLogger requestLogger, + final AuthorizerMapper authorizerMapper, + final AuthConfig authConfig, + final PolicyEnforcer policyEnforcer, + final QueryConfigSnapshot configSnapshot, + final long startMs, + final long startNs + ) + { + this( + conglomerate, + texasRanger, + queryMetricsFactory, + emitter, + requestLogger, + authorizerMapper, + authConfig, + policyEnforcer, + configSnapshot, + Collections.emptyMap(), + startMs, + startNs + ); + } + /** * For callers who have already authorized their query, and where simplicity is desired over flexibility. This method * does it all in one call. Logs and metrics are emitted when the Sequence is either fully iterated or throws an @@ -271,6 +308,9 @@ public void initialize(final Query baseQuery, @Nullable final Set cli public AuthorizationResult authorize(HttpServletRequest req) { transition(State.INITIALIZED, State.AUTHORIZING); + executeNativeQueryLocally = QueryResource.NATIVE_QUERY_ROUTE_LOCAL.equals( + req.getHeader(QueryResource.HEADER_NATIVE_QUERY_ROUTE) + ); final Iterable resourcesToAuthorize = Iterables.concat( Iterables.transform( baseQuery.getDataSource().getTableNames(), @@ -420,13 +460,49 @@ public QueryResponse execute() final ResponseContext responseContext = DirectDruidClient.makeResponseContextForQuery(); @SuppressWarnings("unchecked") - final Sequence res = QueryPlus.wrap((Query) baseQuery) - .withIdentity(authenticationResult.getIdentity()) - .run(texasRanger, responseContext); + final Query query = (Query) baseQuery; + final DataSourceQueryHandler dataSourceQueryHandler = findDataSourceQueryHandler(query.getDataSource()); + final QueryRunner queryRunner = dataSourceQueryHandler == null + ? query.getRunner(texasRanger) + : dataSourceQueryHandler.createRunner( + query, + authenticationResult, + executeNativeQueryLocally + ); + final Sequence res = queryRunner.run( + QueryPlus.wrap(query).withIdentity(authenticationResult.getIdentity()), + responseContext + ); return new QueryResponse<>(res == null ? Sequences.empty() : res, responseContext); } + /** + * Finds a handler for the root datasource or one of its descendants. A handler selected for a descendant receives + * the complete query so it can resolve the matching datasource vertices before normal native execution. + */ + @Nullable + private DataSourceQueryHandler findDataSourceQueryHandler(final DataSource dataSource) + { + final DataSourceQueryHandler directHandler = dataSourceQueryHandlers.get(dataSource.getClass()); + if (directHandler != null) { + return directHandler; + } + + DataSourceQueryHandler descendantHandler = null; + for (final DataSource child : dataSource.getChildren()) { + final DataSourceQueryHandler childHandler = findDataSourceQueryHandler(child); + if (childHandler == null) { + continue; + } + if (descendantHandler != null && descendantHandler != childHandler) { + throw new ISE("Multiple datasource query handlers are required for datasource[%s]", dataSource); + } + descendantHandler = childHandler; + } + return descendantHandler; + } + /** * Emits logs and metrics for this query. *

diff --git a/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java b/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java index d22d24ed3efb..2614301431b0 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java @@ -23,6 +23,7 @@ import org.apache.druid.client.BrokerViewOfBrokerConfig; import org.apache.druid.guice.LazySingleton; import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import org.apache.druid.query.DataSource; import org.apache.druid.query.GenericQueryMetricsFactory; import org.apache.druid.query.QueryConfigProvider; import org.apache.druid.query.QueryRunnerFactoryConglomerate; @@ -34,6 +35,7 @@ import org.apache.druid.server.security.AuthorizerMapper; import javax.annotation.Nullable; +import java.util.Map; @LazySingleton public class QueryLifecycleFactory @@ -48,6 +50,7 @@ public class QueryLifecycleFactory private final AuthConfig authConfig; private final PolicyEnforcer policyEnforcer; private final BrokerViewOfBrokerConfig brokerViewOfBrokerConfig; + private final Map, DataSourceQueryHandler> dataSourceQueryHandlers; @Inject public QueryLifecycleFactory( @@ -60,6 +63,7 @@ public QueryLifecycleFactory( final PolicyEnforcer policyEnforcer, final AuthorizerMapper authorizerMapper, final QueryConfigProvider queryConfigProvider, + final Map, DataSourceQueryHandler> dataSourceQueryHandlers, @Nullable final BrokerViewOfBrokerConfig brokerViewOfBrokerConfig ) { @@ -70,11 +74,40 @@ public QueryLifecycleFactory( this.requestLogger = requestLogger; this.authorizerMapper = authorizerMapper; this.queryConfigProvider = queryConfigProvider; + this.dataSourceQueryHandlers = dataSourceQueryHandlers; this.authConfig = authConfig; this.policyEnforcer = policyEnforcer; this.brokerViewOfBrokerConfig = brokerViewOfBrokerConfig; } + public QueryLifecycleFactory( + final QueryRunnerFactoryConglomerate conglomerate, + final QuerySegmentWalker texasRanger, + final GenericQueryMetricsFactory queryMetricsFactory, + final ServiceEmitter emitter, + final RequestLogger requestLogger, + final AuthConfig authConfig, + final PolicyEnforcer policyEnforcer, + final AuthorizerMapper authorizerMapper, + final QueryConfigProvider queryConfigProvider, + @Nullable final BrokerViewOfBrokerConfig brokerViewOfBrokerConfig + ) + { + this( + conglomerate, + texasRanger, + queryMetricsFactory, + emitter, + requestLogger, + authConfig, + policyEnforcer, + authorizerMapper, + queryConfigProvider, + Map.of(), + brokerViewOfBrokerConfig + ); + } + public QueryLifecycle factorize() { // Read once per query so the whole lifecycle sees one config, even if it is swapped mid-query. @@ -93,6 +126,7 @@ public QueryLifecycle factorize() authConfig, policyEnforcer, configSnapshot, + dataSourceQueryHandlers, System.currentTimeMillis(), System.nanoTime() ); diff --git a/server/src/main/java/org/apache/druid/server/QueryResource.java b/server/src/main/java/org/apache/druid/server/QueryResource.java index 97679d1edcf2..3213e81cf693 100644 --- a/server/src/main/java/org/apache/druid/server/QueryResource.java +++ b/server/src/main/java/org/apache/druid/server/QueryResource.java @@ -86,6 +86,9 @@ public class QueryResource implements QueryCountStatsProvider public static final String HEADER_RESPONSE_CONTEXT = "X-Druid-Response-Context"; public static final String HEADER_IF_NONE_MATCH = "If-None-Match"; public static final String QUERY_ID_RESPONSE_HEADER = "X-Druid-Query-Id"; + /** Selects where a native query sent through a forwarding service should execute. */ + public static final String HEADER_NATIVE_QUERY_ROUTE = "X-Druid-Native-Query-Route"; + public static final String NATIVE_QUERY_ROUTE_LOCAL = "local"; public static final String ERROR_MESSAGE_TRAILER_HEADER = "X-Error-Message"; public static final String RESPONSE_COMPLETE_TRAILER_HEADER = "X-Druid-Response-Complete"; public static final String HEADER_ETAG = "ETag"; @@ -287,7 +290,7 @@ public interface QueryMetricCounter void incrementTimedOut(); } - private Query readQuery( + protected Query readQuery( final HttpServletRequest req, final InputStream in, final ResourceIOReaderWriterFactory.ResourceIOReaderWriter ioReaderWriter diff --git a/server/src/main/java/org/apache/druid/server/http/CoordinatorRedirectInfo.java b/server/src/main/java/org/apache/druid/server/http/CoordinatorRedirectInfo.java index 987ccb94842d..6e7d5c3c1a6e 100644 --- a/server/src/main/java/org/apache/druid/server/http/CoordinatorRedirectInfo.java +++ b/server/src/main/java/org/apache/druid/server/http/CoordinatorRedirectInfo.java @@ -22,6 +22,7 @@ import com.google.common.collect.ImmutableSet; import com.google.inject.Inject; import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.query.SystemTableDataSource; import org.apache.druid.server.coordinator.DruidCoordinator; import java.net.URL; @@ -33,7 +34,9 @@ public class CoordinatorRedirectInfo implements RedirectInfo { private static final Set LOCAL_PATHS = ImmutableSet.of( "/druid/coordinator/v1/leader", - "/druid/coordinator/v1/isLeader" + "/druid/coordinator/v1/isLeader", + "/druid/v2", + "/druid/v2/" ); private final DruidCoordinator coordinator; @@ -47,7 +50,10 @@ public CoordinatorRedirectInfo(DruidCoordinator coordinator) @Override public boolean doLocal(String requestURI) { - return (requestURI != null && LOCAL_PATHS.contains(requestURI)) || coordinator.isLeader(); + return (requestURI != null + && (LOCAL_PATHS.contains(requestURI) + || requestURI.startsWith("/druid/v2/" + SystemTableDataSource.NODE_QUERY_ID_PREFIX))) + || coordinator.isLeader(); } @Override diff --git a/server/src/main/java/org/apache/druid/server/system/SystemTableNotLeaderException.java b/server/src/main/java/org/apache/druid/server/system/SystemTableNotLeaderException.java new file mode 100644 index 000000000000..b66b48f92767 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/SystemTableNotLeaderException.java @@ -0,0 +1,38 @@ +/* + * 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; + +import org.apache.druid.java.util.common.ISE; + +/** + * Signals that a leader-only system table request reached a node that is no longer the leader. + * + *

The Broker may discover a leader and then contact it after leadership has changed. The native-query error + * response preserves this exception's class name, allowing {@code SystemTableQueryClient} to distinguish that race + * from an ordinary query failure, resolve the new leader, and retry once. A generic exception would not provide a + * reliable retry signal and retrying every failure could hide genuine query errors.

+ */ +public class SystemTableNotLeaderException extends ISE +{ + public SystemTableNotLeaderException(final String nodeRole) + { + super("Node role[%s] is not the current leader", nodeRole); + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/handler/SystemTableBrokerQueryHandler.java b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableBrokerQueryHandler.java new file mode 100644 index 000000000000..1a8db1011306 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableBrokerQueryHandler.java @@ -0,0 +1,63 @@ +/* + * 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.handler; + +import com.google.inject.Inject; +import org.apache.druid.query.BadQueryContextException; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.server.DataSourceQueryHandler; +import org.apache.druid.server.security.AuthenticationResult; + +/** Separates original Broker fanout from node-local requests received through the Broker's native endpoint. */ +public class SystemTableBrokerQueryHandler implements DataSourceQueryHandler +{ + private final SystemTableQueryClient systemTableQueryClient; + private final SystemTableQueryHandler localQueryHandler; + + @Inject + public SystemTableBrokerQueryHandler( + final SystemTableQueryClient systemTableQueryClient, + final SystemTableQueryHandler localQueryHandler + ) + { + this.systemTableQueryClient = systemTableQueryClient; + this.localQueryHandler = localQueryHandler; + } + + @Override + public QueryRunner createRunner( + final Query query, + final AuthenticationResult authenticationResult, + final boolean executeLocally + ) + { + if (executeLocally) { + if (!(query.getDataSource() instanceof SystemTableDataSource)) { + throw new BadQueryContextException( + "Local system-table execution requires a SystemTableDataSource at the query root" + ); + } + return localQueryHandler.createRunner(query, authenticationResult, true); + } + return systemTableQueryClient.createRunner(query, authenticationResult, false); + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/handler/SystemTableLeaderLocator.java b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableLeaderLocator.java new file mode 100644 index 000000000000..26f494c9c387 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableLeaderLocator.java @@ -0,0 +1,31 @@ +/* + * 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.handler; + +import com.google.common.util.concurrent.ListenableFuture; + +import java.net.URI; + +/** Resolves the current leader for one Druid node role. */ +@FunctionalInterface +public interface SystemTableLeaderLocator +{ + ListenableFuture findCurrentLeader(); +} diff --git a/server/src/main/java/org/apache/druid/server/system/handler/SystemTableNode.java b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableNode.java new file mode 100644 index 000000000000..23f263eb8d83 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableNode.java @@ -0,0 +1,53 @@ +/* + * 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.handler; + +import org.apache.druid.discovery.DiscoveryDruidNode; +import org.apache.druid.discovery.NodeRole; + +import java.util.LinkedHashSet; +import java.util.Set; + +/** A physical Druid process selected to contribute rows to a native system-table query. */ +class SystemTableNode +{ + private final DiscoveryDruidNode discoveryNode; + private final Set nodeRoles = new LinkedHashSet<>(); + + SystemTableNode(final DiscoveryDruidNode discoveryNode) + { + this.discoveryNode = discoveryNode; + } + + DiscoveryDruidNode getDiscoveryNode() + { + return discoveryNode; + } + + Set getNodeRoles() + { + return nodeRoles; + } + + void addNodeRole(final NodeRole nodeRole) + { + nodeRoles.add(nodeRole); + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/handler/SystemTableNodeLocator.java b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableNodeLocator.java new file mode 100644 index 000000000000..1706860c1407 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableNodeLocator.java @@ -0,0 +1,150 @@ +/* + * 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.handler; + +import com.google.common.collect.Iterables; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.inject.Inject; +import org.apache.druid.client.DirectDruidClient; +import org.apache.druid.client.coordinator.CoordinatorClient; +import org.apache.druid.discovery.DiscoveryDruidNode; +import org.apache.druid.discovery.DruidNodeDiscoveryProvider; +import org.apache.druid.discovery.NodeRole; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryInterruptedException; +import org.apache.druid.query.QueryTimeoutException; +import org.apache.druid.rpc.indexing.OverlordClient; +import org.apache.druid.server.system.table.SystemTableDescriptor; +import org.apache.druid.server.system.table.SystemTableRoutingMode; + +import java.net.URI; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +/** Resolves a descriptor's logical routing policy to concrete Druid processes. */ +public class SystemTableNodeLocator +{ + private final DruidNodeDiscoveryProvider discoveryProvider; + private final Map leaderLocators; + + @Inject + public SystemTableNodeLocator( + final DruidNodeDiscoveryProvider discoveryProvider, + final CoordinatorClient coordinatorClient, + final OverlordClient overlordClient + ) + { + this.discoveryProvider = discoveryProvider; + this.leaderLocators = Map.of( + NodeRole.COORDINATOR, + coordinatorClient::findCurrentLeader, + NodeRole.OVERLORD, + overlordClient::findCurrentLeader + ); + } + + List locate(final SystemTableDescriptor descriptor, final Query query) + { + if (descriptor.getRoutingMode() == SystemTableRoutingMode.ALL_NODES) { + return discoverAllNodes(descriptor); + } + + final NodeRole leaderRole = Iterables.getOnlyElement(descriptor.getNodeRoles()); + final URI leaderUri = findLeader(leaderRole, query); + // Resolve leadership before taking the discovery snapshot. A leader election may complete while discovery still + // contains the previous membership, so taking the snapshot first can reject a valid newly elected leader. + return discoverAllNodes(descriptor).stream() + .filter(node -> sameServer(leaderUri, node.getDiscoveryNode().getDruidNode().getUriToUse())) + .findFirst() + .map(List::of) + .orElseThrow( + () -> new ISE( + "Current leader[%s] for role[%s] is not present in service discovery", + leaderUri, + leaderRole + ) + ); + } + + private List discoverAllNodes(final SystemTableDescriptor descriptor) + { + final Map nodes = new LinkedHashMap<>(); + for (final NodeRole nodeRole : descriptor.getNodeRoles()) { + for (final DiscoveryDruidNode node : discoveryProvider.getForNodeRole(nodeRole).getAllNodes()) { + nodes.computeIfAbsent( + node.getDruidNode().getHostAndPortToUse(), + ignored -> new SystemTableNode(node) + ).addNodeRole(node.getNodeRole()); + } + } + return new ArrayList<>(nodes.values()); + } + + private URI findLeader(final NodeRole nodeRole, final Query query) + { + final SystemTableLeaderLocator leaderLocator = leaderLocators.get(nodeRole); + if (leaderLocator == null) { + throw new ISE("Leader-only system-table routing is not supported for role[%s]", nodeRole); + } + + final ListenableFuture leaderFuture = leaderLocator.findCurrentLeader(); + final long timeLeft = query.context().getLong(DirectDruidClient.QUERY_FAIL_TIME) - System.currentTimeMillis(); + if (timeLeft <= 0) { + leaderFuture.cancel(true); + throw new QueryTimeoutException("Timed out while locating the current leader for role[" + nodeRole + "]"); + } + try { + return leaderFuture.get(timeLeft, TimeUnit.MILLISECONDS); + } + catch (InterruptedException e) { + leaderFuture.cancel(true); + Thread.currentThread().interrupt(); + throw QueryInterruptedException.wrapIfNeeded(e); + } + catch (TimeoutException e) { + leaderFuture.cancel(true); + throw new QueryTimeoutException("Timed out while locating the current leader for role[" + nodeRole + "]"); + } + catch (ExecutionException e) { + throw QueryInterruptedException.wrapIfNeeded(e.getCause() == null ? e : e.getCause()); + } + } + + static boolean sameServer(final URI first, final URI second) + { + return first.getScheme().equalsIgnoreCase(second.getScheme()) + && first.getHost().equalsIgnoreCase(second.getHost()) + && effectivePort(first) == effectivePort(second); + } + + private static int effectivePort(final URI uri) + { + if (uri.getPort() >= 0) { + return uri.getPort(); + } + return "https".equalsIgnoreCase(uri.getScheme()) ? 443 : 80; + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryClient.java b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryClient.java new file mode 100644 index 000000000000..11d4121448dd --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryClient.java @@ -0,0 +1,901 @@ +/* + * 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.handler; + +import com.google.common.base.Preconditions; +import com.google.common.collect.Iterables; +import com.google.inject.Inject; +import io.netty.channel.ChannelException; +import org.apache.druid.client.DirectDruidClient; +import org.apache.druid.client.DirectDruidClientFactory; +import org.apache.druid.client.DruidServer; +import org.apache.druid.discovery.DiscoveryDruidNode; +import org.apache.druid.guice.annotations.Self; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.JodaUtils; +import org.apache.druid.java.util.common.guava.BaseSequence; +import org.apache.druid.java.util.common.guava.LazySequence; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.java.util.common.guava.Sequences; +import org.apache.druid.java.util.common.guava.Yielder; +import org.apache.druid.java.util.common.guava.Yielders; +import org.apache.druid.query.BaseQuery; +import org.apache.druid.query.DataSource; +import org.apache.druid.query.Druids; +import org.apache.druid.query.InlineDataSource; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryContexts; +import org.apache.druid.query.QueryDataSource; +import org.apache.druid.query.QueryException; +import org.apache.druid.query.QueryInterruptedException; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.query.QuerySegmentWalker; +import org.apache.druid.query.QueryTimeoutException; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.query.filter.AndDimFilter; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.operator.OperatorFactory; +import org.apache.druid.query.operator.ScanOperatorFactory; +import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.segment.VirtualColumn; +import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.server.DataSourceQueryHandler; +import org.apache.druid.server.DruidNode; +import org.apache.druid.server.QueryScheduler; +import org.apache.druid.server.coordination.ServerType; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.server.security.Escalator; +import org.apache.druid.server.system.SystemTableNotLeaderException; +import org.apache.druid.server.system.table.SystemTableDescriptor; +import org.apache.druid.server.system.table.SystemTableRoutingMode; +import org.apache.druid.utils.CloseableUtils; + +import javax.annotation.Nullable; +import java.io.Closeable; +import java.io.EOFException; +import java.io.IOException; +import java.net.SocketException; +import java.net.UnknownHostException; +import java.nio.channels.ClosedChannelException; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Optional; +import java.util.UUID; +import java.util.function.Supplier; + +/** + * Recursively replaces native system-table datasources with authorized inline rows fetched from nodes, then + * executes the rebuilt query through the Broker's normal native-query walker. + */ +public class SystemTableQueryClient implements DataSourceQueryHandler +{ + private static final String SYSTEM_TABLE_TIER = "_system"; + + private final SystemTableNodeLocator nodeLocator; + private final DirectDruidClientFactory directDruidClientFactory; + private final QueryScheduler queryScheduler; + private final QuerySegmentWalker querySegmentWalker; + private final Map tableDescriptors; + private final AuthorizerMapper authorizerMapper; + private final SystemTableQueryHandler localQueryHandler; + private final AuthenticationResult escalatedAuthenticationResult; + private final DruidNode selfNode; + + @Inject + public SystemTableQueryClient( + final SystemTableNodeLocator nodeLocator, + final DirectDruidClientFactory directDruidClientFactory, + final QueryScheduler queryScheduler, + final QuerySegmentWalker querySegmentWalker, + final Map tableDescriptors, + final AuthorizerMapper authorizerMapper, + final SystemTableQueryHandler localQueryHandler, + final Escalator escalator, + @Self final DruidNode selfNode + ) + { + this.nodeLocator = nodeLocator; + this.directDruidClientFactory = directDruidClientFactory; + this.queryScheduler = queryScheduler; + this.querySegmentWalker = querySegmentWalker; + this.tableDescriptors = tableDescriptors; + this.authorizerMapper = authorizerMapper; + this.localQueryHandler = localQueryHandler; + this.escalatedAuthenticationResult = escalator.createEscalatedAuthenticationResult(); + this.selfNode = selfNode; + } + + @Override + public QueryRunner createRunner( + final Query query, + final AuthenticationResult authenticationResult, + final boolean executeLocally + ) + { + return (queryPlus, responseContext) -> { + final NodeSequencesCloser nodeSequencesCloser = new NodeSequencesCloser(queryScheduler); + try { + final DataSource resolvedDataSource = resolveDataSource( + query.getDataSource(), + query, + authenticationResult, + responseContext, + nodeSequencesCloser + ); + final Query resolvedQuery = query.withDataSource(resolvedDataSource).withOverriddenContext( + Map.of(QueryContexts.QUERY_RESOURCE_ID, UUID.randomUUID().toString()) + ); + final QueryRunner queryRunner = querySegmentWalker.getQueryRunnerForIntervals( + resolvedQuery, + resolvedQuery.getIntervals() + ); + return Sequences.withBaggage( + queryRunner.run(queryPlus.withQuery(resolvedQuery), responseContext), + nodeSequencesCloser + ); + } + catch (Throwable t) { + throw CloseableUtils.closeAndWrapInCatch(t, nodeSequencesCloser); + } + }; + } + + private DataSource resolveDataSource( + final DataSource dataSource, + final Query owningQuery, + final AuthenticationResult authenticationResult, + final ResponseContext responseContext, + final NodeSequencesCloser nodeSequencesCloser + ) + { + if (dataSource instanceof SystemTableDataSource systemTableDataSource) { + return resolveSystemTableDataSource( + systemTableDataSource, + owningQuery, + authenticationResult, + responseContext, + nodeSequencesCloser + ); + } + + final Query childOwningQuery = dataSource instanceof QueryDataSource queryDataSource + ? queryDataSource.getQuery() + : owningQuery; + final List resolvedChildren = new ArrayList<>(); + for (final DataSource child : dataSource.getChildren()) { + resolvedChildren.add( + resolveDataSource( + child, + childOwningQuery, + authenticationResult, + responseContext, + nodeSequencesCloser + ) + ); + } + return dataSource.withChildren(resolvedChildren); + } + + private InlineDataSource resolveSystemTableDataSource( + final SystemTableDataSource dataSource, + final Query owningQuery, + final AuthenticationResult authenticationResult, + final ResponseContext responseContext, + final NodeSequencesCloser nodeSequencesCloser + ) + { + final SystemTableDescriptor descriptor = tableDescriptors.get(dataSource.getTable()); + if (descriptor == null) { + throw new ISE("No routing descriptor is registered for system table[%s]", dataSource.getTable()); + } + + final ScanQuery nodeQuery = makeNodeQuery(dataSource, descriptor, owningQuery); + final List> nodeRunners = makeNodeRunners( + nodeQuery, + descriptor, + nodeSequencesCloser + ); + if (nodeRunners.isEmpty() && !descriptor.isEmptyDiscoveryAllowed()) { + throw new ISE("No node is available to serve system table[%s]", dataSource.getTable()); + } + + /* + * Keep the scan result transport lazy. The node response is a Sequence, so this lets the local query + * runner consume and process rows while the HTTP response is still arriving instead of materializing every + * node row in two Broker-side lists. The scan transport is intentionally kept behind this adapter so a + * future node-side aggregation mode can use a different result adapter without changing routing or + * authorization orchestration here. + */ + final NodeSequenceCloser nodeSequenceCloser = new NodeSequenceCloser(); + nodeSequencesCloser.add(nodeSequenceCloser); + final Iterable nodeRows = scanNodeRows( + nodeRunners, + nodeQuery, + responseContext, + nodeSequenceCloser + ); + final Iterable authorizedRows = descriptor.getRowAuthorizer().filterAuthorizedRows( + nodeRows, + authenticationResult, + authorizerMapper + ); + final Iterable queryRows = requiresMaterializedRows(owningQuery) + ? materializeRows(authorizedRows) + : authorizedRows; + return InlineDataSource.fromIterable(queryRows, descriptor.getRowSignature()); + } + + private ScanQuery makeNodeQuery( + final SystemTableDataSource dataSource, + final SystemTableDescriptor descriptor, + final Query owningQuery + ) + { + final Map nodeContext = new LinkedHashMap<>(owningQuery.getContext()); + final DimFilter nodeFilter = nodeFilter(dataSource, owningQuery); + nodeContext.put(DirectDruidClient.QUERY_FAIL_TIME, getNodeFailTime(owningQuery)); + // The node endpoint always returns plain ScanResultValue objects. bySegment belongs to the user-facing query and + // would make DirectDruidClient deserialize this private transport using an incompatible result type. + nodeContext.put(QueryContexts.BY_SEGMENT_KEY, false); + return Druids.newScanQueryBuilder() + .dataSource(dataSource) + .eternityInterval() + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .limit(Long.MAX_VALUE) + .filters(nodeFilter) + .virtualColumns(nodeVirtualColumns(dataSource, owningQuery, nodeFilter)) + .columns(descriptor.getRowSignature()) + .context(nodeContext) + .build(); + } + + private List> makeNodeRunners( + final ScanQuery nodeQuery, + final SystemTableDescriptor descriptor, + final NodeSequencesCloser nodeSequencesCloser + ) + { + final List> nodeRunners = new ArrayList<>(); + for (final SystemTableNode node : nodeLocator.locate(descriptor, nodeQuery)) { + nodeRunners.add( + (queryPlus, responseContext) -> recoverNodeFailure( + () -> runNodeQuery(nodeQuery, node, nodeSequencesCloser, queryPlus, responseContext), + node, + descriptor, + descriptor.getRoutingMode() == SystemTableRoutingMode.LEADER_ONLY + ? () -> { + final SystemTableNode currentLeader = Iterables.getOnlyElement( + nodeLocator.locate(descriptor, nodeQuery) + ); + return new NodeQuerySequence( + currentLeader, + runNodeQuery(nodeQuery, currentLeader, nodeSequencesCloser, queryPlus, responseContext) + ); + } + : null + ) + ); + } + return nodeRunners; + } + + private Sequence runNodeQuery( + final ScanQuery nodeQuery, + final SystemTableNode node, + final NodeSequencesCloser nodeSequencesCloser, + final QueryPlus queryPlus, + final ResponseContext responseContext + ) + { + final String nodeResourceId = UUID.randomUUID().toString(); + final String nodeQueryId = SystemTableDataSource.NODE_QUERY_ID_PREFIX + UUID.randomUUID(); + final ScanQuery subNativeQuery = nodeQuery.withOverriddenContext( + Map.of( + BaseQuery.QUERY_ID, + nodeQueryId, + QueryContexts.QUERY_RESOURCE_ID, + nodeResourceId + ) + ); + final QueryRunner nodeRunner; + if (SystemTableNodeLocator.sameServer( + selfNode.getUriToUse(), + node.getDiscoveryNode().getDruidNode().getUriToUse() + )) { + // Invoke the raw local handler rather than the Broker handler. This avoids an HTTP request back to this + // Broker and, more importantly, prevents the local request from entering node fanout recursively. + nodeRunner = localQueryHandler.createRunner(subNativeQuery, escalatedAuthenticationResult, true); + } else { + nodeRunner = directDruidClientFactory.makeDirectClient(toDruidServer(node.getDiscoveryNode())); + nodeSequencesCloser.addQueryId(nodeQueryId); + } + return nodeRunner.run(queryPlus.withQuery(subNativeQuery), responseContext); + } + + /** + * Window operators consume a {@code RowsAndColumns} view of the segment. The current row-based segment can only + * provide that view when its rows are list-backed, so preserve the existing materialized path for those queries. + */ + private static boolean requiresMaterializedRows(final Query query) + { + return query instanceof WindowOperatorQuery; + } + + /** + * Eagerly consumes the lazy, single-pass node response into an {@link ArrayList} without changing its rows. + * + *

This representation change is required by {@link org.apache.druid.segment.InlineSegmentWrangler}: an inline + * datasource backed specifically by an {@code ArrayList} becomes an {@link org.apache.druid.segment.ArrayListSegment}, + * which can provide the {@code RowsAndColumns} view required by window operators. Other query types retain the lazy + * iterable so the Broker can process rows while node responses are still arriving.

+ */ + private static Iterable materializeRows(final Iterable rows) + { + final List materializedRows = new ArrayList<>(); + rows.forEach(materializedRows::add); + return materializedRows; + } + + /** + * Creates a lazy, single-pass view of the rows returned by the node scan runners. + * + *

The view is single-pass because a node response is backed by a live HTTP stream. The normal native query + * path consumes an {@link InlineDataSource} once per query. Re-iterating would otherwise issue duplicate node + * requests or retain all rows to support replay.

+ */ + static Iterable scanNodeRows( + final List> nodeRunners, + final Query nodeQuery, + final ResponseContext responseContext + ) + { + return scanNodeRows( + nodeRunners, + nodeQuery, + responseContext, + new NodeSequenceCloser() + ); + } + + private static Iterable scanNodeRows( + final List> nodeRunners, + final Query nodeQuery, + final ResponseContext responseContext, + final NodeSequenceCloser nodeSequenceCloser + ) + { + final Sequence rows = new LazySequence<>(() -> { + final List> nodeSequences = new ArrayList<>(); + for (final QueryRunner nodeRunner : nodeRunners) { + // DirectDruidClient.run starts its asynchronous HTTP request. Start every request before waiting for any + // response so slow nodes do not serialize the entire system-table query. + nodeSequences.add(nodeRunner.run(QueryPlus.wrap(nodeQuery), responseContext)); + } + // Node scans have no cross-node ordering. Consume them in discovery order so a ready node can yield rows without + // MergeSequence first initializing every other node sequence. The HTTP requests were all started above. + return Sequences.concat(nodeSequences).flatMap(SystemTableQueryClient::rowsFromScanResult); + }); + return sequenceAsIterable(rows, nodeSequenceCloser); + } + + private static Sequence recoverNodeFailure( + final Supplier> sequenceSupplier, + final SystemTableNode node, + final SystemTableDescriptor descriptor, + @Nullable final Supplier leaderRetrySupplier + ) + { + final Sequence sequence; + try { + sequence = sequenceSupplier.get(); + } + catch (Exception failure) { + if (leaderRetrySupplier != null && isLeaderChangeFailure(failure)) { + try { + final NodeQuerySequence retry = leaderRetrySupplier.get(); + return recoveringSequence(retry, descriptor, null); + } + catch (Exception retryFailure) { + return failureSequence(retryFailure, node, descriptor); + } + } + return failureSequence(failure, node, descriptor); + } + return recoveringSequence(new NodeQuerySequence(node, sequence), descriptor, leaderRetrySupplier); + } + + private static Sequence recoveringSequence( + final NodeQuerySequence nodeQuerySequence, + final SystemTableDescriptor descriptor, + @Nullable final Supplier leaderRetrySupplier + ) + { + return new BaseSequence<>( + new BaseSequence.IteratorMaker() + { + @Override + public RecoveringNodeIterator make() + { + return new RecoveringNodeIterator(nodeQuerySequence, descriptor, leaderRetrySupplier); + } + + @Override + public void cleanup(final RecoveringNodeIterator iterator) + { + CloseableUtils.closeAndWrapExceptions(iterator); + } + } + ); + } + + private static Sequence failureSequence( + final Exception failure, + final SystemTableNode node, + final SystemTableDescriptor descriptor + ) + { + if (!isNodeAvailabilityFailure(failure)) { + throw propagate(failure); + } + return descriptor.getNodeFailureRow( + node.getDiscoveryNode().getDruidNode(), + node.getNodeRoles(), + failure + ).map( + row -> Sequences.simple( + List.of( + new ScanResultValue( + null, + descriptor.getRowSignature().getColumnNames(), + List.of((Object) row) + ) + ) + ) + ).orElseThrow(() -> QueryInterruptedException.wrapIfNeeded(failure)); + } + + private static boolean isNodeAvailabilityFailure(final Throwable failure) + { + Throwable cause = failure; + while (cause != null) { + if (cause instanceof SocketException + || cause instanceof UnknownHostException + || cause instanceof EOFException + || cause instanceof ClosedChannelException + || cause instanceof ChannelException) { + return true; + } + cause = cause.getCause(); + } + return false; + } + + private static boolean isLeaderChangeFailure(final Throwable failure) + { + Throwable cause = failure; + while (cause != null) { + if (cause instanceof SystemTableNotLeaderException + || cause instanceof QueryException + && SystemTableNotLeaderException.class.getName().equals(((QueryException) cause).getErrorClass())) { + return true; + } + cause = cause.getCause(); + } + return false; + } + + private static RuntimeException propagate(final Exception failure) + { + if (failure instanceof RuntimeException runtimeException) { + return runtimeException; + } else if (failure instanceof java.util.concurrent.TimeoutException) { + return new QueryTimeoutException(failure.getMessage()); + } else { + if (failure instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } + return QueryInterruptedException.wrapIfNeeded(failure); + } + } + + private static Sequence rowsFromScanResult(final ScanResultValue scanResult) + { + return Sequences.simple((List) scanResult.getEvents()) + .map(event -> event instanceof Object[] ? (Object[]) event : ((List) event).toArray()); + } + + private static long getNodeFailTime(final Query query) + { + final long existingFailTime = query.context().getLong(DirectDruidClient.QUERY_FAIL_TIME, -1L); + if (existingFailTime > 0) { + return existingFailTime; + } else if (query.context().hasTimeout()) { + return System.currentTimeMillis() + query.context().getTimeout(); + } else { + return JodaUtils.MAX_INSTANT; + } + } + + /** + * Adapts a Druid {@link Sequence} to the lazy {@link Iterable} accepted by {@link InlineDataSource}. + * + *

The iterator owns the sequence yielder and closes it when the sequence is exhausted. This is important for + * direct node clients because closing the yielder closes the underlying response stream and allows an early + * terminating Broker query to cancel the node request.

+ */ + private static Iterable sequenceAsIterable( + final Sequence sequence, + final NodeSequenceCloser nodeSequenceCloser + ) + { + return new Iterable<>() + { + private boolean iterated; + + @Override + public Iterator iterator() + { + if (iterated) { + throw new ISE("A native system table result sequence can only be consumed once"); + } + iterated = true; + + return new Iterator<>() + { + @Nullable + private Yielder yielder = Yielders.each(sequence); + private boolean closed; + + { + nodeSequenceCloser.setYielder(yielder); + } + + @Override + public boolean hasNext() + { + if (yielder == null) { + return false; + } + if (yielder.isDone()) { + close(); + return false; + } + return true; + } + + @Override + public T next() + { + if (!hasNext()) { + throw new NoSuchElementException(); + } + final Yielder currentYielder = Preconditions.checkNotNull(yielder); + final T value = currentYielder.get(); + yielder = currentYielder.next(null); + nodeSequenceCloser.setYielder(yielder); + return value; + } + + private void close() + { + if (!closed) { + closed = true; + yielder = null; + nodeSequenceCloser.close(); + } + } + }; + } + }; + } + + private static class NodeSequenceCloser implements Closeable + { + @Nullable + private Yielder yielder; + private boolean closed; + + synchronized void setYielder(final Yielder newYielder) + { + if (closed) { + CloseableUtils.closeAndWrapExceptions(newYielder); + } else { + yielder = newYielder; + } + } + + @Override + public synchronized void close() + { + if (!closed) { + closed = true; + final Yielder currentYielder = yielder; + yielder = null; + CloseableUtils.closeAndWrapExceptions(currentYielder); + } + } + } + + private static class NodeSequencesCloser implements Closeable + { + private final List nodeSequenceClosers = new ArrayList<>(); + private final List nodeQueryIds = new ArrayList<>(); + private final QueryScheduler queryScheduler; + private boolean closed; + + private NodeSequencesCloser(final QueryScheduler queryScheduler) + { + this.queryScheduler = queryScheduler; + } + + synchronized void addQueryId(final String nodeQueryId) + { + if (closed) { + queryScheduler.cancelQuery(nodeQueryId); + } else { + nodeQueryIds.add(nodeQueryId); + } + } + + synchronized void add(final NodeSequenceCloser nodeSequenceCloser) + { + if (closed) { + CloseableUtils.closeAndWrapExceptions(nodeSequenceCloser); + } else { + nodeSequenceClosers.add(nodeSequenceCloser); + } + } + + @Override + public synchronized void close() throws IOException + { + if (!closed) { + closed = true; + try { + CloseableUtils.closeAll(nodeSequenceClosers); + } + finally { + nodeSequenceClosers.clear(); + nodeQueryIds.forEach(queryScheduler::cancelQuery); + nodeQueryIds.clear(); + } + } + } + } + + private static class RecoveringNodeIterator implements Iterator, Closeable + { + private final SystemTableDescriptor descriptor; + @Nullable + private final Supplier leaderRetrySupplier; + + private NodeQuerySequence nodeQuerySequence; + + @Nullable + private Yielder yielder; + @Nullable + private ScanResultValue failureResult; + private boolean started; + private boolean finished; + private boolean retriedLeader; + private boolean emittedResult; + + private RecoveringNodeIterator( + final NodeQuerySequence nodeQuerySequence, + final SystemTableDescriptor descriptor, + @Nullable final Supplier leaderRetrySupplier + ) + { + this.nodeQuerySequence = nodeQuerySequence; + this.descriptor = descriptor; + this.leaderRetrySupplier = leaderRetrySupplier; + } + + @Override + public boolean hasNext() + { + if (failureResult != null) { + return true; + } + if (finished) { + return false; + } + + try { + if (!started) { + started = true; + yielder = Yielders.each(nodeQuerySequence.sequence()); + } + if (Preconditions.checkNotNull(yielder).isDone()) { + finished = true; + close(); + return false; + } + return true; + } + catch (Exception e) { + return recoverOrThrow(e) ? hasNext() : true; + } + } + + @Override + public ScanResultValue next() + { + if (!hasNext()) { + throw new NoSuchElementException(); + } + if (failureResult != null) { + final ScanResultValue result = failureResult; + failureResult = null; + finished = true; + return result; + } + + final Yielder currentYielder = Preconditions.checkNotNull(yielder); + final ScanResultValue result; + try { + result = currentYielder.get(); + emittedResult = true; + } + catch (Exception e) { + // Either the leader retry restarted the sequence or a failure row is pending; next() serves both. + recoverOrThrow(e); + return next(); + } + try { + yielder = currentYielder.next(null); + } + catch (Exception e) { + recoverOrThrow(e); + } + return result; + } + + private boolean recoverOrThrow(final Exception originalFailure) + { + Exception failure = originalFailure; + if (!emittedResult && !retriedLeader && leaderRetrySupplier != null && isLeaderChangeFailure(failure)) { + retriedLeader = true; + try { + close(); + nodeQuerySequence = leaderRetrySupplier.get(); + started = false; + finished = false; + return true; + } + catch (Exception retryFailure) { + failure = retryFailure; + } + } + if (!isNodeAvailabilityFailure(failure)) { + throw CloseableUtils.closeAndWrapInCatch(failure, this); + } + final Optional failureRow = descriptor.getNodeFailureRow( + nodeQuerySequence.node().getDiscoveryNode().getDruidNode(), + nodeQuerySequence.node().getNodeRoles(), + failure + ); + if (failureRow.isEmpty()) { + throw CloseableUtils.closeAndWrapInCatch(failure, this); + } + try { + close(); + } + catch (RuntimeException closeException) { + failure.addSuppressed(closeException); + } + failureResult = new ScanResultValue( + null, + descriptor.getRowSignature().getColumnNames(), + List.of((Object) failureRow.get()) + ); + return false; + } + + @Override + public void close() + { + final Yielder currentYielder = yielder; + yielder = null; + CloseableUtils.closeAndWrapExceptions(currentYielder); + } + } + + private record NodeQuerySequence(SystemTableNode node, Sequence sequence) + { + } + + @Nullable + private static DimFilter nodeFilter(final SystemTableDataSource dataSource, final Query query) + { + // Push into the remote scan only when this query directly owns the system table. A join or other composite filter + // remains on the resolved Broker query, where the normal native query machinery can apply it with full datasource + // semantics; it cannot safely be attributed to an individual remote leaf here. + if (!dataSource.equals(query.getDataSource())) { + return null; + } + + final List filters = new ArrayList<>(); + if (query.getFilter() != null) { + filters.add(query.getFilter()); + } + if (query instanceof WindowOperatorQuery) { + for (final OperatorFactory operator : ((WindowOperatorQuery) query).getLeafOperators()) { + if (operator instanceof ScanOperatorFactory && ((ScanOperatorFactory) operator).getFilter() != null) { + filters.add(((ScanOperatorFactory) operator).getFilter()); + } + } + } + if (filters.isEmpty()) { + return null; + } else if (filters.size() == 1) { + return filters.get(0); + } else { + return new AndDimFilter(filters); + } + } + + private static VirtualColumns nodeVirtualColumns( + final SystemTableDataSource dataSource, + final Query query, + @Nullable final DimFilter nodeFilter + ) + { + if (!dataSource.equals(query.getDataSource()) || nodeFilter == null) { + return VirtualColumns.EMPTY; + } + + final List virtualColumns = new ArrayList<>(List.of(query.getVirtualColumns().getVirtualColumns())); + if (query instanceof WindowOperatorQuery) { + for (final OperatorFactory operator : ((WindowOperatorQuery) query).getLeafOperators()) { + // WindowOperatorQuery stores null rather than empty virtual columns for a leaf scan without any. + if (operator instanceof ScanOperatorFactory + && ((ScanOperatorFactory) operator).getVirtualColumns() != null) { + virtualColumns.addAll(List.of(((ScanOperatorFactory) operator).getVirtualColumns().getVirtualColumns())); + } + } + } + return VirtualColumns.create(virtualColumns); + } + + private static DruidServer toDruidServer(final DiscoveryDruidNode discoveryNode) + { + final DruidNode node = discoveryNode.getDruidNode(); + return new DruidServer( + node.getHostAndPortToUse(), + node.getHostAndPort(), + node.getHostAndTlsPort(), + 0, + null, + ServerType.HISTORICAL, + SYSTEM_TABLE_TIER, + 0 + ); + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryHandler.java b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryHandler.java new file mode 100644 index 000000000000..0214bd10d7cf --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryHandler.java @@ -0,0 +1,167 @@ +/* + * 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.handler; + +import com.google.inject.Inject; +import org.apache.druid.client.DirectDruidClient; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.JodaUtils; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.query.BadQueryContextException; +import org.apache.druid.query.Druids; +import org.apache.druid.query.InlineDataSource; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.scan.ScanQueryEngine; +import org.apache.druid.segment.InlineSegmentWrangler; +import org.apache.druid.segment.Segment; +import org.apache.druid.server.DataSourceQueryHandler; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.server.system.table.SystemTableDataProvider; +import org.apache.druid.server.system.table.SystemTableDescriptor; +import org.apache.druid.server.system.table.SystemTablePushdownFilter; + +import java.util.Map; + +/** Resolves one node-local system table and returns its rows through the standard native Scan stack. */ +public class SystemTableQueryHandler implements DataSourceQueryHandler +{ + private final Map dataSuppliers; + private final Map tableDescriptors; + private final ScanQueryEngine scanQueryEngine; + private final AuthorizerMapper authorizerMapper; + + @Inject + public SystemTableQueryHandler( + final Map dataSuppliers, + final Map tableDescriptors, + final ScanQueryEngine scanQueryEngine, + final AuthorizerMapper authorizerMapper + ) + { + this.dataSuppliers = dataSuppliers; + this.tableDescriptors = tableDescriptors; + this.scanQueryEngine = scanQueryEngine; + this.authorizerMapper = authorizerMapper; + } + + @Override + public QueryRunner createRunner( + final Query query, + final AuthenticationResult requestAuthenticationResult, + final boolean executeLocally + ) + { + if (!(query instanceof ScanQuery)) { + throw new BadQueryContextException( + "Local system-table execution requires a Scan query with a SystemTableDataSource at the query root" + ); + } + if (!(query.getDataSource() instanceof SystemTableDataSource dataSource)) { + throw new BadQueryContextException( + "Local system-table execution requires a Scan query with a SystemTableDataSource at the query root" + ); + } + final SystemTableDataProvider dataSupplier = dataSuppliers.get(dataSource.getTable()); + if (dataSupplier == null) { + throw new ISE("System table[%s] is not served by this node", dataSource.getTable()); + } + final SystemTableDescriptor descriptor = tableDescriptors.get(dataSource.getTable()); + if (descriptor == null) { + throw new ISE("No descriptor is registered for system table[%s]", dataSource.getTable()); + } + + return (queryPlus, responseContext) -> { + final Iterable suppliedRows = () -> dataSupplier.getRows( + SystemTablePushdownFilter.extract(query, dataSupplier.getPushdownFilters()), + requestAuthenticationResult + ).iterator(); + final Iterable authorizedRows = descriptor.getRowAuthorizer().filterAuthorizedRows( + suppliedRows, + requestAuthenticationResult, + authorizerMapper + ); + final ScanQuery resolvedQuery = Druids.ScanQueryBuilder.copy((ScanQuery) query) + .dataSource( + InlineDataSource.fromIterable( + authorizedRows, + descriptor.getRowSignature() + ) + ) + .build(); + final Segment inlineSegment = new InlineSegmentWrangler() + .getSegmentsForIntervals(resolvedQuery.getDataSource(), resolvedQuery.getIntervals()) + .iterator() + .next(); + + return runScan( + scanQueryEngine, + resolvedQuery, + inlineSegment, + queryPlus, + responseContext + ); + }; + } + + @SuppressWarnings("unchecked") + private static Sequence runScan( + final ScanQueryEngine scanQueryEngine, + final ScanQuery query, + final Segment segment, + final QueryPlus queryPlus, + final ResponseContext responseContext + ) + { + ScanQuery.verifyOrderByForNativeExecution(query); + initializeTimeout(query, responseContext); + return (Sequence) (Sequence) scanQueryEngine.process( + query, + segment, + responseContext, + queryPlus.getQueryMetrics() + ); + } + + private static void initializeTimeout(final ScanQuery query, final ResponseContext responseContext) + { + final Long existingTimeoutAt = responseContext.getTimeoutTime(); + if (existingTimeoutAt != null && existingTimeoutAt != 0L) { + return; + } + + final long failTime = query.context().getLong(DirectDruidClient.QUERY_FAIL_TIME, 0L); + final long timeoutAt; + if (failTime > 0L) { + timeoutAt = failTime; + } else if (query.context().hasTimeout()) { + timeoutAt = System.currentTimeMillis() + query.context().getTimeout(); + } else { + timeoutAt = JodaUtils.MAX_INSTANT; + } + responseContext.putTimeoutTime(timeoutAt); + } + +} diff --git a/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryResource.java b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryResource.java new file mode 100644 index 000000000000..4e4f7bb04558 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryResource.java @@ -0,0 +1,86 @@ +/* + * 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.handler; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.inject.Inject; +import org.apache.druid.guice.LazySingleton; +import org.apache.druid.guice.annotations.Json; +import org.apache.druid.query.BadQueryContextException; +import org.apache.druid.query.Query; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.server.QueryLifecycleFactory; +import org.apache.druid.server.QueryResource; +import org.apache.druid.server.QueryResourceQueryResultPusherFactory; +import org.apache.druid.server.QueryScheduler; +import org.apache.druid.server.ResourceIOReaderWriterFactory; +import org.apache.druid.server.initialization.ServerConfig; +import org.apache.druid.server.security.AuthorizerMapper; + +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.Path; +import java.io.IOException; +import java.io.InputStream; + +/** Native query endpoint for nodes that serve only local system-table scans. */ +@LazySingleton +@Path("/druid/v2/") +public class SystemTableQueryResource extends QueryResource +{ + @Inject + public SystemTableQueryResource( + final QueryLifecycleFactory queryLifecycleFactory, + final @Json ObjectMapper jsonMapper, + final QueryScheduler queryScheduler, + final AuthorizerMapper authorizerMapper, + final QueryResourceQueryResultPusherFactory queryResultPusherFactory, + final ResourceIOReaderWriterFactory resourceIOReaderWriterFactory, + final ServerConfig serverConfig + ) + { + super( + queryLifecycleFactory, + jsonMapper, + queryScheduler, + authorizerMapper, + queryResultPusherFactory, + resourceIOReaderWriterFactory, + serverConfig + ); + } + + @Override + protected Query readQuery( + final HttpServletRequest req, + final InputStream in, + final ResourceIOReaderWriterFactory.ResourceIOReaderWriter ioReaderWriter + ) throws IOException + { + final Query query = super.readQuery(req, in, ioReaderWriter); + if (!(query instanceof ScanQuery) + || !(query.getDataSource() instanceof SystemTableDataSource)) { + throw new BadQueryContextException( + "This native query endpoint accepts only local system-table Scan queries" + ); + } + return query; + } +} 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 new file mode 100644 index 000000000000..6e25c527f3ee --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/module/SystemTableModule.java @@ -0,0 +1,56 @@ +/* + * 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.module; + +import com.google.inject.Binder; +import com.google.inject.Module; +import com.google.inject.multibindings.MapBinder; +import org.apache.druid.guice.DruidBinders; +import org.apache.druid.guice.LazySingleton; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.server.system.table.ServerPropertiesTableDataProvider; +import org.apache.druid.server.system.table.ServerPropertiesTableDescriptor; +import org.apache.druid.server.system.table.SystemTableDataProvider; +import org.apache.druid.server.system.table.SystemTableDescriptor; + +/** + * 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.

+ */ +public class SystemTableModule implements Module +{ + @Override + public void configure(final Binder binder) + { + DruidBinders.dataSourceQueryHandlerBinder(binder) + .addBinding(SystemTableDataSource.class) + .toProvider(SystemTableQueryHandlerProvider.class) + .in(LazySingleton.class); + + final MapBinder descriptorBinder = MapBinder.newMapBinder(binder, String.class, SystemTableDescriptor.class); + descriptorBinder.addBinding(ServerPropertiesTableDescriptor.TABLE_NAME) + .toInstance(new ServerPropertiesTableDescriptor()); + final MapBinder dataProviderBinder = MapBinder.newMapBinder(binder, String.class, SystemTableDataProvider.class); + dataProviderBinder.addBinding(ServerPropertiesTableDescriptor.TABLE_NAME) + .to(ServerPropertiesTableDataProvider.class) + .in(LazySingleton.class); + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/module/SystemTableQueryHandlerProvider.java b/server/src/main/java/org/apache/druid/server/system/module/SystemTableQueryHandlerProvider.java new file mode 100644 index 000000000000..6fea25818148 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/module/SystemTableQueryHandlerProvider.java @@ -0,0 +1,56 @@ +/* + * 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.module; + +import com.google.inject.Inject; +import com.google.inject.Provider; +import org.apache.druid.discovery.NodeRole; +import org.apache.druid.guice.annotations.Self; +import org.apache.druid.server.DataSourceQueryHandler; +import org.apache.druid.server.system.handler.SystemTableBrokerQueryHandler; +import org.apache.druid.server.system.handler.SystemTableQueryHandler; + +import java.util.Set; + +/** Selects the system-table query handler for the current server role. */ +public class SystemTableQueryHandlerProvider implements Provider +{ + private final Set nodeRoles; + private final Provider brokerQueryHandler; + private final Provider localQueryHandler; + + @Inject + public SystemTableQueryHandlerProvider( + @Self final Set nodeRoles, + final Provider brokerQueryHandler, + final Provider localQueryHandler + ) + { + this.nodeRoles = nodeRoles; + this.brokerQueryHandler = brokerQueryHandler; + this.localQueryHandler = localQueryHandler; + } + + @Override + public DataSourceQueryHandler get() + { + return nodeRoles.contains(NodeRole.BROKER) ? brokerQueryHandler.get() : localQueryHandler.get(); + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/module/SystemTableQueryResourceModule.java b/server/src/main/java/org/apache/druid/server/system/module/SystemTableQueryResourceModule.java new file mode 100644 index 000000000000..160ae606cee8 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/module/SystemTableQueryResourceModule.java @@ -0,0 +1,43 @@ +/* + * 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.module; + +import com.google.inject.Binder; +import org.apache.druid.guice.Jerseys; +import org.apache.druid.guice.LazySingleton; +import org.apache.druid.guice.LifecycleModule; +import org.apache.druid.initialization.DruidModule; +import org.apache.druid.server.ResponseContextConfig; +import org.apache.druid.server.metrics.QueryCountStatsProvider; +import org.apache.druid.server.system.handler.SystemTableQueryResource; + +/** Registers the restricted native query HTTP resource used by scan-only nodes. */ +public class SystemTableQueryResourceModule implements DruidModule +{ + @Override + public void configure(final Binder binder) + { + binder.bind(ResponseContextConfig.class).toInstance(ResponseContextConfig.newConfig(true)); + binder.bind(SystemTableQueryResource.class).in(LazySingleton.class); + binder.bind(QueryCountStatsProvider.class).to(SystemTableQueryResource.class).in(LazySingleton.class); + Jerseys.addResource(binder, SystemTableQueryResource.class); + LifecycleModule.register(binder, SystemTableQueryResource.class); + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/table/ServerPropertiesTableDataProvider.java b/server/src/main/java/org/apache/druid/server/system/table/ServerPropertiesTableDataProvider.java new file mode 100644 index 000000000000..cd0c2e42a1a6 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/table/ServerPropertiesTableDataProvider.java @@ -0,0 +1,218 @@ +/* + * 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.Maps; +import com.google.inject.Inject; +import org.apache.druid.client.DruidServerConfig; +import org.apache.druid.discovery.NodeRole; +import org.apache.druid.guice.annotations.Self; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.filter.EqualityFilter; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.server.DruidNode; +import org.apache.druid.server.security.Action; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.AuthorizationResult; +import org.apache.druid.server.security.AuthorizationUtils; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.server.security.ForbiddenException; +import org.apache.druid.server.security.Resource; +import org.apache.druid.server.security.ResourceAction; + +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.Properties; +import java.util.Set; +import java.util.TreeMap; +import java.util.stream.Collectors; + +/** Native row supplier for {@code sys.server_properties}. */ +public class ServerPropertiesTableDataProvider implements SystemTableDataProvider +{ + private static final List PUSHDOWN_FILTERS = List.of( + new SystemTablePushdownFilter("server", null), + new SystemTablePushdownFilter("service_name", null) + ); + + private final DruidNode selfNode; + private final Set selfNodeRoles; + private final AuthorizerMapper authorizerMapper; + private final Properties properties; + private final DruidServerConfig druidServerConfig; + + @Inject + public ServerPropertiesTableDataProvider( + @Self final DruidNode selfNode, + @Self final Set selfNodeRoles, + final AuthorizerMapper authorizerMapper, + final Properties properties, + final DruidServerConfig druidServerConfig + ) + { + this.selfNode = selfNode; + this.selfNodeRoles = selfNodeRoles; + this.authorizerMapper = authorizerMapper; + this.properties = properties; + this.druidServerConfig = druidServerConfig; + } + + @Override + public List getPushdownFilters() + { + return PUSHDOWN_FILTERS; + } + + @Override + public Iterable getRows( + final List filters, + final AuthenticationResult internalAuthenticationResult + ) + { + authorizeServerRead(internalAuthenticationResult); + + String serverFilter = null; + String serviceNameFilter = null; + for (final DimFilter filter : filters) { + if (isStringEqualityFilter(filter, "server")) { + serverFilter = getFilterValue(filter); + } else if (isStringEqualityFilter(filter, "service_name")) { + serviceNameFilter = getFilterValue(filter); + } + } + + final String server = selfNode.getHostAndPortToUse(); + if (serverFilter != null && !serverFilter.equals(server)) { + return Collections.emptyList(); + } + if (serviceNameFilter != null && !serviceNameFilter.equals(selfNode.getServiceName())) { + return Collections.emptyList(); + } + + final ServerProperties serverProperties = new ServerProperties(selfNode.getServiceName(), server); + selfNodeRoles.stream() + .map(NodeRole::getJsonName) + .sorted() + .forEach(serverProperties.nodeRoles::add); + + return buildRows(serverProperties); + } + + private void authorizeServerRead(final AuthenticationResult authenticationResult) + { + final AuthorizationResult authorizationResult = AuthorizationUtils.authorizeAllResourceActions( + authenticationResult, + Collections.singletonList(new ResourceAction(Resource.STATE_RESOURCE, Action.READ)), + authorizerMapper + ); + if (!authorizationResult.allowAccessWithNoRestriction()) { + throw new ForbiddenException( + "Insufficient permission to view servers: " + authorizationResult.getErrorMessage() + ); + } + } + + private List buildRows(final ServerProperties serverProperties) + { + final Map localProperties = getProperties(); + final String nodeRoles = serverProperties.nodeRoles.toString(); + if (localProperties.isEmpty()) { + return Collections.singletonList( + row( + serverProperties.server, + serverProperties.serviceName, + nodeRoles, + null, + null, + null + ) + ); + } + + return localProperties.entrySet() + .stream() + .map(entry -> row( + serverProperties.server, + serverProperties.serviceName, + nodeRoles, + entry.getKey(), + entry.getValue(), + null + )) + .collect(Collectors.toList()); + } + + private Object[] row( + final String server, + final String serviceName, + final String nodeRoles, + @Nullable final String property, + @Nullable final String value, + @Nullable final String error + ) + { + return new Object[]{server, serviceName, nodeRoles, property, value, error}; + } + + private Map getProperties() + { + final Map allProperties = Maps.fromProperties(properties); + final Set hiddenProperties = druidServerConfig.getHiddenProperties(); + final Map filteredProperties = new HashMap<>(allProperties); + filteredProperties.keySet().removeIf( + key -> hiddenProperties.stream().anyMatch( + hiddenProperty -> StringUtils.toLowerCase(key).contains(StringUtils.toLowerCase(hiddenProperty)) + ) + ); + return new TreeMap<>(filteredProperties); + } + + private static String getFilterValue(final DimFilter filter) + { + if (filter instanceof SelectorDimFilter) { + return ((SelectorDimFilter) filter).getValue(); + } + return (String) ((EqualityFilter) filter).getMatchValue(); + } + + private static boolean isStringEqualityFilter(final DimFilter filter, final String column) + { + return filter instanceof SelectorDimFilter selector && column.equals(selector.getDimension()) + || filter instanceof EqualityFilter equality && column.equals(equality.getColumn()); + } + + private static class ServerProperties + { + private final String serviceName; + private final String server; + private final List nodeRoles = new ArrayList<>(); + + private ServerProperties(final String serviceName, final String server) + { + this.serviceName = serviceName; + this.server = server; + } + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/table/ServerPropertiesTableDescriptor.java b/server/src/main/java/org/apache/druid/server/system/table/ServerPropertiesTableDescriptor.java new file mode 100644 index 000000000000..530b9c8602ec --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/table/ServerPropertiesTableDescriptor.java @@ -0,0 +1,115 @@ +/* + * 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.DruidNode; +import org.apache.druid.server.security.Action; +import org.apache.druid.server.security.AuthorizationResult; +import org.apache.druid.server.security.AuthorizationUtils; +import org.apache.druid.server.security.ForbiddenException; +import org.apache.druid.server.security.Resource; +import org.apache.druid.server.security.ResourceAction; + +import java.util.Collections; +import java.util.Optional; +import java.util.Set; + +/** Descriptor for the native {@code sys.server_properties} table. */ +public class ServerPropertiesTableDescriptor implements SystemTableDescriptor +{ + public static final String TABLE_NAME = "server_properties"; + public static final RowSignature ROW_SIGNATURE = RowSignature + .builder() + .add("server", ColumnType.STRING) + .add("service_name", ColumnType.STRING) + .add("node_roles", ColumnType.STRING) + .add("property", ColumnType.STRING) + .add("value", ColumnType.STRING) + .add("error_message", ColumnType.STRING) + .build(); + + private static final Set NODE_ROLES = Set.of(NodeRole.values()); + private static final SystemTableRowAuthorizer ROW_AUTHORIZER = (rows, authenticationResult, authorizerMapper) -> { + final AuthorizationResult authorizationResult = AuthorizationUtils.authorizeAllResourceActions( + authenticationResult, + Collections.singletonList(new ResourceAction(Resource.STATE_RESOURCE, Action.READ)), + authorizerMapper + ); + if (!authorizationResult.allowAccessWithNoRestriction()) { + throw new ForbiddenException(authorizationResult.getErrorMessage()); + } + return rows; + }; + + @Override + public String getTableName() + { + return TABLE_NAME; + } + + @Override + public Set getNodeRoles() + { + return NODE_ROLES; + } + + @Override + public RowSignature getRowSignature() + { + return ROW_SIGNATURE; + } + + @Override + public SystemTableRowAuthorizer getRowAuthorizer() + { + return ROW_AUTHORIZER; + } + + @Override + public boolean isEmptyDiscoveryAllowed() + { + return true; + } + + @Override + public Optional getNodeFailureRow( + final DruidNode node, + final Set nodeRoles, + final Exception failure + ) + { + final String errorMessage = failure.getMessage() == null + ? failure.getClass().getSimpleName() + : failure.getMessage(); + return Optional.of( + new Object[]{ + node.getHostAndPortToUse(), + node.getServiceName(), + nodeRoles.stream().map(NodeRole::getJsonName).sorted().toList().toString(), + null, + null, + errorMessage + } + ); + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/table/SystemTableDataProvider.java b/server/src/main/java/org/apache/druid/server/system/table/SystemTableDataProvider.java new file mode 100644 index 000000000000..7f40f0017162 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/table/SystemTableDataProvider.java @@ -0,0 +1,45 @@ +/* + * 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 jakarta.validation.constraints.NotNull; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.server.security.AuthenticationResult; + +import java.util.Collections; +import java.util.List; + +/** + * Supplies storage-prefiltered rows authorized for the internal caller of one native system table. + * The implementation is deployed in related service. + * For example, the data provider of sys.tasks is deployed in overlord module + * */ +public interface SystemTableDataProvider +{ + default List getPushdownFilters() + { + return Collections.emptyList(); + } + + Iterable getRows( + @NotNull List filters, + AuthenticationResult internalAuthenticationResult + ); +} diff --git a/server/src/main/java/org/apache/druid/server/system/table/SystemTableDescriptor.java b/server/src/main/java/org/apache/druid/server/system/table/SystemTableDescriptor.java new file mode 100644 index 000000000000..f9442c51327f --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/table/SystemTableDescriptor.java @@ -0,0 +1,67 @@ +/* + * 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.RowSignature; +import org.apache.druid.server.DruidNode; + +import java.util.Optional; +import java.util.Set; + +/** Describes a native system table and how its rows are distributed across Druid services. */ +public interface SystemTableDescriptor +{ + String getTableName(); + + /** + * Returns the node roles capable of serving this table. {@link #getRoutingMode()} determines whether every node or + * only the leader for the role receives a query. + */ + Set getNodeRoles(); + + default SystemTableRoutingMode getRoutingMode() + { + return SystemTableRoutingMode.ALL_NODES; + } + + RowSignature getRowSignature(); + + SystemTableRowAuthorizer getRowAuthorizer(); + + /** Whether an empty discovery result represents an empty table instead of unavailable infrastructure. */ + default boolean isEmptyDiscoveryAllowed() + { + return false; + } + + /** + * Converts a node failure into a table row when the table supports partial results. An empty result means the + * node failure must fail the query. + */ + default Optional getNodeFailureRow( + final DruidNode node, + final Set nodeRoles, + final Exception failure + ) + { + return Optional.empty(); + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/table/SystemTablePushdownFilter.java b/server/src/main/java/org/apache/druid/server/system/table/SystemTablePushdownFilter.java new file mode 100644 index 000000000000..f4ab7abc8d32 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/table/SystemTablePushdownFilter.java @@ -0,0 +1,278 @@ +/* + * 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.query.Query; +import org.apache.druid.query.filter.AndDimFilter; +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.query.operator.OperatorFactory; +import org.apache.druid.query.operator.ScanOperatorFactory; +import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.query.ordering.StringComparators; +import org.apache.druid.segment.VirtualColumn; +import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.segment.column.ColumnType; + +import javax.annotation.Nullable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Declares a pushdown-capable system-table column and its optional supplier-side name. + */ +public record SystemTablePushdownFilter(String key, @Nullable String value) +{ + public static List extract( + final Query query, + final List pushdownFilters + ) + { + if (pushdownFilters.isEmpty()) { + return Collections.emptyList(); + } + // A virtual column shadows a physical column of the same name, so a filter on it is not a filter on the + // provider's column and must not be pushed down. + final Set virtualColumnNames = getVirtualColumnNames(query); + final Map columnMappings = new HashMap<>(); + for (final SystemTablePushdownFilter filter : pushdownFilters) { + if (!virtualColumnNames.contains(filter.key())) { + columnMappings.put(filter.key(), filter.value() == null ? filter.key() : filter.value()); + } + } + if (columnMappings.isEmpty()) { + return Collections.emptyList(); + } + + final List extracted = new ArrayList<>(); + extractConjuncts(query.getFilter(), columnMappings, extracted); + if (query instanceof WindowOperatorQuery) { + for (final OperatorFactory leafOperator : ((WindowOperatorQuery) query).getLeafOperators()) { + if (leafOperator instanceof ScanOperatorFactory) { + extractConjuncts(((ScanOperatorFactory) leafOperator).getFilter(), columnMappings, extracted); + } + } + } + return Collections.unmodifiableList(extracted); + } + + private static Set getVirtualColumnNames(final Query query) + { + final Set names = new HashSet<>(); + for (final VirtualColumn virtualColumn : query.getVirtualColumns().getVirtualColumns()) { + names.add(virtualColumn.getOutputName()); + } + if (query instanceof WindowOperatorQuery) { + for (final OperatorFactory leafOperator : ((WindowOperatorQuery) query).getLeafOperators()) { + if (leafOperator instanceof ScanOperatorFactory) { + final VirtualColumns leafVirtualColumns = ((ScanOperatorFactory) leafOperator).getVirtualColumns(); + if (leafVirtualColumns != null) { + for (final VirtualColumn virtualColumn : leafVirtualColumns.getVirtualColumns()) { + names.add(virtualColumn.getOutputName()); + } + } + } + } + } + return names; + } + + private static void extractConjuncts( + @Nullable final DimFilter filter, + final Map columnMappings, + final List extracted + ) + { + if (filter instanceof AndDimFilter) { + for (final DimFilter subfilter : ((AndDimFilter) filter).getFields()) { + extractConjuncts(subfilter, columnMappings, extracted); + } + } else if (filter != null) { + final DimFilter rewrittenFilter = rewriteSupportedFilter(filter, columnMappings); + if (rewrittenFilter != null) { + extracted.add(rewrittenFilter); + } + } + } + + @Nullable + private static DimFilter rewriteSupportedFilter( + final DimFilter filter, + final Map columnMappings + ) + { + return switch (filter) { + case SelectorDimFilter selector -> { + final String mappedColumn = columnMappings.get(selector.getDimension()); + yield mappedColumn != null && selector.getValue() != null && selector.getExtractionFn() == null + ? new SelectorDimFilter(mappedColumn, selector.getValue(), null) + : null; + } + case EqualityFilter equality -> { + final String mappedColumn = columnMappings.get(equality.getColumn()); + yield mappedColumn != null + && ColumnType.STRING.equals(equality.getMatchValueType()) + && equality.getMatchValue() instanceof String + ? new EqualityFilter(mappedColumn, ColumnType.STRING, equality.getMatchValue(), null) + : null; + } + case InDimFilter in -> { + final String mappedColumn = columnMappings.get(in.getDimension()); + yield mappedColumn != null + && in.getExtractionFn() == null + && !in.getValues().isEmpty() + && in.getValues().stream().allMatch(Objects::nonNull) + ? new InDimFilter(mappedColumn, in.getValues(), null) + : null; + } + case TypedInFilter in -> { + final String mappedColumn = columnMappings.get(in.getColumn()); + yield mappedColumn != null + && ColumnType.STRING.equals(in.getMatchValueType()) + && !in.getSortedValues().isEmpty() + && in.getSortedValues().stream().allMatch(String.class::isInstance) + ? new TypedInFilter(mappedColumn, ColumnType.STRING, null, in.getSortedValues(), null) + : null; + } + case OrDimFilter or -> { + final List rewrittenFields = new ArrayList<>(); + boolean supported = true; + for (final DimFilter field : or.getFields()) { + final DimFilter rewrittenField = rewriteSupportedFilter(field, columnMappings); + if (rewrittenField == null || !isStringValuesFilter(rewrittenField)) { + supported = false; + break; + } + rewrittenFields.add(rewrittenField); + } + yield supported && !rewrittenFields.isEmpty() && hasSingleColumn(rewrittenFields) + ? new OrDimFilter(rewrittenFields) + : null; + } + case LikeDimFilter like -> { + final String mappedColumn = columnMappings.get(like.getDimension()); + yield mappedColumn != null && like.getExtractionFn() == null && like.getEscape() == null + ? new LikeDimFilter(mappedColumn, like.getPattern(), null, null) + : null; + } + case NotDimFilter not -> { + final DimFilter rewrittenField = rewriteSupportedFilter(not.getField(), columnMappings); + yield rewrittenField != null + && (isStringValuesFilter(rewrittenField) || rewrittenField instanceof LikeDimFilter) + ? new NotDimFilter(rewrittenField) + : null; + } + case BoundDimFilter bound -> { + final String mappedColumn = columnMappings.get(bound.getDimension()); + yield mappedColumn != null + && bound.getExtractionFn() == null + && StringComparators.LEXICOGRAPHIC.equals(bound.getOrdering()) + ? new BoundDimFilter( + mappedColumn, + bound.getLower(), + bound.getUpper(), + bound.isLowerStrict(), + bound.isUpperStrict(), + null, + null, + StringComparators.LEXICOGRAPHIC + ) + : null; + } + case RangeFilter range -> { + final String mappedColumn = columnMappings.get(range.getColumn()); + yield mappedColumn != null + && ColumnType.STRING.equals(range.getMatchValueType()) + && (range.getLower() == null || range.getLower() instanceof String) + && (range.getUpper() == null || range.getUpper() instanceof String) + ? new RangeFilter( + mappedColumn, + ColumnType.STRING, + range.getLower(), + range.getUpper(), + range.isLowerOpen(), + range.isUpperOpen(), + null + ) + : null; + } + default -> null; + }; + } + + private static boolean isStringValuesFilter(final DimFilter filter) + { + return filter instanceof SelectorDimFilter + || filter instanceof EqualityFilter + || filter instanceof InDimFilter + || filter instanceof TypedInFilter + || filter instanceof OrDimFilter; + } + + private static boolean hasSingleColumn(final List filters) + { + final String column = getStringValuesColumn(filters.get(0)); + return filters.stream().allMatch(filter -> column.equals(getStringValuesColumn(filter))); + } + + public static String getStringValuesColumn(final DimFilter filter) + { + return switch (filter) { + case SelectorDimFilter selector -> selector.getDimension(); + case EqualityFilter equality -> equality.getColumn(); + case InDimFilter in -> in.getDimension(); + case TypedInFilter in -> in.getColumn(); + case null, default -> getStringValuesColumn(((OrDimFilter) filter).getFields().get(0)); + }; + } + + /** Returns the string values from a validated string-values pushdown filter. */ + public static Set getStringValues(final DimFilter filter) + { + return switch (filter) { + case SelectorDimFilter selector -> Set.of(selector.getValue()); + case EqualityFilter equality -> Set.of((String) equality.getMatchValue()); + case InDimFilter in -> in.getValues(); + case TypedInFilter in -> in.getSortedValues().stream().map(String.class::cast).collect(Collectors.toSet()); + case null, default -> { + final Set values = new HashSet<>(); + for (final DimFilter field : ((OrDimFilter) filter).getFields()) { + values.addAll(getStringValues(field)); + } + yield values; + } + }; + } +} diff --git a/server/src/main/java/org/apache/druid/server/system/table/SystemTableRoutingMode.java b/server/src/main/java/org/apache/druid/server/system/table/SystemTableRoutingMode.java new file mode 100644 index 000000000000..fa6c1c69ba69 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/table/SystemTableRoutingMode.java @@ -0,0 +1,30 @@ +/* + * 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; + +/** Defines how the Broker selects nodes that contribute rows to a system table. */ +public enum SystemTableRoutingMode +{ + /** Every discovered node for the descriptor's roles contributes an independent set of rows. */ + ALL_NODES, + + /** Only the current leader for the descriptor's single node role contains authoritative rows. */ + LEADER_ONLY +} diff --git a/server/src/main/java/org/apache/druid/server/system/table/SystemTableRowAuthorizer.java b/server/src/main/java/org/apache/druid/server/system/table/SystemTableRowAuthorizer.java new file mode 100644 index 000000000000..658bf5a068f5 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/table/SystemTableRowAuthorizer.java @@ -0,0 +1,33 @@ +/* + * 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.server.security.AuthenticationResult; +import org.apache.druid.server.security.AuthorizerMapper; + +/** Applies the requesting user's authorization to node-supplied system-table rows on the Broker. */ +public interface SystemTableRowAuthorizer +{ + Iterable filterAuthorizedRows( + Iterable rows, + AuthenticationResult authenticationResult, + AuthorizerMapper authorizerMapper + ); +} diff --git a/server/src/test/java/org/apache/druid/client/DirectDruidClientTest.java b/server/src/test/java/org/apache/druid/client/DirectDruidClientTest.java index b4d1e9e1b7b0..5567866adce7 100644 --- a/server/src/test/java/org/apache/druid/client/DirectDruidClientTest.java +++ b/server/src/test/java/org/apache/druid/client/DirectDruidClientTest.java @@ -66,6 +66,7 @@ import org.apache.druid.segment.TestIndex; import org.apache.druid.segment.incremental.IncrementalIndexSchema; import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory; +import org.apache.druid.server.QueryResource; import org.apache.druid.server.QueryStackTests; import org.apache.druid.server.coordination.ServerType; import org.apache.druid.server.coordinator.simulate.BlockingExecutorService; @@ -154,6 +155,10 @@ public void testRun() throws Exception Assertions.assertFalse(requests.isEmpty()); Assertions.assertEquals(url, requests.get(0).getUrl()); Assertions.assertEquals(HttpMethod.POST, requests.get(0).getMethod()); + Assertions.assertEquals( + QueryResource.NATIVE_QUERY_ROUTE_LOCAL, + requests.get(0).getHeaders().get(QueryResource.HEADER_NATIVE_QUERY_ROUTE).iterator().next() + ); Assertions.assertEquals(1, client1.getNumOpenConnections()); // simulate read timeout on second request diff --git a/server/src/test/java/org/apache/druid/guice/QueryRunnerFactoryModuleTest.java b/server/src/test/java/org/apache/druid/guice/QueryRunnerFactoryModuleTest.java new file mode 100644 index 000000000000..4acd9fbd1f4e --- /dev/null +++ b/server/src/test/java/org/apache/druid/guice/QueryRunnerFactoryModuleTest.java @@ -0,0 +1,80 @@ +/* + * 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.guice; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.inject.Guice; +import com.google.inject.Injector; +import com.google.inject.Key; +import com.google.inject.TypeLiteral; +import org.apache.druid.java.util.emitter.service.ServiceEmitter; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryRunnerFactory; +import org.apache.druid.query.QueryToolChest; +import org.apache.druid.query.groupby.GroupByQueryConfig; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.scan.ScanQueryConfig; +import org.apache.druid.server.metrics.NoopServiceEmitter; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.Map; +import java.util.Properties; + +public class QueryRunnerFactoryModuleTest +{ + @Test + public void testRegistersSelectedQueryTypesOnly() + { + final Injector injector = Guice.createInjector( + ImmutableList.of( + new DruidGuiceExtensions(), + new ConfigModule(), + new QueryableModule(), + new QueryRunnerFactoryModule(ImmutableSet.of(ScanQuery.class)), + new LifecycleModule(), + binder -> binder.bind(ServiceEmitter.class).to(NoopServiceEmitter.class), + binder -> binder.bind(Properties.class).toInstance(new Properties()) + ) + ); + + final Map, QueryRunnerFactory> queryRunnerFactories = injector.getInstance( + Key.get(new TypeLiteral<>() {}) + ); + final Map, QueryToolChest> queryToolChests = injector.getInstance( + Key.get(new TypeLiteral<>() {}) + ); + + Assertions.assertEquals(ImmutableSet.of(ScanQuery.class), queryRunnerFactories.keySet()); + Assertions.assertEquals(ImmutableSet.of(ScanQuery.class), queryToolChests.keySet()); + Assertions.assertNotNull(injector.getExistingBinding(Key.get(ScanQueryConfig.class))); + Assertions.assertNull(injector.getExistingBinding(Key.get(GroupByQueryConfig.class))); + } + + @Test + public void testRejectsUnsupportedQueryType() + { + Assertions.assertThrows( + IllegalArgumentException.class, + () -> new QueryRunnerFactoryModule(ImmutableSet.of(Query.class)) + ); + } +} diff --git a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java index 6be0876ecc39..e2db920c8a6a 100644 --- a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java @@ -30,6 +30,7 @@ import com.google.inject.testing.fieldbinder.BoundFieldModule; import org.apache.druid.client.BrokerViewOfBrokerConfig; import org.apache.druid.error.DruidException; +import org.apache.druid.guice.DruidBinders; import org.apache.druid.guice.LazySingleton; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.Intervals; @@ -41,12 +42,14 @@ import org.apache.druid.query.Query; import org.apache.druid.query.QueryConfigProvider; import org.apache.druid.query.QueryContextTest; +import org.apache.druid.query.QueryDataSource; import org.apache.druid.query.QueryMetrics; import org.apache.druid.query.QueryRunner; import org.apache.druid.query.QueryRunnerFactoryConglomerate; import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.QueryToolChest; import org.apache.druid.query.RestrictedDataSource; +import org.apache.druid.query.SystemTableDataSource; import org.apache.druid.query.TableDataSource; import org.apache.druid.query.aggregation.CountAggregatorFactory; import org.apache.druid.query.filter.DimFilter; @@ -58,6 +61,7 @@ import org.apache.druid.query.policy.PolicyEnforcer; import org.apache.druid.query.policy.RestrictAllTablesPolicyEnforcer; import org.apache.druid.query.policy.RowFilterPolicy; +import org.apache.druid.query.scan.ScanQuery; import org.apache.druid.query.timeseries.TimeseriesQuery; import org.apache.druid.server.broker.BrokerDynamicConfig; import org.apache.druid.server.broker.QueryConfigSnapshot; @@ -84,6 +88,7 @@ import java.util.HashMap; import java.util.Map; import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; @LazySingleton public class QueryLifecycleTest @@ -159,7 +164,10 @@ public void setup() injector = Guice.createInjector( BoundFieldModule.of(this), - binder -> binder.bindScope(LazySingleton.class, Scopes.SINGLETON) + binder -> { + binder.bindScope(LazySingleton.class, Scopes.SINGLETON); + DruidBinders.dataSourceQueryHandlerBinder(binder); + } ); } @@ -192,6 +200,7 @@ public void testRunSimple_preauthorizedAsSuperuser() { EasyMock.expect(queryConfig.getContext()).andReturn(ImmutableMap.of()).anyTimes(); EasyMock.expect(authenticationResult.getIdentity()).andReturn(IDENTITY).anyTimes(); + EasyMock.expect(authenticationResult.getAuthorizerName()).andReturn(AUTHORIZER).anyTimes(); EasyMock.expect(conglomerate.getToolChest(EasyMock.anyObject())) .andReturn(toolChest) .once(); @@ -206,6 +215,34 @@ public void testRunSimple_preauthorizedAsSuperuser() lifecycle.runSimple(query, authenticationResult, AuthorizationResult.ALLOW_NO_RESTRICTION); } + /** + * Native query form of: + * + * 
{@code
+   * SELECT COUNT(*)
+   * FROM
+   * (
+   *   SELECT
+   *     task_id,
+   *     COUNT(*) AS task_count
+   *   FROM sys.tasks
+   *   GROUP BY task_id
+   * )
+   * WHERE task_count > 0
+   * }
+ * + * The handler is registered for the leaf datasource and must be found below the query root. + */ + @Test + public void testSystemTableHandlerHandlesQueryDataSource() + { + final Query innerQuery = Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource("tasks")) + .eternityInterval() + .build(); + assertWrappedSystemTableUsesHandler(new QueryDataSource(innerQuery)); + } + @Test public void testRunSimpleUnauthorized() { @@ -917,6 +954,11 @@ public void testRunSimple_queryNotBlocklisted() } private HttpServletRequest mockRequest() + { + return mockRequest(null); + } + + private HttpServletRequest mockRequest(@Nullable final String nativeQueryRoute) { HttpServletRequest request = EasyMock.createNiceMock(HttpServletRequest.class); EasyMock.expect(request.getAttribute(EasyMock.eq(AuthConfig.DRUID_AUTHENTICATION_RESULT))) @@ -925,10 +967,131 @@ private HttpServletRequest mockRequest() .andReturn(null).anyTimes(); EasyMock.expect(request.getAttribute(EasyMock.eq(AuthConfig.DRUID_AUTHORIZATION_CHECKED))) .andReturn(null).anyTimes(); + EasyMock.expect(request.getHeader(QueryResource.HEADER_NATIVE_QUERY_ROUTE)) + .andReturn(nativeQueryRoute).anyTimes(); EasyMock.replay(request); return request; } + /** {@code X-Druid-Native-Query-Route: local} reaches the datasource query handler as server request state. */ + @Test + public void testLocalNativeQueryRouteIsPassedToDataSourceHandler() + { + assertNativeQueryRoute(QueryResource.NATIVE_QUERY_ROUTE_LOCAL, true); + } + + /** The native-query route header is case-sensitive and rejects non-exact values. */ + @Test + public void testNonExactNativeQueryRouteIsNotLocal() + { + assertNativeQueryRoute("LOCAL", false); + } + + private void assertNativeQueryRoute(final String route, final boolean expectedExecuteLocally) + { + final ScanQuery systemTableQuery = Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource("tasks")) + .eternityInterval() + .build(); + final AtomicInteger handlerCalls = new AtomicInteger(); + final DataSourceQueryHandler handler = new DataSourceQueryHandler() + { + @Override + public QueryRunner createRunner( + final Query query, + final AuthenticationResult authenticationResult, + final boolean executeLocally + ) + { + handlerCalls.incrementAndGet(); + Assertions.assertEquals(expectedExecuteLocally, executeLocally); + return (queryPlus, responseContext) -> Sequences.empty(); + } + }; + + EasyMock.expect(queryConfig.getContext()).andReturn(ImmutableMap.of()).anyTimes(); + EasyMock.expect(authenticationResult.getIdentity()).andReturn(IDENTITY).anyTimes(); + EasyMock.expect(authenticationResult.getAuthorizerName()).andReturn(AUTHORIZER).anyTimes(); + EasyMock.expect( + authorizer.authorize( + authenticationResult, + new Resource("sys.tasks", ResourceType.DATASOURCE), + Action.READ + ) + ).andReturn(Access.OK).once(); + EasyMock.expect(conglomerate.getToolChest(EasyMock.anyObject())).andReturn(toolChest).once(); + replayAll(); + + final QueryLifecycle lifecycle = new QueryLifecycle( + conglomerate, + texasRanger, + metricsFactory, + emitter, + requestLogger, + authzMapper, + authConfig, + policyEnforcer, + new QueryConfigSnapshot(ImmutableMap.of(), BrokerDynamicConfig.builder().build()), + Map.of(SystemTableDataSource.class, handler), + System.currentTimeMillis(), + System.nanoTime() + ); + lifecycle.initialize(systemTableQuery); + Assertions.assertTrue( + lifecycle.authorize(mockRequest(route)).allowAccessWithNoRestriction() + ); + lifecycle.execute(); + + Assertions.assertEquals(1, handlerCalls.get()); + } + + private void assertWrappedSystemTableUsesHandler(final DataSource dataSource) + { + final TimeseriesQuery wrappedQuery = Druids.newTimeseriesQueryBuilder() + .dataSource(dataSource) + .intervals(ImmutableList.of(Intervals.ETERNITY)) + .aggregators(new CountAggregatorFactory("count")) + .build(); + final AtomicInteger handlerCalls = new AtomicInteger(); + final DataSourceQueryHandler handler = new DataSourceQueryHandler() + { + @Override + public QueryRunner createRunner( + final Query query, + final AuthenticationResult authenticationResult, + final boolean executeLocally + ) + { + Assertions.assertFalse(executeLocally); + handlerCalls.incrementAndGet(); + return (queryPlus, responseContext) -> Sequences.empty(); + } + }; + + EasyMock.expect(queryConfig.getContext()).andReturn(ImmutableMap.of()).anyTimes(); + EasyMock.expect(authenticationResult.getIdentity()).andReturn(IDENTITY).anyTimes(); + EasyMock.expect(conglomerate.getToolChest(EasyMock.anyObject())).andReturn(toolChest).once(); + replayAll(); + + final QueryLifecycle lifecycle = new QueryLifecycle( + conglomerate, + texasRanger, + metricsFactory, + emitter, + requestLogger, + authzMapper, + authConfig, + policyEnforcer, + new QueryConfigSnapshot(ImmutableMap.of(), BrokerDynamicConfig.builder().build()), + Map.of(SystemTableDataSource.class, handler), + System.currentTimeMillis(), + System.nanoTime() + ); + + lifecycle.runSimple(wrappedQuery, authenticationResult, AuthorizationResult.ALLOW_NO_RESTRICTION); + Assertions.assertEquals(1, handlerCalls.get()); + } + private void replayAll() { EasyMock.replay( diff --git a/server/src/test/java/org/apache/druid/server/QuerySchedulerTest.java b/server/src/test/java/org/apache/druid/server/QuerySchedulerTest.java index f0c75b4d174d..94f3ffac49f6 100644 --- a/server/src/test/java/org/apache/druid/server/QuerySchedulerTest.java +++ b/server/src/test/java/org/apache/druid/server/QuerySchedulerTest.java @@ -27,6 +27,7 @@ import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListeningExecutorService; import com.google.common.util.concurrent.MoreExecutors; +import com.google.common.util.concurrent.SettableFuture; import com.google.inject.Injector; import com.google.inject.Key; import com.google.inject.ProvisionException; @@ -53,6 +54,7 @@ import org.apache.druid.query.QueryPlus; import org.apache.druid.query.QueryRunnerFactory; import org.apache.druid.query.QueryToolChest; +import org.apache.druid.query.SystemTableDataSource; import org.apache.druid.query.aggregation.CountAggregatorFactory; import org.apache.druid.query.dimension.DefaultDimensionSpec; import org.apache.druid.query.groupby.GroupByQuery; @@ -98,6 +100,25 @@ public class QuerySchedulerTest private ListeningExecutorService executorService; private ObservableQueryScheduler scheduler; + /** System-table node queries attach a stable datasource authorization resource for cancellation. */ + @Test + public void testSystemTableCancellationHasAuthorizationResource() + { + final Query query = GroupByQuery.builder() + .setDataSource(new SystemTableDataSource("tasks")) + .setInterval("2000/2001") + .setGranularity(Granularities.ALL) + .setContext(Map.of("queryId", "system-query")) + .build(); + final SettableFuture future = SettableFuture.create(); + + scheduler.registerQueryFuture(query, future); + + Assertions.assertEquals(Set.of("sys.tasks"), scheduler.getQueryDatasources(query.getId())); + scheduler.cancelQuery(query.getId()); + Assertions.assertTrue(future.isCancelled()); + } + @BeforeEach public void setup() { diff --git a/server/src/test/java/org/apache/druid/server/http/CoordinatorRedirectInfoTest.java b/server/src/test/java/org/apache/druid/server/http/CoordinatorRedirectInfoTest.java index c08f87869b7c..cb6997b3cb67 100644 --- a/server/src/test/java/org/apache/druid/server/http/CoordinatorRedirectInfoTest.java +++ b/server/src/test/java/org/apache/druid/server/http/CoordinatorRedirectInfoTest.java @@ -47,6 +47,7 @@ public void testDoLocalWhenLeading() Assertions.assertTrue(coordinatorRedirectInfo.doLocal(null)); Assertions.assertTrue(coordinatorRedirectInfo.doLocal("/druid/coordinator/v1/leader")); Assertions.assertTrue(coordinatorRedirectInfo.doLocal("/druid/coordinator/v1/isLeader")); + Assertions.assertTrue(coordinatorRedirectInfo.doLocal("/druid/v2")); Assertions.assertTrue(coordinatorRedirectInfo.doLocal("/druid/coordinator/v1/other/path")); EasyMock.verify(druidCoordinator); } @@ -59,6 +60,7 @@ public void testDoLocalWhenNotLeading() Assertions.assertFalse(coordinatorRedirectInfo.doLocal(null)); Assertions.assertTrue(coordinatorRedirectInfo.doLocal("/druid/coordinator/v1/leader")); Assertions.assertTrue(coordinatorRedirectInfo.doLocal("/druid/coordinator/v1/isLeader")); + Assertions.assertTrue(coordinatorRedirectInfo.doLocal("/druid/v2")); Assertions.assertFalse(coordinatorRedirectInfo.doLocal("/druid/coordinator/v1/other/path")); EasyMock.verify(druidCoordinator); } diff --git a/server/src/test/java/org/apache/druid/server/system/ServerPropertiesTableDataProviderTest.java b/server/src/test/java/org/apache/druid/server/system/ServerPropertiesTableDataProviderTest.java new file mode 100644 index 000000000000..ab22d8df7ba6 --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/system/ServerPropertiesTableDataProviderTest.java @@ -0,0 +1,153 @@ +/* + * 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; + +import org.apache.druid.client.DruidServerConfig; +import org.apache.druid.discovery.NodeRole; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.server.DruidNode; +import org.apache.druid.server.security.Access; +import org.apache.druid.server.security.AuthConfig; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.Authorizer; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.server.security.ForbiddenException; +import org.apache.druid.server.system.table.ServerPropertiesTableDataProvider; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Properties; +import java.util.Set; + +public class ServerPropertiesTableDataProviderTest +{ + private static final AuthenticationResult AUTHENTICATION_RESULT = + new AuthenticationResult("test-user", AuthConfig.ALLOW_ALL_NAME, null, null); + + @Test + public void testReturnsVisiblePropertiesAndNodeMetadata() + { + final Properties properties = new Properties(); + properties.setProperty("druid.test.visible", "visible"); + properties.setProperty("druid.test.password", "hidden"); + + final ServerPropertiesTableDataProvider supplier = supplier(properties, allowAllAuthorizerMapper()); + final List rows = toRows(supplier.getRows(List.of(), AUTHENTICATION_RESULT)); + + Assertions.assertEquals(1, rows.size()); + Assertions.assertArrayEquals( + new Object[]{ + "localhost:8080", + "overlord", + "[overlord]", + "druid.test.visible", + "visible", + null + }, + rows.get(0) + ); + } + + @Test + public void testAppliesServerAndServiceNameFilters() + { + final ServerPropertiesTableDataProvider supplier = supplier(new Properties(), allowAllAuthorizerMapper()); + final DimFilter wrongServer = new SelectorDimFilter("server", "other:8080", null); + final DimFilter wrongService = new SelectorDimFilter("service_name", "broker", null); + + Assertions.assertTrue(toRows(supplier.getRows(List.of(wrongServer), AUTHENTICATION_RESULT)).isEmpty()); + Assertions.assertTrue(toRows(supplier.getRows(List.of(wrongService), AUTHENTICATION_RESULT)).isEmpty()); + Assertions.assertFalse( + toRows( + supplier.getRows( + List.of(new SelectorDimFilter("server", "localhost:8080", null)), + AUTHENTICATION_RESULT + ) + ).isEmpty() + ); + } + + @Test + public void testReturnsPlaceholderWhenNoPropertiesExist() + { + final ServerPropertiesTableDataProvider supplier = supplier(new Properties(), allowAllAuthorizerMapper()); + final List rows = toRows(supplier.getRows(List.of(), AUTHENTICATION_RESULT)); + + Assertions.assertEquals(1, rows.size()); + Assertions.assertNull(rows.get(0)[3]); + Assertions.assertNull(rows.get(0)[4]); + Assertions.assertNull(rows.get(0)[5]); + } + + @Test + public void testRejectsUnauthorizedRequest() + { + final Authorizer denyAll = (authenticationResult, resource, action) -> Access.DENIED; + final AuthorizerMapper authorizerMapper = new AuthorizerMapper(null) + { + @Override + public Authorizer getAuthorizer(final String name) + { + return denyAll; + } + }; + final ServerPropertiesTableDataProvider supplier = supplier(new Properties(), authorizerMapper); + + Assertions.assertThrows( + ForbiddenException.class, + () -> supplier.getRows(List.of(), AUTHENTICATION_RESULT) + ); + } + + private static ServerPropertiesTableDataProvider supplier( + final Properties properties, + final AuthorizerMapper authorizerMapper + ) + { + return new ServerPropertiesTableDataProvider( + new DruidNode("overlord", "localhost", false, 8080, null, true, false), + Set.of(NodeRole.OVERLORD), + authorizerMapper, + properties, + new DruidServerConfig(null, null) + ); + } + + private static AuthorizerMapper allowAllAuthorizerMapper() + { + return new AuthorizerMapper(null) + { + @Override + public Authorizer getAuthorizer(final String name) + { + return (authenticationResult, resource, action) -> Access.OK; + } + }; + } + + private static List toRows(final Iterable rows) + { + final List result = new java.util.ArrayList<>(); + rows.forEach(result::add); + return result; + } +} diff --git a/server/src/test/java/org/apache/druid/server/system/SystemTablePushdownFilterTest.java b/server/src/test/java/org/apache/druid/server/system/SystemTablePushdownFilterTest.java new file mode 100644 index 000000000000..2e6646479035 --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/system/SystemTablePushdownFilterTest.java @@ -0,0 +1,190 @@ +/* + * 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; + +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.math.expr.ExprMacroTable; +import org.apache.druid.query.Druids; +import org.apache.druid.query.TableDataSource; +import org.apache.druid.query.filter.AndDimFilter; +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.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.operator.ScanOperatorFactory; +import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.query.ordering.StringComparators; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.spec.LegacySegmentSpec; +import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.virtual.ExpressionVirtualColumn; +import org.apache.druid.server.system.table.SystemTablePushdownFilter; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.List; + +public class SystemTablePushdownFilterTest +{ + private static final List PUSHDOWN_FILTERS = List.of( + new SystemTablePushdownFilter("task_id", "id"), + new SystemTablePushdownFilter("created_time", "created_date") + ); + + @Test + public void testExtractsSupportedFilterWithMappedColumn() + { + final List filters = extract( + new EqualityFilter("task_id", ColumnType.STRING, "task-a", null) + ); + + Assertions.assertEquals( + List.of(new EqualityFilter("id", ColumnType.STRING, "task-a", null)), + filters + ); + } + + @Test + public void testExtractsSameColumnOrWithMappedColumn() + { + final List filters = extract( + new OrDimFilter( + new SelectorDimFilter("task_id", "task-a", null), + new SelectorDimFilter("task_id", "task-b", null) + ) + ); + + final OrDimFilter filter = Assertions.assertInstanceOf(OrDimFilter.class, filters.get(0)); + Assertions.assertEquals("id", ((SelectorDimFilter) filter.getFields().get(0)).getDimension()); + Assertions.assertEquals("id", ((SelectorDimFilter) filter.getFields().get(1)).getDimension()); + } + + @Test + public void testExtractsLexicographicRangeWithMappedColumn() + { + final List boundFilters = extract( + new BoundDimFilter( + "created_time", + "2026-01-01", + "2026-02-01", + false, + true, + null, + null, + StringComparators.LEXICOGRAPHIC + ) + ); + final List rangeFilters = extract( + new RangeFilter("task_id", ColumnType.STRING, "a", "z", false, true, null) + ); + + Assertions.assertEquals("created_date", ((BoundDimFilter) boundFilters.get(0)).getDimension()); + Assertions.assertEquals("id", ((RangeFilter) rangeFilters.get(0)).getColumn()); + } + + @Test + public void testDoesNotExtractUnsupportedFilterShapeOrColumn() + { + Assertions.assertTrue(extract(new SelectorDimFilter("datasource", "wikipedia", null)).isEmpty()); + } + + @Test + public void testExtractsLikeAndNotWithMappedColumn() + { + final List likeFilters = extract(new LikeDimFilter("task_id", "%task%", null, null)); + final List notFilters = extract( + new NotDimFilter(new SelectorDimFilter("task_id", "task-a", null)) + ); + + Assertions.assertEquals("id", ((LikeDimFilter) likeFilters.get(0)).getDimension()); + final NotDimFilter notFilter = (NotDimFilter) notFilters.get(0); + Assertions.assertEquals("id", ((SelectorDimFilter) notFilter.getField()).getDimension()); + } + + /** + * A virtual column shadows the physical column of the same name, so its filter must stay in the native scan instead + * of being compared against the provider's column. Filters on columns that are not shadowed are still extracted. + */ + @Test + public void testDoesNotExtractFilterOnColumnShadowedByVirtualColumn() + { + final DimFilter shadowedFilter = new SelectorDimFilter("task_id", "alias", null); + final DimFilter createdTimeFilter = new SelectorDimFilter("created_time", "2026-01-01", null); + final ScanQuery query = Druids.newScanQueryBuilder() + .dataSource(new TableDataSource("test")) + .intervals(new LegacySegmentSpec(Intervals.ETERNITY)) + .virtualColumns(shadowingVirtualColumns()) + .filters(new AndDimFilter(shadowedFilter, createdTimeFilter)) + .build(); + + Assertions.assertEquals( + List.of(new SelectorDimFilter("created_date", "2026-01-01", null)), + SystemTablePushdownFilter.extract(query, PUSHDOWN_FILTERS) + ); + } + + @Test + public void testDoesNotExtractWindowLeafFilterOnColumnShadowedByLeafVirtualColumn() + { + final WindowOperatorQuery query = new WindowOperatorQuery( + new TableDataSource("test"), + new LegacySegmentSpec(Intervals.ETERNITY), + Collections.emptyMap(), + RowSignature.builder().add("task_id", ColumnType.STRING).build(), + Collections.emptyList(), + List.of( + new ScanOperatorFactory( + null, + new SelectorDimFilter("task_id", "alias", null), + null, + null, + shadowingVirtualColumns(), + null + ) + ) + ); + + Assertions.assertTrue(SystemTablePushdownFilter.extract(query, PUSHDOWN_FILTERS).isEmpty()); + } + + private static VirtualColumns shadowingVirtualColumns() + { + return VirtualColumns.create( + new ExpressionVirtualColumn("task_id", "'alias'", ColumnType.STRING, ExprMacroTable.nil()) + ); + } + + private static List extract(final DimFilter filter) + { + final ScanQuery query = Druids.newScanQueryBuilder() + .dataSource(new TableDataSource("test")) + .intervals(new LegacySegmentSpec(Intervals.ETERNITY)) + .filters(filter) + .build(); + return SystemTablePushdownFilter.extract(query, PUSHDOWN_FILTERS); + } +} diff --git a/server/src/test/java/org/apache/druid/server/system/SystemTableQueryHandlerTest.java b/server/src/test/java/org/apache/druid/server/system/SystemTableQueryHandlerTest.java new file mode 100644 index 000000000000..5398b85d3154 --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/system/SystemTableQueryHandlerTest.java @@ -0,0 +1,247 @@ +/* + * 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; + +import org.apache.druid.client.DruidServerConfig; +import org.apache.druid.discovery.NodeRole; +import org.apache.druid.math.expr.ExprMacroTable; +import org.apache.druid.query.BadQueryContextException; +import org.apache.druid.query.Druids; +import org.apache.druid.query.FilteredDataSource; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.scan.ScanQueryEngine; +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.virtual.ExpressionVirtualColumn; +import org.apache.druid.server.DruidNode; +import org.apache.druid.server.security.Access; +import org.apache.druid.server.security.AuthConfig; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.Authorizer; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.server.system.handler.SystemTableQueryHandler; +import org.apache.druid.server.system.table.ServerPropertiesTableDataProvider; +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.SystemTableRowAuthorizer; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; + +public class SystemTableQueryHandlerTest +{ + private static final RowSignature ROW_SIGNATURE = RowSignature.builder() + .add("task_id", ColumnType.STRING) + .add("duration", ColumnType.LONG) + .build(); + + /** A direct local Scan lazily reads provider rows and applies the table descriptor's row authorization. */ + @Test + public void testRunsNodeScanAgainstSuppliedRows() + { + final AtomicInteger getRowsCalls = new AtomicInteger(); + final AtomicInteger authorizationCalls = new AtomicInteger(); + final SystemTableDataProvider supplier = new SystemTableDataProvider() + { + @Override + public Iterable getRows( + final List filters, + final AuthenticationResult authenticationResult + ) + { + getRowsCalls.incrementAndGet(); + return Arrays.asList( + new Object[]{"task-a", 10L}, + new Object[]{"task-b", 20L}, + new Object[]{"task-c", 30L} + ); + } + }; + + final SystemTableDescriptor descriptor = new SystemTableDescriptor() + { + @Override + public String getTableName() + { + return "test"; + } + + @Override + public Set getNodeRoles() + { + return Set.of(); + } + + @Override + public RowSignature getRowSignature() + { + return ROW_SIGNATURE; + } + + @Override + public SystemTableRowAuthorizer getRowAuthorizer() + { + return (rows, authenticationResult, authorizerMapper) -> { + authorizationCalls.incrementAndGet(); + Assertions.assertEquals("alice", authenticationResult.getIdentity()); + return () -> java.util.stream.StreamSupport.stream(rows.spliterator(), false) + .filter(row -> !"task-a".equals(row[0])) + .iterator(); + }; + } + }; + final SystemTableQueryHandler handler = new SystemTableQueryHandler( + Map.of("test", supplier), + Map.of(descriptor.getTableName(), descriptor), + new ScanQueryEngine(), + new AuthorizerMapper(Map.of()) + ); + final ScanQuery query = Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource("test")) + .eternityInterval() + .columns(ROW_SIGNATURE) + .filters(new SelectorDimFilter("task_id", "task-b", null)) + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .build(); + + final QueryRunner runner = handler.createRunner( + query, + new AuthenticationResult("alice", "allow", "external", null), + true + ); + Assertions.assertEquals(0, getRowsCalls.get()); + + final List result = runner.run( + QueryPlus.wrap(query), + ResponseContext.createEmpty() + ).toList(); + + Assertions.assertEquals(1, getRowsCalls.get()); + Assertions.assertEquals(1, authorizationCalls.get()); + Assertions.assertEquals(1, result.size()); + Assertions.assertEquals( + List.of( + List.of("task-b", 20L) + ), + result.get(0).getEvents() + ); + } + + /** + * A virtual column shadows the physical column of the same name. A filter on the virtual {@code server} column must + * be evaluated by the native scan rather than pushed to the provider, which compares it to the physical host. + */ + @Test + public void testFilterOnVirtualColumnShadowingPushdownColumnKeepsMatchingRows() + { + final AuthorizerMapper authorizerMapper = new AuthorizerMapper(null) + { + @Override + public Authorizer getAuthorizer(final String name) + { + return (authenticationResult, resource, action) -> Access.OK; + } + }; + final Properties properties = new Properties(); + properties.setProperty("druid.test.visible", "visible"); + final ServerPropertiesTableDataProvider provider = new ServerPropertiesTableDataProvider( + new DruidNode("overlord", "localhost", false, 8080, null, true, false), + Set.of(NodeRole.OVERLORD), + authorizerMapper, + properties, + new DruidServerConfig(null, null) + ); + final SystemTableQueryHandler handler = new SystemTableQueryHandler( + Map.of(ServerPropertiesTableDescriptor.TABLE_NAME, provider), + Map.of(ServerPropertiesTableDescriptor.TABLE_NAME, new ServerPropertiesTableDescriptor()), + new ScanQueryEngine(), + authorizerMapper + ); + final ScanQuery query = Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource(ServerPropertiesTableDescriptor.TABLE_NAME)) + .eternityInterval() + .virtualColumns( + new ExpressionVirtualColumn( + "server", + "'alias'", + ColumnType.STRING, + ExprMacroTable.nil() + ) + ) + .filters(new SelectorDimFilter("server", "alias", null)) + .columns("property", "value") + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .build(); + + final List result = handler.createRunner( + query, + new AuthenticationResult("alice", AuthConfig.ALLOW_ALL_NAME, null, null), + true + ).run(QueryPlus.wrap(query), ResponseContext.createEmpty()).toList(); + + Assertions.assertEquals(1, result.size()); + Assertions.assertEquals(List.of(List.of("druid.test.visible", "visible")), result.get(0).getEvents()); + } + + /** A node-local handler rejects a composite root instead of casting it to a system-table datasource. */ + @Test + public void testRejectsCompositeDataSource() + { + final SystemTableQueryHandler handler = new SystemTableQueryHandler( + Map.of(), + Map.of(), + new ScanQueryEngine(), + new AuthorizerMapper(Map.of()) + ); + final ScanQuery query = Druids.newScanQueryBuilder() + .dataSource( + FilteredDataSource.create( + new SystemTableDataSource("server_properties"), + new SelectorDimFilter("property", "druid.host", null) + ) + ) + .eternityInterval() + .build(); + + Assertions.assertThrows( + BadQueryContextException.class, + () -> handler.createRunner( + query, + new AuthenticationResult("alice", "allow", "external", null), + true + ) + ); + } + +} diff --git a/server/src/test/java/org/apache/druid/server/system/handler/SystemTableBrokerQueryHandlerTest.java b/server/src/test/java/org/apache/druid/server/system/handler/SystemTableBrokerQueryHandlerTest.java new file mode 100644 index 000000000000..a7c8431f7623 --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/system/handler/SystemTableBrokerQueryHandlerTest.java @@ -0,0 +1,101 @@ +/* + * 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.handler; + +import org.apache.druid.query.BadQueryContextException; +import org.apache.druid.query.Druids; +import org.apache.druid.query.FilteredDataSource; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.server.security.AuthenticationResult; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +public class SystemTableBrokerQueryHandlerTest +{ + private static final AuthenticationResult AUTHENTICATION_RESULT = + new AuthenticationResult("user", "authorizer", "authenticator", null); + + /** {@code X-Druid-Native-Query-Route: local} selects the Broker's local system-table handler. */ + @Test + public void testLocalRouteUsesLocalHandler() + { + final SystemTableQueryClient queryClient = Mockito.mock(SystemTableQueryClient.class); + final SystemTableQueryHandler localHandler = Mockito.mock(SystemTableQueryHandler.class); + final SystemTableBrokerQueryHandler brokerHandler = new SystemTableBrokerQueryHandler(queryClient, localHandler); + final ScanQuery query = query(); + final QueryRunner expectedRunner = Mockito.mock(QueryRunner.class); + Mockito.when(localHandler.createRunner(query, AUTHENTICATION_RESULT, true)).thenReturn(expectedRunner); + + Assertions.assertSame(expectedRunner, brokerHandler.createRunner(query, AUTHENTICATION_RESULT, true)); + Mockito.verifyNoInteractions(queryClient); + } + + /** A request without the local route uses distributed Broker fanout. */ + @Test + public void testDefaultRouteUsesQueryClient() + { + final SystemTableQueryClient queryClient = Mockito.mock(SystemTableQueryClient.class); + final SystemTableQueryHandler localHandler = Mockito.mock(SystemTableQueryHandler.class); + final SystemTableBrokerQueryHandler brokerHandler = new SystemTableBrokerQueryHandler(queryClient, localHandler); + final ScanQuery query = query(); + final QueryRunner expectedRunner = Mockito.mock(QueryRunner.class); + Mockito.when(queryClient.createRunner(query, AUTHENTICATION_RESULT, false)).thenReturn(expectedRunner); + + Assertions.assertSame(expectedRunner, brokerHandler.createRunner(query, AUTHENTICATION_RESULT, false)); + Mockito.verifyNoInteractions(localHandler); + } + + /** A local route cannot bypass recursive resolution for a system table below a composite datasource. */ + @Test + public void testLocalRouteRejectsCompositeDataSource() + { + final SystemTableQueryClient queryClient = Mockito.mock(SystemTableQueryClient.class); + final SystemTableQueryHandler localHandler = Mockito.mock(SystemTableQueryHandler.class); + final SystemTableBrokerQueryHandler brokerHandler = new SystemTableBrokerQueryHandler(queryClient, localHandler); + final ScanQuery query = Druids.newScanQueryBuilder() + .dataSource( + FilteredDataSource.create( + new SystemTableDataSource("server_properties"), + new SelectorDimFilter("property", "druid.host", null) + ) + ) + .eternityInterval() + .build(); + + Assertions.assertThrows( + BadQueryContextException.class, + () -> brokerHandler.createRunner(query, AUTHENTICATION_RESULT, true) + ); + Mockito.verifyNoInteractions(queryClient, localHandler); + } + + private static ScanQuery query() + { + return Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource("server_properties")) + .eternityInterval() + .build(); + } +} diff --git a/server/src/test/java/org/apache/druid/server/system/handler/SystemTableNodeLocatorTest.java b/server/src/test/java/org/apache/druid/server/system/handler/SystemTableNodeLocatorTest.java new file mode 100644 index 000000000000..56655a6043b1 --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/system/handler/SystemTableNodeLocatorTest.java @@ -0,0 +1,198 @@ +/* + * 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.handler; + +import com.google.common.util.concurrent.Futures; +import org.apache.druid.client.DirectDruidClient; +import org.apache.druid.client.coordinator.CoordinatorClient; +import org.apache.druid.discovery.DiscoveryDruidNode; +import org.apache.druid.discovery.DruidNodeDiscovery; +import org.apache.druid.discovery.DruidNodeDiscoveryProvider; +import org.apache.druid.discovery.NodeRole; +import org.apache.druid.query.Druids; +import org.apache.druid.query.Query; +import org.apache.druid.rpc.indexing.OverlordClient; +import org.apache.druid.server.DruidNode; +import org.apache.druid.server.system.table.SystemTableDescriptor; +import org.apache.druid.server.system.table.SystemTableRoutingMode; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.net.URI; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicReference; + +public class SystemTableNodeLocatorTest +{ + /** {@link SystemTableRoutingMode#ALL_NODES} selects every discovered node for the configured role. */ + @Test + public void testAllNodesRouting() + { + final DiscoveryDruidNode first = node(8081); + final DiscoveryDruidNode second = node(8082); + final SystemTableNodeLocator locator = new SystemTableNodeLocator( + discoveryProvider(List.of(first, second)), + Mockito.mock(CoordinatorClient.class), + Mockito.mock(OverlordClient.class) + ); + final SystemTableDescriptor descriptor = descriptor(SystemTableRoutingMode.ALL_NODES); + + final List nodes = locator.locate(descriptor, query()); + + Assertions.assertEquals(2, nodes.size()); + } + + /** {@link SystemTableRoutingMode#LEADER_ONLY} excludes healthy and unavailable Overlord standbys. */ + @Test + public void testLeaderOnlyRouting() + { + final DiscoveryDruidNode standby = node(8081); + final DiscoveryDruidNode leader = node(8082); + final OverlordClient overlordClient = Mockito.mock(OverlordClient.class); + Mockito.when(overlordClient.findCurrentLeader()) + .thenReturn(Futures.immediateFuture(URI.create("http://localhost:8082"))); + final SystemTableNodeLocator locator = new SystemTableNodeLocator( + discoveryProvider(List.of(standby, leader)), + Mockito.mock(CoordinatorClient.class), + overlordClient + ); + + final List nodes = locator.locate( + descriptor(SystemTableRoutingMode.LEADER_ONLY), + query() + ); + + Assertions.assertEquals(1, nodes.size()); + Assertions.assertEquals(leader, nodes.get(0).getDiscoveryNode()); + } + + /** Leader routing is role-driven and therefore also supports a Coordinator locator without special-case logic. */ + @Test + public void testLeaderOnlyRoutingUsesLocatorRegisteredForRole() + { + final DiscoveryDruidNode standby = node(NodeRole.COORDINATOR, 8081); + final DiscoveryDruidNode leader = node(NodeRole.COORDINATOR, 8082); + final CoordinatorClient coordinatorClient = Mockito.mock(CoordinatorClient.class); + Mockito.when(coordinatorClient.findCurrentLeader()) + .thenReturn(Futures.immediateFuture(URI.create("http://localhost:8082"))); + final SystemTableNodeLocator locator = new SystemTableNodeLocator( + discoveryProvider(NodeRole.COORDINATOR, List.of(standby, leader)), + coordinatorClient, + Mockito.mock(OverlordClient.class) + ); + + final List nodes = locator.locate( + descriptor(NodeRole.COORDINATOR, SystemTableRoutingMode.LEADER_ONLY), + query() + ); + + Assertions.assertEquals(1, nodes.size()); + Assertions.assertEquals(leader, nodes.get(0).getDiscoveryNode()); + } + + /** Leader resolution happens before discovery so a newly elected leader is read from a fresh node snapshot. */ + @Test + public void testLeaderOnlyRoutingDiscoversNodesAfterResolvingLeader() + { + final DiscoveryDruidNode leader = node(8082); + final AtomicReference> discoveredNodes = new AtomicReference<>(Collections.emptyList()); + final DruidNodeDiscovery discovery = Mockito.mock(DruidNodeDiscovery.class); + Mockito.when(discovery.getAllNodes()).thenAnswer(ignored -> discoveredNodes.get()); + final DruidNodeDiscoveryProvider provider = Mockito.mock(DruidNodeDiscoveryProvider.class); + Mockito.when(provider.getForNodeRole(NodeRole.OVERLORD)).thenReturn(discovery); + final OverlordClient overlordClient = Mockito.mock(OverlordClient.class); + Mockito.when(overlordClient.findCurrentLeader()).thenAnswer(ignored -> { + discoveredNodes.set(List.of(leader)); + return Futures.immediateFuture(URI.create("http://localhost:8082")); + }); + final SystemTableNodeLocator locator = new SystemTableNodeLocator( + provider, + Mockito.mock(CoordinatorClient.class), + overlordClient + ); + + final List nodes = locator.locate( + descriptor(SystemTableRoutingMode.LEADER_ONLY), + query() + ); + + Assertions.assertEquals(List.of(leader), nodes.stream().map(SystemTableNode::getDiscoveryNode).toList()); + } + + private static SystemTableDescriptor descriptor(final SystemTableRoutingMode routingMode) + { + return descriptor(NodeRole.OVERLORD, routingMode); + } + + private static SystemTableDescriptor descriptor( + final NodeRole nodeRole, + final SystemTableRoutingMode routingMode + ) + { + final SystemTableDescriptor descriptor = Mockito.mock(SystemTableDescriptor.class); + Mockito.when(descriptor.getNodeRoles()).thenReturn(Set.of(nodeRole)); + Mockito.when(descriptor.getRoutingMode()).thenReturn(routingMode); + return descriptor; + } + + private static DruidNodeDiscoveryProvider discoveryProvider(final List nodes) + { + return discoveryProvider(NodeRole.OVERLORD, nodes); + } + + private static DruidNodeDiscoveryProvider discoveryProvider( + final NodeRole nodeRole, + final List nodes + ) + { + final DruidNodeDiscovery discovery = Mockito.mock(DruidNodeDiscovery.class); + Mockito.when(discovery.getAllNodes()).thenReturn(nodes); + final DruidNodeDiscoveryProvider provider = Mockito.mock(DruidNodeDiscoveryProvider.class); + Mockito.when(provider.getForNodeRole(nodeRole)).thenReturn(discovery); + return provider; + } + + private static DiscoveryDruidNode node(final int port) + { + return node(NodeRole.OVERLORD, port); + } + + private static DiscoveryDruidNode node(final NodeRole nodeRole, final int port) + { + return new DiscoveryDruidNode( + new DruidNode(nodeRole.getJsonName(), "localhost", false, port, null, true, false), + nodeRole, + Collections.emptyMap() + ); + } + + private static Query query() + { + return Druids.newScanQueryBuilder() + .dataSource("test") + .eternityInterval() + .context(Map.of(DirectDruidClient.QUERY_FAIL_TIME, Long.MAX_VALUE)) + .build(); + } +} 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 new file mode 100644 index 000000000000..333f7ffef4fc --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/system/handler/SystemTableQueryClientTest.java @@ -0,0 +1,1337 @@ +/* + * 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.handler; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.util.concurrent.SettableFuture; +import org.apache.druid.client.DirectDruidClient; +import org.apache.druid.client.DirectDruidClientFactory; +import org.apache.druid.client.DruidServer; +import org.apache.druid.discovery.DiscoveryDruidNode; +import org.apache.druid.discovery.NodeRole; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.java.util.common.JodaUtils; +import org.apache.druid.java.util.common.guava.LazySequence; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.java.util.common.guava.Sequences; +import org.apache.druid.java.util.common.guava.Yielder; +import org.apache.druid.java.util.common.guava.Yielders; +import org.apache.druid.java.util.http.client.HttpClient; +import org.apache.druid.java.util.http.client.Request; +import org.apache.druid.java.util.http.client.response.HttpResponseHandler; +import org.apache.druid.math.expr.ExprMacroTable; +import org.apache.druid.query.DataSource; +import org.apache.druid.query.Druids; +import org.apache.druid.query.FilteredDataSource; +import org.apache.druid.query.InlineDataSource; +import org.apache.druid.query.JoinAlgorithm; +import org.apache.druid.query.JoinDataSource; +import org.apache.druid.query.QueryCapacityExceededException; +import org.apache.druid.query.QueryContexts; +import org.apache.druid.query.QueryDataSource; +import org.apache.druid.query.QueryException; +import org.apache.druid.query.QueryInterruptedException; +import org.apache.druid.query.QueryPlus; +import org.apache.druid.query.QueryRunner; +import org.apache.druid.query.QueryRunnerTestHelper; +import org.apache.druid.query.QuerySegmentWalker; +import org.apache.druid.query.QueryTimeoutException; +import org.apache.druid.query.QueryUnsupportedException; +import org.apache.druid.query.ResourceLimitExceededException; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.query.context.ResponseContext; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.query.operator.ScanOperatorFactory; +import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.scan.ScanResultValue; +import org.apache.druid.query.spec.LegacySegmentSpec; +import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.join.JoinType; +import org.apache.druid.segment.virtual.ExpressionVirtualColumn; +import org.apache.druid.server.DruidNode; +import org.apache.druid.server.QueryScheduler; +import org.apache.druid.server.QueryStackTests; +import org.apache.druid.server.metrics.NoopServiceEmitter; +import org.apache.druid.server.security.AllowAllAuthorizer; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.server.security.Escalator; +import org.apache.druid.server.security.ForbiddenException; +import org.apache.druid.server.security.NoopEscalator; +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.joda.time.Duration; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.ArgumentMatchers; +import org.mockito.Mockito; + +import javax.annotation.Nullable; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.net.ConnectException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.stream.Stream; + +public class SystemTableQueryClientTest +{ + @RegisterExtension + public static final QueryStackTests.ConglomerateExtension CONGLOMERATE = + new QueryStackTests.ConglomerateExtension(); + + private static final AuthenticationResult AUTHENTICATION_RESULT = + new AuthenticationResult("test", "allow", "test", null); + + @Test + public void testNodeRowsAreLazyAndFlattenScanResultBatches() + { + final AtomicInteger runnerCalls = new AtomicInteger(); + final ScanQuery nodeQuery = Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource("test")) + .eternityInterval() + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .build(); + final QueryRunner nodeRunner = (queryPlus, responseContext) -> { + runnerCalls.incrementAndGet(); + final Sequence result = Sequences.simple( + List.of( + new ScanResultValue( + null, + List.of("task_id"), + List.of(List.of("task-a"), List.of("task-b")) + ), + new ScanResultValue( + null, + List.of("task_id"), + List.of((Object) new Object[]{"task-c"}) + ) + ) + ); + return result; + }; + + final Iterable rows = SystemTableQueryClient.scanNodeRows( + List.of(nodeRunner), + nodeQuery, + ResponseContext.createEmpty() + ); + + Assertions.assertEquals(0, runnerCalls.get()); + + final Iterator iterator = rows.iterator(); + Assertions.assertEquals(1, runnerCalls.get()); + Assertions.assertArrayEquals(new Object[]{"task-a"}, iterator.next()); + + final List remainingRows = new ArrayList<>(); + iterator.forEachRemaining(remainingRows::add); + Assertions.assertEquals(2, remainingRows.size()); + Assertions.assertArrayEquals(new Object[]{"task-b"}, remainingRows.get(0)); + Assertions.assertArrayEquals(new Object[]{"task-c"}, remainingRows.get(1)); + Assertions.assertThrows(IllegalStateException.class, rows::iterator); + } + + /** + * Native query transport for: + * + * 
{@code
+   * SELECT property, value FROM sys.server_properties
+   * }
+ * + * A ready node must yield rows without waiting for a later node's response sequence to initialize. + */ + @Test + public void testReadyNodeDoesNotWaitForLaterNodeInitialization() throws Exception + { + final CompletableFuture secondNodeInitialization = new CompletableFuture<>(); + final CountDownLatch releaseSecondNode = new CountDownLatch(1); + final AtomicReference> rowIterator = new AtomicReference<>(); + final ExecutorService consumerExecutor = Executors.newSingleThreadExecutor(); + final ScanQuery nodeQuery = Druids.newScanQueryBuilder() + .dataSource( + new SystemTableDataSource( + ServerPropertiesTableDescriptor.TABLE_NAME + ) + ) + .eternityInterval() + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .build(); + final QueryRunner readyRunner = (queryPlus, responseContext) -> Sequences.simple( + List.of(scanResult(new Object[]{"first"}, new Object[]{"first-remainder"})) + ); + final QueryRunner delayedRunner = (queryPlus, responseContext) -> new LazySequence<>(() -> { + secondNodeInitialization.complete(null); + try { + releaseSecondNode.await(); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + return Sequences.simple(List.of(scanResult(new Object[]{"second"}))); + }); + final Iterable rows = SystemTableQueryClient.scanNodeRows( + List.of(readyRunner, delayedRunner), + nodeQuery, + ResponseContext.createEmpty() + ); + + try { + final CompletableFuture firstRow = CompletableFuture.supplyAsync( + () -> { + final Iterator iterator = rows.iterator(); + rowIterator.set(iterator); + return iterator.next(); + }, + consumerExecutor + ); + + CompletableFuture.anyOf(firstRow, secondNodeInitialization).get(10, TimeUnit.SECONDS); + Assertions.assertTrue(firstRow.isDone()); + Assertions.assertFalse(secondNodeInitialization.isDone()); + Assertions.assertArrayEquals(new Object[]{"first"}, firstRow.get()); + + releaseSecondNode.countDown(); + final List remainingRows = CompletableFuture.supplyAsync( + () -> { + final List result = new ArrayList<>(); + rowIterator.get().forEachRemaining(result::add); + return result; + }, + consumerExecutor + ).get(10, TimeUnit.SECONDS); + Assertions.assertEquals(2, remainingRows.size()); + Assertions.assertArrayEquals(new Object[]{"first-remainder"}, remainingRows.get(0)); + Assertions.assertArrayEquals(new Object[]{"second"}, remainingRows.get(1)); + } + finally { + releaseSecondNode.countDown(); + consumerExecutor.shutdownNow(); + } + } + + /** A node selected at the Broker's own address executes in-process without creating an HTTP client. */ + @Test + public void testBrokerNodeExecutesLocallyWithoutHttp() + { + final SystemTableDescriptor descriptor = new TestSystemTableDescriptor(); + final DiscoveryDruidNode brokerNode = testNode(); + final SystemTableNodeLocator nodeLocator = Mockito.mock(SystemTableNodeLocator.class); + final SystemTableNode node = new SystemTableNode(brokerNode); + node.addNodeRole(brokerNode.getNodeRole()); + Mockito.when(nodeLocator.locate(Mockito.eq(descriptor), ArgumentMatchers.any())).thenReturn(List.of(node)); + + final DirectDruidClientFactory directClientFactory = Mockito.mock(DirectDruidClientFactory.class); + final SystemTableQueryHandler localQueryHandler = Mockito.mock(SystemTableQueryHandler.class); + final AuthenticationResult escalatedAuthenticationResult = + new AuthenticationResult("system", "allow", "system", null); + final Escalator escalator = Mockito.mock(Escalator.class); + Mockito.when(escalator.createEscalatedAuthenticationResult()).thenReturn(escalatedAuthenticationResult); + Mockito.doAnswer( + ignored -> (QueryRunner) (queryPlus, responseContext) -> + Sequences.simple(List.of(scanResult(new Object[]{"local"}))) + ).when(localQueryHandler).createRunner( + ArgumentMatchers.any(), + Mockito.same(escalatedAuthenticationResult), + Mockito.eq(true) + ); + + final QuerySegmentWalker querySegmentWalker = Mockito.mock(QuerySegmentWalker.class); + Mockito.when(querySegmentWalker.getQueryRunnerForIntervals(ArgumentMatchers.any(), ArgumentMatchers.any())) + .thenAnswer(ignored -> passthroughRunner(descriptor.getRowSignature())); + final SystemTableQueryClient client = new SystemTableQueryClient( + nodeLocator, + directClientFactory, + Mockito.mock(QueryScheduler.class), + querySegmentWalker, + Map.of(descriptor.getTableName(), descriptor), + new AuthorizerMapper(Map.of("allow", new AllowAllAuthorizer(null))), + localQueryHandler, + escalator, + brokerNode.getDruidNode() + ); + final ScanQuery query = query(descriptor, Collections.emptyMap()); + + final List results = client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList(); + + Assertions.assertEquals(1, results.size()); + final List events = (List) results.get(0).getEvents(); + Assertions.assertEquals(1, events.size()); + Assertions.assertArrayEquals(new Object[]{"local"}, (Object[]) events.get(0)); + Mockito.verifyNoInteractions(directClientFactory); + Mockito.verify(localQueryHandler).createRunner( + ArgumentMatchers.any(), + Mockito.same(escalatedAuthenticationResult), + Mockito.eq(true) + ); + } + + /** Delayed node responses verify concurrent fanout through the real {@link DirectDruidClient} transport path. */ + @Test + public void testDirectClientsStartRequestsBeforeWaitingForDelayedResponses() throws Exception + { + final ObjectMapper objectMapper = new DefaultObjectMapper(); + final DelayedHttpClient httpClient = new DelayedHttpClient(2); + final ScheduledExecutorService cancellationExecutor = Executors.newSingleThreadScheduledExecutor(); + final ExecutorService consumerExecutor = Executors.newSingleThreadExecutor(); + try { + final ScanQuery nodeQuery = Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource("test")) + .eternityInterval() + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .context(Map.of(DirectDruidClient.QUERY_FAIL_TIME, Long.MAX_VALUE)) + .build(); + final List> nodeRunners = List.of( + directClient(objectMapper, httpClient, "localhost:8081", cancellationExecutor), + directClient(objectMapper, httpClient, "localhost:8082", cancellationExecutor) + ); + final Iterable rows = SystemTableQueryClient.scanNodeRows( + nodeRunners, + nodeQuery, + DirectDruidClient.makeResponseContextForQuery() + ); + + final CompletableFuture> consumedRows = CompletableFuture.supplyAsync( + () -> { + final List result = new ArrayList<>(); + rows.forEach(result::add); + return result; + }, + consumerExecutor + ); + + Assertions.assertTrue(httpClient.awaitRequests()); + httpClient.respond( + 0, + objectMapper.writeValueAsBytes(List.of(scanResult(new Object[]{"first"}))) + ); + httpClient.respond( + 1, + objectMapper.writeValueAsBytes(List.of(scanResult(new Object[]{"second"}))) + ); + + final List result = consumedRows.get(10, TimeUnit.SECONDS); + Assertions.assertEquals(2, result.size()); + Assertions.assertArrayEquals(new Object[]{"first"}, result.get(0)); + Assertions.assertArrayEquals(new Object[]{"second"}, result.get(1)); + } + finally { + consumerExecutor.shutdownNow(); + cancellationExecutor.shutdownNow(); + } + } + + /** Closing a partial result closes its initialized transport and cancels every concurrently started node query. */ + @Test + public void testClosingPartialResultCleansUpInitializedDirectClientTransport() throws Exception + { + final ObjectMapper objectMapper = new DefaultObjectMapper(); + final DelayedHttpClient httpClient = new DelayedHttpClient(2); + final ScheduledExecutorService cancellationExecutor = Executors.newSingleThreadScheduledExecutor(); + final ExecutorService consumerExecutor = Executors.newSingleThreadExecutor(); + final QueryScheduler queryScheduler = Mockito.mock(QueryScheduler.class); + final SystemTableDescriptor descriptor = new TestSystemTableDescriptor(); + try { + final SystemTableQueryClient client = makeDirectClientBackedClient( + descriptor, + List.of(testNode(8081), testNode(8082)), + objectMapper, + httpClient, + cancellationExecutor, + queryScheduler + ); + final ScanQuery query = query(descriptor, Collections.emptyMap()); + final CompletableFuture> yielderFuture = CompletableFuture.supplyAsync( + () -> Yielders.each( + client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + ), + consumerExecutor + ); + + Assertions.assertTrue(httpClient.awaitRequests()); + final CloseTrackingInputStream firstResponse = httpClient.respond( + 0, + objectMapper.writeValueAsBytes( + List.of(scanResult(new Object[]{"first"}), scanResult(new Object[]{"first-extra"})) + ) + ); + httpClient.respond( + 1, + objectMapper.writeValueAsBytes( + List.of(scanResult(new Object[]{"second"}), scanResult(new Object[]{"second-extra"})) + ) + ); + final Yielder yielder = yielderFuture.get(10, TimeUnit.SECONDS); + Assertions.assertFalse(yielder.isDone()); + + yielder.close(); + + Assertions.assertTrue(firstResponse.isClosed()); + // The second response sequence was deliberately not initialized. Its remote work is stopped by query-id + // cancellation below rather than by opening and closing its response stream during Broker cleanup. + Mockito.verify(queryScheduler, Mockito.times(2)).cancelQuery( + ArgumentMatchers.argThat(queryId -> queryId.startsWith(SystemTableDataSource.NODE_QUERY_ID_PREFIX)) + ); + } + finally { + consumerExecutor.shutdownNow(); + cancellationExecutor.shutdownNow(); + } + } + + /** + * Native query form of: + * + * 
{@code
+   * SELECT COUNT(*)
+   * FROM
+   * (
+   *   SELECT
+   *     task_id,
+   *     COUNT(*) AS task_count
+   *   FROM sys.tasks
+   *   GROUP BY task_id
+   * )
+   * WHERE task_count > 0
+   * }
+ */ + @Test + public void testResolvesSystemTableInsideQueryDataSource() + { + final SystemTableDescriptor descriptor = new TestSystemTableDescriptor(); + final ScanQuery innerQuery = query(descriptor, Collections.emptyMap()); + + assertWrappedSystemTableIsResolved( + descriptor, + new QueryDataSource(innerQuery) + ); + } + + /** + * Native query form of: + * + * 
{@code
+   * SELECT task_id
+   * FROM sys.tasks
+   * WHERE type = 'index_parallel'
+   * }
+ */ + @Test + public void testResolvesSystemTableInsideFilteredDataSource() + { + final SystemTableDescriptor descriptor = new TestSystemTableDescriptor(); + + assertWrappedSystemTableIsResolved( + descriptor, + FilteredDataSource.create( + new SystemTableDataSource(descriptor.getTableName()), + new SelectorDimFilter("type", "index_parallel", null) + ) + ); + } + + /** + * Native query form of: + * + * 
{@code
+   * SELECT l.value
+   * FROM sys.server_properties AS l
+   * JOIN sys.server_properties AS r ON l.property = r.property
+   * WHERE r.value = 'right'   -- or the unprefixed l.value = 'left'
+   * }
+ * + * A root filter belongs to the join result, whether or not it uses the right-side join prefix. It must remain on + * the Broker instead of being copied into either node-local system table scan. + */ + @ParameterizedTest + @ValueSource(strings = {"j.value", "value"}) + public void testJoinRootFilterIsNotPushedIntoSystemTableLeaves(final String filterColumn) + { + final SystemTableDescriptor descriptor = new ServerPropertiesTableDescriptor(); + final List nodeQueries = captureNodeQueries( + descriptor, + selfJoin(descriptor), + new SelectorDimFilter(filterColumn, "match", null) + ); + + Assertions.assertEquals(2, nodeQueries.size()); + Assertions.assertNull(nodeQueries.get(0).getFilter()); + Assertions.assertNull(nodeQueries.get(1).getFilter()); + } + + /** + * Native query form of two filtered subqueries joined together: + * + * 
{@code
+   * SELECT l.value
+   * FROM (SELECT property, value FROM sys.server_properties WHERE value = 'left') AS l
+   * JOIN (SELECT property, value FROM sys.server_properties WHERE value = 'right') AS r
+   *   ON l.property = r.property
+   * }
+ * + * Each subquery filter is owned by its own system table leaf and remains eligible for node pushdown. + */ + @Test + public void testQueryDataSourceFiltersArePushedOnlyIntoTheirOwningLeaves() + { + final SystemTableDescriptor descriptor = new ServerPropertiesTableDescriptor(); + final SelectorDimFilter leftFilter = new SelectorDimFilter("value", "left", null); + final SelectorDimFilter rightFilter = new SelectorDimFilter("value", "right", null); + final DataSource dataSource = JoinDataSource.create( + new QueryDataSource( + Druids.ScanQueryBuilder.copy(query(descriptor, Collections.emptyMap())).filters(leftFilter).build() + ), + new QueryDataSource( + Druids.ScanQueryBuilder.copy(query(descriptor, Collections.emptyMap())).filters(rightFilter).build() + ), + "j.", + "property == \"j.property\"", + JoinType.INNER, + null, + ExprMacroTable.nil(), + null, + JoinAlgorithm.BROADCAST + ); + + final List nodeQueries = captureNodeQueries(descriptor, dataSource, null); + + Assertions.assertEquals(2, nodeQueries.size()); + Assertions.assertEquals(leftFilter, nodeQueries.get(0).getFilter()); + Assertions.assertEquals(rightFilter, nodeQueries.get(1).getFilter()); + } + + /** A filter pushed into a node scan carries the virtual column referenced by that filter. */ + @Test + public void testPushedFilterCarriesVirtualColumns() + { + final SystemTableDescriptor descriptor = new ServerPropertiesTableDescriptor(); + final AtomicReference capturedNodeQuery = new AtomicReference<>(); + final ExpressionVirtualColumn virtualColumn = new ExpressionVirtualColumn( + "upper_property", + "upper(property)", + ColumnType.STRING, + ExprMacroTable.nil() + ); + final SystemTableQueryClient client = makeClient( + descriptor, + List.of(testNode()), + (queryPlus, responseContext) -> { + capturedNodeQuery.set((ScanQuery) queryPlus.getQuery()); + return Sequences.empty(); + } + ); + final ScanQuery query = Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource(descriptor.getTableName())) + .eternityInterval() + .virtualColumns(VirtualColumns.create(virtualColumn)) + .filters(new SelectorDimFilter("upper_property", "MATCH", null)) + .build(); + + client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList(); + + Assertions.assertEquals( + VirtualColumns.create(virtualColumn), + capturedNodeQuery.get().getVirtualColumns() + ); + } + + /** A window leaf filter also carries its leaf virtual columns into the node scan. */ + @Test + public void testWindowLeafFilterCarriesVirtualColumns() + { + final SystemTableDescriptor descriptor = new ServerPropertiesTableDescriptor(); + final AtomicReference capturedNodeQuery = new AtomicReference<>(); + final ExpressionVirtualColumn virtualColumn = new ExpressionVirtualColumn( + "upper_property", + "upper(property)", + ColumnType.STRING, + ExprMacroTable.nil() + ); + final SystemTableQueryClient client = makeClient( + descriptor, + List.of(testNode()), + (queryPlus, responseContext) -> { + capturedNodeQuery.set((ScanQuery) queryPlus.getQuery()); + return Sequences.empty(); + } + ); + final WindowOperatorQuery query = new WindowOperatorQuery( + new SystemTableDataSource(descriptor.getTableName()), + new LegacySegmentSpec(Intervals.ETERNITY), + Collections.emptyMap(), + descriptor.getRowSignature(), + Collections.emptyList(), + List.of( + new ScanOperatorFactory( + null, + new SelectorDimFilter("upper_property", "MATCH", null), + null, + null, + VirtualColumns.create(virtualColumn), + null + ) + ) + ); + + client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList(); + + Assertions.assertEquals( + VirtualColumns.create(virtualColumn), + capturedNodeQuery.get().getVirtualColumns() + ); + } + + /** + * A window query is planned over a Scan subquery, which {@link WindowOperatorQuery} converts into a leaf scan. For a + * Scan without virtual columns the leaf stores null virtual columns; the filter is still pushed into the node scan, + * with no virtual columns. + */ + @Test + public void testWindowLeafFilterWithoutVirtualColumns() + { + final SystemTableDescriptor descriptor = new ServerPropertiesTableDescriptor(); + final AtomicReference capturedNodeQuery = new AtomicReference<>(); + final SystemTableQueryClient client = makeClient( + descriptor, + List.of(testNode()), + (queryPlus, responseContext) -> { + capturedNodeQuery.set((ScanQuery) queryPlus.getQuery()); + return Sequences.empty(); + } + ); + final SelectorDimFilter filter = new SelectorDimFilter("property", "druid.host", null); + final WindowOperatorQuery query = new WindowOperatorQuery( + new QueryDataSource( + Druids.ScanQueryBuilder.copy(query(descriptor, Collections.emptyMap())).filters(filter).build() + ), + new LegacySegmentSpec(Intervals.ETERNITY), + Collections.emptyMap(), + descriptor.getRowSignature(), + Collections.emptyList(), + null + ); + Assertions.assertNull(((ScanOperatorFactory) query.getLeafOperators().get(0)).getVirtualColumns()); + + client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList(); + + Assertions.assertEquals(filter, capturedNodeQuery.get().getFilter()); + Assertions.assertEquals(VirtualColumns.EMPTY, capturedNodeQuery.get().getVirtualColumns()); + } + + /** The private node transport always uses plain scan results even when the user query requests by-segment results. */ + @Test + public void testNodeQueryDisablesBySegmentContext() + { + final SystemTableDescriptor descriptor = new ServerPropertiesTableDescriptor(); + final AtomicReference capturedNodeQuery = new AtomicReference<>(); + final SystemTableQueryClient client = makeClient( + descriptor, + List.of(testNode()), + (queryPlus, responseContext) -> { + capturedNodeQuery.set((ScanQuery) queryPlus.getQuery()); + return Sequences.empty(); + } + ); + final ScanQuery query = query(descriptor, Map.of(QueryContexts.BY_SEGMENT_KEY, true)); + + client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList(); + + Assertions.assertFalse(capturedNodeQuery.get().context().isBySegment()); + } + + /** A node query with {@code timeout = 0} retains Druid's no-timeout semantics. */ + @Test + public void testNoTimeoutNodeQueryDoesNotGetImmediateDeadline() + { + final AtomicReference capturedNodeQuery = new AtomicReference<>(); + final SystemTableDescriptor descriptor = new TestSystemTableDescriptor(); + final SystemTableQueryClient client = makeClient( + descriptor, + List.of(testNode()), + (queryPlus, responseContext) -> { + capturedNodeQuery.set((ScanQuery) queryPlus.getQuery()); + return Sequences.empty(); + } + ); + final ScanQuery query = query( + descriptor, + Map.of( + QueryContexts.TIMEOUT_KEY, + QueryContexts.NO_TIMEOUT, + DirectDruidClient.QUERY_FAIL_TIME, + 0L + ) + ); + + client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList(); + + Assertions.assertEquals( + JodaUtils.MAX_INSTANT, + capturedNodeQuery.get().context().getLong(DirectDruidClient.QUERY_FAIL_TIME) + ); + } + + /** A node made unavailable by a network failure contributes an {@code error_message} row. */ + @Test + public void testServerPropertiesNodeFailureBecomesErrorRow() + { + final ServerPropertiesTableDescriptor descriptor = new ServerPropertiesTableDescriptor(); + final DiscoveryDruidNode healthyNode = testNode(8081); + final DiscoveryDruidNode failedNode = testNode(8082); + final Object[] healthyRow = new Object[]{ + healthyNode.getDruidNode().getHostAndPortToUse(), + healthyNode.getDruidNode().getServiceName(), + "[broker]", + "property", + "value", + null + }; + final AtomicInteger nodeNumber = new AtomicInteger(); + final SystemTableQueryClient client = makeClient( + descriptor, + List.of(healthyNode, failedNode), + ignored -> { + if (nodeNumber.getAndIncrement() == 0) { + return (queryPlus, responseContext) -> Sequences.simple( + List.of( + new ScanResultValue( + null, + descriptor.getRowSignature().getColumnNames(), + List.of((Object) healthyRow) + ) + ) + ); + } + return (queryPlus, responseContext) -> { + throw new QueryInterruptedException(new ConnectException("node unavailable")); + }; + } + ); + final ScanQuery query = query(descriptor, Collections.emptyMap()); + + final List results = client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList(); + + Assertions.assertEquals(2, results.size()); + Assertions.assertArrayEquals(healthyRow, (Object[]) ((List) results.get(0).getEvents()).get(0)); + Assertions.assertArrayEquals( + new Object[]{ + failedNode.getDruidNode().getHostAndPortToUse(), + failedNode.getDruidNode().getServiceName(), + "[broker]", + null, + null, + "node unavailable" + }, + (Object[]) ((List) results.get(1).getEvents()).get(0) + ); + } + + /** No discovered nodes produce an empty {@code sys.server_properties} result, as in the Bindable path. */ + @Test + public void testServerPropertiesWithNoDiscoveredNodesReturnsEmptyResult() + { + final ServerPropertiesTableDescriptor descriptor = new ServerPropertiesTableDescriptor(); + final SystemTableQueryClient client = makeClient( + descriptor, + Collections.emptyList(), + (queryPlus, responseContext) -> Sequences.empty() + ); + final ScanQuery query = query(descriptor, Collections.emptyMap()); + + final List results = client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList(); + + Assertions.assertTrue(results.isEmpty()); + } + + /** + * Timeouts, cancellations, interruptions, and other failures that do not indicate an unreachable node terminate + * {@code sys.server_properties} instead of becoming an error row. + */ + @ParameterizedTest + @MethodSource("nonAvailabilityFailures") + public void testServerPropertiesNonAvailabilityFailureIsPropagated(final RuntimeException failure) + { + assertServerPropertiesQueryFailureIsPropagated(failure); + } + + private static Stream nonAvailabilityFailures() + { + return Stream.of( + new QueryTimeoutException("node timed out"), + new QueryInterruptedException( + QueryException.QUERY_CANCELED_ERROR_CODE, + "node cancelled", + QueryInterruptedException.class.getName(), + "localhost" + ), + new QueryInterruptedException(new InterruptedException()), + new ForbiddenException("forbidden"), + new QueryCapacityExceededException(1), + new ResourceLimitExceededException("resource limit"), + new QueryUnsupportedException("unsupported"), + new QueryInterruptedException(new IOException("malformed response")), + new RuntimeException("unexpected") + ); + } + + /** A timeout raised while reading a node response is also propagated. */ + @Test + public void testServerPropertiesMidStreamTimeoutIsPropagated() + { + final QueryTimeoutException failure = new QueryTimeoutException("node timed out while streaming"); + final ServerPropertiesTableDescriptor descriptor = new ServerPropertiesTableDescriptor(); + final SystemTableQueryClient client = makeClient( + descriptor, + List.of(testNode()), + (queryPlus, responseContext) -> Sequences.simple(List.of(scanResult(new Object[]{"ignored"}))) + .map(result -> { + throw failure; + }) + ); + final ScanQuery query = query(descriptor, Collections.emptyMap()); + + final QueryTimeoutException thrown = Assertions.assertThrows( + QueryTimeoutException.class, + () -> client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList() + ); + Assertions.assertSame(failure, thrown); + } + + /** A leader-only node scan re-resolves the leader once when the selected node loses leadership before emitting rows. */ + @Test + public void testLeaderOnlyQueryRetriesAfterLeadershipChange() + { + final SystemTableDescriptor descriptor = new LeaderOnlySystemTableDescriptor(); + final SystemTableNode oldLeader = systemTableNode(testNode(8081)); + final SystemTableNode newLeader = systemTableNode(testNode(8082)); + final SystemTableNodeLocator nodeLocator = Mockito.mock(SystemTableNodeLocator.class); + Mockito.when(nodeLocator.locate(Mockito.eq(descriptor), ArgumentMatchers.any())) + .thenReturn(List.of(oldLeader), List.of(newLeader)); + final DirectDruidClientFactory directClientFactory = Mockito.mock(DirectDruidClientFactory.class); + Mockito.when(directClientFactory.makeDirectClient(ArgumentMatchers.any())).thenAnswer(invocation -> { + final DruidServer server = invocation.getArgument(0); + @SuppressWarnings("unchecked") + final DirectDruidClient directClient = Mockito.mock(DirectDruidClient.class); + if (server.getHost().endsWith(":8081")) { + Mockito.when(directClient.run(ArgumentMatchers.any(), ArgumentMatchers.any())).thenReturn( + Sequences.simple(List.of(scanResult(taskRow("stale")))).map(ignored -> { + throw new QueryInterruptedException(new SystemTableNotLeaderException("overlord")); + }) + ); + } else { + Mockito.when(directClient.run(ArgumentMatchers.any(), ArgumentMatchers.any())).thenReturn( + Sequences.simple(List.of(scanResult(taskRow("current")))) + ); + } + return directClient; + }); + final QuerySegmentWalker querySegmentWalker = Mockito.mock(QuerySegmentWalker.class); + Mockito.when(querySegmentWalker.getQueryRunnerForIntervals(ArgumentMatchers.any(), ArgumentMatchers.any())) + .thenAnswer(ignored -> passthroughRunner(descriptor.getRowSignature())); + final SystemTableQueryClient client = new SystemTableQueryClient( + nodeLocator, + directClientFactory, + Mockito.mock(QueryScheduler.class), + querySegmentWalker, + Map.of(descriptor.getTableName(), descriptor), + new AuthorizerMapper(Map.of("allow", new AllowAllAuthorizer(null))), + Mockito.mock(SystemTableQueryHandler.class), + NoopEscalator.getInstance(), + nonMatchingSelfNode() + ); + final ScanQuery query = query(descriptor, Collections.emptyMap()); + + final List results = client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList(); + + Assertions.assertEquals(1, results.size()); + Assertions.assertArrayEquals( + taskRow("current"), + (Object[]) ((List) results.get(0).getEvents()).get(0) + ); + Mockito.verify(nodeLocator, Mockito.times(2)).locate(Mockito.eq(descriptor), ArgumentMatchers.any()); + } + + private static void assertServerPropertiesQueryFailureIsPropagated(final RuntimeException failure) + { + final ServerPropertiesTableDescriptor descriptor = new ServerPropertiesTableDescriptor(); + final SystemTableQueryClient client = makeClient( + descriptor, + List.of(testNode()), + (queryPlus, responseContext) -> { + throw failure; + } + ); + final ScanQuery query = query(descriptor, Collections.emptyMap()); + + final RuntimeException thrown = Assertions.assertThrows( + failure.getClass(), + () -> client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList() + ); + Assertions.assertSame(failure, thrown); + } + + private static SystemTableQueryClient makeClient( + final SystemTableDescriptor descriptor, + final List discoveryNodes, + final QueryRunner nodeRunner + ) + { + return makeClient(descriptor, discoveryNodes, ignored -> nodeRunner); + } + + private static SystemTableQueryClient makeDirectClientBackedClient( + final SystemTableDescriptor descriptor, + final List discoveryNodes, + final ObjectMapper objectMapper, + final HttpClient httpClient, + final ScheduledExecutorService cancellationExecutor, + final QueryScheduler queryScheduler + ) + { + final SystemTableNodeLocator nodeLocator = Mockito.mock(SystemTableNodeLocator.class); + Mockito.when(nodeLocator.locate(Mockito.eq(descriptor), ArgumentMatchers.any())).thenReturn( + discoveryNodes.stream().map(discoveryNode -> { + final SystemTableNode systemTableNode = new SystemTableNode(discoveryNode); + systemTableNode.addNodeRole(discoveryNode.getNodeRole()); + return systemTableNode; + }).toList() + ); + final DirectDruidClientFactory directClientFactory = Mockito.mock(DirectDruidClientFactory.class); + Mockito.when(directClientFactory.makeDirectClient(ArgumentMatchers.any())).thenAnswer( + invocation -> directClient( + objectMapper, + httpClient, + ((DruidServer) invocation.getArgument(0)).getHost(), + cancellationExecutor + ) + ); + final QuerySegmentWalker querySegmentWalker = Mockito.mock(QuerySegmentWalker.class); + Mockito.when(querySegmentWalker.getQueryRunnerForIntervals(ArgumentMatchers.any(), ArgumentMatchers.any())) + .thenAnswer(ignored -> passthroughRunner(descriptor.getRowSignature())); + return new SystemTableQueryClient( + nodeLocator, + directClientFactory, + queryScheduler, + querySegmentWalker, + Map.of(descriptor.getTableName(), descriptor), + new AuthorizerMapper(Map.of("allow", new AllowAllAuthorizer(null))), + Mockito.mock(SystemTableQueryHandler.class), + NoopEscalator.getInstance(), + nonMatchingSelfNode() + ); + } + + private static SystemTableQueryClient makeClient( + final SystemTableDescriptor descriptor, + final List discoveryNodes, + final Function> nodeRunnerFactory + ) + { + return makeClient( + descriptor, + discoveryNodes, + nodeRunnerFactory, + passthroughRunner(descriptor.getRowSignature()) + ); + } + + private static SystemTableQueryClient makeClient( + final SystemTableDescriptor descriptor, + final List discoveryNodes, + final Function> nodeRunnerFactory, + final QueryRunner queryRunner + ) + { + final SystemTableNodeLocator nodeLocator = Mockito.mock(SystemTableNodeLocator.class); + Mockito.when(nodeLocator.locate(Mockito.eq(descriptor), ArgumentMatchers.any())).thenReturn( + discoveryNodes.stream().map(discoveryNode -> { + final SystemTableNode systemTableNode = new SystemTableNode(discoveryNode); + systemTableNode.addNodeRole(discoveryNode.getNodeRole()); + return systemTableNode; + }).toList() + ); + + final DirectDruidClientFactory directClientFactory = Mockito.mock(DirectDruidClientFactory.class); + Mockito.when(directClientFactory.makeDirectClient(ArgumentMatchers.any())).thenAnswer(invocation -> { + final QueryRunner nodeRunner = nodeRunnerFactory.apply(invocation.getArgument(0)); + @SuppressWarnings("unchecked") + final DirectDruidClient directClient = Mockito.mock(DirectDruidClient.class); + Mockito.when(directClient.run(ArgumentMatchers.any(), ArgumentMatchers.any())).thenAnswer( + runInvocation -> nodeRunner.run( + runInvocation.getArgument(0), + runInvocation.getArgument(1) + ) + ); + return directClient; + }); + + final QuerySegmentWalker querySegmentWalker = Mockito.mock(QuerySegmentWalker.class); + Mockito.when(querySegmentWalker.getQueryRunnerForIntervals(ArgumentMatchers.any(), ArgumentMatchers.any())) + .thenAnswer(invocation -> queryRunner); + + return new SystemTableQueryClient( + nodeLocator, + directClientFactory, + Mockito.mock(QueryScheduler.class), + querySegmentWalker, + Map.of(descriptor.getTableName(), descriptor), + new AuthorizerMapper(Map.of("allow", new AllowAllAuthorizer(null))), + Mockito.mock(SystemTableQueryHandler.class), + NoopEscalator.getInstance(), + nonMatchingSelfNode() + ); + } + + private static DruidNode nonMatchingSelfNode() + { + return new DruidNode("broker-service", "localhost", false, 9082, null, true, false); + } + + private static DataSource selfJoin(final SystemTableDescriptor descriptor) + { + return JoinDataSource.create( + new SystemTableDataSource(descriptor.getTableName()), + new SystemTableDataSource(descriptor.getTableName()), + "j.", + "property == \"j.property\"", + JoinType.INNER, + null, + ExprMacroTable.nil(), + null, + JoinAlgorithm.BROADCAST + ); + } + + private static List captureNodeQueries( + final SystemTableDescriptor descriptor, + final DataSource dataSource, + @Nullable final DimFilter filter + ) + { + final List nodeQueries = new ArrayList<>(); + final SystemTableNodeLocator nodeLocator = Mockito.mock(SystemTableNodeLocator.class); + final DiscoveryDruidNode discoveryNode = testNode(); + final SystemTableNode systemTableNode = new SystemTableNode(discoveryNode); + systemTableNode.addNodeRole(discoveryNode.getNodeRole()); + Mockito.when(nodeLocator.locate(Mockito.eq(descriptor), ArgumentMatchers.any())).thenAnswer(invocation -> { + nodeQueries.add(invocation.getArgument(1)); + return List.of(systemTableNode); + }); + final DirectDruidClientFactory directClientFactory = Mockito.mock(DirectDruidClientFactory.class); + Mockito.when(directClientFactory.makeDirectClient(ArgumentMatchers.any())) + .thenReturn(Mockito.mock(DirectDruidClient.class)); + final QuerySegmentWalker querySegmentWalker = Mockito.mock(QuerySegmentWalker.class); + Mockito.when(querySegmentWalker.getQueryRunnerForIntervals(ArgumentMatchers.any(), ArgumentMatchers.any())) + .thenReturn((queryPlus, responseContext) -> Sequences.empty()); + final SystemTableQueryClient client = new SystemTableQueryClient( + nodeLocator, + directClientFactory, + Mockito.mock(QueryScheduler.class), + querySegmentWalker, + Map.of(descriptor.getTableName(), descriptor), + new AuthorizerMapper(Map.of("allow", new AllowAllAuthorizer(null))), + Mockito.mock(SystemTableQueryHandler.class), + NoopEscalator.getInstance(), + nonMatchingSelfNode() + ); + final ScanQuery query = Druids.newScanQueryBuilder() + .dataSource(dataSource) + .eternityInterval() + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .filters(filter) + .build(); + + client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList(); + return nodeQueries; + } + + private static void assertWrappedSystemTableIsResolved( + final SystemTableDescriptor descriptor, + final DataSource wrappedDataSource + ) + { + final AtomicReference resolvedDataSource = new AtomicReference<>(); + final QueryRunner resolvedQueryRunner = (queryPlus, responseContext) -> { + resolvedDataSource.set(queryPlus.getQuery().getDataSource()); + return Sequences.empty(); + }; + final SystemTableQueryClient client = makeClient( + descriptor, + List.of(testNode()), + ignored -> (queryPlus, responseContext) -> Sequences.simple(List.of(scanResult(new Object[]{"value"}))), + resolvedQueryRunner + ); + final ScanQuery query = Druids.newScanQueryBuilder() + .dataSource(wrappedDataSource) + .eternityInterval() + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .build(); + + client.createRunner(query, AUTHENTICATION_RESULT, false) + .run(QueryPlus.wrap(query), ResponseContext.createEmpty()) + .toList(); + + Assertions.assertNotNull(resolvedDataSource.get()); + Assertions.assertFalse(containsSystemTableDataSource(resolvedDataSource.get())); + } + + private static boolean containsSystemTableDataSource(final DataSource dataSource) + { + if (dataSource instanceof SystemTableDataSource) { + return true; + } + return dataSource.getChildren().stream().anyMatch(SystemTableQueryClientTest::containsSystemTableDataSource); + } + + private static QueryRunner passthroughRunner(final RowSignature rowSignature) + { + return (queryPlus, responseContext) -> { + final InlineDataSource dataSource = (InlineDataSource) queryPlus.getQuery().getDataSource(); + return Sequences.simple(dataSource.getRows()) + .map(row -> new ScanResultValue(null, rowSignature.getColumnNames(), List.of((Object) row))); + }; + } + + private static ScanQuery query(final SystemTableDescriptor descriptor, final Map context) + { + return Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource(descriptor.getTableName())) + .eternityInterval() + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .context(context) + .build(); + } + + private static ScanResultValue scanResult(final Object[]... rows) + { + return new ScanResultValue(null, List.of("value"), List.of(rows)); + } + + private static Object[] taskRow(final String taskId) + { + return new Object[]{taskId, null, null, null, null, null, null, null, 0L, null, null, -1L, -1L, null}; + } + + private static DirectDruidClient directClient( + final ObjectMapper objectMapper, + final HttpClient httpClient, + final String host, + final ScheduledExecutorService cancellationExecutor + ) + { + return new DirectDruidClient<>( + CONGLOMERATE.getConglomerate(), + QueryRunnerTestHelper.NOOP_QUERYWATCHER, + objectMapper, + httpClient, + "http", + host, + new NoopServiceEmitter(), + cancellationExecutor + ); + } + + private static class DelayedHttpClient implements HttpClient + { + private final List> responses; + private final CountDownLatch requestLatch; + private final AtomicInteger requestNumber = new AtomicInteger(); + + private DelayedHttpClient(final int requestCount) + { + responses = new ArrayList<>(requestCount); + for (int i = 0; i < requestCount; i++) { + responses.add(SettableFuture.create()); + } + requestLatch = new CountDownLatch(requestCount); + } + + private boolean awaitRequests() throws InterruptedException + { + return requestLatch.await(10, TimeUnit.SECONDS); + } + + private CloseTrackingInputStream respond(final int request, final byte[] response) + { + final CloseTrackingInputStream inputStream = new CloseTrackingInputStream(response); + responses.get(request).set(inputStream); + return inputStream; + } + + @Override + @SuppressWarnings("unchecked") + public com.google.common.util.concurrent.ListenableFuture go( + final Request request, + final HttpResponseHandler handler, + final Duration readTimeout + ) + { + final SettableFuture response = responses.get(requestNumber.getAndIncrement()); + requestLatch.countDown(); + return (com.google.common.util.concurrent.ListenableFuture) response; + } + + @Override + public com.google.common.util.concurrent.ListenableFuture go( + final Request request, + final HttpResponseHandler handler + ) + { + throw new UnsupportedOperationException(); + } + } + + private static class CloseTrackingInputStream extends ByteArrayInputStream + { + private final AtomicBoolean closed = new AtomicBoolean(); + + private CloseTrackingInputStream(final byte[] response) + { + super(response); + } + + private boolean isClosed() + { + return closed.get(); + } + + @Override + public void close() throws IOException + { + closed.set(true); + super.close(); + } + } + + private static DiscoveryDruidNode testNode() + { + return testNode(8082); + } + + private static DiscoveryDruidNode testNode(final int port) + { + return new DiscoveryDruidNode( + new DruidNode("broker-service", "localhost", false, port, null, true, false), + NodeRole.BROKER, + Collections.emptyMap() + ); + } + + private static SystemTableNode systemTableNode(final DiscoveryDruidNode discoveryNode) + { + final SystemTableNode node = new SystemTableNode(discoveryNode); + node.addNodeRole(discoveryNode.getNodeRole()); + return node; + } + + private static class TestSystemTableDescriptor implements SystemTableDescriptor + { + private static final RowSignature ROW_SIGNATURE = RowSignature.builder().add("value", null).build(); + + @Override + public String getTableName() + { + return "test"; + } + + @Override + public Set getNodeRoles() + { + return Set.of(NodeRole.BROKER); + } + + @Override + public RowSignature getRowSignature() + { + return ROW_SIGNATURE; + } + + @Override + public org.apache.druid.server.system.table.SystemTableRowAuthorizer getRowAuthorizer() + { + 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/server/src/test/java/org/apache/druid/server/system/handler/SystemTableQueryResourceTest.java b/server/src/test/java/org/apache/druid/server/system/handler/SystemTableQueryResourceTest.java new file mode 100644 index 000000000000..d15ca81510e2 --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/system/handler/SystemTableQueryResourceTest.java @@ -0,0 +1,154 @@ +/* + * 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.handler; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.server.QueryLifecycle; +import org.apache.druid.server.QueryLifecycleFactory; +import org.apache.druid.server.QueryResourceQueryResultPusherFactory; +import org.apache.druid.server.QueryScheduler; +import org.apache.druid.server.ResourceIOReaderWriterFactory; +import org.apache.druid.server.initialization.ServerConfig; +import org.apache.druid.server.mocks.MockHttpServletRequest; +import org.apache.druid.server.security.AuthorizationResult; +import org.apache.druid.server.security.AuthorizerMapper; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentMatchers; +import org.mockito.Mockito; + +import javax.ws.rs.core.MediaType; +import javax.ws.rs.core.Response; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; + +public class SystemTableQueryResourceTest +{ + private static final String INTERVAL = "2000-01-01/2001-01-01"; + + private final ObjectMapper objectMapper = new DefaultObjectMapper(); + private final QueryLifecycleFactory queryLifecycleFactory = Mockito.mock(QueryLifecycleFactory.class); + private final QueryResourceQueryResultPusherFactory resultPusherFactory = + Mockito.mock(QueryResourceQueryResultPusherFactory.class); + private final MockHttpServletRequest request = new MockHttpServletRequest(); + + private SystemTableQueryResource resource; + + @BeforeEach + public void setUp() + { + request.contentType = MediaType.APPLICATION_JSON; + resource = new SystemTableQueryResource( + queryLifecycleFactory, + objectMapper, + Mockito.mock(QueryScheduler.class), + Mockito.mock(AuthorizerMapper.class), + resultPusherFactory, + new ResourceIOReaderWriterFactory(objectMapper, objectMapper), + new ServerConfig() + ); + } + + /** A scan-only node must reject an externally submitted {@code segmentMetadata} native query. */ + @Test + public void testRejectsSegmentMetadataQuery() throws IOException + { + assertRejected( + "{" + + "\"queryType\":\"segmentMetadata\"," + + "\"dataSource\":\"foo\"," + + "\"intervals\":[\"" + INTERVAL + "\"]" + + "}" + ); + } + + /** A scan-only node must reject a regular segment-backed native Scan query. */ + @Test + public void testRejectsTableScanQuery() throws IOException + { + assertRejected( + "{" + + "\"queryType\":\"scan\"," + + "\"dataSource\":\"foo\"," + + "\"intervals\":[\"" + INTERVAL + "\"]" + + "}" + ); + } + + /** A direct local system-table Scan is admitted to normal authentication, authorization, and execution. */ + @Test + public void testAcceptsDirectSystemTableScanQuery() throws IOException + { + final QueryLifecycle queryLifecycle = Mockito.mock(QueryLifecycle.class); + final QueryResourceQueryResultPusherFactory.QueryResourceQueryResultPusher resultPusher = + Mockito.mock(QueryResourceQueryResultPusherFactory.QueryResourceQueryResultPusher.class); + Mockito.when(queryLifecycleFactory.factorize()).thenReturn(queryLifecycle); + Mockito.when(queryLifecycle.threadName(ArgumentMatchers.anyString())).thenReturn("system-table-query-test"); + Mockito.when(queryLifecycle.authorize(request)).thenReturn(AuthorizationResult.ALLOW_NO_RESTRICTION); + Mockito.when(resultPusherFactory.factorize( + ArgumentMatchers.any(), + Mockito.eq(request), + Mockito.eq(queryLifecycle), + ArgumentMatchers.any() + )).thenReturn(resultPusher); + Mockito.when(resultPusher.push()).thenReturn(Response.ok().build()); + + final Response response = post(systemTableScanContext("")); + + Assertions.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + Mockito.verify(queryLifecycle).initialize(ArgumentMatchers.any()); + } + + private void assertRejected(final String queryJson) throws IOException + { + final Response response = post(queryJson); + + Assertions.assertEquals(Response.Status.BAD_REQUEST.getStatusCode(), response.getStatus()); + Assertions.assertTrue( + objectMapper.readTree((byte[]) response.getEntity()) + .get("errorMessage") + .asText() + .contains("local system-table Scan") + ); + Mockito.verifyNoInteractions(queryLifecycleFactory); + } + + private Response post(final String queryJson) throws IOException + { + return resource.doPost( + new ByteArrayInputStream(queryJson.getBytes(StandardCharsets.UTF_8)), + null, + request + ); + } + + private static String systemTableScanContext(final String context) + { + return "{" + + "\"queryType\":\"scan\"," + + "\"dataSource\":{\"type\":\"systemTable\",\"table\":\"server_properties\"}," + + "\"intervals\":[\"" + INTERVAL + "\"]," + + "\"context\":{" + context + "}" + + "}"; + } +} diff --git a/server/src/test/java/org/apache/druid/server/system/module/SystemTableQueryHandlerProviderTest.java b/server/src/test/java/org/apache/druid/server/system/module/SystemTableQueryHandlerProviderTest.java new file mode 100644 index 000000000000..e976337441a1 --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/system/module/SystemTableQueryHandlerProviderTest.java @@ -0,0 +1,59 @@ +/* + * 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.module; + +import com.google.common.collect.ImmutableSet; +import org.apache.druid.discovery.NodeRole; +import org.apache.druid.server.system.handler.SystemTableBrokerQueryHandler; +import org.apache.druid.server.system.handler.SystemTableQueryHandler; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +public class SystemTableQueryHandlerProviderTest +{ + @Test + public void testReturnsBrokerHandlerForBroker() + { + final SystemTableBrokerQueryHandler brokerHandler = Mockito.mock(SystemTableBrokerQueryHandler.class); + final SystemTableQueryHandler localHandler = Mockito.mock(SystemTableQueryHandler.class); + final SystemTableQueryHandlerProvider provider = new SystemTableQueryHandlerProvider( + ImmutableSet.of(NodeRole.BROKER), + () -> brokerHandler, + () -> localHandler + ); + + Assertions.assertSame(brokerHandler, provider.get()); + } + + @Test + public void testReturnsLocalHandlerForNonBroker() + { + final SystemTableBrokerQueryHandler brokerHandler = Mockito.mock(SystemTableBrokerQueryHandler.class); + final SystemTableQueryHandler localHandler = Mockito.mock(SystemTableQueryHandler.class); + final SystemTableQueryHandlerProvider provider = new SystemTableQueryHandlerProvider( + ImmutableSet.of(NodeRole.COORDINATOR), + () -> brokerHandler, + () -> localHandler + ); + + Assertions.assertSame(localHandler, provider.get()); + } +} diff --git a/services/pom.xml b/services/pom.xml index 38eb1d03f171..5bc91d0bc5c6 100644 --- a/services/pom.xml +++ b/services/pom.xml @@ -170,6 +170,10 @@ com.sun.jersey jersey-server + + com.sun.jersey.contribs + jersey-guice + org.roaringbitmap diff --git a/services/src/main/java/org/apache/druid/cli/CliBroker.java b/services/src/main/java/org/apache/druid/cli/CliBroker.java index 4e8d79ce71d7..2723b496b978 100644 --- a/services/src/main/java/org/apache/druid/cli/CliBroker.java +++ b/services/src/main/java/org/apache/druid/cli/CliBroker.java @@ -25,7 +25,6 @@ import com.google.inject.Key; import com.google.inject.Module; import com.google.inject.name.Names; -import com.google.inject.util.Modules; import org.apache.druid.client.BrokerSegmentWatcherConfig; import org.apache.druid.client.BrokerServerView; import org.apache.druid.client.BrokerViewOfBrokerConfig; @@ -45,17 +44,15 @@ import org.apache.druid.client.selector.TierSelectorStrategy; import org.apache.druid.discovery.NodeRole; import org.apache.druid.guice.BrokerProcessingModule; +import org.apache.druid.guice.BrokerQueryResourceModule; import org.apache.druid.guice.BrokerServiceModule; import org.apache.druid.guice.CacheModule; import org.apache.druid.guice.Jerseys; -import org.apache.druid.guice.JoinableFactoryModule; import org.apache.druid.guice.JsonConfigProvider; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.LifecycleModule; import org.apache.druid.guice.ManageLifecycle; -import org.apache.druid.guice.QueryRunnerFactoryModule; -import org.apache.druid.guice.QueryableModule; -import org.apache.druid.guice.SegmentWranglerModule; +import org.apache.druid.guice.NativeQueryEngineModule; import org.apache.druid.guice.ServerTypeConfig; import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.msq.dart.guice.DartControllerMemoryManagementModule; @@ -70,10 +67,8 @@ import org.apache.druid.query.RetryQueryRunnerConfig; import org.apache.druid.query.lookup.LookupModule; import org.apache.druid.server.BrokerDynamicConfigResource; -import org.apache.druid.server.BrokerQueryResource; import org.apache.druid.server.ClientInfoResource; import org.apache.druid.server.ClientQuerySegmentWalker; -import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.SegmentManager; import org.apache.druid.server.SubqueryGuardrailHelper; import org.apache.druid.server.SubqueryGuardrailHelperProvider; @@ -84,7 +79,6 @@ import org.apache.druid.server.http.SegmentListerResource; import org.apache.druid.server.http.SelfDiscoveryResource; import org.apache.druid.server.initialization.jetty.JettyServerInitializer; -import org.apache.druid.server.metrics.QueryCountStatsProvider; import org.apache.druid.server.metrics.SubqueryCountStatsProvider; import org.apache.druid.server.router.TieredBrokerConfig; import org.apache.druid.sql.calcite.schema.MetadataSegmentView; @@ -121,12 +115,13 @@ protected List getModules() { return ImmutableList.of( new BrokerProcessingModule(), - new QueryableModule(), - Modules.override(new QueryRunnerFactoryModule()).with( - overrideBinder -> overrideBinder.bind(QueryConfigProvider.class).to(BrokerViewOfBrokerConfig.class) - ), - new SegmentWranglerModule(), - new JoinableFactoryModule(), + NativeQueryEngineModule.builder() + .withOverrideModule( + overrideBinder -> overrideBinder.bind(QueryConfigProvider.class) + .to(BrokerViewOfBrokerConfig.class) + ) + .withQueryResourceModule(new BrokerQueryResourceModule()) + .build(), new BrokerServiceModule(), binder -> { validateCentralizedDatasourceSchemaConfig(getProperties()); @@ -137,8 +132,6 @@ protected List getModules() binder.bindConstant().annotatedWith(Names.named("servicePort")).to(8082); binder.bindConstant().annotatedWith(Names.named("tlsServicePort")).to(8282); binder.bindConstant().annotatedWith(PruneLoadSpec.class).to(true); - binder.bind(ResponseContextConfig.class).toInstance(ResponseContextConfig.newConfig(false)); - binder.bind(CachingClusteredClient.class).in(LazySingleton.class); LifecycleModule.register(binder, BrokerServerView.class); LifecycleModule.register(binder, MetadataSegmentView.class); @@ -162,17 +155,12 @@ protected List getModules() binder.bind(QuerySegmentWalker.class).to(ClientQuerySegmentWalker.class).in(LazySingleton.class); binder.bind(JettyServerInitializer.class).to(QueryJettyServerInitializer.class).in(LazySingleton.class); - binder.bind(BrokerQueryResource.class).in(LazySingleton.class); - Jerseys.addResource(binder, BrokerQueryResource.class); binder.bind(SubqueryGuardrailHelper.class).toProvider(SubqueryGuardrailHelperProvider.class); - binder.bind(QueryCountStatsProvider.class).to(BrokerQueryResource.class).in(LazySingleton.class); binder.bind(SubqueryCountStatsProvider.class).toInstance(new SubqueryCountStatsProvider()); Jerseys.addResource(binder, BrokerResource.class); Jerseys.addResource(binder, ClientInfoResource.class); Jerseys.addResource(binder, BrokerDynamicConfigResource.class); - LifecycleModule.register(binder, BrokerQueryResource.class); - Jerseys.addResource(binder, HttpServerInventoryViewResource.class); LifecycleModule.register(binder, Server.class); diff --git a/services/src/main/java/org/apache/druid/cli/CliCoordinator.java b/services/src/main/java/org/apache/druid/cli/CliCoordinator.java index 07b8ac6398ab..c2758d6ca00f 100644 --- a/services/src/main/java/org/apache/druid/cli/CliCoordinator.java +++ b/services/src/main/java/org/apache/druid/cli/CliCoordinator.java @@ -49,11 +49,12 @@ import org.apache.druid.guice.ManageLifecycle; import org.apache.druid.guice.MetadataConfigModule; import org.apache.druid.guice.MetadataManagerModule; -import org.apache.druid.guice.QueryableModule; +import org.apache.druid.guice.NativeQueryEngineModule; import org.apache.druid.guice.RegexEngineModule; import org.apache.druid.guice.SegmentSchemaCacheModule; import org.apache.druid.guice.SupervisorCleanupModule; import org.apache.druid.guice.annotations.EscalatedGlobal; +import org.apache.druid.guice.annotations.Global; import org.apache.druid.guice.http.JettyHttpClientModule; import org.apache.druid.java.util.common.IAE; import org.apache.druid.java.util.common.ISE; @@ -69,9 +70,12 @@ import org.apache.druid.msq.guice.MSQDurableStorageModule; import org.apache.druid.msq.guice.MSQExternalDataSourceModule; import org.apache.druid.msq.guice.MSQIndexingModule; +import org.apache.druid.query.DefaultQueryConfig; +import org.apache.druid.query.QueryConfigProvider; import org.apache.druid.query.lookup.LookupSerdeModule; import org.apache.druid.segment.metadata.CoordinatorSegmentMetadataCache; import org.apache.druid.segment.metadata.SegmentMetadataCacheConfig; +import org.apache.druid.server.QuerySchedulerProvider; import org.apache.druid.server.compaction.CompactionStatusTracker; import org.apache.druid.server.coordinator.CloneStatusManager; import org.apache.druid.server.coordinator.CoordinatorConfigManager; @@ -177,8 +181,38 @@ protected List getModules() if (isSegmentSchemaCacheEnabled) { validateCentralizedDatasourceSchemaConfig(properties); + modules.add( + NativeQueryEngineModule.builder() + .scanOnly() + // Segment metadata queries historically use Coordinator-specific scheduler and + // default-query configuration prefixes. Override the generic native-engine bindings + // here to preserve that behavior without coupling it to SegmentSchemaCacheModule. + .withOverrideModule( + new Module() + { + @Override + public void configure(final Binder binder) + { + JsonConfigProvider.bind( + binder, + "druid.coordinator.query.scheduler", + QuerySchedulerProvider.class, + Global.class + ); + JsonConfigProvider.bind( + binder, + "druid.coordinator.query.default", + DefaultQueryConfig.class + ); + binder.bind(QueryConfigProvider.class).to(DefaultQueryConfig.class); + } + } + ) + .build() + ); modules.add(new SegmentSchemaCacheModule()); - modules.add(new QueryableModule()); + } else { + modules.add(NativeQueryEngineModule.builder().scanOnly().build()); } modules.add( @@ -194,7 +228,6 @@ public void configure(Binder binder) binder.bindConstant().annotatedWith(Names.named("tlsServicePort")).to(8281); binder.bind(MetadataStorage.class).toProvider(MetadataStorageProvider.class); - JsonConfigProvider.bind(binder, "druid.manager.lookups", LookupCoordinatorManagerConfig.class); JsonConfigProvider.bind(binder, "druid.coordinator", CoordinatorRunConfig.class); JsonConfigProvider.bind(binder, "druid.coordinator.kill", CoordinatorKillConfigs.class); diff --git a/services/src/main/java/org/apache/druid/cli/CliHistorical.java b/services/src/main/java/org/apache/druid/cli/CliHistorical.java index ae1145f860fd..cb24d8f961ec 100644 --- a/services/src/main/java/org/apache/druid/cli/CliHistorical.java +++ b/services/src/main/java/org/apache/druid/cli/CliHistorical.java @@ -33,13 +33,10 @@ import org.apache.druid.guice.DruidProcessingModule; import org.apache.druid.guice.HistoricalServiceModule; import org.apache.druid.guice.Jerseys; -import org.apache.druid.guice.JoinableFactoryModule; import org.apache.druid.guice.JsonConfigProvider; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.LifecycleModule; -import org.apache.druid.guice.QueryRunnerFactoryModule; -import org.apache.druid.guice.QueryableModule; -import org.apache.druid.guice.SegmentWranglerModule; +import org.apache.druid.guice.NativeQueryEngineModule; import org.apache.druid.guice.ServerTypeConfig; import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.msq.dart.guice.DartWorkerMemoryManagementModule; @@ -49,8 +46,6 @@ import org.apache.druid.msq.guice.MSQIndexingModule; import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.lookup.LookupModule; -import org.apache.druid.server.QueryResource; -import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.SegmentManager; import org.apache.druid.server.ServerManager; import org.apache.druid.server.coordination.SegmentCacheBootstrapper; @@ -59,7 +54,6 @@ import org.apache.druid.server.http.SegmentListerResource; import org.apache.druid.server.http.SelfDiscoveryResource; import org.apache.druid.server.initialization.jetty.JettyServerInitializer; -import org.apache.druid.server.metrics.QueryCountStatsProvider; import org.apache.druid.storage.local.LocalTmpStorageConfig; import org.apache.druid.timeline.PruneLastCompactionState; import org.eclipse.jetty.server.Server; @@ -92,18 +86,13 @@ protected List getModules() { return ImmutableList.of( new DruidProcessingModule(), - new QueryableModule(), - new QueryRunnerFactoryModule(), - new SegmentWranglerModule(), - new JoinableFactoryModule(), + NativeQueryEngineModule.builder().build(), new HistoricalServiceModule(), binder -> { binder.bindConstant().annotatedWith(Names.named("serviceName")).to("druid/historical"); binder.bindConstant().annotatedWith(Names.named("servicePort")).to(8083); binder.bindConstant().annotatedWith(Names.named("tlsServicePort")).to(8283); binder.bindConstant().annotatedWith(PruneLastCompactionState.class).to(true); - binder.bind(ResponseContextConfig.class).toInstance(ResponseContextConfig.newConfig(true)); - LifecycleModule.register(binder, Server.class); binder.bind(ServerManager.class).in(LazySingleton.class); binder.bind(SegmentManager.class).in(LazySingleton.class); @@ -111,11 +100,8 @@ protected List getModules() binder.bind(ServerTypeConfig.class).toInstance(new ServerTypeConfig(ServerType.HISTORICAL)); binder.bind(JettyServerInitializer.class).to(QueryJettyServerInitializer.class).in(LazySingleton.class); - binder.bind(QueryCountStatsProvider.class).to(QueryResource.class); - Jerseys.addResource(binder, QueryResource.class); Jerseys.addResource(binder, SegmentListerResource.class); Jerseys.addResource(binder, HistoricalResource.class); - LifecycleModule.register(binder, QueryResource.class); LifecycleModule.register(binder, SegmentCacheBootstrapper.class); diff --git a/services/src/main/java/org/apache/druid/cli/CliIndexer.java b/services/src/main/java/org/apache/druid/cli/CliIndexer.java index 6ffb17ad637e..2116ebf9dfb0 100644 --- a/services/src/main/java/org/apache/druid/cli/CliIndexer.java +++ b/services/src/main/java/org/apache/druid/cli/CliIndexer.java @@ -40,15 +40,11 @@ import org.apache.druid.guice.IndexingServiceTaskLogsModule; import org.apache.druid.guice.IndexingServiceTuningConfigModule; import org.apache.druid.guice.Jerseys; -import org.apache.druid.guice.JoinableFactoryModule; import org.apache.druid.guice.JsonConfigProvider; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.LifecycleModule; -import org.apache.druid.guice.QueryRunnerFactoryModule; -import org.apache.druid.guice.QueryableModule; -import org.apache.druid.guice.QueryablePeonModule; +import org.apache.druid.guice.NativeQueryEngineModule; import org.apache.druid.guice.RegexEngineModule; -import org.apache.druid.guice.SegmentWranglerModule; import org.apache.druid.guice.ServerTypeConfig; import org.apache.druid.guice.annotations.AttemptId; import org.apache.druid.guice.annotations.Parent; @@ -73,7 +69,6 @@ import org.apache.druid.segment.realtime.appenderator.AppenderatorsManager; import org.apache.druid.segment.realtime.appenderator.UnifiedIndexerAppenderatorsManager; import org.apache.druid.server.DruidNode; -import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.SegmentManager; import org.apache.druid.server.coordination.SegmentCacheBootstrapper; import org.apache.druid.server.coordination.ServerType; @@ -125,10 +120,7 @@ protected List getModules() { return ImmutableList.of( new DruidProcessingModule(), - new QueryableModule(), - new QueryRunnerFactoryModule(), - new SegmentWranglerModule(), - new JoinableFactoryModule(), + NativeQueryEngineModule.builder().build(), new IndexerServiceModule(), new Module() { @@ -140,7 +132,6 @@ public void configure(Binder binder) binder.bindConstant().annotatedWith(Names.named("serviceName")).to("druid/indexer"); binder.bindConstant().annotatedWith(Names.named("servicePort")).to(8091); binder.bindConstant().annotatedWith(Names.named("tlsServicePort")).to(8291); - binder.bind(ResponseContextConfig.class).toInstance(ResponseContextConfig.newConfig(true)); // needed for the CliPeon, not needed for indexer, but have to bind annotation. binder.bindConstant().annotatedWith(AttemptId.class).to(""); @@ -238,7 +229,6 @@ public DataNodeService getDataNodeService(DruidServerConfig serverConfig) new IndexingServiceTaskLogsModule(properties), new IndexingServiceTuningConfigModule(), new InputSourceModule(), - new QueryablePeonModule(), new CliIndexerServerModule(properties), new LookupModule(), new MSQIndexingModule(), diff --git a/services/src/main/java/org/apache/druid/cli/CliMiddleManager.java b/services/src/main/java/org/apache/druid/cli/CliMiddleManager.java index e5cebe3adc39..1e3f9be42fca 100644 --- a/services/src/main/java/org/apache/druid/cli/CliMiddleManager.java +++ b/services/src/main/java/org/apache/druid/cli/CliMiddleManager.java @@ -44,6 +44,7 @@ import org.apache.druid.guice.LifecycleModule; import org.apache.druid.guice.ManageLifecycle; import org.apache.druid.guice.MiddleManagerServiceModule; +import org.apache.druid.guice.NativeQueryEngineModule; import org.apache.druid.guice.PolyBind; import org.apache.druid.guice.RegexEngineModule; import org.apache.druid.guice.annotations.Self; @@ -116,6 +117,7 @@ protected List getModules() { return ImmutableList.of( new MiddleManagerServiceModule(), + NativeQueryEngineModule.builder().scanOnly().build(), new Module() { @Override @@ -127,7 +129,6 @@ public void configure(Binder binder) binder.bindConstant().annotatedWith(Names.named("servicePort")).to(8091); binder.bindConstant().annotatedWith(Names.named("tlsServicePort")).to(8291); binder.bindConstant().annotatedWith(PruneLastCompactionState.class).to(true); - IndexingServiceModuleHelper.configureTaskRunnerConfigs(binder); JsonConfigProvider.bind(binder, "druid.indexer.task", TaskConfig.class); 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 9b81a4a987b0..b6af7205cdd3 100644 --- a/services/src/main/java/org/apache/druid/cli/CliOverlord.java +++ b/services/src/main/java/org/apache/druid/cli/CliOverlord.java @@ -51,6 +51,7 @@ import org.apache.druid.guice.ListProvider; import org.apache.druid.guice.ManageLifecycle; import org.apache.druid.guice.MetadataManagerModule; +import org.apache.druid.guice.NativeQueryEngineModule; import org.apache.druid.guice.PolyBind; import org.apache.druid.guice.RegexEngineModule; import org.apache.druid.guice.SupervisorModule; @@ -203,9 +204,13 @@ public void configure(Properties properties) protected List getModules(final boolean standalone) { - return ImmutableList.of( - new DerbyTaskStorageModule(), - standalone ? new MetadataManagerModule() : binder -> {}, + final ImmutableList.Builder modules = ImmutableList.builder(); + modules.add(new DerbyTaskStorageModule()); + modules.add(standalone ? new MetadataManagerModule() : binder -> {}); + if (standalone) { + modules.add(NativeQueryEngineModule.builder().scanOnly().build()); + } + modules.add( new Module() { @Override @@ -516,6 +521,7 @@ private void configureOverlordWebResources(Binder binder) new MSQExternalDataSourceModule(), new RegexEngineModule() ); + return modules.build(); } /** diff --git a/services/src/main/java/org/apache/druid/cli/CliPeon.java b/services/src/main/java/org/apache/druid/cli/CliPeon.java index 84effc35a6fa..4f5fcc156545 100644 --- a/services/src/main/java/org/apache/druid/cli/CliPeon.java +++ b/services/src/main/java/org/apache/druid/cli/CliPeon.java @@ -50,20 +50,16 @@ import org.apache.druid.guice.IndexingServiceTaskLogsModule; import org.apache.druid.guice.IndexingServiceTuningConfigModule; import org.apache.druid.guice.Jerseys; -import org.apache.druid.guice.JoinableFactoryModule; import org.apache.druid.guice.JsonConfigProvider; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.LifecycleModule; import org.apache.druid.guice.ManageLifecycle; import org.apache.druid.guice.ManageLifecycleServer; +import org.apache.druid.guice.NativeQueryEngineModule; import org.apache.druid.guice.PeonProcessingModule; import org.apache.druid.guice.PeonServerModule; import org.apache.druid.guice.PolyBind; -import org.apache.druid.guice.QueryRunnerFactoryModule; -import org.apache.druid.guice.QueryableModule; -import org.apache.druid.guice.QueryablePeonModule; import org.apache.druid.guice.RegexEngineModule; -import org.apache.druid.guice.SegmentWranglerModule; import org.apache.druid.guice.ServerTypeConfig; import org.apache.druid.guice.annotations.AttemptId; import org.apache.druid.guice.annotations.Json; @@ -118,7 +114,6 @@ import org.apache.druid.segment.realtime.appenderator.AppenderatorsManager; import org.apache.druid.segment.realtime.appenderator.PeonAppenderatorsManager; import org.apache.druid.server.DruidNode; -import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.SegmentManager; import org.apache.druid.server.coordination.BroadcastDatasourceLoadingSpec; import org.apache.druid.server.coordination.SegmentCacheBootstrapper; @@ -214,10 +209,7 @@ protected List getModules() { return ImmutableList.of( new PeonProcessingModule(), - new QueryableModule(), - new QueryRunnerFactoryModule(), - new SegmentWranglerModule(), - new JoinableFactoryModule(), + NativeQueryEngineModule.builder().build(), new IndexingServiceTaskLogsModule(properties), new RegexEngineModule(), new Module() @@ -239,7 +231,6 @@ public void configure(Binder binder) binder.bindConstant().annotatedWith(Names.named("serviceName")).to("druid/peon"); binder.bindConstant().annotatedWith(Names.named("servicePort")).to(0); binder.bindConstant().annotatedWith(Names.named("tlsServicePort")).to(-1); - binder.bind(ResponseContextConfig.class).toInstance(ResponseContextConfig.newConfig(true)); binder.bindConstant().annotatedWith(AttemptId.class).to(attemptId); JsonConfigProvider.bind(binder, "druid.task.executor", DruidNode.class, Parent.class); @@ -342,7 +333,6 @@ public LocalTmpStorageConfig getLocalTmpStorage() return () -> tmpDir; } }, - new QueryablePeonModule(), new PeonServerModule(), new IndexingServiceInputSourceModule(), new IndexingServiceTuningConfigModule(), diff --git a/services/src/main/java/org/apache/druid/cli/CliRouter.java b/services/src/main/java/org/apache/druid/cli/CliRouter.java index da6e4772188f..f501fb6aeb46 100644 --- a/services/src/main/java/org/apache/druid/cli/CliRouter.java +++ b/services/src/main/java/org/apache/druid/cli/CliRouter.java @@ -32,19 +32,17 @@ import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.LifecycleModule; import org.apache.druid.guice.ManageLifecycle; -import org.apache.druid.guice.QueryRunnerFactoryModule; -import org.apache.druid.guice.QueryableModule; +import org.apache.druid.guice.NativeQueryEngineModule; import org.apache.druid.guice.RouterProcessingModule; +import org.apache.druid.guice.RouterQueryResourceModule; import org.apache.druid.guice.http.JettyHttpClientModule; import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.lookup.LookupSerdeModule; -import org.apache.druid.server.AsyncQueryForwardingServlet; import org.apache.druid.server.NoopQuerySegmentWalker; import org.apache.druid.server.http.RouterResource; import org.apache.druid.server.http.SelfDiscoveryResource; import org.apache.druid.server.initialization.jetty.JettyServerInitializer; -import org.apache.druid.server.metrics.QueryCountStatsProvider; import org.apache.druid.server.router.AvaticaConnectionBalancer; import org.apache.druid.server.router.CoordinatorRuleManager; import org.apache.druid.server.router.ManagementProxyConfig; @@ -88,15 +86,15 @@ protected List getModules() { return ImmutableList.of( new RouterProcessingModule(), - new QueryableModule(), - new QueryRunnerFactoryModule(), + NativeQueryEngineModule.builder() + .withQueryResourceModule(new RouterQueryResourceModule()) + .build(), new JettyHttpClientModule("druid.router.http", Router.class), JettyHttpClientModule.global(), binder -> { binder.bindConstant().annotatedWith(Names.named("serviceName")).to("druid/router"); binder.bindConstant().annotatedWith(Names.named("servicePort")).to(8888); binder.bindConstant().annotatedWith(Names.named("tlsServicePort")).to(9088); - JsonConfigProvider.bind(binder, "druid.router", TieredBrokerConfig.class); JsonConfigProvider.bind(binder, "druid.router.avatica.balancer", AvaticaConnectionBalancer.class); JsonConfigProvider.bind(binder, "druid.router.managementProxy", ManagementProxyConfig.class); @@ -112,7 +110,6 @@ protected List getModules() .toProvider(TieredBrokerSelectorStrategiesProvider.class) .in(LazySingleton.class); - binder.bind(QueryCountStatsProvider.class).to(AsyncQueryForwardingServlet.class).in(LazySingleton.class); binder.bind(JettyServerInitializer.class).to(RouterJettyServerInitializer.class).in(LazySingleton.class); Jerseys.addResource(binder, RouterResource.class); diff --git a/services/src/main/java/org/apache/druid/cli/CoordinatorJettyServerInitializer.java b/services/src/main/java/org/apache/druid/cli/CoordinatorJettyServerInitializer.java index ba32c3f666fb..637bbedf042d 100644 --- a/services/src/main/java/org/apache/druid/cli/CoordinatorJettyServerInitializer.java +++ b/services/src/main/java/org/apache/druid/cli/CoordinatorJettyServerInitializer.java @@ -120,6 +120,7 @@ public void initialize(Server server, Injector injector) // Can't use '/*' here because of Guice and Jetty static content conflicts root.addFilter(guiceFilterHolder, "/info/*", null); root.addFilter(guiceFilterHolder, "/druid/coordinator/*", null); + root.addFilter(guiceFilterHolder, "/druid/v2/*", null); if (beOverlord) { root.addFilter(guiceFilterHolder, "/druid/indexer/*", null); } diff --git a/services/src/main/java/org/apache/druid/guice/RouterQueryResourceModule.java b/services/src/main/java/org/apache/druid/guice/RouterQueryResourceModule.java new file mode 100644 index 000000000000..f539fa72f2e9 --- /dev/null +++ b/services/src/main/java/org/apache/druid/guice/RouterQueryResourceModule.java @@ -0,0 +1,40 @@ +/* + * 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.guice; + +import com.google.inject.Binder; +import org.apache.druid.initialization.DruidModule; +import org.apache.druid.server.AsyncQueryForwardingServlet; +import org.apache.druid.server.QueryResource; +import org.apache.druid.server.ResponseContextConfig; +import org.apache.druid.server.metrics.QueryCountStatsProvider; + +/** Registers the Router's local native query resource and forwarding-servlet query metrics. */ +public class RouterQueryResourceModule implements DruidModule +{ + @Override + public void configure(final Binder binder) + { + binder.bind(ResponseContextConfig.class).toInstance(ResponseContextConfig.newConfig(true)); + binder.bind(QueryCountStatsProvider.class).to(AsyncQueryForwardingServlet.class).in(LazySingleton.class); + Jerseys.addResource(binder, QueryResource.class); + LifecycleModule.register(binder, QueryResource.class); + } +} diff --git a/services/src/main/java/org/apache/druid/guice/SegmentSchemaCacheModule.java b/services/src/main/java/org/apache/druid/guice/SegmentSchemaCacheModule.java index 80f85c0d1c41..4ab9406277ca 100644 --- a/services/src/main/java/org/apache/druid/guice/SegmentSchemaCacheModule.java +++ b/services/src/main/java/org/apache/druid/guice/SegmentSchemaCacheModule.java @@ -20,23 +20,14 @@ package org.apache.druid.guice; import com.google.inject.Binder; -import com.google.inject.Key; import com.google.inject.Module; -import com.google.inject.Provides; import com.google.inject.multibindings.MapBinder; +import com.google.inject.multibindings.OptionalBinder; import org.apache.druid.client.InternalQueryConfig; -import org.apache.druid.guice.annotations.Global; -import org.apache.druid.query.DefaultGenericQueryMetricsFactory; -import org.apache.druid.query.DefaultQueryConfig; -import org.apache.druid.query.GenericQueryMetricsFactory; -import org.apache.druid.query.MapQueryToolChestWarehouse; import org.apache.druid.query.Query; -import org.apache.druid.query.QueryConfigProvider; import org.apache.druid.query.QueryRunnerFactory; import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.QueryToolChest; -import org.apache.druid.query.QueryToolChestWarehouse; -import org.apache.druid.query.QueryWatcher; import org.apache.druid.query.RetryQueryRunnerConfig; import org.apache.druid.query.metadata.SegmentMetadataQueryConfig; import org.apache.druid.query.metadata.SegmentMetadataQueryQueryToolChest; @@ -45,8 +36,6 @@ import org.apache.druid.segment.metadata.CentralizedDatasourceSchemaConfig; import org.apache.druid.segment.metadata.CoordinatorSegmentMetadataCache; import org.apache.druid.segment.metadata.SegmentMetadataQuerySegmentWalker; -import org.apache.druid.server.QueryScheduler; -import org.apache.druid.server.QuerySchedulerProvider; /** * Module that binds dependencies required for segment schema management and @@ -61,38 +50,22 @@ public class SegmentSchemaCacheModule implements Module public void configure(Binder binder) { JsonConfigProvider.bind(binder, "druid.coordinator.segmentMetadata", SegmentMetadataQueryConfig.class); - JsonConfigProvider.bind(binder, "druid.coordinator.query.scheduler", QuerySchedulerProvider.class, Global.class); - JsonConfigProvider.bind(binder, "druid.coordinator.query.default", DefaultQueryConfig.class); - binder.bind(QueryConfigProvider.class).to(DefaultQueryConfig.class); JsonConfigProvider.bind(binder, "druid.coordinator.query.retryPolicy", RetryQueryRunnerConfig.class); JsonConfigProvider.bind(binder, "druid.coordinator.internal.query.config", InternalQueryConfig.class); MapBinder, QueryToolChest> toolChests = DruidBinders.queryToolChestBinder(binder); toolChests.addBinding(SegmentMetadataQuery.class).to(SegmentMetadataQueryQueryToolChest.class); binder.bind(SegmentMetadataQueryQueryToolChest.class).in(LazySingleton.class); - binder.bind(QueryToolChestWarehouse.class).to(MapQueryToolChestWarehouse.class); final MapBinder, QueryRunnerFactory> queryFactoryBinder = DruidBinders.queryRunnerFactoryBinder(binder); queryFactoryBinder.addBinding(SegmentMetadataQuery.class).to(SegmentMetadataQueryRunnerFactory.class); - DruidBinders.queryBinder(binder); binder.bind(SegmentMetadataQueryRunnerFactory.class).in(LazySingleton.class); - - binder.bind(GenericQueryMetricsFactory.class).to(DefaultGenericQueryMetricsFactory.class); - - binder.bind(QueryScheduler.class) - .toProvider(Key.get(QuerySchedulerProvider.class, Global.class)) - .in(LazySingleton.class); - binder.bind(QuerySchedulerProvider.class).in(LazySingleton.class); - binder.bind(QuerySegmentWalker.class).to(SegmentMetadataQuerySegmentWalker.class).in(LazySingleton.class); + OptionalBinder.newOptionalBinder(binder, QuerySegmentWalker.class) + .setBinding() + .to(SegmentMetadataQuerySegmentWalker.class) + .in(LazySingleton.class); LifecycleModule.register(binder, CoordinatorSegmentMetadataCache.class); } - - @LazySingleton - @Provides - public QueryWatcher getWatcher(QueryScheduler scheduler) - { - return scheduler; - } } diff --git a/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java b/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java index c26e8eb9af8d..93f74550807a 100644 --- a/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java +++ b/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java @@ -26,6 +26,7 @@ import com.google.common.annotations.VisibleForTesting; import com.google.inject.Inject; import com.google.inject.Provider; +import com.sun.jersey.guice.spi.container.servlet.GuiceContainer; import org.apache.calcite.avatica.remote.ProtobufTranslation; import org.apache.calcite.avatica.remote.ProtobufTranslationImpl; import org.apache.calcite.avatica.remote.Service; @@ -46,6 +47,7 @@ import org.apache.druid.query.QueryInterruptedException; import org.apache.druid.query.QueryMetrics; import org.apache.druid.query.QueryToolChestWarehouse; +import org.apache.druid.query.SystemTableDataSource; import org.apache.druid.server.initialization.ServerConfig; import org.apache.druid.server.initialization.jetty.HttpException; import org.apache.druid.server.initialization.jetty.ResponseIdentityHeaderHandler; @@ -72,12 +74,18 @@ import org.eclipse.jetty.http.HttpMethod; import javax.annotation.Nullable; +import javax.servlet.ReadListener; import javax.servlet.ServletException; +import javax.servlet.ServletInputStream; import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletRequestWrapper; import javax.servlet.http.HttpServletResponse; import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response.Status; +import java.io.ByteArrayInputStream; import java.io.IOException; +import java.util.Collections; +import java.util.Enumeration; import java.util.HashMap; import java.util.Map; import java.util.Properties; @@ -103,6 +111,7 @@ public class AsyncQueryForwardingServlet extends AsyncProxyServlet implements Qu private static final String AVATICA_QUERY_ATTRIBUTE = "org.apache.druid.proxy.avaticaQuery"; private static final String SQL_QUERY_ATTRIBUTE = "org.apache.druid.proxy.sqlQuery"; private static final String OBJECTMAPPER_ATTRIBUTE = "org.apache.druid.proxy.objectMapper"; + private static final String NODE_LOCAL_ATTRIBUTE = "org.apache.druid.proxy.nodeLocalSystemQuery"; private static final String PROPERTY_SQL_ENABLE = "druid.router.sql.enable"; private static final String PROPERTY_SQL_ENABLE_DEFAULT = "false"; @@ -145,6 +154,8 @@ void handleException(HttpServletResponse response, ObjectMapper objectMapper, Ex private final ProtobufTranslation protobufTranslation; private final ServerConfig serverConfig; + private GuiceContainer localQueryContainer; + private final boolean routeSqlByStrategy; private HttpClient broadcastClient; @@ -182,6 +193,12 @@ public AsyncQueryForwardingServlet( this.serverConfig = serverConfig; } + @Inject(optional = true) + public void setLocalQueryContainer(final GuiceContainer localQueryContainer) + { + this.localQueryContainer = localQueryContainer; + } + @Override public void init() throws ServletException { @@ -224,7 +241,8 @@ protected void service(HttpServletRequest request, HttpServletResponse response) // The Router does not have the ability to look inside SQL queries and route them intelligently, so just treat // them as a generic request. - final boolean isNativeQueryEndpoint = requestURI.startsWith("/druid/v2") && !requestURI.startsWith("/druid/v2/sql"); + final boolean isNativeQueryEndpoint = requestURI.startsWith("/druid/v2") + && !requestURI.startsWith("/druid/v2/sql"); final boolean isSqlQueryEndpoint = requestURI.startsWith("/druid/v2/sql"); final boolean isAvaticaJson = requestURI.startsWith("/druid/v2/sql/avatica"); @@ -247,6 +265,9 @@ protected void service(HttpServletRequest request, HttpServletResponse response) byte[] requestBytes = objectMapper.writeValueAsBytes(requestMap); request.setAttribute(AVATICA_QUERY_ATTRIBUTE, requestBytes); LOG.debug("Forwarding JDBC connection [%s] to broker [%s]", connectionId, targetServer.getHost()); + } else if (HttpMethod.DELETE.is(method) && isNativeQueryEndpoint && isLocalNativeQueryRoute(request)) { + dispatchNodeLocalCancellation(request, response); + return; } else if (HttpMethod.DELETE.is(method)) { // query cancellation request targetServer = hostFinder.pickDefaultServer(); @@ -257,6 +278,10 @@ protected void service(HttpServletRequest request, HttpServletResponse response) try { Query inputQuery = objectMapper.readValue(request.getInputStream(), Query.class); if (inputQuery != null) { + if (isLocalSystemTableQuery(request, inputQuery)) { + dispatchNodeLocalQuery(request, response, inputQuery, objectMapper); + return; + } targetServer = hostFinder.pickServer(inputQuery); if (inputQuery.getId() == null) { inputQuery = inputQuery.withId(UUID.randomUUID().toString()); @@ -307,6 +332,45 @@ protected void service(HttpServletRequest request, HttpServletResponse response) doService(request, response); } + private void dispatchNodeLocalQuery( + final HttpServletRequest request, + final HttpServletResponse response, + final Query query, + final ObjectMapper objectMapper + ) throws ServletException, IOException + { + request.setAttribute(NODE_LOCAL_ATTRIBUTE, true); + if (localQueryContainer == null) { + throw new IAE("Router local query container is not available"); + } + localQueryContainer.service(new CachedBodyRequest(request, objectMapper.writeValueAsBytes(query)), response); + } + + private void dispatchNodeLocalCancellation( + final HttpServletRequest request, + final HttpServletResponse response + ) throws ServletException, IOException + { + request.setAttribute(NODE_LOCAL_ATTRIBUTE, true); + if (localQueryContainer == null) { + throw new IAE("Router local query container is not available"); + } + localQueryContainer.service(new NodeLocalRequest(request), response); + } + + private static boolean isLocalSystemTableQuery(final HttpServletRequest request, final Query query) + { + return query.getDataSource() instanceof SystemTableDataSource + && isLocalNativeQueryRoute(request); + } + + private static boolean isLocalNativeQueryRoute(final HttpServletRequest request) + { + return QueryResource.NATIVE_QUERY_ROUTE_LOCAL.equals( + request.getHeader(QueryResource.HEADER_NATIVE_QUERY_ROUTE) + ); + } + /** * Rebuilds the {@link SqlQuery} object with sqlQueryId and queryId context parameters if not present * @@ -457,10 +521,116 @@ protected void doService( HttpServletResponse response ) throws ServletException, IOException { + if (Boolean.TRUE.equals(request.getAttribute(NODE_LOCAL_ATTRIBUTE))) { + throw new IAE("Node-local system table queries must not be proxied"); + } // Just call the superclass service method. Overridden in tests. super.service(request, response); } + private static class NodeLocalRequest extends HttpServletRequestWrapper + { + NodeLocalRequest(final HttpServletRequest request) + { + super(request); + } + + @Override + public String getServletPath() + { + return ""; + } + + @Override + public String getPathInfo() + { + return getRequestURI(); + } + } + + private static class CachedBodyRequest extends NodeLocalRequest + { + private final byte[] body; + + CachedBodyRequest(final HttpServletRequest request, final byte[] body) + { + super(request); + this.body = body; + } + + @Override + public int getContentLength() + { + return body.length; + } + + @Override + public long getContentLengthLong() + { + return body.length; + } + + @Override + public String getHeader(final String name) + { + if (HttpHeader.CONTENT_LENGTH.asString().equalsIgnoreCase(name)) { + return String.valueOf(body.length); + } + return super.getHeader(name); + } + + @Override + public Enumeration getHeaders(final String name) + { + if (HttpHeader.CONTENT_LENGTH.asString().equalsIgnoreCase(name)) { + return Collections.enumeration(Collections.singleton(String.valueOf(body.length))); + } + return super.getHeaders(name); + } + + @Override + public int getIntHeader(final String name) + { + if (HttpHeader.CONTENT_LENGTH.asString().equalsIgnoreCase(name)) { + return body.length; + } + return super.getIntHeader(name); + } + + @Override + public ServletInputStream getInputStream() + { + final ByteArrayInputStream input = new ByteArrayInputStream(body); + return new ServletInputStream() + { + @Override + public boolean isFinished() + { + return input.available() == 0; + } + + @Override + public boolean isReady() + { + return true; + } + + @Override + public void setReadListener(final ReadListener readListener) + { + // Synchronous in-memory request body. + } + + @Override + public int read() + { + return input.read(); + } + }; + } + + } + @Override protected void sendProxyRequest( HttpServletRequest clientRequest, diff --git a/services/src/test/java/org/apache/druid/cli/CliCoordinatorTest.java b/services/src/test/java/org/apache/druid/cli/CliCoordinatorTest.java index b5fe3e34eea9..e2ae659f584d 100644 --- a/services/src/test/java/org/apache/druid/cli/CliCoordinatorTest.java +++ b/services/src/test/java/org/apache/druid/cli/CliCoordinatorTest.java @@ -26,15 +26,26 @@ import com.google.inject.TypeLiteral; import jakarta.validation.Validation; import jakarta.validation.Validator; -import org.apache.druid.discovery.NodeRole; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.LifecycleModule; +import org.apache.druid.guice.annotations.JSR311Resource; import org.apache.druid.jackson.JacksonModule; +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryConfigProvider; +import org.apache.druid.query.QueryRunnerFactory; +import org.apache.druid.query.QuerySegmentWalker; +import org.apache.druid.query.metadata.metadata.SegmentMetadataQuery; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.segment.metadata.SegmentMetadataQuerySegmentWalker; +import org.apache.druid.server.NoopQuerySegmentWalker; +import org.apache.druid.server.QueryResource; import org.apache.druid.server.initialization.jetty.JettyBindings; +import org.apache.druid.server.system.handler.SystemTableQueryResource; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import java.util.Arrays; +import java.util.Map; import java.util.Properties; import java.util.Set; @@ -91,6 +102,49 @@ public void testLeaderEndpointsExcludedFromQos() ); } + @Test + public void testCoordinatorAsOverlordWithCentralizedDatasourceSchema() + { + final Properties properties = new Properties(); + properties.setProperty("druid.coordinator.asOverlord.enabled", "true"); + properties.setProperty("druid.centralizedDatasourceSchema.enabled", "true"); + properties.setProperty("druid.coordinator.query.default.context.testKey", "testValue"); + + final Injector injector = makeCoordinatorInjector(properties); + + Assertions.assertNotNull(injector.getInstance(SystemTableQueryResource.class)); + Assertions.assertFalse(jerseyResources(injector).contains(QueryResource.class)); + Assertions.assertTrue(jerseyResources(injector).contains(SystemTableQueryResource.class)); + Assertions.assertInstanceOf( + SegmentMetadataQuerySegmentWalker.class, + injector.getInstance(QuerySegmentWalker.class) + ); + Assertions.assertTrue(queryRunnerFactories(injector).containsKey(ScanQuery.class)); + Assertions.assertTrue(queryRunnerFactories(injector).containsKey(SegmentMetadataQuery.class)); + Assertions.assertEquals("testValue", injector.getInstance(QueryConfigProvider.class).getContext().get("testKey")); + } + + @Test + public void testCoordinatorUsesRestrictedQueryResourceWithoutCentralizedDatasourceSchema() + { + final Injector injector = makeCoordinatorInjector(new Properties()); + + Assertions.assertNotNull(injector.getInstance(SystemTableQueryResource.class)); + Assertions.assertInstanceOf(NoopQuerySegmentWalker.class, injector.getInstance(QuerySegmentWalker.class)); + Assertions.assertTrue(queryRunnerFactories(injector).containsKey(ScanQuery.class)); + Assertions.assertFalse(queryRunnerFactories(injector).containsKey(SegmentMetadataQuery.class)); + } + + private static Map, QueryRunnerFactory> queryRunnerFactories(final Injector injector) + { + return injector.getInstance(Key.get(new TypeLiteral<>() {})); + } + + private static Set> jerseyResources(final Injector injector) + { + return injector.getInstance(Key.get(new TypeLiteral<>() {}, JSR311Resource.class)); + } + private static boolean hasCoordinatorQosFilter(Set qosFilters) { return qosFilters.stream() @@ -116,6 +170,6 @@ private static Injector makeCoordinatorInjector(final Properties props) final CliCoordinator coordinator = new CliCoordinator(); baseInjector.injectMembers(coordinator); - return coordinator.makeInjector(Set.of(NodeRole.COORDINATOR)); + return coordinator.makeInjector(coordinator.getNodeRoles(props)); } } diff --git a/services/src/test/java/org/apache/druid/cli/CliPeonTest.java b/services/src/test/java/org/apache/druid/cli/CliPeonTest.java index 7f9a8a6f6079..b0a13daa4171 100644 --- a/services/src/test/java/org/apache/druid/cli/CliPeonTest.java +++ b/services/src/test/java/org/apache/druid/cli/CliPeonTest.java @@ -24,8 +24,10 @@ import com.google.common.collect.ImmutableMap; import com.google.inject.Guice; import com.google.inject.Injector; +import com.google.inject.Key; import com.google.inject.Module; import com.google.inject.Scopes; +import com.google.inject.TypeLiteral; import jakarta.validation.Validation; import jakarta.validation.Validator; import org.apache.commons.io.FileUtils; @@ -74,9 +76,13 @@ import org.apache.druid.segment.TestHelper; import org.apache.druid.segment.TestIndex; import org.apache.druid.segment.indexing.DataSchema; +import org.apache.druid.server.QueryResource; import org.apache.druid.server.coordination.BroadcastDatasourceLoadingSpec; import org.apache.druid.server.lookup.cache.LookupLoadingSpec; import org.apache.druid.server.metrics.LoadSpecHolder; +import org.apache.druid.server.system.handler.SystemTableQueryResource; +import org.apache.druid.server.system.table.ServerPropertiesTableDescriptor; +import org.apache.druid.server.system.table.SystemTableDataProvider; import org.apache.druid.storage.local.LocalTmpStorageConfig; import org.apache.druid.testing.TemporaryFolderExtension; import org.joda.time.Duration; @@ -109,6 +115,20 @@ public class CliPeonTest .id("compact_test_taskid") .inputSpec(new CompactionIntervalSpec(Intervals.of("2020/2021"), null)); + @Test + public void testPeonNativeServerPropertiesInjection() throws IOException + { + final Injector peonInjector = makePeonInjector(NoopTask.create(), new Properties()); + final Map dataProviders = peonInjector.getInstance( + Key.get(new TypeLiteral<>() {}) + ); + + Assertions.assertNotNull(peonInjector.getInstance(QueryResource.class)); + // Peons use the normal query resource to serve native queries over both segments and system tables. + Assertions.assertNull(peonInjector.getExistingBinding(Key.get(SystemTableQueryResource.class))); + Assertions.assertTrue(dataProviders.containsKey(ServerPropertiesTableDescriptor.TABLE_NAME)); + } + @Test public void testCliPeonK8sMode() throws IOException { diff --git a/services/src/test/java/org/apache/druid/cli/MainTest.java b/services/src/test/java/org/apache/druid/cli/MainTest.java index 33fb605d1f60..4ac8f2a035a3 100644 --- a/services/src/test/java/org/apache/druid/cli/MainTest.java +++ b/services/src/test/java/org/apache/druid/cli/MainTest.java @@ -20,11 +20,18 @@ package org.apache.druid.cli; import com.google.inject.Injector; +import com.google.inject.Key; +import com.google.inject.TypeLiteral; import org.apache.druid.guice.GuiceInjectors; +import org.apache.druid.server.system.handler.SystemTableQueryResource; +import org.apache.druid.server.system.table.ServerPropertiesTableDescriptor; +import org.apache.druid.server.system.table.SystemTableDataProvider; import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; +import java.util.Map; import java.util.Properties; import java.util.stream.Stream; @@ -51,4 +58,27 @@ public void testSimpleInjection(ServerRunnable runnable) injector.injectMembers(runnable); Assertions.assertNotNull(runnable.makeInjector(runnable.getNodeRoles(new Properties()))); } + + @Test + public void testMiddleManagerNativeServerPropertiesInjection() + { + final CliMiddleManager middleManager = new CliMiddleManager(); + final Injector injector = GuiceInjectors.makeStartupInjector(); + injector.injectMembers(middleManager); + + final Injector middleManagerInjector = middleManager.makeInjector( + middleManager.getNodeRoles(new Properties()) + ); + + // Middle Managers expose /druid/v2 through the resource restricted to internal system-table scans. + Assertions.assertNotNull(middleManagerInjector.getInstance(SystemTableQueryResource.class)); + Assertions.assertTrue( + systemTableDataProviders(middleManagerInjector).containsKey(ServerPropertiesTableDescriptor.TABLE_NAME) + ); + } + + private static Map systemTableDataProviders(final Injector injector) + { + return injector.getInstance(Key.get(new TypeLiteral<>() {})); + } } diff --git a/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java b/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java index 5b339b8d6212..c352708f5b03 100644 --- a/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java +++ b/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java @@ -31,6 +31,7 @@ import com.google.inject.Key; import com.google.inject.Module; import com.google.inject.servlet.GuiceFilter; +import com.sun.jersey.guice.spi.container.servlet.GuiceContainer; import org.apache.calcite.avatica.Meta; import org.apache.calcite.avatica.remote.Service; import org.apache.commons.io.IOUtils; @@ -60,6 +61,7 @@ import org.apache.druid.query.Query; import org.apache.druid.query.QueryException; import org.apache.druid.query.QueryMetrics; +import org.apache.druid.query.SystemTableDataSource; import org.apache.druid.query.aggregation.FilteredAggregatorFactory; import org.apache.druid.query.aggregation.any.StringAnyAggregatorFactory; import org.apache.druid.query.filter.SelectorDimFilter; @@ -96,6 +98,7 @@ import org.eclipse.jetty.ee8.servlet.ServletHolder; import org.eclipse.jetty.http.HttpField; import org.eclipse.jetty.http.HttpFields; +import org.eclipse.jetty.http.HttpHeader; import org.eclipse.jetty.http.HttpVersion; import org.eclipse.jetty.io.EofException; import org.eclipse.jetty.server.Handler; @@ -134,6 +137,7 @@ import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; import java.util.zip.Deflater; @@ -248,6 +252,168 @@ public void testDeleteBroadcast() throws Exception latch.await(); } + /** {@code X-Druid-Native-Query-Route: local} dispatches a user's cancellation to the Router-local query resource. */ + @Test + public void testLocalNativeQueryRouteDispatchesDeleteLocally() throws Exception + { + assertNativeQueryRouteForDelete(QueryResource.NATIVE_QUERY_ROUTE_LOCAL, true); + } + + /** A non-exact native-query route value follows the normal cancellation forwarding path. */ + @Test + public void testNonExactNativeQueryRouteDoesNotDispatchDeleteLocally() throws Exception + { + assertNativeQueryRouteForDelete("LOCAL", false); + } + + /** A Router-local system table query reports the length of its reserialized body, not the original request body. */ + @Test + public void testLocalNativeQueryRouteUsesReserializedBodyLength() throws Exception + { + final ObjectMapper objectMapper = TestHelper.makeJsonMapper(); + final Query query = Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource("server_properties")) + .eternityInterval() + .build(); + final byte[] originalBody = StringUtils.toUtf8(" " + objectMapper.writeValueAsString(query) + " "); + final ByteArrayInputStream input = new ByteArrayInputStream(originalBody); + final ServletInputStream requestInputStream = new ServletInputStream() + { + @Override + public boolean isFinished() + { + return input.available() == 0; + } + + @Override + public boolean isReady() + { + return true; + } + + @Override + public void setReadListener(final ReadListener readListener) + { + // Synchronous in-memory request body. + } + + @Override + public int read() + { + return input.read(); + } + }; + + final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet( + new MapQueryToolChestWarehouse(ImmutableMap.of()), + objectMapper, + TestHelper.makeSmileMapper(), + Mockito.mock(QueryHostFinder.class), + null, + null, + NoopServiceEmitter.instance(), + NoopRequestLogger.instance(), + new DefaultGenericQueryMetricsFactory(), + new AuthenticatorMapper(ImmutableMap.of()), + new Properties(), + new ServerConfig() + ); + final GuiceContainer localQueryContainer = Mockito.mock(GuiceContainer.class); + servlet.setLocalQueryContainer(localQueryContainer); + + final HttpServletRequest request = Mockito.mock(HttpServletRequest.class); + final HttpServletResponse response = Mockito.mock(HttpServletResponse.class); + Mockito.when(request.getContentType()).thenReturn(MediaType.APPLICATION_JSON); + Mockito.when(request.getRequestURI()).thenReturn("/druid/v2"); + Mockito.when(request.getMethod()).thenReturn("POST"); + Mockito.when(request.getInputStream()).thenReturn(requestInputStream); + Mockito.when(request.getContentLength()).thenReturn(originalBody.length); + Mockito.when(request.getContentLengthLong()).thenReturn((long) originalBody.length); + Mockito.when(request.getHeader(HttpHeader.CONTENT_LENGTH.asString())) + .thenReturn(String.valueOf(originalBody.length)); + Mockito.when(request.getHeader(QueryResource.HEADER_NATIVE_QUERY_ROUTE)) + .thenReturn(QueryResource.NATIVE_QUERY_ROUTE_LOCAL); + + servlet.service(request, response); + + final ArgumentCaptor localRequest = ArgumentCaptor.forClass(HttpServletRequest.class); + Mockito.verify(localQueryContainer).service(localRequest.capture(), Mockito.same(response)); + final HttpServletRequest reserializedRequest = localRequest.getValue(); + final byte[] reserializedBody = IOUtils.toByteArray(reserializedRequest.getInputStream()); + Assertions.assertNotEquals(originalBody.length, reserializedBody.length); + Assertions.assertEquals(reserializedBody.length, reserializedRequest.getContentLength()); + Assertions.assertEquals(reserializedBody.length, reserializedRequest.getContentLengthLong()); + Assertions.assertEquals( + String.valueOf(reserializedBody.length), + reserializedRequest.getHeader(HttpHeader.CONTENT_LENGTH.asString()) + ); + Assertions.assertEquals( + Collections.singletonList(String.valueOf(reserializedBody.length)), + Collections.list(reserializedRequest.getHeaders(HttpHeader.CONTENT_LENGTH.asString())) + ); + Assertions.assertEquals( + reserializedBody.length, + reserializedRequest.getIntHeader(HttpHeader.CONTENT_LENGTH.asString()) + ); + } + + private void assertNativeQueryRouteForDelete(final String route, final boolean expectedLocal) throws Exception + { + final QueryHostFinder hostFinder = Mockito.mock(QueryHostFinder.class); + final org.apache.druid.client.selector.Server defaultServer = + new TestServer("http", "localhost", 8082); + Mockito.when(hostFinder.pickDefaultServer()).thenReturn(defaultServer); + Mockito.when(hostFinder.getAllServers()).thenReturn(List.of(defaultServer)); + + final AtomicInteger proxyCalls = new AtomicInteger(); + final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet( + new MapQueryToolChestWarehouse(ImmutableMap.of()), + TestHelper.makeJsonMapper(), + TestHelper.makeSmileMapper(), + hostFinder, + null, + null, + NoopServiceEmitter.instance(), + NoopRequestLogger.instance(), + new DefaultGenericQueryMetricsFactory(), + new AuthenticatorMapper(ImmutableMap.of()), + new Properties(), + new ServerConfig() + ) + { + @Override + protected void doService(final HttpServletRequest request, final HttpServletResponse response) + { + proxyCalls.incrementAndGet(); + } + }; + final GuiceContainer localQueryContainer = Mockito.mock(GuiceContainer.class); + servlet.setLocalQueryContainer(localQueryContainer); + + final HttpServletRequest request = Mockito.mock(HttpServletRequest.class); + final HttpServletResponse response = Mockito.mock(HttpServletResponse.class); + Mockito.when(request.getContentType()).thenReturn(MediaType.APPLICATION_JSON); + Mockito.when(request.getRequestURI()) + .thenReturn("/druid/v2/" + SystemTableDataSource.NODE_QUERY_ID_PREFIX + "query-id"); + Mockito.when(request.getMethod()).thenReturn("DELETE"); + Mockito.when(request.getHeader(QueryResource.HEADER_NATIVE_QUERY_ROUTE)) + .thenReturn(route); + Mockito.when(request.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)) + .thenReturn(new AuthenticationResult("alice", "allow", "external", null)); + + servlet.service(request, response); + + if (expectedLocal) { + final ArgumentCaptor localRequest = ArgumentCaptor.forClass(HttpServletRequest.class); + Mockito.verify(localQueryContainer).service(localRequest.capture(), Mockito.same(response)); + Assertions.assertEquals("", localRequest.getValue().getServletPath()); + Assertions.assertEquals(request.getRequestURI(), localRequest.getValue().getPathInfo()); + } else { + Mockito.verifyNoInteractions(localQueryContainer); + } + Assertions.assertEquals(expectedLocal ? 0 : 1, proxyCalls.get()); + } + @Test public void testSqlQueryProxy() throws Exception { diff --git a/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedClusterApis.java b/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedClusterApis.java index a4fb736ee649..fe0805dbf541 100644 --- a/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedClusterApis.java +++ b/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedClusterApis.java @@ -150,6 +150,15 @@ public String runSql(String sql, Object... args) return client.runSql(sql, args); } + public String runSql( + final String sql, + final Map context, + final Object... args + ) + { + return client.runSql(sql, context, args); + } + /** * Runs the given SQL query for a datasource and verifies the result. * diff --git a/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedDruidCluster.java b/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedDruidCluster.java index cb428d302765..6db8482236d9 100644 --- a/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedDruidCluster.java +++ b/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedDruidCluster.java @@ -33,6 +33,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Properties; import java.util.stream.Collectors; @@ -347,6 +348,11 @@ public String runSql(String sql, Object... args) return clusterApis.runSql(sql, args); } + public String runSql(final String sql, final Map context, final Object... args) + { + return clusterApis.runSql(sql, context, args); + } + EmbeddedDruidServer anyServer() { return servers.get(0); diff --git a/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedServiceClient.java b/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedServiceClient.java index 932b82cb41a3..b31da1eb200a 100644 --- a/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedServiceClient.java +++ b/services/src/test/java/org/apache/druid/testing/embedded/EmbeddedServiceClient.java @@ -53,6 +53,7 @@ import org.apache.druid.sql.http.ResultFormat; import javax.annotation.Nullable; +import java.util.Map; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.function.Function; @@ -232,6 +233,15 @@ private BrokerClient createBrokerClientForBroker(EmbeddedBroker targetBroker) * @return The result of the SQL as a single CSV string. */ public String runSql(String sql, Object... args) + { + return runSql(sql, null, args); + } + + public String runSql( + final String sql, + @Nullable final Map context, + final Object... args + ) { try { return onAnyBroker( @@ -242,7 +252,7 @@ public String runSql(String sql, Object... args) false, false, false, - null, + context, null ) ) diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java index 923bc1d198a2..25242b25cb0b 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java @@ -138,6 +138,9 @@ public class PlannerContext public static final String CTX_SQL_USE_GRANULARITY = "sqlUseGranularity"; public static final boolean DEFAULT_SQL_USE_GRANULARITY = true; + public static final String CTX_USE_NATIVE_QUERY_FOR_SYSTEM_TABLES = "useNativeQueryForSystemTables"; + public static final boolean DEFAULT_USE_NATIVE_QUERY_FOR_SYSTEM_TABLES = false; + // DataContext keys public static final String DATA_CTX_AUTHENTICATION_RESULT = "authenticationResult"; @@ -424,6 +427,14 @@ public QueryContext queryContext() return QueryContext.of(queryContext); } + public boolean useNativeQueryForSystemTables() + { + return queryContext().getBoolean( + CTX_USE_NATIVE_QUERY_FOR_SYSTEM_TABLES, + DEFAULT_USE_NATIVE_QUERY_FOR_SYSTEM_TABLES + ); + } + public boolean isStringifyArrays() { return stringifyArrays; diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java index 77a5d01dcaea..09e26bdd3570 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java @@ -77,6 +77,7 @@ import org.apache.druid.sql.calcite.rel.logical.DruidLogicalNode; import org.apache.druid.sql.calcite.run.EngineFeature; import org.apache.druid.sql.calcite.run.QueryMaker; +import org.apache.druid.sql.calcite.schema.SystemSchema; import org.apache.druid.sql.calcite.table.DruidTable; import org.apache.druid.sql.hook.DruidHook; import org.apache.druid.utils.Throwables; @@ -184,7 +185,7 @@ private static RelDataType getExplainStructType(RelDataTypeFactory typeFactory) public PlannerResult plan() { prepare(); - final Set bindableTables = getBindableTables(rootQueryRel.rel); + final Set bindableTables = getBindableTables(rootQueryRel.rel, handlerContext.plannerContext()); // the planner's type factory is not available until after parsing rexBuilder = new RexBuilder(handlerContext.planner().getTypeFactory()); @@ -259,7 +260,10 @@ public ExplainAttributes explainAttributes() ); } - private static Set getBindableTables(final RelNode relNode) + private static Set getBindableTables( + final RelNode relNode, + final PlannerContext plannerContext + ) { class HasBindableVisitor extends RelVisitor { @@ -270,8 +274,10 @@ public void visit(RelNode node, int ordinal, RelNode parent) { if (node instanceof TableScan) { RelOptTable table = node.getTable(); + // This is the routing point that decides whether a system table uses Bindable or native planning. if ((table.unwrap(ScannableTable.class) != null || table.unwrap(ProjectableFilterableTable.class) != null) - && table.unwrap(DruidTable.class) == null) { + && table.unwrap(DruidTable.class) == null + && !SystemSchema.canUseNativeSystemTable(table, plannerContext)) { found.add(table); return; } @@ -605,12 +611,20 @@ protected PlannerResult planWithDruidConvention() throws ValidationException .stream() .filter(action -> action.getAction() == Action.READ) .collect(Collectors.toSet()); + final Set authorizationTrackedDataSourceNames = new HashSet<>(druidRel.getDataSourceNames()); + if (plannerContext.useNativeQueryForSystemTables() + && !plannerContext.getPlannerConfig().isAuthorizeSystemTablesDirectly()) { + // Native system tables preserve the traditional row-level authorization model. They therefore have no + // SQL resource action unless direct system-table authorization is enabled, and must not make this + // datasource-resource sanity check fail. + authorizationTrackedDataSourceNames.removeAll(getNativeSystemTableNames(rootQueryRel.rel)); + } Preconditions.checkState( - readResourceActions.isEmpty() == druidRel.getDataSourceNames().isEmpty() + readResourceActions.isEmpty() == authorizationTrackedDataSourceNames.isEmpty() // The resources found in the plannerContext can be less than the datasources in // the query plan, because the query planner can eliminate empty tables by replacing // them with InlineDataSource of empty rows. - || readResourceActions.size() >= druidRel.getDataSourceNames().size(), + || readResourceActions.size() >= authorizationTrackedDataSourceNames.size(), "Authorization sanity check failed" ); @@ -619,6 +633,30 @@ protected PlannerResult planWithDruidConvention() throws ValidationException } } + /** + * Returns the datasource names of the native system tables scanned by the logical plan. Reading them from the + * logical plan avoids generating the native query just for the authorization sanity check. + */ + private static Set getNativeSystemTableNames(final RelNode relNode) + { + final Set names = new HashSet<>(); + new RelVisitor() + { + @Override + public void visit(RelNode node, int ordinal, RelNode parent) + { + if (node instanceof TableScan) { + final DruidTable nativeTable = SystemSchema.getNativeSystemTable(node.getTable()); + if (nativeTable != null) { + names.addAll(nativeTable.getDataSource().getTableNames()); + } + } + super.visit(node, ordinal, parent); + } + }.go(relNode); + return names; + } + /** * This method wraps the root with a {@link LogicalSort} that applies a limit (no ordering change). If the outer rel * is already a {@link Sort}, we can merge our outerLimit into it, similar to what is going on in diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rel/logical/DruidTableScan.java b/sql/src/main/java/org/apache/druid/sql/calcite/rel/logical/DruidTableScan.java index 45d97b04f32c..059170a7c85c 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/rel/logical/DruidTableScan.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/rel/logical/DruidTableScan.java @@ -35,6 +35,7 @@ import org.apache.calcite.schema.Table; import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.calcite.planner.querygen.SourceDescProducer; +import org.apache.druid.sql.calcite.schema.SystemSchema; import org.apache.druid.sql.calcite.table.DruidTable; import java.util.List; @@ -107,7 +108,11 @@ public SourceDesc getSourceDesc(PlannerContext plannerContext, List private DruidTable getDruidTable() { final RelOptTable table = getTable(); - final DruidTable druidTable = table.unwrap(DruidTable.class); + DruidTable druidTable = table.unwrap(DruidTable.class); + if (druidTable == null) { + // QueryHandler has already selected native planning; this decoupled stage only needs the native representation. + druidTable = SystemSchema.getNativeSystemTable(table); + } Preconditions.checkNotNull(druidTable, "DruidTable may not be null"); return druidTable; } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidTableScanRule.java b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidTableScanRule.java index 11533f7a2754..5c94738c07cf 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidTableScanRule.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidTableScanRule.java @@ -25,6 +25,7 @@ import org.apache.calcite.rel.logical.LogicalTableScan; import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.calcite.rel.DruidQueryRel; +import org.apache.druid.sql.calcite.schema.SystemSchema; import org.apache.druid.sql.calcite.table.DruidTable; public class DruidTableScanRule extends RelOptRule @@ -42,7 +43,11 @@ public void onMatch(final RelOptRuleCall call) { final LogicalTableScan scan = call.rel(0); final RelOptTable table = scan.getTable(); - final DruidTable druidTable = table.unwrap(DruidTable.class); + DruidTable druidTable = table.unwrap(DruidTable.class); + if (druidTable == null) { + // QueryHandler has already selected native planning, so only resolve the advertised native representation here. + druidTable = SystemSchema.getNativeSystemTable(table); + } if (druidTable != null) { call.transformTo( DruidQueryRel.scanTable(scan, table, druidTable, plannerContext) diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rule/logical/DruidBindableTableScanRule.java b/sql/src/main/java/org/apache/druid/sql/calcite/rule/logical/DruidBindableTableScanRule.java new file mode 100644 index 000000000000..ffcec29f16fc --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/calcite/rule/logical/DruidBindableTableScanRule.java @@ -0,0 +1,107 @@ +/* + * 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.rule.logical; + +import org.apache.calcite.interpreter.BindableConvention; +import org.apache.calcite.interpreter.Bindables; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.convert.ConverterRule; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexUtil; +import org.apache.druid.sql.calcite.rel.logical.DruidFilter; +import org.apache.druid.sql.calcite.rel.logical.DruidLogicalConvention; +import org.apache.druid.sql.calcite.rel.logical.DruidProject; +import org.apache.druid.sql.calcite.rel.logical.DruidTableScan; +import org.apache.druid.sql.calcite.schema.SystemSchema; + +import java.util.List; +import java.util.stream.Collectors; + +/// Converts a native system table that Calcite has represented as a [Bindables.BindableTableScan]. Calcite may push +/// filters and projections into this node for a `ProjectableFilterableTable`; they must be restored as Druid logical +/// nodes so native query generation can process them. +/// +/// This rule is used only by the **DECOUPLED** planner. The coupled planner converts a regular table scan through +/// `DruidTableScanRule` before Calcite produces a `BindableTableScan`. +/// +/// ``` +/// BindableTableScan(filters, projects) +/// -> DruidProject +/// `- DruidFilter +/// `- DruidTableScan +/// ``` +/// +/// The project and filter nodes are omitted when the Bindable scan contains an identity projection or no filters. +public class DruidBindableTableScanRule extends ConverterRule +{ + public DruidBindableTableScanRule() + { + super( + Config.INSTANCE.withConversion( + Bindables.BindableTableScan.class, + BindableConvention.INSTANCE, + DruidLogicalConvention.instance(), + DruidBindableTableScanRule.class.getSimpleName() + ) + ); + } + + @Override + public boolean matches(final RelOptRuleCall call) + { + final Bindables.BindableTableScan scan = call.rel(0); + // QueryHandler has already selected native planning; only native-capable system tables can use this conversion. + return SystemSchema.getNativeSystemTable(scan.getTable()) != null; + } + + @Override + public RelNode convert(final RelNode rel) + { + final Bindables.BindableTableScan bindableScan = (Bindables.BindableTableScan) rel; + RelNode current = new DruidTableScan( + bindableScan.getCluster(), + bindableScan.getTraitSet().replace(DruidLogicalConvention.instance()), + bindableScan.getTable() + ); + + if (!bindableScan.filters.isEmpty()) { + current = new DruidFilter( + bindableScan.getCluster(), + current.getTraitSet(), + current, + RexUtil.composeConjunction(bindableScan.getCluster().getRexBuilder(), bindableScan.filters) + ); + } + + if (!bindableScan.projects.equals(bindableScan.identity())) { + final RexBuilder rexBuilder = bindableScan.getCluster().getRexBuilder(); + final RelNode projectInput = current; + final List projects = bindableScan.projects + .stream() + .map(index -> rexBuilder.makeInputRef(projectInput, index)) + .collect(Collectors.toList()); + current = DruidProject.create(current, projects, bindableScan.getRowType()); + } + + return current; + } +} diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rule/logical/DruidLogicalRules.java b/sql/src/main/java/org/apache/druid/sql/calcite/rule/logical/DruidLogicalRules.java index 7d53aeb6fe35..0a8697219ebd 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/rule/logical/DruidLogicalRules.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/rule/logical/DruidLogicalRules.java @@ -57,6 +57,7 @@ public List rules() DruidLogicalConvention.instance(), DruidTableScanRule.class.getSimpleName() ), + new DruidBindableTableScanRule(), new DruidAggregateRule( LogicalAggregate.class, Convention.NONE, diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeServerPropertiesTable.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeServerPropertiesTable.java new file mode 100644 index 000000000000..57c4d926f156 --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeServerPropertiesTable.java @@ -0,0 +1,69 @@ +/* + * 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.sql.calcite.table.DruidTable; + +/** Native-query representation of {@code sys.server_properties}. */ +public class NativeServerPropertiesTable extends DruidTable +{ + private static final DataSource DATA_SOURCE = new SystemTableDataSource("server_properties"); + + NativeServerPropertiesTable() + { + super(SystemServerPropertiesTable.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/NativeSystemTable.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeSystemTable.java new file mode 100644 index 000000000000..cc98482a51a4 --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/NativeSystemTable.java @@ -0,0 +1,38 @@ +/* + * 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.druid.sql.calcite.table.DruidTable; + +/** + * Capability implemented by a traditional system table that also has a native-query representation. + * The capability is discovered from the table resolved through {@link SystemSchemaProvider}, so native planning + * inherits the provider's table-visibility authorization. + */ +interface NativeSystemTable +{ + /** + * Returns the representation used by the SQL planner after native system-table planning has been selected. + * The returned table supplies the native {@code DataSource} and row signature needed to translate the Calcite + * relational plan into a native Druid query. It does not read the system-table rows itself; those rows are supplied + * by the corresponding component-side system-table data provider when the native query executes. + */ + DruidTable asNativeTable(); +} 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 885fc1ca3522..555f2d771fc8 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 @@ -33,6 +33,7 @@ import org.apache.calcite.linq4j.Enumerable; import org.apache.calcite.linq4j.Enumerator; import org.apache.calcite.linq4j.Linq4j; +import org.apache.calcite.plan.RelOptTable; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.rex.RexNode; @@ -74,7 +75,10 @@ import org.apache.druid.server.security.ResourceAction; import org.apache.druid.server.security.ResourceType; import org.apache.druid.sql.calcite.planner.PlannerConfig; +import org.apache.druid.sql.calcite.planner.PlannerContext; +import org.apache.druid.sql.calcite.run.NativeSqlEngine; import org.apache.druid.sql.calcite.run.SqlEngine; +import org.apache.druid.sql.calcite.table.DruidTable; import org.apache.druid.sql.calcite.table.RowSignatures; import org.apache.druid.sql.http.GetQueriesResponse; import org.apache.druid.sql.http.QueryInfo; @@ -132,6 +136,27 @@ public class SystemSchema extends AbstractTableSchema private static final long IS_OVERSHADOWED_FALSE = 0L; private static final long IS_OVERSHADOWED_TRUE = 1L; + public static boolean canUseNativeSystemTable( + final RelOptTable table, + final PlannerContext plannerContext + ) + { + return plannerContext.useNativeQueryForSystemTables() + && NativeSqlEngine.NAME.equals(plannerContext.getEngine().name()) + && table.unwrap(NativeSystemTable.class) != null; + } + + /** + * Returns the native representation advertised by a system table resolved through {@link SystemSchemaProvider}. + * Eligibility for native planning must be checked with {@link #canUseNativeSystemTable} before calling this method. + */ + @Nullable + public static DruidTable getNativeSystemTable(final RelOptTable table) + { + final NativeSystemTable nativeSystemTable = table.unwrap(NativeSystemTable.class); + return nativeSystemTable == null ? null : nativeSystemTable.asNativeTable(); + } + static final RowSignature SEGMENTS_SIGNATURE = RowSignature .builder() .add("segment_id", ColumnType.STRING) diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemServerPropertiesTable.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemServerPropertiesTable.java index bf8e74c1050d..6849a2bf7ac4 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemServerPropertiesTable.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemServerPropertiesTable.java @@ -39,11 +39,11 @@ import org.apache.druid.java.util.http.client.Request; import org.apache.druid.java.util.http.client.response.StringFullResponseHandler; import org.apache.druid.java.util.http.client.response.StringFullResponseHolder; -import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.column.RowSignature; import org.apache.druid.server.DruidNode; import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.server.system.table.ServerPropertiesTableDescriptor; import org.apache.druid.sql.calcite.table.RowSignatures; import javax.annotation.Nullable; @@ -66,21 +66,13 @@ * that server would have multiple values in the column {@code node_roles} rather than duplicating all the * rows. */ -public class SystemServerPropertiesTable extends AbstractTable implements ProjectableFilterableTable +public class SystemServerPropertiesTable extends AbstractTable implements ProjectableFilterableTable, NativeSystemTable { private static final Logger log = new Logger(SystemServerPropertiesTable.class); - public static final String TABLE_NAME = "server_properties"; + public static final String TABLE_NAME = ServerPropertiesTableDescriptor.TABLE_NAME; - static final RowSignature ROW_SIGNATURE = RowSignature - .builder() - .add("server", ColumnType.STRING) - .add("service_name", ColumnType.STRING) - .add("node_roles", ColumnType.STRING) - .add("property", ColumnType.STRING) - .add("value", ColumnType.STRING) - .add("error_message", ColumnType.STRING) - .build(); + static final RowSignature ROW_SIGNATURE = ServerPropertiesTableDescriptor.ROW_SIGNATURE; private static final int SERVER_INDEX = ROW_SIGNATURE.indexOf("server"); private static final int SERVICE_NAME_INDEX = ROW_SIGNATURE.indexOf("service_name"); @@ -122,6 +114,12 @@ public Schema.TableType getJdbcTableType() return Schema.TableType.SYSTEM_TABLE; } + @Override + public NativeServerPropertiesTable asNativeTable() + { + return new NativeServerPropertiesTable(); + } + @Override public Enumerable scan( final DataContext root, diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSysQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSysQueryTest.java index ab0efe32d09c..586ce1f75b00 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSysQueryTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSysQueryTest.java @@ -19,14 +19,55 @@ package org.apache.druid.sql.calcite; +import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.ImmutableList; +import com.google.inject.Inject; +import org.apache.druid.server.QueryLifecycleFactory; +import org.apache.druid.sql.SqlToolbox; import org.apache.druid.sql.calcite.NotYetSupported.Modes; import org.apache.druid.sql.calcite.NotYetSupported.NotYetSupportedProcessor; +import org.apache.druid.sql.calcite.run.NativeSqlEngine; +import org.apache.druid.sql.calcite.run.SqlEngine; +import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; +@SqlTestFrameworkConfig.ComponentSupplier(CalciteSysQueryTest.BindableComponentSupplier.class) public class CalciteSysQueryTest extends BaseCalciteQueryTest { + public static class BindableComponentSupplier extends StandardComponentSupplier + { + public BindableComponentSupplier(final TempDirProducer tempDirProducer) + { + super(tempDirProducer); + } + + @Override + public Class getSqlEngineClass() + { + return BindableTestSqlEngine.class; + } + } + + public static class BindableTestSqlEngine extends NativeSqlEngine + { + @Inject + public BindableTestSqlEngine( + final QueryLifecycleFactory queryLifecycleFactory, + final ObjectMapper jsonMapper, + final SqlToolbox toolbox + ) + { + super(queryLifecycleFactory, jsonMapper, toolbox); + } + + @Override + public String name() + { + return "bindable-test"; + } + } + @RegisterExtension NotYetSupportedProcessor notYetSupportedProcessor = new NotYetSupportedProcessor(NotYetSupported.Scope.BINDABLE); diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/rule/logical/DruidBindableTableScanRuleTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/rule/logical/DruidBindableTableScanRuleTest.java new file mode 100644 index 000000000000..93a5467f63b4 --- /dev/null +++ b/sql/src/test/java/org/apache/druid/sql/calcite/rule/logical/DruidBindableTableScanRuleTest.java @@ -0,0 +1,170 @@ +/* + * 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.rule.logical; + +import org.apache.calcite.interpreter.Bindables; +import org.apache.calcite.plan.ConventionTraitDef; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.plan.RelOptTable; +import org.apache.calcite.plan.volcano.VolcanoPlanner; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.schema.ProjectableFilterableTable; +import org.apache.calcite.schema.Table; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.druid.sql.calcite.planner.DruidTypeSystem; +import org.apache.druid.sql.calcite.rel.logical.DruidFilter; +import org.apache.druid.sql.calcite.rel.logical.DruidProject; +import org.apache.druid.sql.calcite.rel.logical.DruidTableScan; +import org.apache.druid.sql.calcite.schema.SystemServerPropertiesTable; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.util.List; +import java.util.Optional; + +public class DruidBindableTableScanRuleTest +{ + private RelOptCluster cluster; + private RelDataType rowType; + private DruidBindableTableScanRule rule; + + @BeforeEach + public void setUp() + { + final RexBuilder rexBuilder = new RexBuilder(DruidTypeSystem.TYPE_FACTORY); + final VolcanoPlanner planner = new VolcanoPlanner(); + planner.addRelTraitDef(ConventionTraitDef.INSTANCE); + cluster = RelOptCluster.create(planner, rexBuilder); + rowType = DruidTypeSystem.TYPE_FACTORY.builder() + .add("server", SqlTypeName.VARCHAR) + .add("service_name", SqlTypeName.VARCHAR) + .add("node_roles", SqlTypeName.VARCHAR) + .add("property", SqlTypeName.VARCHAR) + .add("value", SqlTypeName.VARCHAR) + .add("error_message", SqlTypeName.VARCHAR) + .build(); + rule = new DruidBindableTableScanRule(); + } + + @Test + public void testMatchesNativeSystemTableOnly() + { + final Bindables.BindableTableScan nativeScan = createScan( + new SystemServerPropertiesTable(null, null, null, null, null), + List.of(), + identityProjects() + ); + final Bindables.BindableTableScan nonNativeScan = createScan( + Mockito.mock(ProjectableFilterableTable.class), + List.of(), + identityProjects() + ); + final RelOptRuleCall call = Mockito.mock(RelOptRuleCall.class); + + Mockito.when(call.rel(0)).thenReturn(nativeScan, nonNativeScan); + + Assertions.assertTrue(rule.matches(call)); + Assertions.assertFalse(rule.matches(call)); + } + + @Test + public void testConvertPreservesFilter() + { + final RexNode filter = equalsLiteral(3, "some.property"); + final Bindables.BindableTableScan bindableScan = createScan( + new SystemServerPropertiesTable(null, null, null, null, null), + List.of(filter), + identityProjects() + ); + + final DruidFilter converted = Assertions.assertInstanceOf(DruidFilter.class, rule.convert(bindableScan)); + Assertions.assertEquals(filter, converted.getCondition()); + Assertions.assertInstanceOf(DruidTableScan.class, converted.getInput()); + } + + @Test + public void testConvertPreservesFilterAndProjection() + { + final RexNode filter = equalsLiteral(3, "some.property"); + final Bindables.BindableTableScan bindableScan = createScan( + new SystemServerPropertiesTable(null, null, null, null, null), + List.of(filter), + List.of(3, 1) + ); + + final DruidProject converted = Assertions.assertInstanceOf(DruidProject.class, rule.convert(bindableScan)); + final DruidFilter convertedFilter = Assertions.assertInstanceOf(DruidFilter.class, converted.getInput()); + Assertions.assertEquals(filter, convertedFilter.getCondition()); + Assertions.assertInstanceOf(DruidTableScan.class, convertedFilter.getInput()); + Assertions.assertEquals(bindableScan.getRowType(), converted.getRowType()); + Assertions.assertEquals( + List.of(3, 1), + converted.getProjects().stream().map(project -> ((RexInputRef) project).getIndex()).toList() + ); + } + + private Bindables.BindableTableScan createScan( + final Table table, + final List filters, + final List projects + ) + { + return Bindables.BindableTableScan.create(cluster, createRelOptTable(table), filters, projects); + } + + private RelOptTable createRelOptTable(final Table table) + { + final RelOptTable relOptTable = Mockito.mock(RelOptTable.class); + Mockito.when(relOptTable.getRowType()).thenReturn(rowType); + Mockito.when(relOptTable.getQualifiedName()).thenReturn(List.of("sys", "test")); + Mockito.when(relOptTable.maybeUnwrap(ProjectableFilterableTable.class)).thenReturn( + Optional.of((ProjectableFilterableTable) table) + ); + Mockito.when(relOptTable.unwrap(Mockito.any())).thenAnswer( + invocation -> { + final Class requestedClass = invocation.getArgument(0); + return requestedClass.isInstance(table) ? table : null; + } + ); + return relOptTable; + } + + private RexNode equalsLiteral(final int column, final String value) + { + final RexBuilder rexBuilder = cluster.getRexBuilder(); + return rexBuilder.makeCall( + SqlStdOperatorTable.EQUALS, + rexBuilder.makeInputRef(rowType.getFieldList().get(column).getType(), column), + rexBuilder.makeLiteral(value) + ); + } + + private List identityProjects() + { + return List.of(0, 1, 2, 3, 4, 5); + } +} 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 new file mode 100644 index 000000000000..795208f8dfd8 --- /dev/null +++ b/sql/src/test/java/org/apache/druid/sql/calcite/schema/SystemTableDataProviderTest.java @@ -0,0 +1,94 @@ +/* + * 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.schema.Schema; +import org.apache.druid.query.SystemTableDataSource; +import org.apache.druid.sql.calcite.planner.PlannerContext; +import org.apache.druid.sql.calcite.run.NativeSqlEngine; +import org.apache.druid.sql.calcite.run.SqlEngine; +import org.easymock.EasyMock; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class SystemTableDataProviderTest +{ + @Test + public void testNativeTablesExposeSystemMetadata() + { + final NativeServerPropertiesTable serverProperties = new NativeServerPropertiesTable(); + Assertions.assertEquals( + "server_properties", + ((SystemTableDataSource) serverProperties.getDataSource()).getTable() + ); + Assertions.assertFalse(serverProperties.isJoinable()); + Assertions.assertFalse(serverProperties.isBroadcast()); + Assertions.assertEquals(Schema.TableType.SYSTEM_TABLE, serverProperties.getJdbcTableType()); + } + + @Test + public void testSystemSchemaSelectsNativeTablesOnlyWhenEnabledForNativeEngine() + { + final RelOptTable table = EasyMock.createMock(RelOptTable.class); + final SqlEngine nativeEngine = EasyMock.createMock(SqlEngine.class); + final SqlEngine bindableEngine = EasyMock.createMock(SqlEngine.class); + final PlannerContext disabledContext = EasyMock.createMock(PlannerContext.class); + final PlannerContext bindableContext = EasyMock.createMock(PlannerContext.class); + final PlannerContext enabledContext = EasyMock.createMock(PlannerContext.class); + EasyMock.expect(nativeEngine.name()).andReturn(NativeSqlEngine.NAME).anyTimes(); + EasyMock.expect(bindableEngine.name()).andReturn("bindable").anyTimes(); + EasyMock.expect(disabledContext.useNativeQueryForSystemTables()).andReturn(false).once(); + EasyMock.expect(bindableContext.useNativeQueryForSystemTables()).andReturn(true).once(); + 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.replay(table, nativeEngine, bindableEngine, disabledContext, bindableContext, enabledContext); + + Assertions.assertFalse(SystemSchema.canUseNativeSystemTable(table, disabledContext)); + Assertions.assertFalse(SystemSchema.canUseNativeSystemTable(table, bindableContext)); + Assertions.assertTrue(SystemSchema.canUseNativeSystemTable(table, enabledContext)); + + EasyMock.verify(table, nativeEngine, bindableEngine, disabledContext, bindableContext, enabledContext); + } + + @Test + public void testSystemSchemaGetsNativeRepresentation() + { + final RelOptTable table = EasyMock.createMock(RelOptTable.class); + EasyMock.expect(table.unwrap(NativeSystemTable.class)).andReturn(NativeServerPropertiesTable::new).once(); + EasyMock.replay(table); + + Assertions.assertInstanceOf(NativeServerPropertiesTable.class, SystemSchema.getNativeSystemTable(table)); + EasyMock.verify(table); + } + + @Test + public void testSystemSchemaRejectsTableWithoutNativeCapability() + { + final RelOptTable table = EasyMock.createMock(RelOptTable.class); + EasyMock.expect(table.unwrap(NativeSystemTable.class)).andReturn(null).once(); + EasyMock.replay(table); + + Assertions.assertNull(SystemSchema.getNativeSystemTable(table)); + EasyMock.verify(table); + } +} diff --git a/sql/src/test/java/org/apache/druid/sql/guice/SqlModuleTest.java b/sql/src/test/java/org/apache/druid/sql/guice/SqlModuleTest.java index 8a4b1908467e..aa41bf662fc1 100644 --- a/sql/src/test/java/org/apache/druid/sql/guice/SqlModuleTest.java +++ b/sql/src/test/java/org/apache/druid/sql/guice/SqlModuleTest.java @@ -33,6 +33,7 @@ import org.apache.druid.client.coordinator.NoopCoordinatorClient; import org.apache.druid.discovery.DruidNodeDiscoveryProvider; import org.apache.druid.guice.CatalogCoreModule; +import org.apache.druid.guice.DruidBinders; import org.apache.druid.guice.DruidGuiceExtensions; import org.apache.druid.guice.JsonConfigurator; import org.apache.druid.guice.LazySingleton; @@ -195,6 +196,7 @@ private Injector makeInjectorWithProperties(final Properties props) new AuthenticatorMapperModule(), new CatalogCoreModule(), binder -> { + DruidBinders.dataSourceQueryHandlerBinder(binder); binder.bind(Validator.class).toInstance(Validation.buildDefaultValidatorFactory().getValidator()); binder.bind(JsonConfigurator.class).in(LazySingleton.class); binder.bind(Properties.class).toInstance(props);