Add client tags to directory context properties

presto-hive/src/main/java/com/facebook/presto/hive/HiveUtil.java

@@ -171,6 +171,7 @@
 import static java.lang.Long.parseLong;
 import static java.lang.Short.parseShort;
 import static java.lang.String.format;
+import static java.lang.String.join;
 import static java.math.BigDecimal.ROUND_UNNECESSARY;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Objects.requireNonNull;
@@ -193,6 +194,8 @@ public final class HiveUtil
     public static final String PRESTO_CLIENT_INFO = "presto_client_info";
     public static final String PRESTO_USER_NAME = "presto_user_name";
     public static final String PRESTO_METASTORE_HEADER = "presto_metastore_header";
+    public static final String PRESTO_CLIENT_TAGS = "presto_client_tags";
+    public static final String CLIENT_TAGS_DELIMITER = ",";
 
     private static final Pattern DEFAULT_HIVE_COLUMN_NAME_PATTERN = Pattern.compile("_col\\d+");
 
@@ -1300,6 +1303,9 @@ public static Map<String, String> buildDirectoryContextProperties(ConnectorSessi
         session.getClientInfo().ifPresent(clientInfo -> directoryContextProperties.put(PRESTO_CLIENT_INFO, clientInfo));
         getMetastoreHeaders(session).ifPresent(metastoreHeaders -> directoryContextProperties.put(PRESTO_METASTORE_HEADER, metastoreHeaders));
         directoryContextProperties.put(PRESTO_USER_NAME, session.getUser());
+        if (!session.getClientTags().isEmpty()) {
+            directoryContextProperties.put(PRESTO_CLIENT_TAGS, join(CLIENT_TAGS_DELIMITER, session.getClientTags()));
+        }
         return directoryContextProperties.build();
     }
 }

presto-hive/src/test/java/com/facebook/presto/hive/HiveTestUtils.java

@@ -103,7 +103,9 @@ private HiveTestUtils()
     public static final JsonCodec<PartitionUpdate> PARTITION_UPDATE_CODEC = jsonCodec(PartitionUpdate.class);
     public static final SmileCodec<PartitionUpdate> PARTITION_UPDATE_SMILE_CODEC = smileCodec(PartitionUpdate.class);
 
-    public static final ConnectorSession SESSION = new TestingConnectorSession(getAllSessionProperties(new HiveClientConfig(), new HiveCommonClientConfig()));
+    public static final Set<String> TEST_CLIENT_TAGS = ImmutableSet.of("TAG1", "TAG2");
+
+    public static final ConnectorSession SESSION = new TestingConnectorSession(getAllSessionProperties(new HiveClientConfig(), new HiveCommonClientConfig()), TEST_CLIENT_TAGS);
     public static final MetadataManager METADATA = MetadataManager.createTestMetadataManager();
 
     public static final FunctionAndTypeManager FUNCTION_AND_TYPE_MANAGER = METADATA.getFunctionAndTypeManager();

presto-hive/src/test/java/com/facebook/presto/hive/TestHiveUtil.java

@@ -37,6 +37,7 @@
 import java.time.ZoneId;
 import java.util.AbstractList;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -49,8 +50,10 @@
 import static com.facebook.presto.common.type.IntegerType.INTEGER;
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
 import static com.facebook.presto.hive.HiveTestUtils.SESSION;
+import static com.facebook.presto.hive.HiveUtil.CLIENT_TAGS_DELIMITER;
 import static com.facebook.presto.hive.HiveUtil.CUSTOM_FILE_SPLIT_CLASS_KEY;
 import static com.facebook.presto.hive.HiveUtil.PRESTO_CLIENT_INFO;
+import static com.facebook.presto.hive.HiveUtil.PRESTO_CLIENT_TAGS;
 import static com.facebook.presto.hive.HiveUtil.PRESTO_METASTORE_HEADER;
 import static com.facebook.presto.hive.HiveUtil.PRESTO_QUERY_ID;
 import static com.facebook.presto.hive.HiveUtil.PRESTO_QUERY_SOURCE;
@@ -68,6 +71,7 @@
 import static com.facebook.presto.hive.util.HudiRealtimeSplitConverter.HUDI_DELTA_FILEPATHS_KEY;
 import static com.facebook.presto.hive.util.HudiRealtimeSplitConverter.HUDI_MAX_COMMIT_TIME_KEY;
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_CLASS;
 import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_FORMAT;
 import static org.apache.hadoop.hive.serde.serdeConstants.SERIALIZATION_LIB;
@@ -192,6 +196,7 @@ public void testBuildDirectoryContextProperties()
         assertEquals(Optional.ofNullable(additionalProperties.get(PRESTO_CLIENT_INFO)), SESSION.getClientInfo());
         assertEquals(additionalProperties.get(PRESTO_USER_NAME), SESSION.getUser());
         assertEquals(Optional.ofNullable(additionalProperties.get(PRESTO_METASTORE_HEADER)), getMetastoreHeaders(SESSION));
+        assertEquals(Arrays.stream(additionalProperties.get(PRESTO_CLIENT_TAGS).split(CLIENT_TAGS_DELIMITER)).collect(toImmutableSet()), SESSION.getClientTags());
     }
 
     @Test

presto-main/src/main/java/com/facebook/presto/testing/TestingConnectorSession.java

@@ -68,6 +68,11 @@ public TestingConnectorSession(List<PropertyMetadata<?>> properties)
         this("user", Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, ImmutableMap.of(), new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), ImmutableSet.of(), Optional.empty(), ImmutableMap.of());
     }
 
+    public TestingConnectorSession(List<PropertyMetadata<?>> properties, Set<String> clientTags)
+    {
+        this("user", Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, ImmutableMap.of(), new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), clientTags, Optional.empty(), ImmutableMap.of());
+    }
+
     public TestingConnectorSession(List<PropertyMetadata<?>> properties, Map<String, Object> propertyValues)
     {
         this("user", Optional.of("test"), Optional.empty(), UTC_KEY, ENGLISH, System.currentTimeMillis(), properties, propertyValues, new FeaturesConfig().isLegacyTimestamp(), Optional.empty(), ImmutableSet.of(), Optional.empty(), ImmutableMap.of());