Add system tables support with registry and broker wiring

pinot-broker/pom.xml

@@ -33,6 +33,10 @@
     <pinot.root>${basedir}/..</pinot.root>
   </properties>
   <dependencies>
+    <dependency>
+      <groupId>org.apache.pinot</groupId>
+      <artifactId>pinot-java-client</artifactId>
+    </dependency>
     <dependency>
       <groupId>org.apache.pinot</groupId>
       <artifactId>pinot-query-runtime</artifactId>
@@ -90,6 +94,11 @@
       <artifactId>pinot-yammer</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.pinot</groupId>
+      <artifactId>pinot-system-table</artifactId>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>org.testng</groupId>
       <artifactId>testng</artifactId>

pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/SystemTableDataTableResource.java

@@ -0,0 +1,84 @@
+/**
+ * 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.pinot.broker.api.resources;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import javax.inject.Inject;
+import javax.ws.rs.POST;
+import javax.ws.rs.Path;
+import javax.ws.rs.Produces;
+import javax.ws.rs.core.Context;
+import javax.ws.rs.core.HttpHeaders;
+import javax.ws.rs.core.MediaType;
+import javax.ws.rs.core.Response;
+import org.apache.pinot.broker.requesthandler.SystemTableBrokerRequestHandler;
+import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.datatable.DataTableImplV4;
+import org.apache.pinot.core.auth.ManualAuthorization;
+import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.trace.RequestScope;
+import org.apache.pinot.spi.trace.Tracing;
+import org.apache.pinot.spi.utils.JsonUtils;
+
+import static org.apache.pinot.spi.utils.CommonConstants.Broker.Request.SQL;
+
+
+/**
+ * Internal endpoint used for broker-to-broker scatter-gather for system tables.
+ * <p>
+ * Returns a serialized {@link DataTable} payload (application/octet-stream) that represents the local broker's shard
+ * of a system table query.
+ */
+@Path("/")
+public class SystemTableDataTableResource {
+  @Inject
+  private SystemTableBrokerRequestHandler _systemTableBrokerRequestHandler;
+
+  @POST
+  @Produces(MediaType.APPLICATION_OCTET_STREAM)
+  @Path("query/systemTable/datatable")
+  @ManualAuthorization
+  public Response processSystemTableDataTable(String requestBody,
+      @Context org.glassfish.grizzly.http.server.Request requestContext, @Context HttpHeaders httpHeaders) {
+    DataTable dataTable;
+    try {
+      JsonNode requestJson = JsonUtils.stringToJsonNode(requestBody);
+      if (requestJson == null || !requestJson.isObject() || !requestJson.has(SQL)) {
+        dataTable = new DataTableImplV4();
+        dataTable.addException(QueryErrorCode.JSON_PARSING, "Payload is missing the query string field 'sql'");
+      } else {
+        try (RequestScope requestScope = Tracing.getTracer().createRequestScope()) {
+          requestScope.setRequestArrivalTimeMillis(System.currentTimeMillis());
+          dataTable = _systemTableBrokerRequestHandler.handleSystemTableDataTableRequest(requestJson,
+              PinotClientRequest.makeHttpIdentity(requestContext), requestScope, httpHeaders);
+        }
+      }
+    } catch (Exception e) {
+      dataTable = new DataTableImplV4();
+      dataTable.addException(QueryErrorCode.QUERY_EXECUTION, e.getMessage());
+    }
+
+    try {
+      return Response.ok(dataTable.toBytes(), MediaType.APPLICATION_OCTET_STREAM).build();
+    } catch (Exception e) {
+      // As a last resort, return an empty body; the caller will treat this as a failure.
+      return Response.ok(new byte[0], MediaType.APPLICATION_OCTET_STREAM).build();
+    }
+  }
+}

pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerAdminApiApplication.java

@@ -34,6 +34,7 @@
 import org.apache.helix.HelixManager;
 import org.apache.pinot.broker.queryquota.QueryQuotaManager;
 import org.apache.pinot.broker.requesthandler.BrokerRequestHandler;
+import org.apache.pinot.broker.requesthandler.SystemTableBrokerRequestHandler;
 import org.apache.pinot.broker.routing.manager.BrokerRoutingManager;
 import org.apache.pinot.common.audit.AuditLogFilter;
 import org.apache.pinot.common.cursors.AbstractResponseStore;
