From 8df5a63bb727e51993088b61eb3f42808ddfd784 Mon Sep 17 00:00:00 2001
From: Marius Grama <findinpath@gmail.com>
Date: Wed, 25 Oct 2023 17:30:42 +0200
Subject: [PATCH] Test access operations for `CREATE OR REPLACE TABLE`

---
 .../iceberg/TestIcebergFileOperations.java    | 49 +++++++++++++++++++
 .../TestIcebergMetastoreAccessOperations.java | 38 ++++++++++++++
 2 files changed, 87 insertions(+)

diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergFileOperations.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergFileOperations.java
index 85e39e03ebe8..8b2b79cc6de2 100644
--- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergFileOperations.java
+++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergFileOperations.java
@@ -129,6 +129,26 @@ public void testCreateTable()
                         .build());
     }
 
+    @Test
+    public void testCreateOrReplaceTable()
+    {
+        assertFileSystemAccesses("CREATE OR REPLACE TABLE test_create_or_replace (id VARCHAR, age INT)",
+                ImmutableMultiset.<FileOperation>builder()
+                        .add(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE))
+                        .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH))
+                        .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM))
+                        .add(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE))
+                        .build());
+        assertFileSystemAccesses("CREATE OR REPLACE TABLE test_create_or_replace (id VARCHAR, age INT)",
+                ImmutableMultiset.<FileOperation>builder()
+                        .add(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE))
+                        .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM))
+                        .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH))
+                        .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM))
+                        .add(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE))
+                        .build());
+    }
+
     @Test
     public void testCreateTableAsSelect()
     {
@@ -157,6 +177,35 @@ public void testCreateTableAsSelect()
                         .build());
     }
 
+    @Test
+    public void testCreateOrReplaceTableAsSelect()
+    {
+        assertFileSystemAccesses(
+                "CREATE OR REPLACE TABLE test_create_or_replace_as_select AS SELECT 1 col_name",
+                ImmutableMultiset.<FileOperation>builder()
+                        .addCopies(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE), 2)
+                        .add(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM))
+                        .add(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH))
+                        .add(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM))
+                        .add(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE))
+                        .add(new FileOperation(MANIFEST, OUTPUT_FILE_CREATE_OR_OVERWRITE))
+                        .add(new FileOperation(STATS, OUTPUT_FILE_CREATE))
+                        .build());
+
+        assertFileSystemAccesses(
+                "CREATE OR REPLACE TABLE test_create_or_replace_as_select AS SELECT 1 col_name",
+                ImmutableMultiset.<FileOperation>builder()
+                        .addCopies(new FileOperation(METADATA_JSON, OUTPUT_FILE_CREATE), 2)
+                        .addCopies(new FileOperation(METADATA_JSON, INPUT_FILE_NEW_STREAM), 2)
+                        .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_NEW_STREAM), 2)
+                        .addCopies(new FileOperation(SNAPSHOT, INPUT_FILE_GET_LENGTH), 2)
+                        .add(new FileOperation(SNAPSHOT, OUTPUT_FILE_CREATE_OR_OVERWRITE))
+                        .add(new FileOperation(MANIFEST, OUTPUT_FILE_CREATE_OR_OVERWRITE))
+                        .add(new FileOperation(MANIFEST, INPUT_FILE_NEW_STREAM))
+                        .add(new FileOperation(STATS, OUTPUT_FILE_CREATE))
+                        .build());
+    }
+
     @Test
     public void testSelect()
     {
diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetastoreAccessOperations.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetastoreAccessOperations.java
index 58beead75d31..2d46297e5b2b 100644
--- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetastoreAccessOperations.java
+++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergMetastoreAccessOperations.java
@@ -108,6 +108,23 @@ public void testCreateTable()
                         .build());
     }
 
+    @Test
+    public void testCreateOrReplaceTable()
+    {
+        assertMetastoreInvocations("CREATE OR REPLACE TABLE test_create_or_replace (id VARCHAR, age INT)",
+                ImmutableMultiset.builder()
+                        .add(CREATE_TABLE)
+                        .add(GET_DATABASE)
+                        .add(GET_TABLE)
+                        .build());
+        assertMetastoreInvocations("CREATE OR REPLACE TABLE test_create_or_replace (id VARCHAR, age INT)",
+                ImmutableMultiset.builder()
+                        .add(GET_DATABASE)
+                        .add(REPLACE_TABLE)
+                        .add(GET_TABLE)
+                        .build());
+    }
+
     @Test
     public void testCreateTableAsSelect()
     {
@@ -131,6 +148,27 @@ public void testCreateTableAsSelect()
                         .build());
     }
 
+    @Test
+    public void testCreateOrReplaceTableAsSelect()
+    {
+        assertMetastoreInvocations(
+                "CREATE OR REPLACE TABLE test_cortas AS SELECT 1 AS age",
+                ImmutableMultiset.builder()
+                        .add(GET_DATABASE)
+                        .add(CREATE_TABLE)
+                        .addCopies(GET_TABLE, 4)
+                        .add(REPLACE_TABLE)
+                        .build());
+
+        assertMetastoreInvocations(
+                "CREATE OR REPLACE TABLE test_cortas AS SELECT 1 AS age",
+                ImmutableMultiset.builder()
+                        .add(GET_DATABASE)
+                        .addCopies(GET_TABLE, 3)
+                        .addCopies(REPLACE_TABLE, 2)
+                        .build());
+    }
+
     @Test
     public void testSelect()
     {