[trinodb#6] Add HDFS Caller Context

rice668 · Mar 28, 2023 · e51cf96 · e51cf96
1 parent 7d3d8fc
commit e51cf96
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 9 deletions.
diff --git a/lib/trino-filesystem/src/main/java/io/trino/filesystem/hdfs/HdfsFileSystemFactory.java b/lib/trino-filesystem/src/main/java/io/trino/filesystem/hdfs/HdfsFileSystemFactory.java
@@ -17,6 +17,7 @@
 import io.trino.filesystem.TrinoFileSystemFactory;
 import io.trino.hdfs.HdfsContext;
 import io.trino.hdfs.HdfsEnvironment;
+import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.security.ConnectorIdentity;
 
 import javax.inject.Inject;
@@ -39,4 +40,10 @@ public TrinoFileSystem create(ConnectorIdentity identity)
     {
         return new HdfsFileSystem(environment, new HdfsContext(identity));
     }
+
+    @Override
+    public TrinoFileSystem create(ConnectorSession session)
+    {
+        return new HdfsFileSystem(environment, new HdfsContext(session));
+    }
 }
diff --git a/lib/trino-hdfs/src/main/java/io/trino/hdfs/HdfsContext.java b/lib/trino-hdfs/src/main/java/io/trino/hdfs/HdfsContext.java
@@ -16,29 +16,39 @@
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.security.ConnectorIdentity;
 
+import java.util.Optional;
+
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static java.util.Objects.requireNonNull;
 
 public class HdfsContext
 {
     private final ConnectorIdentity identity;
+    private final Optional<String> queryId;
 
     public HdfsContext(ConnectorIdentity identity)
     {
         this.identity = requireNonNull(identity, "identity is null");
+        this.queryId = Optional.empty();
     }
 
     public HdfsContext(ConnectorSession session)
     {
         requireNonNull(session, "session is null");
         this.identity = requireNonNull(session.getIdentity(), "session.getIdentity() is null");
+        this.queryId = Optional.of(session.getQueryId());
     }
 
     public ConnectorIdentity getIdentity()
     {
         return identity;
     }
 
+    public Optional<String> getQueryId()
+    {
+        return queryId;
+    }
+
     @Override
     public String toString()
     {

diff --git a/lib/trino-hdfs/src/main/java/io/trino/hdfs/HdfsEnvironment.java b/lib/trino-hdfs/src/main/java/io/trino/hdfs/HdfsEnvironment.java
@@ -16,12 +16,14 @@
 import io.trino.hadoop.HadoopNative;
 import io.trino.hdfs.authentication.GenericExceptionAction;
 import io.trino.hdfs.authentication.HdfsAuthentication;
+import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.security.ConnectorIdentity;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.FileSystemManager;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.ipc.CallerContext;
 
 import javax.inject.Inject;
 
@@ -64,9 +66,26 @@ public Configuration getConfiguration(HdfsContext context, Path path)
     public FileSystem getFileSystem(HdfsContext context, Path path)
             throws IOException
     {
+        if (context.getQueryId().isPresent()) {
+            setCallerContext(context.getIdentity(), context.getQueryId().get());
+        }
         return getFileSystem(context.getIdentity(), path, getConfiguration(context, path));
     }
 
+    public FileSystem getFileSystem(ConnectorSession session, Path path, Configuration configuration)
+            throws IOException
+    {
+        setCallerContext(session.getIdentity(), session.getQueryId());
+        return getFileSystem(session.getIdentity(), path, configuration);
+    }
+
+    private void setCallerContext(ConnectorIdentity identity, String queryId)
+    {
+        String callerContent = String.format("JobId:olap_trino_%s$User:%s", queryId, identity.getUser());
+        CallerContext callerContext = new CallerContext.Builder(callerContent).build();
+        CallerContext.setCurrent(callerContext);
+    }
+
     public FileSystem getFileSystem(ConnectorIdentity identity, Path path, Configuration configuration)
             throws IOException
     {

diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/orc/OrcPageSourceFactory.java
@@ -47,7 +47,6 @@
 import io.trino.spi.connector.EmptyPageSource;
 import io.trino.spi.predicate.Domain;
 import io.trino.spi.predicate.TupleDomain;
-import io.trino.spi.security.ConnectorIdentity;
 import io.trino.spi.type.Type;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FSDataInputStream;
@@ -192,7 +191,7 @@ public Optional<ReaderPageSource> createPageSource(
 
         ConnectorPageSource orcPageSource = createOrcPageSource(
                 hdfsEnvironment,
-                session.getIdentity(),
+                session,
                 configuration,
                 path,
                 start,
@@ -224,7 +223,7 @@ public Optional<ReaderPageSource> createPageSource(
 
     private ConnectorPageSource createOrcPageSource(
             HdfsEnvironment hdfsEnvironment,
-            ConnectorIdentity identity,
+            ConnectorSession session,
             Configuration configuration,
             Path path,
             long start,
@@ -252,8 +251,8 @@ private ConnectorPageSource createOrcPageSource(
 
         boolean originalFilesPresent = acidInfo.isPresent() && !acidInfo.get().getOriginalFiles().isEmpty();
         try {
-            FileSystem fileSystem = hdfsEnvironment.getFileSystem(identity, path, configuration);
-            FSDataInputStream inputStream = hdfsEnvironment.doAs(identity, () -> fileSystem.open(path));
+            FileSystem fileSystem = hdfsEnvironment.getFileSystem(session, path, configuration);
+            FSDataInputStream inputStream = hdfsEnvironment.doAs(session.getIdentity(), () -> fileSystem.open(path));
             orcDataSource = new HdfsOrcDataSource(
                     new OrcDataSourceId(path.toString()),
                     estimatedFileSize,
@@ -397,8 +396,8 @@ else if (column.getBaseHiveColumnIndex() < fileColumns.size()) {
             Optional<OrcDeletedRows> deletedRows = acidInfo.map(info ->
                     new OrcDeletedRows(
                             path.getName(),
-                            new OrcDeleteDeltaPageSourceFactory(options, identity, configuration, hdfsEnvironment, stats),
-                            identity,
+                            new OrcDeleteDeltaPageSourceFactory(options, session.getIdentity(), configuration, hdfsEnvironment, stats),
+                            session.getIdentity(),
                             configuration,
                             hdfsEnvironment,
                             info,
@@ -412,7 +411,7 @@ else if (column.getBaseHiveColumnIndex() < fileColumns.size()) {
                             acidInfo.get().getOriginalFiles(),
                             path,
                             hdfsEnvironment,
-                            identity,
+                            session.getIdentity(),
                             options,
                             configuration,
                             stats));

diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/rcfile/RcFilePageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/rcfile/RcFilePageSourceFactory.java
@@ -152,7 +152,7 @@ else if (deserializerClassName.equals(ColumnarSerDe.class.getName())) {
 
         RcFileDataSource dataSource;
         try {
-            FileSystem fileSystem = hdfsEnvironment.getFileSystem(session.getIdentity(), path, configuration);
+            FileSystem fileSystem = hdfsEnvironment.getFileSystem(session, path, configuration);
             FSDataInputStream inputStream = hdfsEnvironment.doAs(session.getIdentity(), () -> fileSystem.open(path));
             if (estimatedFileSize < BUFFER_SIZE.toBytes()) {
                 //  Handle potentially imprecise file lengths by reading the footer