[CQ] Check READ_DATA permissions during continuous query execution

Author: JackieTien97

integration-test/src/test/java/org/apache/iotdb/db/it/cq/IoTDBCQExecIT.java

@@ -18,10 +18,18 @@
  */
 package org.apache.iotdb.db.it.cq;
 
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.client.ClientPoolFactory;
+import org.apache.iotdb.commons.client.IClientManager;
+import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient;
 import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper;
 import org.apache.iotdb.it.framework.IoTDBTestRunner;
 import org.apache.iotdb.itbase.category.ClusterIT;
 import org.apache.iotdb.itbase.category.LocalStandaloneIT;
+import org.apache.iotdb.mpp.rpc.thrift.TExecuteCQ;
+import org.apache.iotdb.rpc.TSStatusCode;
 
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -31,11 +39,16 @@
 
 import java.sql.Connection;
 import java.sql.ResultSet;
+import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.ArrayDeque;
+import java.util.Deque;
 import java.util.concurrent.TimeUnit;
 
+import static org.apache.iotdb.db.it.utils.TestUtils.resultSetEqualTest;
 import static org.apache.iotdb.itbase.constant.TestConstant.TIMESTAMP_STR;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 @RunWith(IoTDBTestRunner.class)
@@ -52,6 +65,205 @@ public static void tearDown() throws Exception {
     EnvFactory.getEnv().cleanClusterEnvironment();
   }
 
