[WIP] Add key-pair auth to Iceberg Snowflake catalog

.github/workflows/ci.yml

@@ -716,6 +716,7 @@ jobs:
           ABFS_ACCESS_KEY: ${{ secrets.AZURE_ABFS_HIERARCHICAL_ACCESS_KEY }}
           SNOWFLAKE_USER: ${{ vars.SNOWFLAKE_USER }}
           SNOWFLAKE_PASSWORD: ${{ secrets.SNOWFLAKE_PASSWORD }}
+          SNOWFLAKE_KEY: ${{ secrets.SNOWFLAKE_KEY }}
           SNOWFLAKE_URL: ${{ vars.SNOWFLAKE_URL }}
           SNOWFLAKE_DATABASE: ${{ vars.SNOWFLAKE_DATABASE }}
           SNOWFLAKE_CATALOG_SCHEMA: ${{ vars.SNOWFLAKE_CATALOG_SCHEMA }}
@@ -737,6 +738,7 @@ jobs:
             -Dtesting.azure-abfs-access-key="${ABFS_ACCESS_KEY}" \
             -Dtesting.snowflake.catalog.user="${SNOWFLAKE_USER}" \
             -Dtesting.snowflake.catalog.password="${SNOWFLAKE_PASSWORD}" \
+            -Dtesting.snowflake.catalog.key="${SNOWFLAKE_KEY}" \
             -Dtesting.snowflake.catalog.account-url="${SNOWFLAKE_URL}" \
             -Dtesting.snowflake.catalog.database="${SNOWFLAKE_DATABASE}" \
             -Dtesting.snowflake.catalog.schema="${SNOWFLAKE_CATALOG_SCHEMA}" \

docs/src/main/sphinx/object-storage/metastores.md

@@ -675,7 +675,10 @@ properties:
 * - `iceberg.snowflake-catalog.user`
   - Snowflake user (required).
 * - `iceberg.snowflake-catalog.password`
-  - Snowflake password (required).
+  - Snowflake password (deprecated).
+* - `iceberg.snowflake-catalog.key`
+  - Snowflake key for key-pair authentication (required). 
+    Set either key or password, but not both.
 * - `iceberg.snowflake-catalog.database`
   - Snowflake database name (required).
 * - `iceberg.snowflake-catalog.role`
@@ -687,7 +690,7 @@ connector.name=iceberg
 iceberg.catalog.type=snowflake
 iceberg.snowflake-catalog.account-uri=jdbc:snowflake://example1234567890.snowflakecomputing.com
 iceberg.snowflake-catalog.user=user
