Introduce DeltaLakeFileSystemFactory and VendedCredentialsProvider in Delta Lake

Author: chenjian2664

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/AbstractDeltaLakePageSink.java

@@ -23,10 +23,10 @@
 import io.airlift.slice.Slice;
 import io.trino.filesystem.Location;
 import io.trino.filesystem.TrinoFileSystem;
-import io.trino.filesystem.TrinoFileSystemFactory;
 import io.trino.metastore.Partitions;
 import io.trino.parquet.writer.ParquetWriterOptions;
 import io.trino.plugin.deltalake.DataFileInfo.DataFileType;
+import io.trino.plugin.deltalake.metastore.VendedCredentialsHandle;
 import io.trino.plugin.deltalake.util.DeltaLakeWriteUtils;
 import io.trino.plugin.hive.parquet.ParquetFileWriter;
 import io.trino.spi.Page;
@@ -118,10 +118,11 @@ public AbstractDeltaLakePageSink(
             List<DeltaLakeColumnHandle> inputColumns,
             List<String> originalPartitionColumns,
             PageIndexerFactory pageIndexerFactory,
-            TrinoFileSystemFactory fileSystemFactory,
+            DeltaLakeFileSystemFactory fileSystemFactory,
             int maxOpenWriters,
             JsonCodec<DataFileInfo> dataFileInfoCodec,
             Location tableLocation,
+            VendedCredentialsHandle credentialsHandle,
             Location outputPathDirectory,
             ConnectorSession session,
             DeltaLakeWriterStats stats,
@@ -133,7 +134,7 @@ public AbstractDeltaLakePageSink(
 
         requireNonNull(pageIndexerFactory, "pageIndexerFactory is null");
 
-        this.fileSystem = requireNonNull(fileSystemFactory, "fileSystemFactory is null").create(session);
+        this.fileSystem = requireNonNull(fileSystemFactory, "fileSystemFactory is null").create(session, credentialsHandle);
         this.maxOpenWriters = maxOpenWriters;
         this.dataFileInfoCodec = requireNonNull(dataFileInfoCodec, "dataFileInfoCodec is null");
         this.parquetSchemaMapping = requireNonNull(parquetSchemaMapping, "parquetSchemaMapping is null");

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/BaseTransactionsTable.java

@@ -16,7 +16,6 @@
 import com.google.common.collect.ImmutableList;
 import io.airlift.units.DataSize;
 import io.trino.filesystem.TrinoFileSystem;
-import io.trino.filesystem.TrinoFileSystemFactory;
 import io.trino.plugin.deltalake.metastore.DeltaMetastoreTable;
 import io.trino.plugin.deltalake.transactionlog.TableSnapshot;
 import io.trino.plugin.deltalake.transactionlog.Transaction;
@@ -53,13 +52,13 @@ public abstract class BaseTransactionsTable
         implements SystemTable
 {
     private final DeltaMetastoreTable table;
-    private final TrinoFileSystemFactory fileSystemFactory;
+    private final DeltaLakeFileSystemFactory fileSystemFactory;
     private final TransactionLogAccess transactionLogAccess;
     private final ConnectorTableMetadata tableMetadata;
 
     public BaseTransactionsTable(
             DeltaMetastoreTable table,
-            TrinoFileSystemFactory fileSystemFactory,
+            DeltaLakeFileSystemFactory fileSystemFactory,
             TransactionLogAccess transactionLogAccess,
             TypeManager typeManager,
             ConnectorTableMetadata tableMetadata)
@@ -86,12 +85,13 @@ public ConnectorTableMetadata getTableMetadata()
     @Override
     public ConnectorPageSource pageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain<Integer> constraint)
     {
+        TrinoFileSystem fileSystem = fileSystemFactory.create(session, table);
         long snapshotVersion;
         try {
             // Verify the transaction log is readable
             TableSnapshot tableSnapshot = transactionLogAccess.loadSnapshot(session, table, Optional.empty());
             snapshotVersion = tableSnapshot.getVersion();
-            transactionLogAccess.getMetadataEntry(session, tableSnapshot);
+            transactionLogAccess.getMetadataEntry(session, fileSystem, tableSnapshot);
         }
         catch (IOException e) {
             throw new TrinoException(DeltaLakeErrorCode.DELTA_LAKE_INVALID_SCHEMA, "Unable to load table metadata from location: " + table.location(), e);
@@ -137,7 +137,6 @@ public ConnectorPageSource pageSource(ConnectorTransactionHandle transactionHand
             endVersionInclusive = Optional.of(snapshotVersion);
         }
 
-        TrinoFileSystem fileSystem = fileSystemFactory.create(session);
         PageListBuilder pagesBuilder = PageListBuilder.forTable(tableMetadata);
         try {
             List<Transaction> transactions = loadNewTailBackward(fileSystem, table.location(), startVersionExclusive, endVersionInclusive.get()).reversed();

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/CorruptedDeltaLakeTableHandle.java

@@ -13,15 +13,21 @@
  */
 package io.trino.plugin.deltalake;
 
+import io.trino.plugin.deltalake.metastore.VendedCredentials;
+import io.trino.plugin.deltalake.metastore.VendedCredentialsHandle;
 import io.trino.spi.TrinoException;
 import io.trino.spi.connector.SchemaTableName;
 
+import java.util.Optional;
+
 import static java.util.Objects.requireNonNull;
 
 public record CorruptedDeltaLakeTableHandle(
         SchemaTableName schemaTableName,
+        boolean catalogOwned,
         boolean managed,
         String location,
+        Optional<VendedCredentials> vendedCredentials,
         TrinoException originalException)
         implements LocatedTableHandle
 {
@@ -37,4 +43,10 @@ public TrinoException createException()
         // Original exception originates from a different place. Create a new exception not to confuse reader with a stacktrace not matching call site.
         return new TrinoException(originalException::getErrorCode, originalException.getMessage(), originalException);
     }
+
+    @Override
+    public VendedCredentialsHandle toCredentialsHandle()
+    {
+        return new VendedCredentialsHandle(catalogOwned, managed, location, vendedCredentials);
+    }
 }

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DefaultDeltaLakeFileSystemFactory.java

@@ -0,0 +1,53 @@
+/*
+ * 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.deltalake;
+
+import com.google.inject.Inject;
+import io.trino.filesystem.TrinoFileSystem;
+import io.trino.filesystem.TrinoFileSystemFactory;
+import io.trino.plugin.deltalake.metastore.VendedCredentialsHandle;
+import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.security.ConnectorIdentity;
+
+import static java.util.Objects.requireNonNull;
+
+public class DefaultDeltaLakeFileSystemFactory
+        implements DeltaLakeFileSystemFactory
+{
+    private final TrinoFileSystemFactory fileSystemFactory;
+
+    @Inject
+    public DefaultDeltaLakeFileSystemFactory(TrinoFileSystemFactory fileSystemFactory)
+    {
+        this.fileSystemFactory = requireNonNull(fileSystemFactory, "fileSystemFactory is null");
+    }
+
+    @Override
+    public TrinoFileSystem create(ConnectorSession session, VendedCredentialsHandle table)
+    {
+        return fileSystemFactory.create(session.getIdentity());
+    }
+
+    @Override
+    public TrinoFileSystem create(ConnectorSession session, String tableLocation)
+    {
+        return fileSystemFactory.create(session.getIdentity());
+    }
+
+    @Override
+    public TrinoFileSystem create(ConnectorIdentity identity)
+    {
+        throw new UnsupportedOperationException();
+    }
+}

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeCdfPageSink.java

@@ -15,7 +15,7 @@
 
 import io.airlift.json.JsonCodec;
 import io.trino.filesystem.Location;
-import io.trino.filesystem.TrinoFileSystemFactory;
+import io.trino.plugin.deltalake.metastore.VendedCredentialsHandle;
 import io.trino.spi.PageIndexerFactory;
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.type.TypeOperators;
@@ -35,10 +35,11 @@ public DeltaLakeCdfPageSink(
             List<DeltaLakeColumnHandle> inputColumns,
             List<String> originalPartitionColumns,
             PageIndexerFactory pageIndexerFactory,
-            TrinoFileSystemFactory fileSystemFactory,
+            DeltaLakeFileSystemFactory fileSystemFactory,
             int maxOpenWriters,
             JsonCodec<DataFileInfo> dataFileInfoCodec,
             Location tableLocation,
+            VendedCredentialsHandle credentialsHandle,
             Location outputPath,
             ConnectorSession session,
             DeltaLakeWriterStats stats,
@@ -54,6 +55,7 @@ public DeltaLakeCdfPageSink(
                 maxOpenWriters,
                 dataFileInfoCodec,
                 tableLocation,
+                credentialsHandle,
                 outputPath,
                 session,
                 stats,

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeFileSystemFactory.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.deltalake;
+
+import io.trino.filesystem.TrinoFileSystem;
+import io.trino.filesystem.TrinoFileSystemFactory;
+import io.trino.plugin.deltalake.metastore.DeltaMetastoreTable;
+import io.trino.plugin.deltalake.metastore.VendedCredentialsHandle;
+import io.trino.spi.connector.ConnectorSession;
+
+public interface DeltaLakeFileSystemFactory
+        extends TrinoFileSystemFactory
+{
+    default TrinoFileSystem create(ConnectorSession session, DeltaLakeTableHandle table)
+    {
+        return create(session, table.toCredentialsHandle());
+    }
+
+    default TrinoFileSystem create(ConnectorSession session, DeltaMetastoreTable table)
+    {
+        return create(session, new VendedCredentialsHandle(table.catalogOwned(), table.managed(), table.location(), table.vendedCredentials()));
+    }
+
+    TrinoFileSystem create(ConnectorSession session, VendedCredentialsHandle table);
+
+    /**
+     * For external table create/write using location
+     */
+    TrinoFileSystem create(ConnectorSession session, String tableLocation);
+
+    @Override
+    default TrinoFileSystem create(ConnectorSession session)
+    {
+        throw new UnsupportedOperationException();
+    }
+}

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeHistoryTable.java

@@ -15,7 +15,6 @@
 
 import com.google.common.collect.ImmutableList;
 import io.trino.filesystem.TrinoFileSystem;
-import io.trino.filesystem.TrinoFileSystemFactory;
 import io.trino.plugin.deltalake.metastore.DeltaMetastoreTable;
 import io.trino.plugin.deltalake.transactionlog.CommitInfoEntry;
 import io.trino.plugin.deltalake.transactionlog.DeltaLakeTransactionLogEntry;
@@ -46,7 +45,7 @@ public class DeltaLakeHistoryTable
 {
     public DeltaLakeHistoryTable(
             DeltaMetastoreTable table,
-            TrinoFileSystemFactory fileSystemFactory,
+            DeltaLakeFileSystemFactory fileSystemFactory,
             TransactionLogAccess transactionLogAccess,
             TypeManager typeManager)
     {

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeInsertTableHandle.java

@@ -14,6 +14,7 @@
 package io.trino.plugin.deltalake;
 
 import com.google.common.collect.ImmutableList;
+import io.trino.plugin.deltalake.metastore.VendedCredentialsHandle;
 import io.trino.plugin.deltalake.transactionlog.MetadataEntry;
 import io.trino.plugin.deltalake.transactionlog.ProtocolEntry;
 import io.trino.spi.connector.ConnectorInsertTableHandle;
@@ -30,7 +31,8 @@ public record DeltaLakeInsertTableHandle(
         ProtocolEntry protocolEntry,
         List<DeltaLakeColumnHandle> inputColumns,
         long readVersion,
-        boolean retriesEnabled)
+        boolean retriesEnabled,
+        VendedCredentialsHandle credentialsHandle)
         implements ConnectorInsertTableHandle
 {
     public DeltaLakeInsertTableHandle
@@ -40,6 +42,7 @@ public record DeltaLakeInsertTableHandle(
         requireNonNull(protocolEntry, "protocolEntry is null");
         inputColumns = ImmutableList.copyOf(requireNonNull(inputColumns, "inputColumns is null"));
         requireNonNull(location, "location is null");
+        requireNonNull(credentialsHandle, "credentialsHandle is null");
     }
 
     @Override

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMergeSink.java

@@ -21,7 +21,6 @@
 import io.airlift.slice.Slices;
 import io.trino.filesystem.Location;
 import io.trino.filesystem.TrinoFileSystem;
-import io.trino.filesystem.TrinoFileSystemFactory;
 import io.trino.filesystem.TrinoInputFile;
 import io.trino.parquet.ParquetDataSource;
 import io.trino.parquet.ParquetReaderOptions;
@@ -31,6 +30,7 @@
 import io.trino.parquet.writer.ParquetWriterOptions;
 import io.trino.plugin.base.metrics.FileFormatDataSourceStats;
 import io.trino.plugin.deltalake.delete.RoaringBitmapArray;
+import io.trino.plugin.deltalake.metastore.VendedCredentialsHandle;
 import io.trino.plugin.deltalake.transactionlog.DeletionVectorEntry;
 import io.trino.plugin.hive.parquet.ParquetFileWriter;
 import io.trino.plugin.hive.parquet.ParquetPageSourceFactory;
@@ -139,14 +139,15 @@ public class DeltaLakeMergeSink
 
     public DeltaLakeMergeSink(
             TypeOperators typeOperators,
-            TrinoFileSystemFactory fileSystemFactory,
+            DeltaLakeFileSystemFactory fileSystemFactory,
             ConnectorSession session,
             DateTimeZone parquetDateTimeZone,
             String trinoVersion,
             JsonCodec<DataFileInfo> dataFileInfoCodec,
             JsonCodec<DeltaLakeMergeResult> mergeResultJsonCodec,
             DeltaLakeWriterStats writerStats,
             Location rootTableLocation,
+            VendedCredentialsHandle credentialsHandle,
             ConnectorPageSink insertPageSink,
             List<DeltaLakeColumnHandle> tableColumns,
             int domainCompactionThreshold,
@@ -162,7 +163,7 @@ public DeltaLakeMergeSink(
     {
         this.typeOperators = requireNonNull(typeOperators, "typeOperators is null");
         this.session = requireNonNull(session, "session is null");
-        this.fileSystem = fileSystemFactory.create(session);
+        this.fileSystem = fileSystemFactory.create(session, credentialsHandle);
         this.parquetDateTimeZone = requireNonNull(parquetDateTimeZone, "parquetDateTimeZone is null");
         this.trinoVersion = requireNonNull(trinoVersion, "trinoVersion is null");
         this.dataFileInfoCodec = requireNonNull(dataFileInfoCodec, "dataFileInfoCodec is null");