@@ -78,6 +79,7 @@ public class BrokerAdminApiApplication extends ResourceConfig {
   public BrokerAdminApiApplication(BrokerRoutingManager routingManager, BrokerRequestHandler brokerRequestHandler,
       BrokerMetrics brokerMetrics, PinotConfiguration brokerConf, SqlQueryExecutor sqlQueryExecutor,
       ServerRoutingStatsManager serverRoutingStatsManager, AccessControlFactory accessFactory,
+      SystemTableBrokerRequestHandler systemTableBrokerRequestHandler,
       HelixManager helixManager, QueryQuotaManager queryQuotaManager, ThreadAccountant threadAccountant,
       AbstractResponseStore responseStore) {
     _brokerResourcePackages = brokerConf.getProperty(CommonConstants.Broker.BROKER_RESOURCE_PACKAGES,
@@ -108,6 +110,7 @@ protected void configure() {
         bind(sqlQueryExecutor).to(SqlQueryExecutor.class);
         bind(routingManager).to(BrokerRoutingManager.class);
         bind(brokerRequestHandler).to(BrokerRequestHandler.class);
+        bind(systemTableBrokerRequestHandler).to(SystemTableBrokerRequestHandler.class);
         bind(brokerMetrics).to(BrokerMetrics.class);
         String loggerRootDir = brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_LOGGER_ROOT_DIR);
         if (loggerRootDir != null) {

pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java

@@ -55,6 +55,7 @@
 import org.apache.pinot.broker.requesthandler.MultiStageBrokerRequestHandler;
 import org.apache.pinot.broker.requesthandler.MultiStageQueryThrottler;
 import org.apache.pinot.broker.requesthandler.SingleConnectionBrokerRequestHandler;
+import org.apache.pinot.broker.requesthandler.SystemTableBrokerRequestHandler;
 import org.apache.pinot.broker.requesthandler.TimeSeriesRequestHandler;
 import org.apache.pinot.broker.routing.manager.BrokerRoutingManager;
 import org.apache.pinot.common.Utils;
@@ -73,6 +74,7 @@
 import org.apache.pinot.common.metrics.BrokerMeter;
 import org.apache.pinot.common.metrics.BrokerMetrics;
 import org.apache.pinot.common.metrics.BrokerTimer;
+import org.apache.pinot.common.systemtable.SystemTableRegistry;
 import org.apache.pinot.common.utils.PinotAppConfigs;
 import org.apache.pinot.common.utils.ServiceStartableUtils;
 import org.apache.pinot.common.utils.ServiceStatus;
@@ -157,6 +159,7 @@ public abstract class BaseBrokerStarter implements ServiceStartable {
   protected BrokerMetrics _brokerMetrics;
   protected BrokerRoutingManager _routingManager;
   protected AccessControlFactory _accessControlFactory;
+  protected SystemTableBrokerRequestHandler _systemTableBrokerRequestHandler;
   protected BrokerRequestHandler _brokerRequestHandler;
   protected SqlQueryExecutor _sqlQueryExecutor;
   protected BrokerAdminApiApplication _brokerAdminApplication;
@@ -353,6 +356,7 @@ public void start()
     boolean caseInsensitive =
         _brokerConf.getProperty(Helix.ENABLE_CASE_INSENSITIVE_KEY, Helix.DEFAULT_ENABLE_CASE_INSENSITIVE);
     _tableCache = new ZkTableCache(_propertyStore, caseInsensitive);
+    SystemTableRegistry.init(_tableCache, _helixAdmin, _clusterName, _brokerConf);
 
     LOGGER.info("Initializing Broker Event Listener Factory");
     BrokerQueryEventListenerFactory.init(_brokerConf.subset(Broker.EVENT_LISTENER_CONFIG_PREFIX));
@@ -460,9 +464,13 @@ public void start()
     _responseStore.init(responseStoreConfiguration.subset(_responseStore.getType()), _hostname, _port, brokerId,
         _brokerMetrics, expirationTime);
 
+    _systemTableBrokerRequestHandler =
+        new SystemTableBrokerRequestHandler(_brokerConf, brokerId, requestIdGenerator, _routingManager,
+            _accessControlFactory, _queryQuotaManager, _tableCache, _threadAccountant, multiClusterRoutingContext,
+            _spectatorHelixManager);
     _brokerRequestHandler =
-        new BrokerRequestHandlerDelegate(singleStageBrokerRequestHandler, multiStageBrokerRequestHandler,
-            timeSeriesRequestHandler, _responseStore);
+        new BrokerRequestHandlerDelegate(singleStageBrokerRequestHandler, _systemTableBrokerRequestHandler,
+            multiStageBrokerRequestHandler, timeSeriesRequestHandler, _responseStore);
     _brokerRequestHandler.start();
 
     String controllerUrl = _brokerConf.getProperty(Broker.CONTROLLER_URL);
@@ -775,6 +783,11 @@ public void stop() {
     } catch (IOException e) {
       LOGGER.error("Caught exception when shutting down PinotFsFactory", e);
     }
+    try {
+      SystemTableRegistry.close();
+    } catch (Exception e) {
+      LOGGER.warn("Failed to close system table registry cleanly", e);
+    }
 
     LOGGER.info("Disconnecting spectator Helix manager");
     _spectatorHelixManager.disconnect();
@@ -821,7 +834,8 @@ public BrokerRequestHandler getBrokerRequestHandler() {
   protected BrokerAdminApiApplication createBrokerAdminApp() {
     BrokerAdminApiApplication brokerAdminApiApplication =
         new BrokerAdminApiApplication(_routingManager, _brokerRequestHandler, _brokerMetrics, _brokerConf,
-            _sqlQueryExecutor, _serverRoutingStatsManager, _accessControlFactory, _spectatorHelixManager,
+            _sqlQueryExecutor, _serverRoutingStatsManager, _accessControlFactory, _systemTableBrokerRequestHandler,
+            _spectatorHelixManager,
             _queryQuotaManager, _threadAccountant, _responseStore);
     brokerAdminApiApplication.register(
         new AuditServiceBinder(_defaultClusterConfigChangeHandler, getServiceRole(), _brokerMetrics));

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestHandlerDelegate.java

@@ -43,27 +43,29 @@
 
 /**
  * {@code BrokerRequestHandlerDelegate} delegates the inbound broker request to one of the enabled
- * {@link BrokerRequestHandler} based on the requested handle type.
- *
- * {@see: @CommonConstant
+ * {@link BrokerRequestHandler} implementations based on the request type.
  */
 public class BrokerRequestHandlerDelegate implements BrokerRequestHandler {
   private final BaseSingleStageBrokerRequestHandler _singleStageBrokerRequestHandler;
+  private final SystemTableBrokerRequestHandler _systemTableBrokerRequestHandler;
   private final MultiStageBrokerRequestHandler _multiStageBrokerRequestHandler;
   private final TimeSeriesRequestHandler _timeSeriesRequestHandler;
   private final AbstractResponseStore _responseStore;
 
   public BrokerRequestHandlerDelegate(BaseSingleStageBrokerRequestHandler singleStageBrokerRequestHandler,
+      SystemTableBrokerRequestHandler systemTableBrokerRequestHandler,
       @Nullable MultiStageBrokerRequestHandler multiStageBrokerRequestHandler,
       @Nullable TimeSeriesRequestHandler timeSeriesRequestHandler, AbstractResponseStore responseStore) {
     _singleStageBrokerRequestHandler = singleStageBrokerRequestHandler;
+    _systemTableBrokerRequestHandler = systemTableBrokerRequestHandler;
     _multiStageBrokerRequestHandler = multiStageBrokerRequestHandler;
     _timeSeriesRequestHandler = timeSeriesRequestHandler;
     _responseStore = responseStore;
   }
 
   @Override
   public void start() {
+    _systemTableBrokerRequestHandler.start();
     _singleStageBrokerRequestHandler.start();
     if (_multiStageBrokerRequestHandler != null) {
       _multiStageBrokerRequestHandler.start();
@@ -75,6 +77,7 @@ public void start() {
 
   @Override
   public void shutDown() {
+    _systemTableBrokerRequestHandler.shutDown();
     _singleStageBrokerRequestHandler.shutDown();
     if (_multiStageBrokerRequestHandler != null) {
       _multiStageBrokerRequestHandler.shutDown();
@@ -108,11 +111,19 @@ public BrokerResponse handleRequest(JsonNode request, @Nullable SqlNodeAndOption
 
     BaseBrokerRequestHandler requestHandler = _singleStageBrokerRequestHandler;
     if (QueryOptionsUtils.isUseMultistageEngine(sqlNodeAndOptions.getOptions())) {
+      if (isSystemTableQuery(sqlNodeAndOptions)) {
+        requestContext.setErrorCode(QueryErrorCode.QUERY_VALIDATION);
+        return new BrokerResponseNative(QueryErrorCode.QUERY_VALIDATION,
+            "System tables require the single-stage query engine. Remove the `useMultistageEngine=true` query option "
+                + "to execute this query.");
+      }
       if (_multiStageBrokerRequestHandler != null) {
         requestHandler = _multiStageBrokerRequestHandler;
       } else {
         return new BrokerResponseNative(QueryErrorCode.INTERNAL, "V2 Multi-Stage query engine not enabled.");
       }
+    } else if (isSystemTableQuery(sqlNodeAndOptions)) {
+      requestHandler = _systemTableBrokerRequestHandler;
     }
 
     BrokerResponse response = requestHandler.handleRequest(request, sqlNodeAndOptions, requesterIdentity,
@@ -160,7 +171,7 @@ public boolean cancelQuery(long queryId, int timeoutMs, Executor executor, HttpC
       throws Exception {
     if (_multiStageBrokerRequestHandler != null && _multiStageBrokerRequestHandler.cancelQuery(
         queryId, timeoutMs, executor, connMgr, serverResponses)) {
-        return true;
+      return true;
     }
     return _singleStageBrokerRequestHandler.cancelQuery(queryId, timeoutMs, executor, connMgr, serverResponses);
   }
@@ -188,6 +199,10 @@ public OptionalLong getRequestIdByClientId(String clientQueryId) {
     return _singleStageBrokerRequestHandler.getRequestIdByClientId(clientQueryId);
   }
 
+  private static boolean isSystemTableQuery(SqlNodeAndOptions sqlNodeAndOptions) {
+    return SystemTableQueryDetector.containsSystemTableReference(sqlNodeAndOptions.getSqlNode());
+  }
+
   private CursorResponse getCursorResponse(Integer numRows, BrokerResponse response)
       throws Exception {
     if (numRows == null) {