-iceberg.snowflake-catalog.password=secret
+iceberg.snowflake-catalog.key=xxxx
 iceberg.snowflake-catalog.database=db
 ```
 

plugin/trino-iceberg/pom.xml

@@ -759,7 +759,9 @@
                                 <exclude>**/TestIcebergAbfsConnectorSmokeTest.java</exclude>
                                 <exclude>**/Test*FailureRecoveryTest.java</exclude>
                                 <exclude>**/TestIcebergSnowflakeCatalogConnectorSmokeTest.java</exclude>
+                                <exclude>**/TestIcebergSnowflakeCatalogWithKeyPairConnectorSmokeTest.java</exclude>
                                 <exclude>**/TestTrinoSnowflakeCatalog.java</exclude>
+                                <exclude>**/TestTrinoSnowflakeCatalogWithKeyPair.java</exclude>
                             </excludes>
                         </configuration>
                     </plugin>
@@ -824,7 +826,9 @@
                                 <include>**/TestIcebergGcsConnectorSmokeTest.java</include>
                                 <include>**/TestIcebergAbfsConnectorSmokeTest.java</include>
                                 <include>**/TestIcebergSnowflakeCatalogConnectorSmokeTest.java</include>
+                                <include>**/TestIcebergSnowflakeCatalogWithKeyPairConnectorSmokeTest.java</include>
                                 <include>**/TestTrinoSnowflakeCatalog.java</include>
+                                <include>**/TestTrinoSnowflakeCatalogWithKeyPair.java</include>
                             </includes>
                         </configuration>
                     </plugin>

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/snowflake/IcebergSnowflakeCatalogConfig.java

@@ -29,7 +29,8 @@ public class IcebergSnowflakeCatalogConfig
 {
     private URI uri;
     private String user;
-    private String password;
+    private Optional<String> password = Optional.empty();
+    private Optional<String> key = Optional.empty();
     private String database;
     private Optional<String> role = Optional.empty();
 
@@ -41,6 +42,12 @@ public boolean isUrlValid()
         return driver.acceptsURL(uri.toString());
     }
 
+    @AssertTrue(message = "Either iceberg.snowflake-catalog.password or iceberg.snowflake-catalog.key must be set, but not both")
+    public boolean isAuthenticationMethodSet()
+    {
+        return key.isPresent() != password.isPresent();
+    }
+
     @NotNull
     public URI getUri()
     {
@@ -69,18 +76,33 @@ public IcebergSnowflakeCatalogConfig setUser(String user)
         return this;
     }
 
-    @NotNull
-    public String getPassword()
+    @Deprecated
+    public Optional<String> getPassword()
     {
         return password;
     }
 
+    @Deprecated
     @Config("iceberg.snowflake-catalog.password")
     @ConfigDescription("Password for Snowflake")
     @ConfigSecuritySensitive
     public IcebergSnowflakeCatalogConfig setPassword(String password)
     {
-        this.password = password;
+        this.password = Optional.ofNullable(password);
+        return this;
+    }
+
+    public Optional<String> getKey()
+    {
+        return key;
+    }
+
+    @Config("iceberg.snowflake-catalog.key")
+    @ConfigDescription("The base64 encoded private key for key-pair authentication")
+    @ConfigSecuritySensitive
+    public IcebergSnowflakeCatalogConfig setKey(String key)
+    {
+        this.key = Optional.ofNullable(key);
         return this;
     }
 

plugin/trino-iceberg/src/main/java/org/apache/iceberg/snowflake/TrinoIcebergSnowflakeCatalogFactory.java

@@ -69,6 +69,7 @@ public TrinoIcebergSnowflakeCatalogFactory(
                 snowflakeCatalogConfig.getUri(),
                 snowflakeCatalogConfig.getUser(),
                 snowflakeCatalogConfig.getPassword(),
+                snowflakeCatalogConfig.getKey(),
                 snowflakeCatalogConfig.getRole());
         this.snowflakeDatabase = snowflakeCatalogConfig.getDatabase();
         this.snowflakeConnectionPool = new JdbcClientPool(snowflakeCatalogConfig.getUri().toString(), snowflakeDriverProperties);
@@ -94,7 +95,7 @@ public synchronized TrinoCatalog create(ConnectorIdentity identity)
         return new TrinoSnowflakeCatalog(icebergSnowflakeCatalog, catalogName, typeManager, fileSystemFactory, tableOperationsProvider, snowflakeDatabase);
     }
 
-    public static Map<String, String> getSnowflakeDriverProperties(URI snowflakeUri, String snowflakeUser, String snowflakePassword, Optional<String> snowflakeRole)
+    public static Map<String, String> getSnowflakeDriverProperties(URI snowflakeUri, String snowflakeUser, Optional<String> snowflakePassword, Optional<String> snowflakeKey, Optional<String> snowflakeRole)
     {
         // Below property values are copied from https://github.com/apache/iceberg/blob/apache-iceberg-1.5.0/snowflake/src/main/java/org/apache/iceberg/snowflake/SnowflakeCatalog.java#L122-L129
 
@@ -106,14 +107,15 @@ public static Map<String, String> getSnowflakeDriverProperties(URI snowflakeUri,
         ImmutableMap.Builder<String, String> properties = ImmutableMap.builder();
         properties
                 .put(PROPERTY_PREFIX + "user", snowflakeUser)
-                .put(PROPERTY_PREFIX + "password", snowflakePassword)
                 .put("uri", snowflakeUri.toString())
                 .put(PROPERTY_PREFIX + "JDBC_QUERY_RESULT_FORMAT", "JSON")
                 // Populate application identifier in jdbc client
                 .put(PROPERTY_PREFIX + JDBC_APPLICATION_PROPERTY, uniqueAppIdentifier)
                 // Adds application identifier to the user agent header of the JDBC requests.
                 .put(PROPERTY_PREFIX + JDBC_USER_AGENT_SUFFIX_PROPERTY, userAgentSuffix);
         snowflakeRole.ifPresent(role -> properties.put(PROPERTY_PREFIX + "role", role));
+        snowflakePassword.ifPresent(password -> properties.put(PROPERTY_PREFIX + "password", password));
+        snowflakeKey.ifPresent(key -> properties.put(PROPERTY_PREFIX + "private_key_base64", key));
 
         return properties.buildOrThrow();
     }

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/IcebergQueryRunner.java

@@ -488,7 +488,7 @@ public static void main(String[] args)
                             .put("iceberg.file-format", "PARQUET")
                             .put("iceberg.snowflake-catalog.account-uri", requiredNonEmptySystemProperty("testing.snowflake.catalog.account-url"))
                             .put("iceberg.snowflake-catalog.user", requiredNonEmptySystemProperty("testing.snowflake.catalog.user"))
-                            .put("iceberg.snowflake-catalog.password", requiredNonEmptySystemProperty("testing.snowflake.catalog.password"))
+                            .put("iceberg.snowflake-catalog.key", requiredNonEmptySystemProperty("testing.snowflake.catalog.key"))
                             .put("iceberg.snowflake-catalog.database", requiredNonEmptySystemProperty("testing.snowflake.catalog.database"))
                             .buildOrThrow())
                     .setSchemaInitializer(

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergPlugin.java

@@ -374,7 +374,7 @@ public void testSnowflakeCatalog()
                                 "iceberg.catalog.type", "snowflake",
                                 "iceberg.snowflake-catalog.account-uri", "jdbc:snowflake://sample.url",
                                 "iceberg.snowflake-catalog.user", "user",
-                                "iceberg.snowflake-catalog.password", "password",
+                                "iceberg.snowflake-catalog.key", "key",
                                 "iceberg.snowflake-catalog.database", "database"),
                         new TestingConnectorContext())
                 .shutdown();

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/snowflake/README.md

@@ -6,7 +6,7 @@ For testing purpose, some of the TPCH tables are created through tests by copyin
 
 Iceberg Snowflake catalog creation steps are taken from [official Iceberg Snowflake catalog guide](https://quickstarts.snowflake.com/guide/getting_started_iceberg_tables/index.html?_ga=2.81143831.1840596713.1702066099-371582571.1622222327#2)
 
-Below are the steps to create an external volume as outline in [external volume creation official guide](https://docs.snowflake.com/en/user-guide/tables-iceberg-configure-external-volume#label-tables-iceberg-configure-external-volume-s3-create-role)
+Below are the steps to create an external volume as outlined in [external volume creation official guide](https://docs.snowflake.com/en/user-guide/tables-iceberg-configure-external-volume#label-tables-iceberg-configure-external-volume-s3-create-role)
 
 1. Create an S3 bucket.  
 

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/snowflake/TestIcebergSnowflakeCatalogConfig.java

@@ -33,6 +33,7 @@
         assertRecordedDefaults(recordDefaults(IcebergSnowflakeCatalogConfig.class)
                 .setUser(null)
                 .setPassword(null)
+                .setKey(null)
                 .setDatabase(null)
                 .setUri(null)
                 .setRole(null));
@@ -43,6 +44,7 @@
     {
         Map<String, String> properties = ImmutableMap.<String, String>builder()
                 .put("iceberg.snowflake-catalog.password", "password")
+                .put("iceberg.snowflake-catalog.key", "key")
                 .put("iceberg.snowflake-catalog.user", "user")
                 .put("iceberg.snowflake-catalog.role", "role")
                 .put("iceberg.snowflake-catalog.account-uri", "jdbc:snowflake://sample.url")
@@ -51,12 +53,13 @@
 
         IcebergSnowflakeCatalogConfig expected = new IcebergSnowflakeCatalogConfig()
                 .setPassword("password")
+                .setKey("key")
                 .setUser("user")
                 .setRole("role")
                 .setUri(URI.create("jdbc:snowflake://sample.url"))
                 .setDatabase("database");
 
         assertFullMapping(properties, expected);
     }
 
     @Test
@@ -65,10 +68,47 @@
     {
         IcebergSnowflakeCatalogConfig config = new IcebergSnowflakeCatalogConfig()
                 .setPassword("password")
+                .setKey("key")
                 .setUser("user")
                 .setRole("role")
                 .setUri(URI.create("foobar"))
                 .setDatabase("database");
         assertThat(config.isUrlValid()).isFalse();
     }
+
+    @Test
+    public void testInvalidAuthenticationMethod()
+    {
+        IcebergSnowflakeCatalogConfig configWithBothKeyAndPassword = new IcebergSnowflakeCatalogConfig()
+                .setPassword("password")
+                .setKey("key")
+                .setUser("user")
+                .setRole("role")
+                .setUri(URI.create("foobar"))
+                .setDatabase("database");
+        assertThat(configWithBothKeyAndPassword.isAuthenticationMethodSet()).isFalse();
+
+        IcebergSnowflakeCatalogConfig configWithNeitherKeyNorPassword = new IcebergSnowflakeCatalogConfig()
+                .setUser("user")
+                .setRole("role")
+                .setUri(URI.create("foobar"))
+                .setDatabase("database");
+        assertThat(configWithNeitherKeyNorPassword.isAuthenticationMethodSet()).isFalse();
+
+        IcebergSnowflakeCatalogConfig configWithPassword = new IcebergSnowflakeCatalogConfig()
+                .setPassword("password")
+                .setUser("user")
+                .setRole("role")
+                .setUri(URI.create("foobar"))
+                .setDatabase("database");
+        assertThat(configWithPassword.isAuthenticationMethodSet()).isTrue();
+
+        IcebergSnowflakeCatalogConfig configWithKey = new IcebergSnowflakeCatalogConfig()
+                .setKey("key")
+                .setUser("user")
+                .setRole("role")
+                .setUri(URI.create("foobar"))
+                .setDatabase("database");
+        assertThat(configWithKey.isAuthenticationMethodSet()).isTrue();
+    }
 }

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/snowflake/TestIcebergSnowflakeCatalogConnectorSmokeTest.java

@@ -98,7 +98,20 @@ REGIONKEY NUMBER(38,0),
                     .formatted(TpchTable.REGION.getTableName(), TpchTable.REGION.getTableName()));
         }
 
-        ImmutableMap<String, String> properties = ImmutableMap.<String, String>builder()
+        ImmutableMap<String, String> properties = createIcebergProperties();
+
+        return IcebergQueryRunner.builder(SNOWFLAKE_TEST_SCHEMA.toLowerCase(ENGLISH))
+                .setIcebergProperties(properties)
+                .setSchemaInitializer(
+                        SchemaInitializer.builder()
+                                .withSchemaName(SNOWFLAKE_TEST_SCHEMA.toLowerCase(ENGLISH))
+                                .build())
+                .build();
+    }
+
+    protected ImmutableMap<String, String> createIcebergProperties()
+    {
+        return ImmutableMap.<String, String>builder()
                 .put("fs.native-s3.enabled", "true")
                 .put("s3.aws-access-key", S3_ACCESS_KEY)
                 .put("s3.aws-secret-key", S3_SECRET_KEY)
@@ -111,14 +124,6 @@ REGIONKEY NUMBER(38,0),
                 .put("iceberg.snowflake-catalog.user", SNOWFLAKE_USER)
                 .put("iceberg.snowflake-catalog.password", SNOWFLAKE_PASSWORD)
                 .buildOrThrow();
-
-        return IcebergQueryRunner.builder(SNOWFLAKE_TEST_SCHEMA.toLowerCase(ENGLISH))
-                .setIcebergProperties(properties)
-                .setSchemaInitializer(
-                        SchemaInitializer.builder()
-                                .withSchemaName(SNOWFLAKE_TEST_SCHEMA.toLowerCase(ENGLISH))
-                                .build())
-                .build();
     }
 
     @Override

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/snowflake/TestIcebergSnowflakeCatalogWithKeyPairConnectorSmokeTest.java

@@ -0,0 +1,47 @@
+/*
+ * Licensed 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 io.trino.plugin.iceberg.catalog.snowflake;
+
+import com.google.common.collect.ImmutableMap;
+import org.junit.jupiter.api.TestInstance;
+
+import static io.trino.plugin.iceberg.catalog.snowflake.TestingSnowflakeServer.SNOWFLAKE_JDBC_URI;
+import static io.trino.plugin.iceberg.catalog.snowflake.TestingSnowflakeServer.SNOWFLAKE_KEY;
+import static io.trino.plugin.iceberg.catalog.snowflake.TestingSnowflakeServer.SNOWFLAKE_ROLE;
+import static io.trino.plugin.iceberg.catalog.snowflake.TestingSnowflakeServer.SNOWFLAKE_TEST_DATABASE;
+import static io.trino.plugin.iceberg.catalog.snowflake.TestingSnowflakeServer.SNOWFLAKE_USER;
+import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS;
+
+@TestInstance(PER_CLASS)
+public class TestIcebergSnowflakeCatalogWithKeyPairConnectorSmokeTest
+        extends TestIcebergSnowflakeCatalogConnectorSmokeTest
+{
+    @Override
+    protected ImmutableMap<String, String> createIcebergProperties()
+    {
+        return ImmutableMap.<String, String>builder()
+                .put("fs.native-s3.enabled", "true")
+                .put("s3.aws-access-key", S3_ACCESS_KEY)
+                .put("s3.aws-secret-key", S3_SECRET_KEY)
+                .put("s3.region", S3_REGION)
+                .put("iceberg.file-format", "parquet") // only Parquet is supported
+                .put("iceberg.catalog.type", "snowflake")
+                .put("iceberg.snowflake-catalog.role", SNOWFLAKE_ROLE)
+                .put("iceberg.snowflake-catalog.database", SNOWFLAKE_TEST_DATABASE)
+                .put("iceberg.snowflake-catalog.account-uri", SNOWFLAKE_JDBC_URI)
+                .put("iceberg.snowflake-catalog.user", SNOWFLAKE_USER)
+                .put("iceberg.snowflake-catalog.key", SNOWFLAKE_KEY)
+                .buildOrThrow();
+    }
+}

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/snowflake/TestingSnowflakeServer.java

@@ -39,6 +39,7 @@ public enum TableType
     public static final String SNOWFLAKE_JDBC_URI = requiredNonEmptySystemProperty("testing.snowflake.catalog.account-url");
     public static final String SNOWFLAKE_USER = requiredNonEmptySystemProperty("testing.snowflake.catalog.user");
     public static final String SNOWFLAKE_PASSWORD = requiredNonEmptySystemProperty("testing.snowflake.catalog.password");
+    public static final String SNOWFLAKE_KEY = requiredNonEmptySystemProperty("testing.snowflake.catalog.key");
     public static final String SNOWFLAKE_ROLE = requiredNonEmptySystemProperty("testing.snowflake.catalog.role");
     public static final String SNOWFLAKE_WAREHOUSE = requiredNonEmptySystemProperty("testing.snowflake.catalog.warehouse");
     public static final String SNOWFLAKE_TEST_DATABASE = requiredNonEmptySystemProperty("testing.snowflake.catalog.database");
@@ -71,7 +72,7 @@ private <T> Optional<T> executeOnDatabaseWithResultSetOperator(String schema, Op
     {
         Properties properties = new Properties();
         properties.put("user", SNOWFLAKE_USER);
-        properties.put("password", SNOWFLAKE_PASSWORD);
+        properties.put("private_key_base64", SNOWFLAKE_KEY);
         properties.put("role", SNOWFLAKE_ROLE);
         properties.put("warehouse", SNOWFLAKE_WAREHOUSE);
         properties.put("db", SNOWFLAKE_TEST_DATABASE);