[Pipe] Refresh credential user IDs and retain privilege metadata across retries

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java

@@ -556,6 +556,9 @@ public boolean internallyDecreaseResourceReferenceCount(final String holderMessa
               holderMessage),
           e);
     } finally {
+      // Retries must repeat the privilege check, so release its metadata only with the event.
+      tableNames = null;
+      treeSchemaMap = null;
       if (Objects.nonNull(pipeName)) {
         try {
           PipeDataNodeSinglePipeMetrics.getInstance()
@@ -932,10 +935,6 @@ else if (Objects.nonNull(treeSchemaMap)) {
                   e.getMessage());
       LOGGER.warn(errorMsg, e);
       throw new PipeException(errorMsg, e);
-    } finally {
-      // GC useless
-      tableNames = null;
-      treeSchemaMap = null;
     }
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TableConfigTaskVisitor.java

@@ -1370,6 +1370,15 @@ public static void checkAndEnrichSourceUser(
                   .FAILED_TO_S_PIPE_S_IN_IOTDB_SOURCE_PASSWORD_MUST_BE_SET_WHEN_THE_USERNAME_IS_SPECIFIED,
               isAlter ? DataNodeQueryMessages.ALTER : DataNodeQueryMessages.CREATE,
               pipeName));
+    } else {
+      updatePipeUserIdFromUsername(
+          replacedSourceAttributes,
+          sourceParameters,
+          PipeSourceConstant.SOURCE_IOTDB_USER_ID,
+          PipeSourceConstant.EXTRACTOR_IOTDB_USER_KEY,
+          PipeSourceConstant.SOURCE_IOTDB_USER_KEY,
+          PipeSourceConstant.EXTRACTOR_IOTDB_USERNAME_KEY,
+          PipeSourceConstant.SOURCE_IOTDB_USERNAME_KEY);
     }
   }
 
@@ -1442,6 +1451,15 @@ public static void checkAndEnrichSinkUser(
                   .FAILED_TO_S_PIPE_S_IN_WRITE_BACK_SINK_PASSWORD_MUST_BE_SET_WHEN_THE_USERNAME_IS,
               isAlter ? DataNodeQueryMessages.ALTER : DataNodeQueryMessages.CREATE,
               pipeName));
+    } else {
+      updatePipeUserIdFromUsername(
+          connectorAttributes,
+          connectorParameters,
+          PipeSinkConstant.SINK_IOTDB_USER_ID,
+          PipeSinkConstant.CONNECTOR_IOTDB_USER_KEY,
+          PipeSinkConstant.SINK_IOTDB_USER_KEY,
+          PipeSinkConstant.CONNECTOR_IOTDB_USERNAME_KEY,
+          PipeSinkConstant.SINK_IOTDB_USERNAME_KEY);
     }
   }
 
@@ -1510,6 +1528,14 @@ public static void markSourceAuthenticationAsExplicitIfNecessary(
         PipeSourceConstant.SOURCE_IOTDB_PASSWORD_KEY)) {
       sourceAttributes.put(
           SystemConstant.SOURCE_AUTHENTICATION_INJECTED_KEY, Boolean.FALSE.toString());
+      updatePipeUserIdFromUsername(
+          sourceAttributes,
+          sourceParameters,
+          PipeSourceConstant.SOURCE_IOTDB_USER_ID,
+          PipeSourceConstant.EXTRACTOR_IOTDB_USER_KEY,
+          PipeSourceConstant.SOURCE_IOTDB_USER_KEY,
+          PipeSourceConstant.EXTRACTOR_IOTDB_USERNAME_KEY,
+          PipeSourceConstant.SOURCE_IOTDB_USERNAME_KEY);
     }
   }
 
@@ -1524,6 +1550,26 @@ public static void markSinkAuthenticationAsExplicitIfNecessary(
         PipeSinkConstant.CONNECTOR_IOTDB_PASSWORD_KEY,
         PipeSinkConstant.SINK_IOTDB_PASSWORD_KEY)) {
       sinkAttributes.put(SystemConstant.SINK_AUTHENTICATION_INJECTED_KEY, Boolean.FALSE.toString());
+      updatePipeUserIdFromUsername(
+          sinkAttributes,
+          sinkParameters,
+          PipeSinkConstant.SINK_IOTDB_USER_ID,
+          PipeSinkConstant.CONNECTOR_IOTDB_USER_KEY,
+          PipeSinkConstant.SINK_IOTDB_USER_KEY,
+          PipeSinkConstant.CONNECTOR_IOTDB_USERNAME_KEY,
+          PipeSinkConstant.SINK_IOTDB_USERNAME_KEY);
+    }
+  }
+
+  private static void updatePipeUserIdFromUsername(
+      final Map<String, String> attributes,
+      final PipeParameters parameters,
+      final String userIdAttributeKey,
+      final String... usernameAttributeKeys) {
+    final String username = parameters.getStringByKeys(usernameAttributeKeys);
+    if (username != null) {
+      attributes.put(
+          userIdAttributeKey, String.valueOf(AuthorityChecker.getUserId(username).orElse(-1L)));
     }
   }
 

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/PipeTsFileInsertionEventTest.java

