Optimize load logic

integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBLoadTsFileAuthIT.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.iotdb.db.it;
+
+import org.apache.iotdb.commons.auth.entity.PrivilegeType;
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.it.utils.TsFileGenerator;
+import org.apache.iotdb.itbase.category.ClusterIT;
+import org.apache.iotdb.itbase.category.LocalStandaloneIT;
+import org.apache.iotdb.jdbc.IoTDBSQLException;
+
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.external.commons.io.FileUtils;
+import org.apache.tsfile.file.metadata.enums.TSEncoding;
+import org.apache.tsfile.write.schema.IMeasurementSchema;
+import org.apache.tsfile.write.schema.MeasurementSchema;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.nio.file.Files;
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.Statement;
+import java.util.Collections;
+
+import static org.apache.iotdb.db.it.utils.TestUtils.assertNonQueryTestFail;
+import static org.apache.iotdb.db.it.utils.TestUtils.createUser;
+import static org.apache.iotdb.db.it.utils.TestUtils.executeNonQuery;
+import static org.apache.iotdb.db.it.utils.TestUtils.grantUserSeriesPrivilege;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({LocalStandaloneIT.class, ClusterIT.class})
+public class IoTDBLoadTsFileAuthIT {
+  private static final long PARTITION_INTERVAL = 10 * 1000L;
+  private static final String DATABASE = "root.load_auth";
+  private static final String DEVICE = DATABASE + ".d1";
+  private static final IMeasurementSchema MEASUREMENT =
+      new MeasurementSchema("s1", TSDataType.INT32, TSEncoding.RLE);
+  private static final String NO_WRITE_USER = "load_no_write_user";
+  private static final String WRITE_USER = "load_write_user";
+  private static final String OTHER_PATH_WRITE_USER = "load_other_path_write_user";
+  private static final String PASSWORD = "test123123456";
+
+  private static File tmpDir;
+
+  @BeforeClass
+  public static void setUp() throws Exception {
+    tmpDir = new File(Files.createTempDirectory("load-auth").toUri());
+    EnvFactory.getEnv().getConfig().getCommonConfig().setTimePartitionInterval(PARTITION_INTERVAL);
+    EnvFactory.getEnv().getConfig().getCommonConfig().setEnforceStrongPassword(false);
+    EnvFactory.getEnv().getConfig().getCommonConfig().setAutoCreateSchemaEnabled(false);
+    EnvFactory.getEnv()
+        .getConfig()
+        .getDataNodeConfig()
+        .setLoadTsFileAnalyzeSchemaMemorySizeInBytes(10 * 1024L);
+
+    EnvFactory.getEnv().initClusterEnvironment();
+  }
+
+  @AfterClass
+  public static void tearDown() throws Exception {
+    deleteDatabase();
+    EnvFactory.getEnv().cleanClusterEnvironment();
+    FileUtils.deleteDirectory(tmpDir);
+  }
+
+  @After
+  public void cleanData() throws Exception {
+    deleteDatabase();
+  }
+
+  @Test
+  public void testLoadWithoutSchemaCheckStillChecksWriteDataPermission() throws Exception {
+    final File tsFile = new File(tmpDir, "1-0-0-0.tsfile");
+    prepareSchemaAndTsFile(tsFile);
+    createUser(NO_WRITE_USER, PASSWORD);
+
+    assertNonQueryTestFail(
+        String.format("load \"%s\" with ('database-level'='2', 'verify'='false')", tsFile),
+        "No permissions for this operation, please add privilege WRITE_DATA",
+        NO_WRITE_USER,
+        PASSWORD);
+  }
+
+  @Test
+  public void testLoadWithoutSchemaCheckAllowsUserWithWriteDataPermission() throws Exception {
+    final File tsFile = new File(tmpDir, "2-0-0-0.tsfile");
+    prepareSchemaAndTsFile(tsFile);
+    createUser(WRITE_USER, PASSWORD);
+    grantUserSeriesPrivilege(WRITE_USER, PrivilegeType.WRITE_DATA, DATABASE + ".**");
+
+    executeNonQuery(
+        String.format("load \"%s\" with ('database-level'='2', 'verify'='false')", tsFile),
+        WRITE_USER,
+        PASSWORD);
+
+    try (final Connection connection = EnvFactory.getEnv().getConnection();
+        final Statement statement = connection.createStatement();
+        final ResultSet resultSet = statement.executeQuery("select count(s1) from " + DEVICE)) {
+      Assert.assertTrue(resultSet.next());
+      Assert.assertEquals(10, resultSet.getLong(1));
+    }
+  }
+
+  @Test
+  public void testLoadWithoutSchemaCheckRejectsUserWithOtherPathWriteDataPermission()
+      throws Exception {
+    final File tsFile = new File(tmpDir, "3-0-0-0.tsfile");
+    prepareSchemaAndTsFile(tsFile);
+    createUser(OTHER_PATH_WRITE_USER, PASSWORD);
+    grantUserSeriesPrivilege(OTHER_PATH_WRITE_USER, PrivilegeType.WRITE_DATA, "root.other.**");
+
+    assertNonQueryTestFail(
+        String.format("load \"%s\" with ('database-level'='2', 'verify'='false')", tsFile),
+        "No permissions for this operation, please add privilege WRITE_DATA",
+        OTHER_PATH_WRITE_USER,
+        PASSWORD);
+  }
+
+  private static void prepareSchemaAndTsFile(final File tsFile) throws Exception {
+    try (final Connection connection = EnvFactory.getEnv().getConnection();
+        final Statement statement = connection.createStatement()) {
+      statement.execute("create database " + DATABASE);
+      statement.execute(
+          String.format(
+              "create timeseries %s.%s %s",
+              DEVICE, MEASUREMENT.getMeasurementName(), MEASUREMENT.getType()));
+    }
+
+    try (final TsFileGenerator generator = new TsFileGenerator(tsFile)) {
+      generator.registerTimeseries(DEVICE, Collections.singletonList(MEASUREMENT));
+      generator.generateData(DEVICE, 10, PARTITION_INTERVAL / 10, false);
+    }
+  }
+
+  private static void deleteDatabase() throws Exception {
+    try (final Connection connection = EnvFactory.getEnv().getConnection();
+        final Statement statement = connection.createStatement()) {
+      statement.execute("delete database " + DATABASE);
+    } catch (final IoTDBSQLException ignored) {
+    }
+  }
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java

@@ -528,6 +528,8 @@ private void doAnalyzeSingleTreeFile(
 
       if (isAutoCreateSchemaOrVerifySchemaEnabled) {
         getOrCreateTreeSchemaVerifier().autoCreateAndVerify(reader, device2TimeseriesMetadata);
+      } else {
+        getOrCreateTreeSchemaVerifier().checkWritePermission(device2TimeseriesMetadata);
       }
 
       // TODO: how to get the correct write point count when

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/TreeSchemaAutoCreatorAndVerifier.java

@@ -154,20 +154,7 @@ public void autoCreateAndVerify(
           // not a timeseries, skip
         } else {
           // check WRITE_DATA permission of timeseries
-          long startTime = System.nanoTime();
-          try {
-            UserEntity userEntity = loadTsFileAnalyzer.context.getSession().getUserEntity();
-            TSStatus status =
-                AuthorityChecker.getAccessControl()
-                    .checkFullPathWriteDataPermission(
-                        userEntity, device, timeseriesMetadata.getMeasurementId());
-            if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
-              throw new AuthException(
-                  TSStatusCode.representOf(status.getCode()), status.getMessage());
-            }
-          } finally {
-            PerformanceOverviewMetrics.getInstance().recordAuthCost(System.nanoTime() - startTime);
-          }
+          checkWritePermission(device, timeseriesMetadata.getMeasurementId());
           final Pair<CompressionType, TSEncoding> compressionEncodingPair =
               reader.readTimeseriesCompressionTypeAndEncoding(timeseriesMetadata);
           schemaCache.addTimeSeries(
@@ -191,6 +178,63 @@ public void autoCreateAndVerify(
     }
   }
 
+  public void checkWritePermission(
+      Map<IDeviceID, List<TimeseriesMetadata>> device2TimeseriesMetadataList) throws AuthException {
+    for (final Map.Entry<IDeviceID, List<TimeseriesMetadata>> entry :
+        device2TimeseriesMetadataList.entrySet()) {
+      final IDeviceID device = entry.getKey();
+
+      try {
+        if (schemaCache.isDeviceDeletedByMods(device)) {
+          continue;
+        }
+      } catch (IllegalPathException e) {
+        LOGGER.warn(
+            DataNodeQueryMessages
+                .FAILED_TO_CHECK_IF_DEVICE_ARG_IS_DELETED_BY_MODS_WILL_SEE_IT_AS_NOT_DELETED,
+            device,
+            e);
+      }
+
+      for (final TimeseriesMetadata timeseriesMetadata : entry.getValue()) {
+        try {
+          if (schemaCache.isTimeSeriesDeletedByMods(device, timeseriesMetadata)) {
+            continue;
+          }
+        } catch (IllegalPathException e) {
+          if (!timeseriesMetadata.getMeasurementId().isEmpty()) {
+            LOGGER.warn(
+                DataNodeQueryMessages
+                    .FAILED_TO_CHECK_IF_DEVICE_ARG_TIMESERIES_ARG_IS_DELETED_BY_MODS_WILL_SEE_IT_AS_NOT,
+                device,
+                timeseriesMetadata.getMeasurementId(),
+                e);
+          }
+        }
+
+        if (!TSDataType.VECTOR.equals(timeseriesMetadata.getTsDataType())) {
+          checkWritePermission(device, timeseriesMetadata.getMeasurementId());
+        }
+      }
+    }
+  }
+
+  private void checkWritePermission(final IDeviceID device, final String measurementId)
+      throws AuthException {
+    final long startTime = System.nanoTime();
+    try {
+      UserEntity userEntity = loadTsFileAnalyzer.context.getSession().getUserEntity();
+      TSStatus status =
+          AuthorityChecker.getAccessControl()
+              .checkFullPathWriteDataPermission(userEntity, device, measurementId);
+      if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+        throw new AuthException(TSStatusCode.representOf(status.getCode()), status.getMessage());
+      }
+    } finally {
+      PerformanceOverviewMetrics.getInstance().recordAuthCost(System.nanoTime() - startTime);
+    }
+  }
+
   /**
    * This can only be invoked after all timeseries in the current tsfile have been processed.
    * Otherwise, the isAligned status may be wrong.