Skip to content

Commit

Permalink
Enrich MetastoreContext with clientTags
Browse files Browse the repository at this point in the history
  • Loading branch information
@konjac-h @jainxrohit
konjac-h authored and jainxrohit committed May 17, 2024
1 parent ea64563 commit 08eed5a
Show file tree
Hide file tree
Showing 40 changed files with 531 additions and 100 deletions.
1 change: 1 addition & 0 deletions presto-delta/src/main/java/com/facebook/presto/delta/DeltaMetadata.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -368,6 +368,7 @@ private MetastoreContext metastoreContext(ConnectorSession session)
session.getIdentity(),
session.getQueryId(),
session.getClientInfo(),
session.getClientTags(),
session.getSource(),
Optional.empty(),
false,
Expand Down
1 change: 1 addition & 0 deletions ...astore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -865,6 +865,7 @@ private <T> KeyAndContext<T> getCachingKey(MetastoreContext context, T key)
context.getUsername(),
context.getQueryId(),
context.getClientInfo(),
context.getClientTags(),
context.getSource(),
true,
context.getMetastoreHeaders(),
Expand Down
16 changes: 14 additions & 2 deletions presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/MetastoreContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@

import java.util.Objects;
import java.util.Optional;
import java.util.Set;

import static com.google.common.base.MoreObjects.toStringHelper;
import static java.util.Objects.requireNonNull;
Expand All @@ -31,6 +32,7 @@ public class MetastoreContext
private final String username;
private final String queryId;
private final Optional<String> clientInfo;
private final Set<String> clientTags;
private final Optional<String> source;
private final boolean impersonationEnabled;
private final Optional<String> metastoreHeaders;
Expand All @@ -47,34 +49,37 @@ public MetastoreContext(
ConnectorIdentity identity,
String queryId,
Optional<String> clientInfo,
Set<String> clientTags,
Optional<String> source,
Optional<String> metastoreHeaders,
boolean userDefinedTypeEncodingEnabled,
ColumnConverterProvider columnConverterProvider,
WarningCollector warningCollector,
RuntimeStats runtimeStats)
{
this(requireNonNull(identity, "identity is null").getUser(), queryId, clientInfo, source, metastoreHeaders, userDefinedTypeEncodingEnabled, columnConverterProvider, warningCollector, runtimeStats);
this(requireNonNull(identity, "identity is null").getUser(), queryId, clientInfo, clientTags, source, metastoreHeaders, userDefinedTypeEncodingEnabled, columnConverterProvider, warningCollector, runtimeStats);
}

public MetastoreContext(
String username,
String queryId,
Optional<String> clientInfo,
Set<String> clientTags,
Optional<String> source,
Optional<String> metastoreHeaders,
boolean userDefinedTypeEncodingEnabled,
ColumnConverterProvider columnConverterProvider,
WarningCollector warningCollector,
RuntimeStats runtimeStats)
{
this(username, queryId, clientInfo, source, false, metastoreHeaders, userDefinedTypeEncodingEnabled, columnConverterProvider, warningCollector, runtimeStats);
this(username, queryId, clientInfo, clientTags, source, false, metastoreHeaders, userDefinedTypeEncodingEnabled, columnConverterProvider, warningCollector, runtimeStats);
}

