Repository: apache/iceberg (language: python)

Issue:
Flink: refactor sink tests to reduce the number of combinations with parameterized tests
E.g. it is not really necessary to test different file format for distribution mode (hash, range). I found this when trying to add new tests for the new range distribution feature. We often run 20+ combinations for every test/method in some classes.


Discussion / hints:
@rodmeneses: Could you please take a look. I think it will greatly affect your testing for Sink V2
thanks @rodmeneses and @pvary for the review

Fix patch:
diff --git a/flink/v1.18/flink/src/main/java/org/apache/iceberg/flink/FlinkCatalogFactory.java b/flink/v1.18/flink/src/main/java/org/apache/iceberg/flink/FlinkCatalogFactory.java
index 1453753849ec..9b0c7a938920 100644
--- a/flink/v1.18/flink/src/main/java/org/apache/iceberg/flink/FlinkCatalogFactory.java
+++ b/flink/v1.18/flink/src/main/java/org/apache/iceberg/flink/FlinkCatalogFactory.java
@@ -70,6 +70,7 @@ public class FlinkCatalogFactory implements CatalogFactory {
   public static final String HADOOP_CONF_DIR = "hadoop-conf-dir";
   public static final String DEFAULT_DATABASE = "default-database";
   public static final String DEFAULT_DATABASE_NAME = "default";
+  public static final String DEFAULT_CATALOG_NAME = "default_catalog";
   public static final String BASE_NAMESPACE = "base-namespace";
 
   public static final String TYPE = "type";


Test patch:
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/CatalogTestBase.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/CatalogTestBase.java
index 91ed3c4adea3..062ff68d5d85 100644
--- a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/CatalogTestBase.java
+++ b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/CatalogTestBase.java
@@ -118,26 +118,4 @@ protected String getFullQualifiedTableName(String tableName) {
   static String getURI(HiveConf conf) {
     return conf.get(HiveConf.ConfVars.METASTOREURIS.varname);
   }
-
-  static String toWithClause(Map<String, String> props) {
-    StringBuilder builder = new StringBuilder();
-    builder.append("(");
-    int propCount = 0;
-    for (Map.Entry<String, String> entry : props.entrySet()) {
-      if (propCount > 0) {
-        builder.append(",");
-      }
-      builder
-          .append("'")
-          .append(entry.getKey())
-          .append("'")
-          .append("=")
-          .append("'")
-          .append(entry.getValue())
-          .append("'");
-      propCount++;
-    }
-    builder.append(")");
-    return builder.toString();
-  }
 }
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/SqlBase.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/SqlBase.java
new file mode 100644
index 000000000000..9411ea4f7d71
--- /dev/null
+++ b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/SqlBase.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.iceberg.flink;
+
+import static org.apache.iceberg.flink.FlinkCatalogFactory.DEFAULT_CATALOG_NAME;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.List;
+import java.util.Map;
+import org.apache.flink.table.api.TableEnvironment;
+import org.apache.flink.table.api.TableResult;
+import org.apache.flink.types.Row;
+import org.apache.flink.util.CloseableIterator;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+
+public abstract class SqlBase {
+  protected abstract TableEnvironment getTableEnv();
+
+  protected static TableResult exec(TableEnvironment env, String query, Object... args) {
+    return env.executeSql(String.format(query, args));
+  }
+
+  protected TableResult exec(String query, Object... args) {
+    return exec(getTableEnv(), query, args);
+  }
+
+  protected List<Row> sql(String query, Object... args) {
+    TableResult tableResult = exec(query, args);
+    try (CloseableIterator<Row> iter = tableResult.collect()) {
+      return Lists.newArrayList(iter);
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to collect table result", e);
+    }
+  }
+
+  protected void assertSameElements(Iterable<Row> expected, Iterable<Row> actual) {
+    assertThat(actual).isNotNull().containsExactlyInAnyOrderElementsOf(expected);
+  }
+
+  protected void assertSameElements(String message, Iterable<Row> expected, Iterable<Row> actual) {
+    assertThat(actual).isNotNull().as(message).containsExactlyInAnyOrderElementsOf(expected);
+  }
+
+  /**
+   * We can not drop currently used catalog after FLINK-29677, so we have make sure that we do not
+   * use the current catalog before dropping it. This method switches to the 'default_catalog' and
+   * drops the one requested.
+   *
+   * @param catalogName The catalog to drop
+   * @param ifExists If we should use the 'IF EXISTS' when dropping the catalog
+   */
+  protected void dropCatalog(String catalogName, boolean ifExists) {
+    sql("USE CATALOG %s", DEFAULT_CATALOG_NAME);
+    sql("DROP CATALOG %s %s", ifExists ? "IF EXISTS" : "", catalogName);
+  }
+
+  /**
+   * We can not drop currently used database after FLINK-33226, so we have make sure that we do not
+   * use the current database before dropping it. This method switches to the default database in
+   * the default catalog, and then it and drops the one requested.
+   *
+   * @param database The database to drop
+   * @param ifExists If we should use the 'IF EXISTS' when dropping the database
+   */
+  protected void dropDatabase(String database, boolean ifExists) {
+    String currentCatalog = getTableEnv().getCurrentCatalog();
+    sql("USE CATALOG %s", DEFAULT_CATALOG_NAME);
+    sql("USE %s", getTableEnv().listDatabases()[0]);
+    sql("USE CATALOG %s", currentCatalog);
+    sql("DROP DATABASE %s %s", ifExists ? "IF EXISTS" : "", database);
+  }
+
+  protected static String toWithClause(Map<String, String> props) {
+    StringBuilder builder = new StringBuilder();
+    builder.append("(");
+    int propCount = 0;
+    for (Map.Entry<String, String> entry : props.entrySet()) {
+      if (propCount > 0) {
+        builder.append(",");
+      }
+      builder
+          .append("'")
+          .append(entry.getKey())
+          .append("'")
+          .append("=")
+          .append("'")
+          .append(entry.getValue())
+          .append("'");
+      propCount++;
+    }
+    builder.append(")");
+    return builder.toString();
+  }
+}
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestBase.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestBase.java
index a74226092f38..401960c3591b 100644
--- a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestBase.java
+++ b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestBase.java
@@ -18,6 +18,7 @@
  */
 package org.apache.iceberg.flink;
 
+import static org.apache.iceberg.flink.FlinkCatalogFactory.DEFAULT_CATALOG_NAME;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import java.nio.file.Path;
@@ -26,7 +27,6 @@
 import org.apache.flink.table.api.TableEnvironment;
 import org.apache.flink.table.api.TableResult;
 import org.apache.flink.test.junit5.MiniClusterExtension;
-import org.apache.flink.test.util.TestBaseUtils;
 import org.apache.flink.types.Row;
 import org.apache.flink.util.CloseableIterator;
 import org.apache.hadoop.hive.conf.HiveConf;
@@ -40,7 +40,7 @@
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.api.io.TempDir;
 
-public abstract class TestBase extends TestBaseUtils {
+public abstract class TestBase extends SqlBase {
 
   @RegisterExtension
   public static MiniClusterExtension miniClusterExtension =
@@ -71,6 +71,7 @@ public static void stopMetastore() throws Exception {
     TestBase.catalog = null;
   }
 
+  @Override
   protected TableEnvironment getTableEnv() {
     if (tEnv == null) {
       synchronized (this) {
@@ -122,7 +123,23 @@ protected void assertSameElements(String message, Iterable<Row> expected, Iterab
    * @param ifExists If we should use the 'IF EXISTS' when dropping the catalog
    */
   protected void dropCatalog(String catalogName, boolean ifExists) {
-    sql("USE CATALOG default_catalog");
+    sql("USE CATALOG %s", DEFAULT_CATALOG_NAME);
     sql("DROP CATALOG %s %s", ifExists ? "IF EXISTS" : "", catalogName);
   }
+
+  /**
+   * We can not drop currently used database after FLINK-33226, so we have make sure that we do not
+   * use the current database before dropping it. This method switches to the default database in
+   * the default catalog, and then it and drops the one requested.
+   *
+   * @param database The database to drop
+   * @param ifExists If we should use the 'IF EXISTS' when dropping the database
+   */
+  protected void dropDatabase(String database, boolean ifExists) {
+    String currentCatalog = getTableEnv().getCurrentCatalog();
+    sql("USE CATALOG %s", DEFAULT_CATALOG_NAME);
+    sql("USE %s", getTableEnv().listDatabases()[0]);
+    sql("USE CATALOG %s", currentCatalog);
+    sql("DROP DATABASE %s %s", ifExists ? "IF EXISTS" : "", database);
+  }
 }
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestFlinkTableSink.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestFlinkTableSink.java
index b73300e3f170..2978a92945a2 100644
--- a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestFlinkTableSink.java
+++ b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestFlinkTableSink.java
@@ -18,36 +18,21 @@
  */
 package org.apache.iceberg.flink;
 
-import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assumptions.assumeThat;
 
-import java.util.Collections;
 import java.util.List;
-import java.util.Map;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
-import org.apache.flink.api.dag.Transformation;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 import org.apache.flink.table.api.EnvironmentSettings;
 import org.apache.flink.table.api.Expressions;
 import org.apache.flink.table.api.TableEnvironment;
 import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
-import org.apache.flink.table.api.internal.TableEnvironmentImpl;
-import org.apache.flink.table.operations.ModifyOperation;
-import org.apache.flink.table.planner.delegation.PlannerBase;
-import org.apache.flink.types.Row;
-import org.apache.iceberg.DataFile;
-import org.apache.iceberg.DistributionMode;
 import org.apache.iceberg.FileFormat;
 import org.apache.iceberg.Parameter;
 import org.apache.iceberg.Parameters;
 import org.apache.iceberg.Table;
-import org.apache.iceberg.TableProperties;
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.apache.iceberg.flink.source.BoundedTableFactory;
-import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
-import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -122,7 +107,7 @@ public void before() {
   @AfterEach
   public void clean() {
     sql("DROP TABLE IF EXISTS %s.%s", flinkDatabase, TABLE_NAME);
-    sql("DROP DATABASE IF EXISTS %s", flinkDatabase);
+    dropDatabase(flinkDatabase, true);
     BoundedTableFactory.clearDataSets();
     super.clean();
   }
@@ -169,39 +154,6 @@ public void testOverwriteTable() throws Exception {
         icebergTable, Lists.newArrayList(SimpleDataUtil.createRecord(2, "b")));
   }
 
-  @TestTemplate
-  public void testWriteParallelism() throws Exception {
-    List<Row> dataSet =
-        IntStream.range(1, 1000)
-            .mapToObj(i -> ImmutableList.of(Row.of(i, "aaa"), Row.of(i, "bbb"), Row.of(i, "ccc")))
-            .flatMap(List::stream)
-            .collect(Collectors.toList());
-    String dataId = BoundedTableFactory.registerDataSet(ImmutableList.of(dataSet));
-    sql(
-        "CREATE TABLE %s(id INT NOT NULL, data STRING NOT NULL)"
-            + " WITH ('connector'='BoundedSource', 'data-id'='%s')",
-        SOURCE_TABLE, dataId);
-
-    PlannerBase planner = (PlannerBase) ((TableEnvironmentImpl) getTableEnv()).getPlanner();
-    String insertSQL =
-        String.format(
-            "INSERT INTO %s /*+ OPTIONS('write-parallelism'='1') */ SELECT * FROM %s",
-            TABLE_NAME, SOURCE_TABLE);
-    ModifyOperation operation = (ModifyOperation) planner.getParser().parse(insertSQL).get(0);
-    Transformation<?> dummySink = planner.translate(Collections.singletonList(operation)).get(0);
-    Transformation<?> committer = dummySink.getInputs().get(0);
-    Transformation<?> writer = committer.getInputs().get(0);
-
-    assertThat(writer.getParallelism()).as("Should have the expected 1 parallelism.").isEqualTo(1);
-    writer
-        .getInputs()
-        .forEach(
-            input ->
-                assertThat(input.getParallelism())
-                    .as("Should have the expected parallelism.")
-                    .isEqualTo(isStreamingJob ? 2 : 4));
-  }
-
   @TestTemplate
   public void testReplacePartitions() throws Exception {
     assumeThat(isStreamingJob)
@@ -289,70 +241,4 @@ public void testInsertIntoPartition() throws Exception {
       sql("DROP TABLE IF EXISTS %s.%s", flinkDatabase, tableName);
     }
   }
-
-  @TestTemplate
-  public void testHashDistributeMode() throws Exception {
-    String tableName = "test_hash_distribution_mode";
-    Map<String, String> tableProps =
-        ImmutableMap.of(
-            "write.format.default",
-            format.name(),
-            TableProperties.WRITE_DISTRIBUTION_MODE,
-            DistributionMode.HASH.modeName());
-
-    // Initialize a BoundedSource table to precisely emit those rows in only one checkpoint.
-    List<Row> dataSet =
-        IntStream.range(1, 1000)
-            .mapToObj(i -> ImmutableList.of(Row.of(i, "aaa"), Row.of(i, "bbb"), Row.of(i, "ccc")))
-            .flatMap(List::stream)
-            .collect(Collectors.toList());
-    String dataId = BoundedTableFactory.registerDataSet(ImmutableList.of(dataSet));
-    sql(
-        "CREATE TABLE %s(id INT NOT NULL, data STRING NOT NULL)"
-            + " WITH ('connector'='BoundedSource', 'data-id'='%s')",
-        SOURCE_TABLE, dataId);
-
-    assertThat(sql("SELECT * FROM %s", SOURCE_TABLE))
-        .as("Should have the expected rows in source table.")
-        .containsExactlyInAnyOrderElementsOf(dataSet);
-
-    sql(
-        "CREATE TABLE %s(id INT, data VARCHAR) PARTITIONED BY (data) WITH %s",
-        tableName, toWithClause(tableProps));
-
-    try {
-      // Insert data set.
-      sql("INSERT INTO %s SELECT * FROM %s", tableName, SOURCE_TABLE);
-
-      assertThat(sql("SELECT * FROM %s", tableName))
-          .as("Should have the expected rows in sink table.")
-          .containsExactlyInAnyOrderElementsOf(dataSet);
-
-      // Sometimes we will have more than one checkpoint if we pass the auto checkpoint interval,
-      // thus producing multiple snapshots.  Here we assert that each snapshot has only 1 file per
-      // partition.
-      Table table = validationCatalog.loadTable(TableIdentifier.of(icebergNamespace, tableName));
-      Map<Long, List<DataFile>> snapshotToDataFiles = SimpleDataUtil.snapshotToDataFiles(table);
-      for (List<DataFile> dataFiles : snapshotToDataFiles.values()) {
-        if (dataFiles.isEmpty()) {
-          continue;
-        }
-
-        assertThat(
-                SimpleDataUtil.matchingPartitions(
-                    dataFiles, table.spec(), ImmutableMap.of("data", "aaa")))
-            .hasSize(1);
-        assertThat(
-                SimpleDataUtil.matchingPartitions(
-                    dataFiles, table.spec(), ImmutableMap.of("data", "bbb")))
-            .hasSize(1);
-        assertThat(
-                SimpleDataUtil.matchingPartitions(
-                    dataFiles, table.spec(), ImmutableMap.of("data", "ccc")))
-            .hasSize(1);
-      }
-    } finally {
-      sql("DROP TABLE IF EXISTS %s.%s", flinkDatabase, tableName);
-    }
-  }
 }
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestFlinkTableSinkExtended.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestFlinkTableSinkExtended.java
new file mode 100644
index 000000000000..b63547d433a4
--- /dev/null
+++ b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestFlinkTableSinkExtended.java
@@ -0,0 +1,336 @@
+/*
+ * 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.iceberg.flink;
+
+import static org.apache.iceberg.flink.FlinkCatalogFactory.ICEBERG_CATALOG_TYPE_HADOOP;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assumptions.assumeThat;
+
+import java.io.File;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.apache.flink.api.dag.Transformation;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.api.EnvironmentSettings;
+import org.apache.flink.table.api.TableEnvironment;
+import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
+import org.apache.flink.table.api.internal.TableEnvironmentImpl;
+import org.apache.flink.table.operations.ModifyOperation;
+import org.apache.flink.table.planner.delegation.PlannerBase;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.iceberg.CatalogProperties;
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.Parameter;
+import org.apache.iceberg.ParameterizedTestExtension;
+import org.apache.iceberg.Parameters;
+import org.apache.iceberg.Snapshot;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.catalog.Namespace;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.flink.source.BoundedTableFactory;
+import org.apache.iceberg.hadoop.HadoopCatalog;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestTemplate;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.api.io.TempDir;
+
+/**
+ * This class tests the more extended features of Flink sink. Extract them separately since it is
+ * unnecessary to test all the parameters combinations in {@link TestFlinkTableSink}, like catalog
+ * types, namespaces, file format, streaming/batch. Those combinations explode exponentially. Each
+ * test method in {@link TestFlinkTableSink} runs 21 combinations, which are expensive and slow.
+ */
+@ExtendWith(ParameterizedTestExtension.class)
+public class TestFlinkTableSinkExtended extends SqlBase {
+  protected static final String CATALOG = "testhadoop";
+  protected static final String DATABASE = "db";
+  protected static final String TABLE = "tbl";
+
+  @RegisterExtension
+  public static MiniClusterExtension miniClusterResource =
+      MiniFlinkClusterExtension.createWithClassloaderCheckDisabled();
+
+  private static final String SOURCE_TABLE = "default_catalog.default_database.bounded_source";
+  private static final String FLINK_DATABASE = CATALOG + "." + DATABASE;
+  private static final Namespace ICEBERG_NAMESPACE = Namespace.of(new String[] {DATABASE});
+
+  @TempDir protected File warehouseRoot;
+
+  protected HadoopCatalog catalog = null;
+
+  private TableEnvironment tEnv;
+
+  @Parameter protected boolean isStreamingJob;
+
+  @Parameters(name = "isStreamingJob={0}")
+  protected static List<Object[]> parameters() {
+    return Arrays.asList(new Boolean[] {true}, new Boolean[] {false});
+  }
+
+  protected synchronized TableEnvironment getTableEnv() {
+    if (tEnv == null) {
+      EnvironmentSettings.Builder settingsBuilder = EnvironmentSettings.newInstance();
+      if (isStreamingJob) {
+        settingsBuilder.inStreamingMode();
+        StreamExecutionEnvironment env =
+            StreamExecutionEnvironment.getExecutionEnvironment(
+                MiniFlinkClusterExtension.DISABLE_CLASSLOADER_CHECK_CONFIG);
+        env.enableCheckpointing(400);
+        env.setMaxParallelism(2);
+        env.setParallelism(2);
+        tEnv = StreamTableEnvironment.create(env, settingsBuilder.build());
+      } else {
+        settingsBuilder.inBatchMode();
+        tEnv = TableEnvironment.create(settingsBuilder.build());
+      }
+    }
+    return tEnv;
+  }
+
+  @BeforeEach
+  public void before() {
+    String warehouseLocation = "file:" + warehouseRoot.getPath();
+    this.catalog = new HadoopCatalog(new Configuration(), warehouseLocation);
+    Map<String, String> config = Maps.newHashMap();
+    config.put("type", "iceberg");
+    config.put(FlinkCatalogFactory.ICEBERG_CATALOG_TYPE, ICEBERG_CATALOG_TYPE_HADOOP);
+    config.put(CatalogProperties.WAREHOUSE_LOCATION, warehouseLocation);
+    sql("CREATE CATALOG %s WITH %s", CATALOG, toWithClause(config));
+
+    sql("CREATE DATABASE %s", FLINK_DATABASE);
+    sql("USE CATALOG %s", CATALOG);
+    sql("USE %s", DATABASE);
+    sql(
+        "CREATE TABLE %s (id int, data varchar) with ('write.format.default'='%s')",
+        TABLE, FileFormat.PARQUET.name());
+  }
+
+  @AfterEach
+  public void clean() throws Exception {
+    sql("DROP TABLE IF EXISTS %s.%s", FLINK_DATABASE, TABLE);
+    dropDatabase(FLINK_DATABASE, true);
+    BoundedTableFactory.clearDataSets();
+
+    dropCatalog(CATALOG, true);
+    catalog.close();
+  }
+
+  @TestTemplate
+  public void testWriteParallelism() {
+    List<Row> dataSet =
+        IntStream.range(1, 1000)
+            .mapToObj(i -> ImmutableList.of(Row.of(i, "aaa"), Row.of(i, "bbb"), Row.of(i, "ccc")))
+            .flatMap(List::stream)
+            .collect(Collectors.toList());
+    String dataId = BoundedTableFactory.registerDataSet(ImmutableList.of(dataSet));
+    sql(
+        "CREATE TABLE %s(id INT NOT NULL, data STRING NOT NULL)"
+            + " WITH ('connector'='BoundedSource', 'data-id'='%s')",
+        SOURCE_TABLE, dataId);
+
+    PlannerBase planner = (PlannerBase) ((TableEnvironmentImpl) getTableEnv()).getPlanner();
+    String insertSQL =
+        String.format(
+            "INSERT INTO %s /*+ OPTIONS('write-parallelism'='1') */ SELECT * FROM %s",
+            TABLE, SOURCE_TABLE);
+    ModifyOperation operation = (ModifyOperation) planner.getParser().parse(insertSQL).get(0);
+    Transformation<?> dummySink = planner.translate(Collections.singletonList(operation)).get(0);
+    Transformation<?> committer = dummySink.getInputs().get(0);
+    Transformation<?> writer = committer.getInputs().get(0);
+
+    assertThat(writer.getParallelism()).as("Should have the expected 1 parallelism.").isEqualTo(1);
+    writer
+        .getInputs()
+        .forEach(
+            input ->
+                assertThat(input.getParallelism())
+                    .as("Should have the expected parallelism.")
+                    .isEqualTo(isStreamingJob ? 2 : 4));
+  }
+
+  @TestTemplate
+  public void testHashDistributeMode() throws Exception {
+    // Initialize a BoundedSource table to precisely emit those rows in only one checkpoint.
+    List<Row> dataSet =
+        IntStream.range(1, 1000)
+            .mapToObj(i -> ImmutableList.of(Row.of(i, "aaa"), Row.of(i, "bbb"), Row.of(i, "ccc")))
+            .flatMap(List::stream)
+            .collect(Collectors.toList());
+    String dataId = BoundedTableFactory.registerDataSet(ImmutableList.of(dataSet));
+    sql(
+        "CREATE TABLE %s(id INT NOT NULL, data STRING NOT NULL)"
+            + " WITH ('connector'='BoundedSource', 'data-id'='%s')",
+        SOURCE_TABLE, dataId);
+
+    assertThat(sql("SELECT * FROM %s", SOURCE_TABLE))
+        .as("Should have the expected rows in source table.")
+        .containsExactlyInAnyOrderElementsOf(dataSet);
+
+    Map<String, String> tableProps =
+        ImmutableMap.of(
+            "write.format.default",
+            FileFormat.PARQUET.name(),
+            TableProperties.WRITE_DISTRIBUTION_MODE,
+            DistributionMode.HASH.modeName());
+
+    String tableName = "test_hash_distribution_mode";
+    sql(
+        "CREATE TABLE %s(id INT, data VARCHAR) PARTITIONED BY (data) WITH %s",
+        tableName, toWithClause(tableProps));
+
+    try {
+      // Insert data set.
+      sql("INSERT INTO %s SELECT * FROM %s", tableName, SOURCE_TABLE);
+
+      assertThat(sql("SELECT * FROM %s", tableName))
+          .as("Should have the expected rows in sink table.")
+          .containsExactlyInAnyOrderElementsOf(dataSet);
+
+      // Sometimes we will have more than one checkpoint if we pass the auto checkpoint interval,
+      // thus producing multiple snapshots.  Here we assert that each snapshot has only 1 file per
+      // partition.
+      Table table = catalog.loadTable(TableIdentifier.of(ICEBERG_NAMESPACE, tableName));
+      Map<Long, List<DataFile>> snapshotToDataFiles = SimpleDataUtil.snapshotToDataFiles(table);
+      for (List<DataFile> dataFiles : snapshotToDataFiles.values()) {
+        if (dataFiles.isEmpty()) {
+          continue;
+        }
+
+        assertThat(
+                SimpleDataUtil.matchingPartitions(
+                    dataFiles, table.spec(), ImmutableMap.of("data", "aaa")))
+            .hasSize(1);
+        assertThat(
+                SimpleDataUtil.matchingPartitions(
+                    dataFiles, table.spec(), ImmutableMap.of("data", "bbb")))
+            .hasSize(1);
+        assertThat(
+                SimpleDataUtil.matchingPartitions(
+                    dataFiles, table.spec(), ImmutableMap.of("data", "ccc")))
+            .hasSize(1);
+      }
+    } finally {
+      sql("DROP TABLE IF EXISTS %s.%s", FLINK_DATABASE, tableName);
+    }
+  }
+
+  @TestTemplate
+  public void testRangeDistributionPartitionColumn() {
+    // Range partitioner currently only works with streaming writes (with checkpoints)
+    assumeThat(isStreamingJob).isTrue();
+
+    // Initialize a BoundedSource table to precisely emit those rows in only one checkpoint.
+    List<List<Row>> rowsPerCheckpoint =
+        IntStream.range(1, 6)
+            .mapToObj(
+                checkpointId -> {
+                  List<Row> charRows = Lists.newArrayList();
+                  // emit 26x10 rows for each checkpoint cycle
+                  for (int i = 0; i < 10; ++i) {
+                    for (char c = 'a'; c <= 'z'; c++) {
+                      charRows.add(Row.of(c - 'a', String.valueOf(c)));
+                    }
+                  }
+                  return charRows;
+                })
+            .collect(Collectors.toList());
+    List<Row> flattenedRows =
+        rowsPerCheckpoint.stream().flatMap(List::stream).collect(Collectors.toList());
+
+    String dataId = BoundedTableFactory.registerDataSet(rowsPerCheckpoint);
+    sql(
+        "CREATE TABLE %s(id INT NOT NULL, data STRING NOT NULL)"
+            + " WITH ('connector'='BoundedSource', 'data-id'='%s')",
+        SOURCE_TABLE, dataId);
+
+    assertThat(sql("SELECT * FROM %s", SOURCE_TABLE))
+        .as("Should have the expected rows in source table.")
+        .containsExactlyInAnyOrderElementsOf(flattenedRows);
+
+    Map<String, String> tableProps =
+        ImmutableMap.of(
+            "write.format.default",
+            FileFormat.PARQUET.name(),
+            TableProperties.WRITE_DISTRIBUTION_MODE,
+            DistributionMode.RANGE.modeName());
+
+    String tableName = "test_hash_distribution_mode";
+    sql(
+        "CREATE TABLE %s(id INT, data VARCHAR) PARTITIONED BY (data) WITH %s",
+        tableName, toWithClause(tableProps));
+
+    try {
+      // Insert data set.
+      sql("INSERT INTO %s SELECT * FROM %s", tableName, SOURCE_TABLE);
+
+      assertThat(sql("SELECT * FROM %s", tableName))
+          .as("Should have the expected rows in sink table.")
+          .containsExactlyInAnyOrderElementsOf(flattenedRows);
+
+      Table table = catalog.loadTable(TableIdentifier.of(ICEBERG_NAMESPACE, tableName));
+      // ordered in reverse timeline from the newest snapshot to the oldest snapshot
+      List<Snapshot> snapshots = Lists.newArrayList(table.snapshots().iterator());
+      // only keep the snapshots with added data files
+      snapshots =
+          snapshots.stream()
+              .filter(snapshot -> snapshot.addedDataFiles(table.io()).iterator().hasNext())
+              .collect(Collectors.toList());
+
+      // Sometimes we will have more checkpoints than the bounded source if we pass the
+      // auto checkpoint interval. Thus producing multiple snapshots.
+      assertThat(snapshots).hasSizeGreaterThanOrEqualTo(5);
+
+      // It takes 2 checkpoint cycle for statistics collection and application
+      // of the globally aggregated statistics in the range partitioner.
+      // The last two checkpoints should have range shuffle applied
+      List<Snapshot> rangePartitionedCycles =
+          snapshots.subList(snapshots.size() - 2, snapshots.size());
+
+      for (Snapshot snapshot : rangePartitionedCycles) {
+        List<DataFile> addedDataFiles =
+            Lists.newArrayList(snapshot.addedDataFiles(table.io()).iterator());
+        // range partition results in each partition only assigned to one writer task
+        // maybe less than 26 partitions as BoundedSource doesn't always precisely
+        // control the checkpoint boundary.
+        // It is hard to precisely control the test condition in SQL tests.
+        // Here only minimal safe assertions are applied to avoid flakiness.
+        // If there are no shuffling, the number of data files could be as high as
+        // 26 * 4 as the default parallelism is set to 4 for the mini cluster.
+        assertThat(addedDataFiles).hasSizeLessThanOrEqualTo(26);
+      }
+    } finally {
+      sql("DROP TABLE IF EXISTS %s.%s", FLINK_DATABASE, tableName);
+    }
+  }
+}
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestIcebergConnector.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestIcebergConnector.java
index b709c0058f7d..47f5485df879 100644
--- a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestIcebergConnector.java
+++ b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/TestIcebergConnector.java
@@ -248,20 +248,12 @@ public void testCreateTableUnderDefaultDatabase() {
   public void testCatalogDatabaseConflictWithFlinkDatabase() {
     sql("CREATE DATABASE IF NOT EXISTS `%s`", databaseName());
     sql("USE `%s`", databaseName());
-
-    try {
-      testCreateConnectorTable();
-      // Ensure that the table was created under the specific database.
-      assertThatThrownBy(
-              () -> sql("CREATE TABLE `default_catalog`.`%s`.`%s`", databaseName(), TABLE_NAME))
-          .isInstanceOf(org.apache.flink.table.api.TableException.class)
-          .hasMessageStartingWith("Could not execute CreateTable in path");
-    } finally {
-      sql("DROP TABLE IF EXISTS `%s`.`%s`", databaseName(), TABLE_NAME);
-      if (!isDefaultDatabaseName()) {
-        sql("DROP DATABASE `%s`", databaseName());
-      }
-    }
+    testCreateConnectorTable();
+    // Ensure that the table was created under the specific database.
+    assertThatThrownBy(
+            () -> sql("CREATE TABLE `default_catalog`.`%s`.`%s`", databaseName(), TABLE_NAME))
+        .isInstanceOf(org.apache.flink.table.api.TableException.class)
+        .hasMessageStartingWith("Could not execute CreateTable in path");
   }
 
   @TestTemplate
@@ -327,10 +319,6 @@ private String databaseName() {
     return properties.getOrDefault("catalog-database", "default_database");
   }
 
-  private String toWithClause(Map<String, String> props) {
-    return CatalogTestBase.toWithClause(props);
-  }
-
   private String createWarehouse() {
     try {
       return String.format(
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSink.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSink.java
index 61ab087f2ca3..b778037c559c 100644
--- a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSink.java
+++ b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSink.java
@@ -18,20 +18,11 @@
  */
 package org.apache.iceberg.flink.sink;
 
-import static org.apache.iceberg.flink.TestFixtures.DATABASE;
-import static org.assertj.core.api.Assertions.assertThat;
-import static org.assertj.core.api.Assertions.assertThatThrownBy;
-
 import java.io.IOException;
-import java.util.Collections;
 import java.util.List;
-import java.util.Map;
-import org.apache.flink.configuration.Configuration;
 import org.apache.flink.streaming.api.datastream.DataStream;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
-import org.apache.flink.table.api.TableSchema;
 import org.apache.flink.table.data.RowData;
-import org.apache.flink.test.junit5.MiniClusterExtension;
 import org.apache.flink.types.Row;
 import org.apache.iceberg.DistributionMode;
 import org.apache.iceberg.FileFormat;
@@ -39,37 +30,19 @@
 import org.apache.iceberg.ParameterizedTestExtension;
 import org.apache.iceberg.Parameters;
 import org.apache.iceberg.PartitionSpec;
-import org.apache.iceberg.Table;
 import org.apache.iceberg.TableProperties;
-import org.apache.iceberg.catalog.TableIdentifier;
-import org.apache.iceberg.flink.FlinkWriteOptions;
-import org.apache.iceberg.flink.HadoopCatalogExtension;
 import org.apache.iceberg.flink.MiniFlinkClusterExtension;
 import org.apache.iceberg.flink.SimpleDataUtil;
-import org.apache.iceberg.flink.TableLoader;
 import org.apache.iceberg.flink.TestFixtures;
 import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
-import org.apache.iceberg.relocated.com.google.common.collect.Maps;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.extension.ExtendWith;
-import org.junit.jupiter.api.extension.RegisterExtension;
 
 @ExtendWith(ParameterizedTestExtension.class)
 public class TestFlinkIcebergSink extends TestFlinkIcebergSinkBase {
-
-  @RegisterExtension
-  public static final MiniClusterExtension MINI_CLUSTER_EXTENSION =
-      MiniFlinkClusterExtension.createWithClassloaderCheckDisabled();
-
-  @RegisterExtension
-  private static final HadoopCatalogExtension CATALOG_EXTENSION =
-      new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
-
-  private TableLoader tableLoader;
-
   @Parameter(index = 0)
   private FileFormat format;
 
@@ -99,7 +72,7 @@ public static Object[][] parameters() {
 
   @BeforeEach
   public void before() throws IOException {
-    table =
+    this.table =
         CATALOG_EXTENSION
             .catalog()
             .createTable(
@@ -110,14 +83,14 @@ public void before() throws IOException {
                     : PartitionSpec.unpartitioned(),
                 ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
 
-    env =
+    this.env =
         StreamExecutionEnvironment.getExecutionEnvironment(
                 MiniFlinkClusterExtension.DISABLE_CLASSLOADER_CHECK_CONFIG)
             .enableCheckpointing(100)
             .setParallelism(parallelism)
             .setMaxParallelism(parallelism);
 
-    tableLoader = CATALOG_EXTENSION.tableLoader();
+    this.tableLoader = CATALOG_EXTENSION.tableLoader();
   }
 
   @TestTemplate
@@ -140,246 +113,13 @@ public void testWriteRowData() throws Exception {
     SimpleDataUtil.assertTableRows(table, convertToRowData(rows));
   }
 
-  private void testWriteRow(TableSchema tableSchema, DistributionMode distributionMode)
-      throws Exception {
-    List<Row> rows = createRows("");
-    DataStream<Row> dataStream = env.addSource(createBoundedSource(rows), ROW_TYPE_INFO);
-
-    FlinkSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
-        .table(table)
-        .tableLoader(tableLoader)
-        .tableSchema(tableSchema)
-        .writeParallelism(parallelism)
-        .distributionMode(distributionMode)
-        .append();
-
-    // Execute the program.
-    env.execute("Test Iceberg DataStream.");
-
-    SimpleDataUtil.assertTableRows(table, convertToRowData(rows));
-  }
-
-  private int partitionFiles(String partition) throws IOException {
-    return SimpleDataUtil.partitionDataFiles(table, ImmutableMap.of("data", partition)).size();
-  }
-
   @TestTemplate
   public void testWriteRow() throws Exception {
-    testWriteRow(null, DistributionMode.NONE);
+    testWriteRow(parallelism, null, DistributionMode.NONE);
   }
 
   @TestTemplate
   public void testWriteRowWithTableSchema() throws Exception {
-    testWriteRow(SimpleDataUtil.FLINK_SCHEMA, DistributionMode.NONE);
-  }
-
-  @TestTemplate
-  public void testJobNoneDistributeMode() throws Exception {
-    table
-        .updateProperties()
-        .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
-        .commit();
-
-    testWriteRow(null, DistributionMode.NONE);
-
-    if (parallelism > 1) {
-      if (partitioned) {
-        int files = partitionFiles("aaa") + partitionFiles("bbb") + partitionFiles("ccc");
-        assertThat(files).isGreaterThan(3);
-      }
-    }
-  }
-
-  @TestTemplate
-  public void testJobHashDistributionMode() {
-    table
-        .updateProperties()
-        .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
-        .commit();
-
-    assertThatThrownBy(() -> testWriteRow(null, DistributionMode.RANGE))
-        .isInstanceOf(IllegalArgumentException.class)
-        .hasMessage("Flink does not support 'range' write distribution mode now.");
-  }
-
-  @TestTemplate
-  public void testJobNullDistributionMode() throws Exception {
-    table
-        .updateProperties()
-        .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
-        .commit();
-
-    testWriteRow(null, null);
-
-    if (partitioned) {
-      assertThat(partitionFiles("aaa")).isEqualTo(1);
-      assertThat(partitionFiles("bbb")).isEqualTo(1);
-      assertThat(partitionFiles("ccc")).isEqualTo(1);
-    }
-  }
-
-  @TestTemplate
-  public void testPartitionWriteMode() throws Exception {
-    testWriteRow(null, DistributionMode.HASH);
-    if (partitioned) {
-      assertThat(partitionFiles("aaa")).isEqualTo(1);
-      assertThat(partitionFiles("bbb")).isEqualTo(1);
-      assertThat(partitionFiles("ccc")).isEqualTo(1);
-    }
-  }
-
-  @TestTemplate
-  public void testShuffleByPartitionWithSchema() throws Exception {
-    testWriteRow(SimpleDataUtil.FLINK_SCHEMA, DistributionMode.HASH);
-    if (partitioned) {
-      assertThat(partitionFiles("aaa")).isEqualTo(1);
-      assertThat(partitionFiles("bbb")).isEqualTo(1);
-      assertThat(partitionFiles("ccc")).isEqualTo(1);
-    }
-  }
-
-  @TestTemplate
-  public void testTwoSinksInDisjointedDAG() throws Exception {
-    Map<String, String> props = ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name());
-
-    Table leftTable =
-        CATALOG_EXTENSION
-            .catalog()
-            .createTable(
-                TableIdentifier.of("left"),
-                SimpleDataUtil.SCHEMA,
-                partitioned
-                    ? PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).identity("data").build()
-                    : PartitionSpec.unpartitioned(),
-                props);
-    TableLoader leftTableLoader =
-        TableLoader.fromCatalog(CATALOG_EXTENSION.catalogLoader(), TableIdentifier.of("left"));
-
-    Table rightTable =
-        CATALOG_EXTENSION
-            .catalog()
-            .createTable(
-                TableIdentifier.of("right"),
-                SimpleDataUtil.SCHEMA,
-                partitioned
-                    ? PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).identity("data").build()
-                    : PartitionSpec.unpartitioned(),
-                props);
-    TableLoader rightTableLoader =
-        TableLoader.fromCatalog(CATALOG_EXTENSION.catalogLoader(), TableIdentifier.of("right"));
-
-    env =
-        StreamExecutionEnvironment.getExecutionEnvironment(
-                MiniFlinkClusterExtension.DISABLE_CLASSLOADER_CHECK_CONFIG)
-            .enableCheckpointing(100)
-            .setParallelism(parallelism)
-            .setMaxParallelism(parallelism);
-    env.getConfig().disableAutoGeneratedUIDs();
-
-    List<Row> leftRows = createRows("left-");
-    DataStream<Row> leftStream =
-        env.fromCollection(leftRows, ROW_TYPE_INFO)
-            .name("leftCustomSource")
-            .uid("leftCustomSource");
-    FlinkSink.forRow(leftStream, SimpleDataUtil.FLINK_SCHEMA)
-        .table(leftTable)
-        .tableLoader(leftTableLoader)
-        .tableSchema(SimpleDataUtil.FLINK_SCHEMA)
-        .distributionMode(DistributionMode.NONE)
-        .uidPrefix("leftIcebergSink")
-        .append();
-
-    List<Row> rightRows = createRows("right-");
-    DataStream<Row> rightStream =
-        env.fromCollection(rightRows, ROW_TYPE_INFO)
-            .name("rightCustomSource")
-            .uid("rightCustomSource");
-    FlinkSink.forRow(rightStream, SimpleDataUtil.FLINK_SCHEMA)
-        .table(rightTable)
-        .tableLoader(rightTableLoader)
-        .tableSchema(SimpleDataUtil.FLINK_SCHEMA)
-        .writeParallelism(parallelism)
-        .distributionMode(DistributionMode.HASH)
-        .uidPrefix("rightIcebergSink")
-        .setSnapshotProperty("flink.test", TestFlinkIcebergSink.class.getName())
-        .setSnapshotProperties(Collections.singletonMap("direction", "rightTable"))
-        .append();
-
-    // Execute the program.
-    env.execute("Test Iceberg DataStream.");
-
-    SimpleDataUtil.assertTableRows(leftTable, convertToRowData(leftRows));
-    SimpleDataUtil.assertTableRows(rightTable, convertToRowData(rightRows));
-
-    leftTable.refresh();
-    assertThat(leftTable.currentSnapshot().summary()).doesNotContainKeys("flink.test", "direction");
-    rightTable.refresh();
-    assertThat(rightTable.currentSnapshot().summary())
-        .containsEntry("flink.test", TestFlinkIcebergSink.class.getName())
-        .containsEntry("direction", "rightTable");
-  }
-
-  @TestTemplate
-  public void testOverrideWriteConfigWithUnknownDistributionMode() {
-    Map<String, String> newProps = Maps.newHashMap();
-    newProps.put(FlinkWriteOptions.DISTRIBUTION_MODE.key(), "UNRECOGNIZED");
-
-    List<Row> rows = createRows("");
-    DataStream<Row> dataStream = env.addSource(createBoundedSource(rows), ROW_TYPE_INFO);
-
-    FlinkSink.Builder builder =
-        FlinkSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
-            .table(table)
-            .tableLoader(tableLoader)
-            .writeParallelism(parallelism)
-            .setAll(newProps);
-
-    assertThatThrownBy(builder::append)
-        .isInstanceOf(IllegalArgumentException.class)
-        .hasMessage("Invalid distribution mode: UNRECOGNIZED");
-  }
-
-  @TestTemplate
-  public void testOverrideWriteConfigWithUnknownFileFormat() {
-    Map<String, String> newProps = Maps.newHashMap();
-    newProps.put(FlinkWriteOptions.WRITE_FORMAT.key(), "UNRECOGNIZED");
-
-    List<Row> rows = createRows("");
-    DataStream<Row> dataStream = env.addSource(createBoundedSource(rows), ROW_TYPE_INFO);
-
-    FlinkSink.Builder builder =
-        FlinkSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
-            .table(table)
-            .tableLoader(tableLoader)
-            .writeParallelism(parallelism)
-            .setAll(newProps);
-
-    assertThatThrownBy(builder::append)
-        .isInstanceOf(IllegalArgumentException.class)
-        .hasMessage("Invalid file format: UNRECOGNIZED");
-  }
-
-  @TestTemplate
-  public void testWriteRowWithTableRefreshInterval() throws Exception {
-    List<Row> rows = Lists.newArrayList(Row.of(1, "hello"), Row.of(2, "world"), Row.of(3, "foo"));
-    DataStream<RowData> dataStream =
-        env.addSource(createBoundedSource(rows), ROW_TYPE_INFO)
-            .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
-
-    Configuration flinkConf = new Configuration();
-    flinkConf.setString(FlinkWriteOptions.TABLE_REFRESH_INTERVAL.key(), "100ms");
-
-    FlinkSink.forRowData(dataStream)
-        .table(table)
-        .tableLoader(tableLoader)
-        .flinkConf(flinkConf)
-        .writeParallelism(parallelism)
-        .append();
-
-    // Execute the program.
-    env.execute("Test Iceberg DataStream");
-
-    // Assert the iceberg table's records.
-    SimpleDataUtil.assertTableRows(table, convertToRowData(rows));
+    testWriteRow(parallelism, SimpleDataUtil.FLINK_SCHEMA, DistributionMode.NONE);
   }
 }
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkBase.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkBase.java
index b38aa6b50ce6..9ce36cc1e8d0 100644
--- a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkBase.java
+++ b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkBase.java
@@ -18,29 +18,52 @@
  */
 package org.apache.iceberg.flink.sink;
 
+import static org.apache.iceberg.flink.TestFixtures.DATABASE;
+
+import java.io.IOException;
 import java.util.List;
 import java.util.stream.Collectors;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.streaming.api.datastream.DataStream;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.api.TableSchema;
 import org.apache.flink.table.data.RowData;
 import org.apache.flink.table.data.util.DataFormatConverters;
+import org.apache.flink.test.junit5.MiniClusterExtension;
 import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
 import org.apache.iceberg.Table;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.MiniFlinkClusterExtension;
 import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
 import org.apache.iceberg.flink.source.BoundedTestSource;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 public class TestFlinkIcebergSinkBase {
 
-  protected Table table;
-  protected StreamExecutionEnvironment env;
+  @RegisterExtension
+  public static MiniClusterExtension miniClusterResource =
+      MiniFlinkClusterExtension.createWithClassloaderCheckDisabled();
+
+  @RegisterExtension
+  protected static final HadoopCatalogExtension CATALOG_EXTENSION =
+      new HadoopCatalogExtension(DATABASE, TestFixtures.TABLE);
+
   protected static final TypeInformation<Row> ROW_TYPE_INFO =
       new RowTypeInfo(SimpleDataUtil.FLINK_SCHEMA.getFieldTypes());
 
   protected static final DataFormatConverters.RowConverter CONVERTER =
       new DataFormatConverters.RowConverter(SimpleDataUtil.FLINK_SCHEMA.getFieldDataTypes());
 
+  protected TableLoader tableLoader;
+  protected Table table;
+  protected StreamExecutionEnvironment env;
+
   protected BoundedTestSource<Row> createBoundedSource(List<Row> rows) {
     return new BoundedTestSource<>(rows.toArray(new Row[0]));
   }
@@ -61,4 +84,28 @@ protected List<Row> createRows(String prefix) {
   protected List<RowData> convertToRowData(List<Row> rows) {
     return rows.stream().map(CONVERTER::toInternal).collect(Collectors.toList());
   }
+
+  protected void testWriteRow(
+      int writerParallelism, TableSchema tableSchema, DistributionMode distributionMode)
+      throws Exception {
+    List<Row> rows = createRows("");
+    DataStream<Row> dataStream = env.addSource(createBoundedSource(rows), ROW_TYPE_INFO);
+
+    FlinkSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+        .table(table)
+        .tableLoader(tableLoader)
+        .tableSchema(tableSchema)
+        .writeParallelism(writerParallelism)
+        .distributionMode(distributionMode)
+        .append();
+
+    // Execute the program.
+    env.execute("Test Iceberg DataStream.");
+
+    SimpleDataUtil.assertTableRows(table, convertToRowData(rows));
+  }
+
+  protected int partitionFiles(String partition) throws IOException {
+    return SimpleDataUtil.partitionDataFiles(table, ImmutableMap.of("data", partition)).size();
+  }
 }
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkDistributionMode.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkDistributionMode.java
new file mode 100644
index 000000000000..75e397d3f203
--- /dev/null
+++ b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkDistributionMode.java
@@ -0,0 +1,180 @@
+/*
+ * 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.iceberg.flink.sink;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.test.junit5.MiniClusterExtension;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.Parameter;
+import org.apache.iceberg.ParameterizedTestExtension;
+import org.apache.iceberg.Parameters;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.flink.FlinkWriteOptions;
+import org.apache.iceberg.flink.HadoopCatalogExtension;
+import org.apache.iceberg.flink.MiniFlinkClusterExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestTemplate;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/**
+ * This tests the distribution mode of Flink sink. Extract them separately since it is unnecessary
+ * to test different file formats (Avro, Orc, Parquet) like in {@link TestFlinkIcebergSink}.
+ * Removing the file format dimension reduces the number of combinations from 12 to 4, which helps
+ * reduce test run time.
+ */
+@ExtendWith(ParameterizedTestExtension.class)
+public class TestFlinkIcebergSinkDistributionMode extends TestFlinkIcebergSinkBase {
+
+  @RegisterExtension
+  public static MiniClusterExtension miniClusterResource =
+      MiniFlinkClusterExtension.createWithClassloaderCheckDisabled();
+
+  @RegisterExtension
+  private static final HadoopCatalogExtension CATALOG_EXTENSION =
+      new HadoopCatalogExtension(TestFixtures.DATABASE, TestFixtures.TABLE);
+
+  private final FileFormat format = FileFormat.PARQUET;
+
+  @Parameter(index = 0)
+  private int parallelism;
+
+  @Parameter(index = 1)
+  private boolean partitioned;
+
+  @Parameters(name = "parallelism = {0}, partitioned = {1}")
+  public static Object[][] parameters() {
+    return new Object[][] {
+      {1, true},
+      {1, false},
+      {2, true},
+      {2, false}
+    };
+  }
+
+  @BeforeEach
+  public void before() throws IOException {
+    this.table =
+        CATALOG_EXTENSION
+            .catalog()
+            .createTable(
+                TestFixtures.TABLE_IDENTIFIER,
+                SimpleDataUtil.SCHEMA,
+                partitioned
+                    ? PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).identity("data").build()
+                    : PartitionSpec.unpartitioned(),
+                ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+
+    this.env =
+        StreamExecutionEnvironment.getExecutionEnvironment(
+                MiniFlinkClusterExtension.DISABLE_CLASSLOADER_CHECK_CONFIG)
+            .enableCheckpointing(100)
+            .setParallelism(parallelism)
+            .setMaxParallelism(parallelism);
+
+    this.tableLoader = CATALOG_EXTENSION.tableLoader();
+  }
+
+  @TestTemplate
+  public void testShuffleByPartitionWithSchema() throws Exception {
+    testWriteRow(parallelism, SimpleDataUtil.FLINK_SCHEMA, DistributionMode.HASH);
+    if (partitioned) {
+      assertThat(partitionFiles("aaa")).isEqualTo(1);
+      assertThat(partitionFiles("bbb")).isEqualTo(1);
+      assertThat(partitionFiles("ccc")).isEqualTo(1);
+    }
+  }
+
+  @TestTemplate
+  public void testJobNoneDistributeMode() throws Exception {
+    table
+        .updateProperties()
+        .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
+        .commit();
+
+    testWriteRow(parallelism, null, DistributionMode.NONE);
+
+    if (parallelism > 1) {
+      if (partitioned) {
+        int files = partitionFiles("aaa") + partitionFiles("bbb") + partitionFiles("ccc");
+        assertThat(files).isGreaterThan(3);
+      }
+    }
+  }
+
+  @TestTemplate
+  public void testJobNullDistributionMode() throws Exception {
+    table
+        .updateProperties()
+        .set(TableProperties.WRITE_DISTRIBUTION_MODE, DistributionMode.HASH.modeName())
+        .commit();
+
+    testWriteRow(parallelism, null, null);
+
+    if (partitioned) {
+      assertThat(partitionFiles("aaa")).isEqualTo(1);
+      assertThat(partitionFiles("bbb")).isEqualTo(1);
+      assertThat(partitionFiles("ccc")).isEqualTo(1);
+    }
+  }
+
+  @TestTemplate
+  public void testPartitionWriteMode() throws Exception {
+    testWriteRow(parallelism, null, DistributionMode.HASH);
+    if (partitioned) {
+      assertThat(partitionFiles("aaa")).isEqualTo(1);
+      assertThat(partitionFiles("bbb")).isEqualTo(1);
+      assertThat(partitionFiles("ccc")).isEqualTo(1);
+    }
+  }
+
+  @TestTemplate
+  public void testOverrideWriteConfigWithUnknownDistributionMode() {
+    Map<String, String> newProps = Maps.newHashMap();
+    newProps.put(FlinkWriteOptions.DISTRIBUTION_MODE.key(), "UNRECOGNIZED");
+
+    List<Row> rows = createRows("");
+    DataStream<Row> dataStream = env.addSource(createBoundedSource(rows), ROW_TYPE_INFO);
+
+    FlinkSink.Builder builder =
+        FlinkSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+            .table(table)
+            .tableLoader(tableLoader)
+            .writeParallelism(parallelism)
+            .setAll(newProps);
+
+    assertThatThrownBy(builder::append)
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("Invalid distribution mode: UNRECOGNIZED");
+  }
+}
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkExtended.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkExtended.java
new file mode 100644
index 000000000000..36a59b20431c
--- /dev/null
+++ b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/sink/TestFlinkIcebergSinkExtended.java
@@ -0,0 +1,208 @@
+/*
+ * 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.iceberg.flink.sink;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.types.Row;
+import org.apache.iceberg.DistributionMode;
+import org.apache.iceberg.FileFormat;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.flink.FlinkWriteOptions;
+import org.apache.iceberg.flink.MiniFlinkClusterExtension;
+import org.apache.iceberg.flink.SimpleDataUtil;
+import org.apache.iceberg.flink.TableLoader;
+import org.apache.iceberg.flink.TestFixtures;
+import org.apache.iceberg.flink.util.FlinkCompatibilityUtil;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/**
+ * This class tests the more extended features of Flink sink. Extract them separately since it is
+ * unnecessary to test all the parameters combinations in {@link TestFlinkIcebergSink}. Each test
+ * method in {@link TestFlinkIcebergSink} runs 12 combinations, which are expensive and slow.
+ */
+public class TestFlinkIcebergSinkExtended extends TestFlinkIcebergSinkBase {
+  private final boolean partitioned = true;
+  private final int parallelism = 2;
+  private final FileFormat format = FileFormat.PARQUET;
+
+  @BeforeEach
+  public void before() throws IOException {
+    this.table =
+        CATALOG_EXTENSION
+            .catalog()
+            .createTable(
+                TestFixtures.TABLE_IDENTIFIER,
+                SimpleDataUtil.SCHEMA,
+                partitioned
+                    ? PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).identity("data").build()
+                    : PartitionSpec.unpartitioned(),
+                ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name()));
+
+    this.env =
+        StreamExecutionEnvironment.getExecutionEnvironment(
+                MiniFlinkClusterExtension.DISABLE_CLASSLOADER_CHECK_CONFIG)
+            .enableCheckpointing(100)
+            .setParallelism(parallelism)
+            .setMaxParallelism(parallelism);
+
+    this.tableLoader = CATALOG_EXTENSION.tableLoader();
+  }
+
+  @Test
+  public void testTwoSinksInDisjointedDAG() throws Exception {
+    Map<String, String> props = ImmutableMap.of(TableProperties.DEFAULT_FILE_FORMAT, format.name());
+
+    Table leftTable =
+        CATALOG_EXTENSION
+            .catalog()
+            .createTable(
+                TableIdentifier.of("left"),
+                SimpleDataUtil.SCHEMA,
+                partitioned
+                    ? PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).identity("data").build()
+                    : PartitionSpec.unpartitioned(),
+                props);
+    TableLoader leftTableLoader =
+        TableLoader.fromCatalog(CATALOG_EXTENSION.catalogLoader(), TableIdentifier.of("left"));
+
+    Table rightTable =
+        CATALOG_EXTENSION
+            .catalog()
+            .createTable(
+                TableIdentifier.of("right"),
+                SimpleDataUtil.SCHEMA,
+                partitioned
+                    ? PartitionSpec.builderFor(SimpleDataUtil.SCHEMA).identity("data").build()
+                    : PartitionSpec.unpartitioned(),
+                props);
+    TableLoader rightTableLoader =
+        TableLoader.fromCatalog(CATALOG_EXTENSION.catalogLoader(), TableIdentifier.of("right"));
+
+    env =
+        StreamExecutionEnvironment.getExecutionEnvironment(
+                MiniFlinkClusterExtension.DISABLE_CLASSLOADER_CHECK_CONFIG)
+            .enableCheckpointing(100)
+            .setParallelism(parallelism)
+            .setMaxParallelism(parallelism);
+    env.getConfig().disableAutoGeneratedUIDs();
+
+    List<Row> leftRows = createRows("left-");
+    DataStream<Row> leftStream =
+        env.fromCollection(leftRows, ROW_TYPE_INFO)
+            .name("leftCustomSource")
+            .uid("leftCustomSource");
+    FlinkSink.forRow(leftStream, SimpleDataUtil.FLINK_SCHEMA)
+        .table(leftTable)
+        .tableLoader(leftTableLoader)
+        .tableSchema(SimpleDataUtil.FLINK_SCHEMA)
+        .distributionMode(DistributionMode.NONE)
+        .uidPrefix("leftIcebergSink")
+        .append();
+
+    List<Row> rightRows = createRows("right-");
+    DataStream<Row> rightStream =
+        env.fromCollection(rightRows, ROW_TYPE_INFO)
+            .name("rightCustomSource")
+            .uid("rightCustomSource");
+    FlinkSink.forRow(rightStream, SimpleDataUtil.FLINK_SCHEMA)
+        .table(rightTable)
+        .tableLoader(rightTableLoader)
+        .tableSchema(SimpleDataUtil.FLINK_SCHEMA)
+        .writeParallelism(parallelism)
+        .distributionMode(DistributionMode.HASH)
+        .uidPrefix("rightIcebergSink")
+        .setSnapshotProperty("flink.test", TestFlinkIcebergSink.class.getName())
+        .setSnapshotProperties(Collections.singletonMap("direction", "rightTable"))
+        .append();
+
+    // Execute the program.
+    env.execute("Test Iceberg DataStream.");
+
+    SimpleDataUtil.assertTableRows(leftTable, convertToRowData(leftRows));
+    SimpleDataUtil.assertTableRows(rightTable, convertToRowData(rightRows));
+
+    leftTable.refresh();
+    assertThat(leftTable.currentSnapshot().summary()).doesNotContainKeys("flink.test", "direction");
+    rightTable.refresh();
+    assertThat(rightTable.currentSnapshot().summary())
+        .containsEntry("flink.test", TestFlinkIcebergSink.class.getName())
+        .containsEntry("direction", "rightTable");
+  }
+
+  @Test
+  public void testOverrideWriteConfigWithUnknownFileFormat() {
+    Map<String, String> newProps = Maps.newHashMap();
+    newProps.put(FlinkWriteOptions.WRITE_FORMAT.key(), "UNRECOGNIZED");
+
+    List<Row> rows = createRows("");
+    DataStream<Row> dataStream = env.addSource(createBoundedSource(rows), ROW_TYPE_INFO);
+
+    FlinkSink.Builder builder =
+        FlinkSink.forRow(dataStream, SimpleDataUtil.FLINK_SCHEMA)
+            .table(table)
+            .tableLoader(tableLoader)
+            .writeParallelism(parallelism)
+            .setAll(newProps);
+
+    assertThatThrownBy(builder::append)
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("Invalid file format: UNRECOGNIZED");
+  }
+
+  @Test
+  public void testWriteRowWithTableRefreshInterval() throws Exception {
+    List<Row> rows = Lists.newArrayList(Row.of(1, "hello"), Row.of(2, "world"), Row.of(3, "foo"));
+    DataStream<RowData> dataStream =
+        env.addSource(createBoundedSource(rows), ROW_TYPE_INFO)
+            .map(CONVERTER::toInternal, FlinkCompatibilityUtil.toTypeInfo(SimpleDataUtil.ROW_TYPE));
+
+    Configuration flinkConf = new Configuration();
+    flinkConf.setString(FlinkWriteOptions.TABLE_REFRESH_INTERVAL.key(), "100ms");
+
+    FlinkSink.forRowData(dataStream)
+        .table(table)
+        .tableLoader(tableLoader)
+        .flinkConf(flinkConf)
+        .writeParallelism(parallelism)
+        .append();
+
+    // Execute the program.
+    env.execute("Test Iceberg DataStream");
+
+    // Assert the iceberg table's records.
+    SimpleDataUtil.assertTableRows(table, convertToRowData(rows));
+  }
+}
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/source/TestIcebergSpeculativeExecutionSupport.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/source/TestIcebergSpeculativeExecutionSupport.java
index 95d0b90b6ca9..51f9025b4159 100644
--- a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/source/TestIcebergSpeculativeExecutionSupport.java
+++ b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/source/TestIcebergSpeculativeExecutionSupport.java
@@ -103,7 +103,7 @@ public void before() throws IOException {
   public void after() {
     sql("DROP TABLE IF EXISTS %s.%s", DATABASE_NAME, INPUT_TABLE_NAME);
     sql("DROP TABLE IF EXISTS %s.%s", DATABASE_NAME, OUTPUT_TABLE_NAME);
-    sql("DROP DATABASE %s", DATABASE_NAME);
+    dropDatabase(DATABASE_NAME, true);
     dropCatalog(CATALOG_NAME, true);
   }