+  @Test
+  public void testCQExecutionWithReadDataPermission() throws Exception {
+    testCQExecutionWithReadDataPermission(false, false);
+  }
+
+  @Test
+  public void testCQExecutionWithReadDataPermissionAlignByDevice() throws Exception {
+    testCQExecutionWithReadDataPermission(true, false);
+  }
+
+  @Test
+  public void testCQExecutionWithExplicitTargetReadDataPermission() throws Exception {
+    testCQExecutionWithReadDataPermission(false, true);
+  }
+
+  @Test
+  public void testCQExecutionWithExplicitTargetReadDataPermissionAlignByDevice() throws Exception {
+    testCQExecutionWithReadDataPermission(true, true);
+  }
+
+  private void testCQExecutionWithReadDataPermission(boolean alignByDevice, boolean explicitTargets)
+      throws Exception {
+    String cqId =
+        (explicitTargets ? "cq_exp_auth" : "cq_read_auth") + (alignByDevice ? "_aligned" : "");
+    String database = "root." + cqId;
+    String username = cqId + "_user";
+    String password = "CQRead!Test2026";
+    String firstTarget = explicitTargets ? "m1" : "s1";
+    String secondTarget = explicitTargets ? "m2" : "s2";
+    String queryBody =
+        "SELECT "
+            + (explicitTargets ? "s1, s2" : "*")
+            + " INTO "
+            + database
+            + ".target("
+            + (explicitTargets ? "m1, m2" : "::")
+            + ") FROM "
+            + database
+            + ".source"
+            + (alignByDevice ? " ALIGN BY DEVICE" : "");
+    String resultQuery =
+        "SELECT " + firstTarget + ", " + secondTarget + " FROM " + database + ".target";
+    String resultHeader =
+        "Time,"
+            + database
+            + ".target."
+            + firstTarget
+            + ","
+            + database
+            + ".target."
+            + secondTarget
+            + ",";
+
+    IClientManager<TEndPoint, SyncDataNodeInternalServiceClient> clients =
+        new IClientManager.Factory<TEndPoint, SyncDataNodeInternalServiceClient>()
+            .createClientManager(
+                new ClientPoolFactory.SyncDataNodeInternalServiceClientPoolFactory());
+    try (Connection connection = EnvFactory.getEnv().getConnection();
+        Statement statement = connection.createStatement();
+        CQTestCleanup cleanup = new CQTestCleanup(statement)) {
+      statement.execute("CREATE USER " + username + " '" + password + "'");
+      cleanup.add("DROP USER " + username);
+      statement.execute("CREATE DATABASE " + database);
+      cleanup.add("DELETE DATABASE " + database);
+      statement.execute("CREATE TIMESERIES " + database + ".source.s1 WITH DATATYPE=INT64");
+      statement.execute("CREATE TIMESERIES " + database + ".source.s2 WITH DATATYPE=INT64");
+      statement.execute(
+          "CREATE TIMESERIES " + database + ".target." + firstTarget + " WITH DATATYPE=INT64");
+      statement.execute(
+          "CREATE TIMESERIES " + database + ".target." + secondTarget + " WITH DATATYPE=INT64");
+      statement.execute(
+          "INSERT INTO "
+              + database
+              + ".source(time, s1, s2) VALUES (1, 11, 21), (1001, 12, 22), (2001, 13, 23), (3001, 14, 24)");
+      statement.execute("GRANT SYSTEM ON root.** TO USER " + username);
+      statement.execute("GRANT WRITE_DATA ON " + database + ".target.** TO USER " + username);
+      statement.execute("GRANT READ_DATA ON " + database + ".source.s1 TO USER " + username);
+
+      try (Connection userConnection = EnvFactory.getEnv().getConnection(username, password);
+          Statement userStatement = userConnection.createStatement()) {
+        // Keep the scheduler idle and execute deterministic windows through its DataNode RPC.
+        userStatement.execute(
+            "CREATE CQ "
+                + cqId
+                + " RESAMPLE EVERY 1s BOUNDARY "
+                + (System.currentTimeMillis() + TimeUnit.DAYS.toMillis(1))
+                + " BEGIN "
+                + queryBody
+                + " END");
+        cleanup.add("DROP CQ " + cqId);
+      }
+
+      TExecuteCQ request = new TExecuteCQ(queryBody, 0, 1000, 10000, "UTC", cqId, username);
+      String[] expectedRows;
+      if (explicitTargets) {
+        String columnMismatch =
+            "the number of source columns and the number of target paths should be the same";
+        executeCQOnAllDataNodes(clients, request, TSStatusCode.SEMANTIC_ERROR, columnMismatch);
+        resultSetEqualTest(resultQuery, resultHeader, new String[0]);
+
+        // Explicit targets retain SELECT INTO's positional mapping. Restore access and retry the
+        // same window, as the CQ scheduler does after an execution failure.
+        statement.execute("GRANT READ_DATA ON " + database + ".source.s2 TO USER " + username);
+        executeCQOnAllDataNodes(clients, request, TSStatusCode.SUCCESS_STATUS, null);
+        resultSetEqualTest(resultQuery, resultHeader, new String[] {"1,11,21,"});
+
+        statement.execute("REVOKE READ_DATA ON " + database + ".source.s2 FROM USER " + username);
+        request.setStartTime(1000).setEndTime(2000);
+        executeCQOnAllDataNodes(clients, request, TSStatusCode.SEMANTIC_ERROR, columnMismatch);
+        resultSetEqualTest(resultQuery, resultHeader, new String[] {"1,11,21,"});
+
+        statement.execute("GRANT READ_DATA ON " + database + ".source.s2 TO USER " + username);
+        executeCQOnAllDataNodes(clients, request, TSStatusCode.SUCCESS_STATUS, null);
+        expectedRows = new String[] {"1,11,21,", "1001,12,22,"};
+      } else {
+        executeCQOnAllDataNodes(clients, request, TSStatusCode.SUCCESS_STATUS, null);
+        // The wildcard must not copy s2, for which the CQ owner has no READ_DATA privilege.
+        resultSetEqualTest(resultQuery, resultHeader, new String[] {"1,11,null,"});
+
+        statement.execute("REVOKE READ_DATA ON " + database + ".source.s1 FROM USER " + username);
+        request.setStartTime(1000).setEndTime(2000);
+        executeCQOnAllDataNodes(clients, request, TSStatusCode.SUCCESS_STATUS, null);
+        resultSetEqualTest(resultQuery, resultHeader, new String[] {"1,11,null,"});
+
+        statement.execute("GRANT READ_DATA ON " + database + ".source.** TO USER " + username);
+        request.setStartTime(2000).setEndTime(3000);
+        executeCQOnAllDataNodes(clients, request, TSStatusCode.SUCCESS_STATUS, null);
+        expectedRows = new String[] {"1,11,null,", "2001,13,23,"};
+      }
+      resultSetEqualTest(resultQuery, resultHeader, expectedRows);
+
+      statement.execute("DROP USER " + username);
+      cleanup.remove("DROP USER " + username);
+      request.setStartTime(3000).setEndTime(4000);
+      executeCQOnAllDataNodes(clients, request, TSStatusCode.EXECUTE_STATEMENT_ERROR, username);
+      resultSetEqualTest(resultQuery, resultHeader, expectedRows);
+    } finally {
+      clients.close();
+    }
+  }
+
+  private void executeCQOnAllDataNodes(
+      IClientManager<TEndPoint, SyncDataNodeInternalServiceClient> clients,
+      TExecuteCQ request,
+      TSStatusCode expectedStatus,
+      String expectedMessage)
+      throws Exception {
+    for (DataNodeWrapper dataNode : EnvFactory.getEnv().getDataNodeWrapperList()) {
+      try (SyncDataNodeInternalServiceClient client =
+          clients.borrowClient(
+              new TEndPoint(dataNode.getInternalAddress(), dataNode.getInternalPort()))) {
+        TSStatus status = client.executeCQ(request);
+        assertEquals(status.toString(), expectedStatus.getStatusCode(), status.getCode());
+        if (expectedMessage != null) {
+          assertTrue(status.toString(), status.isSetMessage());
+          assertTrue(status.getMessage(), status.getMessage().contains(expectedMessage));
+        }
+      }
+    }
+  }
+
+  private static class CQTestCleanup implements AutoCloseable {
+    private final Statement statement;
+    private final Deque<String> statements = new ArrayDeque<>();
+
+    private CQTestCleanup(Statement statement) {
+      this.statement = statement;
+    }
+
+    private void add(String sql) {
+      statements.push(sql);
+    }
+
+    private void remove(String sql) {
+      statements.remove(sql);
+    }
+
+    @Override
+    public void close() throws SQLException {
+      // Attempt every cleanup. Try-with-resources preserves the original test failure and adds
+      // cleanup failures as suppressed exceptions.
+      SQLException failure = null;
+      for (String sql : statements) {
+        try {
+          statement.execute(sql);
+        } catch (SQLException e) {
+          if (failure == null) {
+            failure = e;
+          } else {
+            failure.addSuppressed(e);
+          }
+        }
+      }
+      if (failure != null) {
+        throw failure;
+      }
+    }
+  }
+
   @Test
   public void testCQExecution1() {
     String insertTemplate =

iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java

@@ -21,6 +21,9 @@
 
 /** Compile-time i18n constants for DataNode misc subsystems (English). */
 public final class DataNodeMiscMessages {
+  public static final String LOG_FAILED_TO_AUTHORIZE_CQ_ARG_FOR_USER_ARG_ARG_B436A81F =
+      "Failed to authorize CQ {} for user {}: {}";
+
   public static final String MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237 =
       "The requested query does not belong to the current session.";
 

iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java

@@ -21,6 +21,9 @@
 
 /** 编译时国际化常量 - DataNode 杂项子系统（中文）。 */
 public final class DataNodeMiscMessages {
+  public static final String LOG_FAILED_TO_AUTHORIZE_CQ_ARG_FOR_USER_ARG_ARG_B436A81F =
+      "CQ {} 的用户 {} 鉴权失败：{}";
+
   public static final String MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237 =
       "请求的查询不属于当前会话。";
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java

@@ -187,6 +187,7 @@
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.schema.RollbackTableDevicesBlackListNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.schema.TableAttributeColumnDropNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.schema.TableSchemaQueryWriteVisitor;
+import org.apache.iotdb.db.queryengine.plan.relational.security.TreeAccessCheckContext;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.DeleteDevice;
 import org.apache.iotdb.db.queryengine.plan.statement.component.WhereCondition;
 import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement;
@@ -1875,17 +1876,16 @@ public TSStatus executeCQ(TExecuteCQ req) {
 
     SESSION_MANAGER.registerSession(session);
 
-    SESSION_MANAGER.supplySession(
-        session,
-        // TODO: User the real userId
-        -1,
-        req.getUsername(),
-        ZoneId.of(req.getZoneId()),
-        ClientVersion.V_1_0);
-
     String executedSQL = req.queryBody;
 
     try {
+      SESSION_MANAGER.supplySession(
+          session,
+          AuthorityChecker.getUserId(req.getUsername()).orElse(-1L),
+          req.getUsername(),
+          ZoneId.of(req.getZoneId()),
+          ClientVersion.V_1_0);
+
       QueryStatement s =
           (QueryStatement) StatementGenerator.createStatement(req.queryBody, session.getZoneId());
       if (s == null) {
@@ -1918,6 +1918,22 @@ public TSStatus executeCQ(TExecuteCQ req) {
       }
       executedSQL = String.join(" ", s.constructFormattedSQL().split("\n")).replaceAll(" +", " ");
 
+      // Refresh the owner's READ_DATA scope for each CQ execution.
+      TSStatus status =
+          AuthorityChecker.checkAuthority(
+              s,
+              new TreeAccessCheckContext(
+                      session.getUserId(), session.getUsername(), session.getClientAddress())
+                  .setSqlString(executedSQL));
+      if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+        LOGGER.warn(
+            DataNodeMiscMessages.LOG_FAILED_TO_AUTHORIZE_CQ_ARG_FOR_USER_ARG_ARG_B436A81F,
+            req.getCqId(),
+            req.getUsername(),
+            status);
+        return status;
+      }
+
       long queryId =
           SESSION_MANAGER.requestQueryId(session, SESSION_MANAGER.requestStatementId(session));
       // Create and cache dataset

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java

@@ -1395,7 +1395,7 @@ public TSStatus visitQuery(QueryStatement statement, TreeAccessCheckContext cont
       AUDIT_LOGGER.recordObjectAuthenticationAuditLog(
           context.setResult(false),
           () -> statement.getPaths().stream().distinct().collect(Collectors.toList()).toString());
-      return new TSStatus(e.getCode().getStatusCode());
+      return new TSStatus(e.getCode().getStatusCode()).setMessage(e.getMessage());
     }
     AUDIT_LOGGER.recordObjectAuthenticationAuditLog(
         context.setResult(true),