From b7b2b069f0cbfb3d73b9404ca98cd7735dd130c9 Mon Sep 17 00:00:00 2001
From: caican00 <imcaican@gmail.com>
Date: Thu, 21 Mar 2024 20:47:36 +0800
Subject: [PATCH] [#2587] feat(spark-connector): support iceberg
 metadataColumnsspo

---
 .../integration/test/spark/SparkCommonIT.java |  2 +-
 .../spark/iceberg/SparkIcebergCatalogIT.java  | 24 +++++++++++++++++++
 .../connector/iceberg/SparkIcebergTable.java  |  9 ++++++-
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/spark/SparkCommonIT.java b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/spark/SparkCommonIT.java
index 6b735affd69..4741b18e9ee 100644
--- a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/spark/SparkCommonIT.java
+++ b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/spark/SparkCommonIT.java
@@ -529,7 +529,7 @@ private List<SparkColumnInfo> getSimpleTableColumn() {
 
   // Helper method to create a simple table, and could use corresponding
   // getSimpleTableColumn to check table column.
-  private void createSimpleTable(String identifier) {
+  protected void createSimpleTable(String identifier) {
     String createTableSql = getCreateSimpleTableString(identifier);
     sql(createTableSql);
   }
diff --git a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/spark/iceberg/SparkIcebergCatalogIT.java b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/spark/iceberg/SparkIcebergCatalogIT.java
index 7a7fa5e16cc..fba1ecc37c5 100644
--- a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/spark/iceberg/SparkIcebergCatalogIT.java
+++ b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/spark/iceberg/SparkIcebergCatalogIT.java
@@ -5,7 +5,15 @@
 package com.datastrato.gravitino.integration.test.spark.iceberg;
 
 import com.datastrato.gravitino.integration.test.spark.SparkCommonIT;
+import com.datastrato.gravitino.spark.connector.iceberg.SparkIcebergTable;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.catalyst.analysis.ResolvedTable;
+import org.apache.spark.sql.catalyst.plans.logical.CommandResult;
+import org.apache.spark.sql.catalyst.plans.logical.DescribeRelation;
+import org.apache.spark.sql.connector.catalog.Table;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
 
 @Tag("gravitino-docker-it")
@@ -21,4 +29,20 @@ protected String getCatalogName() {
   protected String getProvider() {
     return "lakehouse-iceberg";
   }
+
+  // TODO
+  @Test
+  void testMetadataColumns() {
+    String tableName = "test_metadata_columns";
+    dropTableIfExists(tableName);
+    createSimpleTable(tableName);
+
+    Dataset ds = getSparkSession().sql("DESC TABLE EXTENDED " + tableName);
+    CommandResult result = (CommandResult) ds.logicalPlan();
+    DescribeRelation relation = (DescribeRelation) result.commandLogicalPlan();
+    ResolvedTable table = (ResolvedTable) relation.child();
+    Table table1 = table.table();
+    Assertions.assertTrue(table1 instanceof SparkIcebergTable);
+    SparkIcebergTable icebergTable = (SparkIcebergTable) table1;
+  }
 }
diff --git a/spark-connector/src/main/java/com/datastrato/gravitino/spark/connector/iceberg/SparkIcebergTable.java b/spark-connector/src/main/java/com/datastrato/gravitino/spark/connector/iceberg/SparkIcebergTable.java
index fe085826ed0..ee9237a5500 100644
--- a/spark-connector/src/main/java/com/datastrato/gravitino/spark/connector/iceberg/SparkIcebergTable.java
+++ b/spark-connector/src/main/java/com/datastrato/gravitino/spark/connector/iceberg/SparkIcebergTable.java
@@ -9,9 +9,11 @@
 import com.datastrato.gravitino.spark.connector.PropertiesConverter;
 import com.datastrato.gravitino.spark.connector.table.SparkBaseTable;
 import org.apache.spark.sql.connector.catalog.Identifier;
+import org.apache.spark.sql.connector.catalog.MetadataColumn;
+import org.apache.spark.sql.connector.catalog.SupportsMetadataColumns;
 import org.apache.spark.sql.connector.catalog.TableCatalog;
 
-public class SparkIcebergTable extends SparkBaseTable {
+public class SparkIcebergTable extends SparkBaseTable implements SupportsMetadataColumns {
 
   public SparkIcebergTable(
       Identifier identifier,
@@ -20,4 +22,9 @@ public SparkIcebergTable(
       PropertiesConverter propertiesConverter) {
     super(identifier, gravitinoTable, sparkCatalog, propertiesConverter);
   }
+
+  @Override
+  public MetadataColumn[] metadataColumns() {
+    return ((SupportsMetadataColumns) getSparkTable()).metadataColumns();
+  }
 }