@@ -80,7 +80,8 @@ public void testAuthCheck() throws Exception {
             TsFileNameGenerator.generateNewTsFilePath(
                 TestConstant.BASE_OUTPUT_PATH + IoTDBConstant.SEQUENCE_FOLDER_NAME, 1, 1, 1, 1));
     try {
-      AuthorityChecker.setAccessControl(new TestAccessControl());
+      final TestAccessControl accessControl = Mockito.spy(new TestAccessControl());
+      AuthorityChecker.setAccessControl(accessControl);
 
       final TsFileResource resource = new TsFileResource(file);
       try (CompactionTestFileWriter writer = new CompactionTestFileWriter(resource)) {
@@ -126,6 +127,18 @@ public void testAuthCheck() throws Exception {
               Long.MIN_VALUE,
               Long.MAX_VALUE);
       Assert.assertThrows(AccessDeniedException.class, tableEvent::throwIfNoPrivilege);
+      // A denied event must remain retryable until SELECT is granted.
+      Assert.assertThrows(AccessDeniedException.class, tableEvent::throwIfNoPrivilege);
+      Mockito.doReturn(true)
+          .when(accessControl)
+          .checkCanSelectFromTable4Pipe(Mockito.anyString(), Mockito.any(), Mockito.any());
+      tableEvent.throwIfNoPrivilege();
+      // Processing can fail after a successful check, so the next attempt must check again.
+      tableEvent.throwIfNoPrivilege();
+      Mockito.doReturn(false)
+          .when(accessControl)
+          .checkCanSelectFromTable4Pipe(Mockito.anyString(), Mockito.any(), Mockito.any());
+      Assert.assertThrows(AccessDeniedException.class, tableEvent::throwIfNoPrivilege);
       tableEvent.close();
 
       final PipeTsFileInsertionEvent treeEvent =
@@ -157,6 +170,8 @@ public void testAuthCheck() throws Exception {
 
       treeEvent.setTreeSchemaMap(Collections.singletonMap(deviceID, new String[] {"s0", "s1"}));
       Assert.assertThrows(AccessDeniedException.class, treeEvent::throwIfNoPrivilege);
+      // Retrying a realtime event must not fall back to the historical coarse filter.
+      Assert.assertThrows(AccessDeniedException.class, treeEvent::throwIfNoPrivilege);
 
       treeEvent.close();
     } finally {

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/config/PipeAuthenticationIdentityTest.java

@@ -0,0 +1,131 @@
+/*
+ * 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.iotdb.db.queryengine.plan.execution.config;
+
+import org.apache.iotdb.commons.audit.UserEntity;
+import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant;
+import org.apache.iotdb.commons.pipe.config.constant.PipeSourceConstant;
+import org.apache.iotdb.commons.pipe.config.constant.SystemConstant;
+import org.apache.iotdb.db.auth.AuthorityChecker;
+import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.when;
+import static org.powermock.api.mockito.PowerMockito.mockStatic;
+
+@PowerMockIgnore({"com.sun.org.apache.xerces.*", "javax.xml.*", "org.xml.*", "javax.management.*"})
+@RunWith(PowerMockRunner.class)
+@PrepareForTest(AuthorityChecker.class)
+public class PipeAuthenticationIdentityTest {
+
+  @Test
+  public void testModifiedCredentialsReplaceInjectedSuperUserId() {
+    mockStatic(AuthorityChecker.class);
+    when(AuthorityChecker.getUserId("thulab")).thenReturn(Optional.of(10001L));
+
+    final Map<String, String> sourceAttributes = new HashMap<>();
+    TableConfigTaskVisitor.checkAndEnrichSourceUser(
+        "pipe", sourceAttributes, new UserEntity(0L, "root", "host"), false);
+
+    final Map<String, String> sourceChanges = new HashMap<>();
+    sourceChanges.put("username", "thulab");
+    sourceChanges.put("password", "password");
+    TableConfigTaskVisitor.markSourceAuthenticationAsExplicitIfNecessary(sourceChanges);
+    new PipeParameters(sourceAttributes)
+        .addOrReplaceEquivalentAttributes(new PipeParameters(sourceChanges));
+
+    final Map<String, String> sinkAttributes = new HashMap<>();
+    sinkAttributes.put(PipeSinkConstant.SINK_KEY, "write-back-sink");
+    TableConfigTaskVisitor.checkAndEnrichSinkUser(
+        "pipe", sinkAttributes, new UserEntity(0L, "root", "host"), false);
+
+    final Map<String, String> sinkChanges = new HashMap<>();
+    sinkChanges.put(PipeSinkConstant.CONNECTOR_IOTDB_USER_KEY, "thulab");
+    sinkChanges.put(PipeSinkConstant.CONNECTOR_IOTDB_PASSWORD_KEY, "password");
+    TableConfigTaskVisitor.markSinkAuthenticationAsExplicitIfNecessary(sinkChanges);
+    new PipeParameters(sinkAttributes)
+        .addOrReplaceEquivalentAttributes(new PipeParameters(sinkChanges));
+
+    assertEquals("10001", sourceAttributes.get(PipeSourceConstant.SOURCE_IOTDB_USER_ID));
+    assertEquals(
+        Boolean.FALSE.toString(),
+        sourceAttributes.get(SystemConstant.SOURCE_AUTHENTICATION_INJECTED_KEY));
+    assertEquals("10001", sinkAttributes.get(PipeSinkConstant.SINK_IOTDB_USER_ID));
+    assertEquals(
+        Boolean.FALSE.toString(),
+        sinkAttributes.get(SystemConstant.SINK_AUTHENTICATION_INJECTED_KEY));
+  }
+
+  @Test
+  public void testExplicitCredentialsUseTheirOwnUserIdOnCreateAndReplace() {
+    mockStatic(AuthorityChecker.class);
+    when(AuthorityChecker.getUserId("thulab")).thenReturn(Optional.of(10001L));
+
+    for (final boolean isAlter : new boolean[] {false, true}) {
+      final Map<String, String> sourceAttributes = new HashMap<>();
+      sourceAttributes.put(PipeSourceConstant.EXTRACTOR_IOTDB_USER_KEY, "thulab");
+      sourceAttributes.put(PipeSourceConstant.EXTRACTOR_IOTDB_PASSWORD_KEY, "password");
+      TableConfigTaskVisitor.checkAndEnrichSourceUser(
+          "pipe", sourceAttributes, new UserEntity(0L, "root", "host"), isAlter);
+
+      final Map<String, String> sinkAttributes = new HashMap<>();
+      sinkAttributes.put(PipeSinkConstant.SINK_KEY, "write-back-sink");
+      sinkAttributes.put(PipeSinkConstant.SINK_IOTDB_USERNAME_KEY, "thulab");
+      sinkAttributes.put(PipeSinkConstant.SINK_IOTDB_PASSWORD_KEY, "password");
+      TableConfigTaskVisitor.checkAndEnrichSinkUser(
+          "pipe", sinkAttributes, new UserEntity(0L, "root", "host"), isAlter);
+
+      assertEquals("10001", sourceAttributes.get(PipeSourceConstant.SOURCE_IOTDB_USER_ID));
+      assertEquals("10001", sinkAttributes.get(PipeSinkConstant.SINK_IOTDB_USER_ID));
+    }
+  }
+
+  @Test
+  public void testPasswordOnlyModificationPreservesExistingUserId() {
+    final Map<String, String> sourceAttributes = new HashMap<>();
+    sourceAttributes.put(PipeSourceConstant.SOURCE_IOTDB_USER_ID, "10001");
+    final Map<String, String> sourceChanges = new HashMap<>();
+    sourceChanges.put("password", "new-password");
+    TableConfigTaskVisitor.markSourceAuthenticationAsExplicitIfNecessary(sourceChanges);
+    new PipeParameters(sourceAttributes)
+        .addOrReplaceEquivalentAttributes(new PipeParameters(sourceChanges));
+
+    final Map<String, String> sinkAttributes = new HashMap<>();
+    sinkAttributes.put(PipeSinkConstant.SINK_IOTDB_USER_ID, "10001");
+    final Map<String, String> sinkChanges = new HashMap<>();
+    sinkChanges.put("password", "new-password");
+    TableConfigTaskVisitor.markSinkAuthenticationAsExplicitIfNecessary(sinkChanges);
+    new PipeParameters(sinkAttributes)
+        .addOrReplaceEquivalentAttributes(new PipeParameters(sinkChanges));
+
+    assertEquals("10001", sourceAttributes.get(PipeSourceConstant.SOURCE_IOTDB_USER_ID));
+    assertEquals("10001", sinkAttributes.get(PipeSinkConstant.SINK_IOTDB_USER_ID));
+  }
+}