Added load table method override

prestodb · Mar 15, 2024 · 89a9221 · 89a9221
1 parent 3065dbb
commit 89a9221
Showing 1 changed file with 49 additions and 0 deletions.
diff --git a/...to-iceberg/src/test/java/com/facebook/presto/iceberg/jdbc/TestIcebergDistributedJdbc.java b/...to-iceberg/src/test/java/com/facebook/presto/iceberg/jdbc/TestIcebergDistributedJdbc.java
@@ -13,11 +13,25 @@
  */
 package com.facebook.presto.iceberg.jdbc;
 
+import com.facebook.presto.hive.gcs.HiveGcsConfig;
+import com.facebook.presto.hive.gcs.HiveGcsConfigurationInitializer;
+import com.facebook.presto.hive.s3.HiveS3Config;
+import com.facebook.presto.hive.s3.PrestoS3ConfigurationUpdater;
+import com.facebook.presto.iceberg.IcebergCatalogName;
+import com.facebook.presto.iceberg.IcebergConfig;
 import com.facebook.presto.iceberg.IcebergDistributedTestBase;
 import com.facebook.presto.iceberg.IcebergQueryRunner;
+import com.facebook.presto.iceberg.IcebergResourceFactory;
+import com.facebook.presto.iceberg.IcebergUtil;
+import com.facebook.presto.iceberg.nessie.NessieConfig;
+import com.facebook.presto.metadata.CatalogManager;
+import com.facebook.presto.spi.ConnectorId;
+import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.testing.containers.JdbcContainer;
 import com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.Table;
+import org.jetbrains.annotations.NotNull;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
@@ -26,6 +40,7 @@
 import java.util.Map;
 
 import static com.facebook.presto.iceberg.CatalogType.JDBC;
+import static com.facebook.presto.iceberg.IcebergQueryRunner.ICEBERG_CATALOG;
 import static com.facebook.presto.iceberg.jdbc.JdbcTestUtil.jdbcConnectorProperties;
 
 @Test
@@ -47,6 +62,40 @@ protected Map<String, String> getProperties()
                 "user", jdbcContainer.getDefaultUser(), "password", jdbcContainer.getDefaultPassword());
     }
 
+    @Override
+    protected Table loadTable(String tableName)
+    {
+        CatalogManager catalogManager = getDistributedQueryRunner().getCoordinator().getCatalogManager();
+        ConnectorId connectorId = catalogManager.getCatalog(ICEBERG_CATALOG).get().getConnectorId();
+        IcebergConfig icebergConfig = new IcebergConfig();
+        icebergConfig.setCatalogType(JDBC);
+        icebergConfig.setCatalogWarehouse(getCatalogDirectory().getPath());
+
+        final IcebergResourceFactory resourceFactory = getIcebergResourceFactory(icebergConfig);
+        return IcebergUtil.getNativeIcebergTable(
+                resourceFactory,
+                getQueryRunner().getDefaultSession().toConnectorSession(connectorId),
+                SchemaTableName.valueOf("tpch." + tableName));
+    }
+
+    @NotNull
+    private IcebergResourceFactory getIcebergResourceFactory(IcebergConfig icebergConfig)
+    {
+        JdbcConfig jdbcConfig = new JdbcConfig();
+        jdbcConfig.setServerUrl(jdbcContainer.getJdbcURI());
+        jdbcConfig.setUser(jdbcContainer.getDefaultUser());
+        jdbcConfig.setPassword(jdbcContainer.getDefaultPassword());
+        jdbcConfig.setDriverClass(jdbcContainer.getDefaultDriverClass());
+
+        IcebergResourceFactory resourceFactory = new IcebergResourceFactory(icebergConfig,
+                new IcebergCatalogName(ICEBERG_CATALOG),
+                new NessieConfig(),
+                jdbcConfig,
+                new PrestoS3ConfigurationUpdater(new HiveS3Config()),
+                new HiveGcsConfigurationInitializer(new HiveGcsConfig()));
+        return resourceFactory;
+    }
+
     @Override
     protected boolean supportsViews()
     {