public MetastoreContext(
String username,
String queryId,
Optional<String> clientInfo,
Set<String> clientTags,
Optional<String> source,
boolean impersonationEnabled,
Optional<String> metastoreHeaders,
Expand All @@ -86,6 +91,7 @@ public MetastoreContext(
this.username = requireNonNull(username, "username is null");
this.queryId = requireNonNull(queryId, "queryId is null");
this.clientInfo = requireNonNull(clientInfo, "clientInfo is null");
this.clientTags = requireNonNull(clientTags, "clientTags is null");
this.source = requireNonNull(source, "source is null");
this.columnConverterProvider = requireNonNull(columnConverterProvider, "columnConverterProvider is null");
this.impersonationEnabled = impersonationEnabled;
Expand Down Expand Up @@ -126,6 +132,11 @@ public Optional<String> getClientInfo()
return clientInfo;
}

public Set<String> getClientTags()
{
return clientTags;
}

public Optional<String> getSource()
{
return source;
Expand Down Expand Up @@ -163,6 +174,7 @@ public String toString()
.add("username", username)
.add("queryId", queryId)
.add("clientInfo", clientInfo.orElse(""))
.add("clientTags", clientTags)
.add("source", source.orElse(""))
.add("impersonationEnabled", Boolean.toString(impersonationEnabled))
.add("userDefinedTypeEncodingEnabled", Boolean.toString(userDefinedTypeEncodingEnabled))
Expand Down
28 changes: 26 additions & 2 deletions ...tore/src/main/java/com/facebook/presto/hive/metastore/SemiTransactionalHiveMetastore.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
Expand Down Expand Up @@ -516,7 +517,17 @@ public synchronized void finishInsertIntoExistingTable(
setShared();
SchemaTableName schemaTableName = new SchemaTableName(databaseName, tableName);
Action<TableAndMore> oldTableAction = tableActions.get(schemaTableName);
MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), columnConverterProvider, session.getWarningCollector(), session.getRuntimeStats());
MetastoreContext metastoreContext = new MetastoreContext(
session.getIdentity(),
session.getQueryId(),
session.getClientInfo(),
session.getClientTags(),
session.getSource(),
getMetastoreHeaders(session),
isUserDefinedTypeEncodingEnabled(session),
columnConverterProvider,
session.getWarningCollector(),
session.getRuntimeStats());
if (oldTableAction == null || oldTableAction.getData().getTable().getTableType().equals(TEMPORARY_TABLE)) {
Table table = getTable(metastoreContext, databaseName, tableName)
.orElseThrow(() -> new TableNotFoundException(schemaTableName));
Expand Down Expand Up @@ -560,6 +571,7 @@ public synchronized void truncateUnpartitionedTable(ConnectorSession session, St
session.getIdentity(),
session.getQueryId(),
session.getClientInfo(),
session.getClientTags(),
session.getSource(),
getMetastoreHeaders(session),
isUserDefinedTypeEncodingEnabled(session),
Expand Down Expand Up @@ -909,7 +921,17 @@ public synchronized void finishInsertIntoExistingPartition(
SchemaTableName schemaTableName = new SchemaTableName(databaseName, tableName);
Map<List<String>, Action<PartitionAndMore>> partitionActionsOfTable = partitionActions.computeIfAbsent(schemaTableName, k -> new HashMap<>());
Action<PartitionAndMore> oldPartitionAction = partitionActionsOfTable.get(partitionValues);
MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), columnConverterProvider, session.getWarningCollector(), session.getRuntimeStats());
MetastoreContext metastoreContext = new MetastoreContext(
session.getIdentity(),
session.getQueryId(),
session.getClientInfo(),
session.getClientTags(),
session.getSource(),
getMetastoreHeaders(session),
isUserDefinedTypeEncodingEnabled(session),
columnConverterProvider,
session.getWarningCollector(),
session.getRuntimeStats());

if (oldPartitionAction == null) {
Partition partition = delegate.getPartition(metastoreContext, databaseName, tableName, partitionValues)
Expand Down Expand Up @@ -1149,6 +1171,7 @@ private ConnectorCommitHandle commitShared()
hdfsContext.getIdentity(),
hdfsContext.getQueryId().orElse(""),
hdfsContext.getClientInfo(),
hdfsContext.getClientTags().orElse(Collections.emptySet()),
hdfsContext.getSource(),
hdfsContext.getSession().flatMap(MetastoreUtil::getMetastoreHeaders),
hdfsContext.getSession().map(MetastoreUtil::isUserDefinedTypeEncodingEnabled).orElse(false),
Expand Down Expand Up @@ -1182,6 +1205,7 @@ private ConnectorCommitHandle commitShared()
hdfsContext.getIdentity(),
hdfsContext.getQueryId().orElse(""),
hdfsContext.getClientInfo(),
hdfsContext.getClientTags().orElse(Collections.emptySet()),
hdfsContext.getSource(),
hdfsContext.getSession().flatMap(MetastoreUtil::getMetastoreHeaders),
hdfsContext.getSession().map(MetastoreUtil::isUserDefinedTypeEncodingEnabled).orElse(false),
Expand Down
3 changes: 2 additions & 1 deletion ...tore/src/test/java/com/facebook/presto/hive/metastore/thrift/MockHiveMetastoreClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
import org.apache.hadoop.hive.metastore.api.UnlockRequest;
import org.apache.thrift.TException;

import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
Expand All @@ -75,7 +76,7 @@ public class MockHiveMetastoreClient
public static final List<SQLUniqueConstraint> TEST_UNIQUE_CONSTRAINT = ImmutableList.of(new SQLUniqueConstraint("", TEST_DATABASE, TEST_TABLE_WITH_CONSTRAINTS, "c2", 1, "uk", true, false, true));
public static final List<SQLNotNullConstraint> TEST_NOT_NULL_CONSTRAINT = ImmutableList.of(new SQLNotNullConstraint("", TEST_DATABASE, TEST_TABLE_WITH_CONSTRAINTS, "c3", "nn", true, true, true));
public static final String TEST_TOKEN = "token";
public static final MetastoreContext TEST_METASTORE_CONTEXT = new MetastoreContext("test_user", "test_queryId", Optional.empty(), Optional.empty(), Optional.empty(), false, HiveColumnConverterProvider.DEFAULT_COLUMN_CONVERTER_PROVIDER, WarningCollector.NOOP, new RuntimeStats());
public static final MetastoreContext TEST_METASTORE_CONTEXT = new MetastoreContext("test_user", "test_queryId", Optional.empty(), Collections.emptySet(), Optional.empty(), Optional.empty(), false, HiveColumnConverterProvider.DEFAULT_COLUMN_CONVERTER_PROVIDER, WarningCollector.NOOP, new RuntimeStats());
public static final String TEST_PARTITION1 = "key=testpartition1";
public static final String TEST_PARTITION2 = "key=testpartition2";
public static final List<String> TEST_PARTITION_VALUES1 = ImmutableList.of("testpartition1");
Expand Down
15 changes: 14 additions & 1 deletion presto-hive/src/main/java/com/facebook/presto/hive/CreateEmptyPartitionProcedure.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,20 @@ private void doCreateEmptyPartition(ConnectorSession session, String schema, Str
.map(String.class::cast)
.collect(toImmutableList());

if (metastore.getPartition(new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), false, HiveColumnConverterProvider.DEFAULT_COLUMN_CONVERTER_PROVIDER, session.getWarningCollector(), session.getRuntimeStats()), schema, table, partitionStringValues).isPresent()) {
if (metastore.getPartition(new MetastoreContext(
session.getIdentity(),
session.getQueryId(),
session.getClientInfo(),
session.getClientTags(),
session.getSource(),
getMetastoreHeaders(session),
false,
HiveColumnConverterProvider.DEFAULT_COLUMN_CONVERTER_PROVIDER,
session.getWarningCollector(),
session.getRuntimeStats()),
schema,
table,
partitionStringValues).isPresent()) {
throw new PrestoException(ALREADY_EXISTS, "Partition already exists");
}
String partitionName = FileUtils.makePartName(actualPartitionColumnNames, partitionStringValues);
Expand Down
12 changes: 11 additions & 1 deletion presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1773,7 +1773,17 @@ public static boolean shouldCreateFilesForMissingBuckets(Table table, ConnectorS

private MetastoreContext getMetastoreContext(ConnectorSession session)
{
return new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector(), session.getRuntimeStats());
return new MetastoreContext(
session.getIdentity(),
session.getQueryId(),
session.getClientInfo(),
session.getClientTags(),
session.getSource(),
getMetastoreHeaders(session),
isUserDefinedTypeEncodingEnabled(session),
metastore.getColumnConverterProvider(),
session.getWarningCollector(),
session.getRuntimeStats());
}

private Column columnHandleToColumn(ConnectorSession session, HiveColumnHandle handle)
Expand Down
2 changes: 1 addition & 1 deletion presto-hive/src/main/java/com/facebook/presto/hive/HivePageSinkProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -179,7 +179,7 @@ private ConnectorPageSink createPageSink(HiveWritableTableHandle handle, boolean
new HivePageSinkMetadataProvider(
handle.getPageSinkMetadata(),
memoizeMetastore(metastore, metastoreImpersonationEnabled, perTransactionMetastoreCacheMaximumSize, metastorePartitionCacheMaxColumnCount),
new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), columnConverterProvider, session.getWarningCollector(), session.getRuntimeStats())),
new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getClientTags(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), columnConverterProvider, session.getWarningCollector(), session.getRuntimeStats())),
typeManager,
hdfsEnvironment,
pageSorter,
Expand Down
8 changes: 4 additions & 4 deletions presto-hive/src/main/java/com/facebook/presto/hive/HivePartitionManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,7 @@ private Map<Column, Domain> createPartitionPredicates(
if (domains.isPresent()) {
Map<ColumnHandle, Domain> columnHandleDomainMap = domains.get();
ImmutableMap.Builder<Column, Domain> partitionPredicateBuilder = ImmutableMap.builder();
MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector(), session.getRuntimeStats());
MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getClientTags(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector(), session.getRuntimeStats());
for (HiveColumnHandle partitionColumn : partitionColumns) {
Column key = new Column(
partitionColumn.getName(),
Expand Down Expand Up @@ -422,7 +422,7 @@ private Optional<HivePartition> parseValuesAndFilterPartition(

private Table getTable(ConnectorSession session, SemiTransactionalHiveMetastore metastore, HiveTableHandle hiveTableHandle, boolean offlineDataDebugModeEnabled)
{
MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector(), session.getRuntimeStats());
MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getClientTags(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector(), session.getRuntimeStats());
Optional<Table> target = metastore.getTable(context, hiveTableHandle);
if (!target.isPresent()) {
throw new TableNotFoundException(hiveTableHandle.getSchemaTableName());
Expand All @@ -442,7 +442,7 @@ private List<PartitionNameWithVersion> getFilteredPartitionNames(ConnectorSessio
return ImmutableList.of();
}

MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector(), session.getRuntimeStats());
MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getClientTags(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector(), session.getRuntimeStats());
// fetch the partition names
return metastore.getPartitionNamesByFilter(context, hiveTableHandle, partitionPredicates)
.orElseThrow(() -> new TableNotFoundException(hiveTableHandle.getSchemaTableName()));
Expand All @@ -454,7 +454,7 @@ private List<PartitionNameWithVersion> getAllPartitionNames(ConnectorSession ses
return ImmutableList.of();
}
// fetch the partition names
MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector(), session.getRuntimeStats());
MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getClientTags(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector(), session.getRuntimeStats());
return metastore.getPartitionNames(context, hiveTableHandle)
.orElseThrow(() -> new TableNotFoundException(hiveTableHandle.getSchemaTableName()));
}
Expand Down
Loading

0 comments on commit 08eed5a

Please sign in to comment.