Skip to content

Commit

Permalink
Add WarningCollector to MetastoreContext
Browse files Browse the repository at this point in the history
  • Loading branch information
NikhilCollooru committed Jan 18, 2024
1 parent d11d79f commit 496e98c
Show file tree
Hide file tree
Showing 37 changed files with 171 additions and 109 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -370,7 +370,8 @@ private MetastoreContext metastoreContext(ConnectorSession session)
session.getSource(),
Optional.empty(),
false,
DEFAULT_COLUMN_CONVERTER_PROVIDER);
DEFAULT_COLUMN_CONVERTER_PROVIDER,
session.getWarningCollector());
}

private void checkConnectorId(DeltaTableHandle tableHandle)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1096,7 +1096,7 @@ public String toString()

private <T> KeyAndContext<T> getCachingKey(MetastoreContext context, T key)
{
MetastoreContext metastoreContext = metastoreImpersonationEnabled ? new MetastoreContext(context.getUsername(), context.getQueryId(), context.getClientInfo(), context.getSource(), true, context.getMetastoreHeaders(), context.isUserDefinedTypeEncodingEnabled(), context.getColumnConverterProvider()) : context;
MetastoreContext metastoreContext = metastoreImpersonationEnabled ? new MetastoreContext(context.getUsername(), context.getQueryId(), context.getClientInfo(), context.getSource(), true, context.getMetastoreHeaders(), context.isUserDefinedTypeEncodingEnabled(), context.getColumnConverterProvider(), context.getWarningCollector()) : context;
return new KeyAndContext<>(metastoreContext, key);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import com.facebook.presto.hive.ColumnConverter;
import com.facebook.presto.hive.ColumnConverterProvider;
import com.facebook.presto.hive.HiveColumnConverterProvider;
import com.facebook.presto.spi.WarningCollector;
import com.facebook.presto.spi.security.ConnectorIdentity;

import java.util.Objects;
Expand All @@ -38,18 +39,44 @@ public class MetastoreContext
// a new MetastoreContext from either an existing MetastoreContext or callers
// that only have a handle to the provider (e.g. SemiTransactionalHiveMetastore)
private final ColumnConverterProvider columnConverterProvider;

public MetastoreContext(ConnectorIdentity identity, String queryId, Optional<String> clientInfo, Optional<String> source, Optional<String> metastoreHeaders, boolean userDefinedTypeEncodingEnabled, ColumnConverterProvider columnConverterProvider)
private final WarningCollector warningCollector;

public MetastoreContext(
ConnectorIdentity identity,
String queryId,
Optional<String> clientInfo,
Optional<String> source,
Optional<String> metastoreHeaders,
boolean userDefinedTypeEncodingEnabled,
ColumnConverterProvider columnConverterProvider,
WarningCollector warningCollector)
{
this(requireNonNull(identity, "identity is null").getUser(), queryId, clientInfo, source, metastoreHeaders, userDefinedTypeEncodingEnabled, columnConverterProvider);
this(requireNonNull(identity, "identity is null").getUser(), queryId, clientInfo, source, metastoreHeaders, userDefinedTypeEncodingEnabled, columnConverterProvider, warningCollector);
}

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

public MetastoreContext(String username, String queryId, Optional<String> clientInfo, Optional<String> source, boolean impersonationEnabled, Optional<String> metastoreHeaders, boolean userDefinedTypeEncodingEnabled, ColumnConverterProvider columnConverterProvider)
public MetastoreContext(
String username,
String queryId,
Optional<String> clientInfo,
Optional<String> source,
boolean impersonationEnabled,
Optional<String> metastoreHeaders,
boolean userDefinedTypeEncodingEnabled,
ColumnConverterProvider columnConverterProvider,
WarningCollector warningCollector)
{
this.username = requireNonNull(username, "username is null");
this.queryId = requireNonNull(queryId, "queryId is null");
Expand All @@ -65,6 +92,7 @@ public MetastoreContext(String username, String queryId, Optional<String> client
else {
this.columnConverter = requireNonNull(HiveColumnConverterProvider.DEFAULT_COLUMN_CONVERTER, "columnConverter is null");
}
this.warningCollector = requireNonNull(warningCollector, "warningCollector is null");
}

public ColumnConverterProvider getColumnConverterProvider()
Expand Down Expand Up @@ -112,6 +140,11 @@ public Optional<String> getMetastoreHeaders()
return metastoreHeaders;
}

public WarningCollector getWarningCollector()
{
return warningCollector;
}

@Override
public String toString()
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@
import static com.facebook.presto.spi.StandardErrorCode.ALREADY_EXISTS;
import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
import static com.facebook.presto.spi.StandardErrorCode.TRANSACTION_CONFLICT;
import static com.facebook.presto.spi.WarningCollector.NOOP;
import static com.facebook.presto.spi.security.PrincipalType.USER;
import static com.google.common.base.MoreObjects.toStringHelper;
import static com.google.common.base.Preconditions.checkArgument;
Expand Down Expand Up @@ -512,7 +513,7 @@ 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);
MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), columnConverterProvider, session.getWarningCollector());
if (oldTableAction == null || oldTableAction.getData().getTable().getTableType().equals(TEMPORARY_TABLE)) {
Table table = getTable(metastoreContext, databaseName, tableName)
.orElseThrow(() -> new TableNotFoundException(schemaTableName));
Expand Down Expand Up @@ -558,7 +559,8 @@ public synchronized void truncateUnpartitionedTable(ConnectorSession session, St
session.getSource(),
getMetastoreHeaders(session),
isUserDefinedTypeEncodingEnabled(session),
columnConverterProvider),
columnConverterProvider,
session.getWarningCollector()),
databaseName,
tableName);
SchemaTableName schemaTableName = new SchemaTableName(databaseName, tableName);
Expand Down Expand Up @@ -902,7 +904,7 @@ 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);
MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), columnConverterProvider, session.getWarningCollector());

if (oldPartitionAction == null) {
Partition partition = delegate.getPartition(metastoreContext, databaseName, tableName, partitionValues)
Expand Down Expand Up @@ -1129,7 +1131,8 @@ private ConnectorCommitHandle commitShared()
hdfsContext.getSource(),
hdfsContext.getSession().flatMap(MetastoreUtil::getMetastoreHeaders),
hdfsContext.getSession().map(MetastoreUtil::isUserDefinedTypeEncodingEnabled).orElse(false),
columnConverterProvider);
columnConverterProvider,
hdfsContext.getSession().map(ConnectorSession::getWarningCollector).orElse(NOOP));
switch (action.getType()) {
case DROP:
committer.prepareDropTable(metastoreContext, schemaTableName);
Expand Down Expand Up @@ -1160,7 +1163,8 @@ private ConnectorCommitHandle commitShared()
hdfsContext.getSource(),
hdfsContext.getSession().flatMap(MetastoreUtil::getMetastoreHeaders),
hdfsContext.getSession().map(MetastoreUtil::isUserDefinedTypeEncodingEnabled).orElse(false),
columnConverterProvider);
columnConverterProvider,
hdfsContext.getSession().map(ConnectorSession::getWarningCollector).orElse(NOOP));
switch (action.getType()) {
case DROP:
committer.prepareDropPartition(metastoreContext, schemaTableName, partitionValues);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import com.facebook.presto.hive.metastore.Column;
import com.facebook.presto.hive.metastore.MetastoreContext;
import com.facebook.presto.hive.metastore.PartitionNameWithVersion;
import com.facebook.presto.spi.WarningCollector;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
Expand Down Expand Up @@ -70,7 +71,7 @@ public class MockHiveMetastoreClient
public static final List<SQLPrimaryKey> TEST_PRIMARY_KEY = ImmutableList.of(new SQLPrimaryKey(TEST_DATABASE, TEST_TABLE_WITH_CONSTRAINTS, "c1", 0, "", true, false, true));
public static final List<SQLUniqueConstraint> TEST_UNIQUE_CONSTRAINT = ImmutableList.of(new SQLUniqueConstraint("", TEST_DATABASE, TEST_TABLE_WITH_CONSTRAINTS, "c2", 1, "", true, false, 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);
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);
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
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ 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), schema, table, partitionStringValues).isPresent()) {
if (metastore.getPartition(new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), false, HiveColumnConverterProvider.DEFAULT_COLUMN_CONVERTER_PROVIDER, session.getWarningCollector()), schema, table, partitionStringValues).isPresent()) {
throw new PrestoException(ALREADY_EXISTS, "Partition already exists");
}
String partitionName = FileUtils.makePartName(actualPartitionColumnNames, partitionStringValues);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1753,7 +1753,7 @@ 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());
return new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector());
}

private Column columnHandleToColumn(ConnectorSession session, HiveColumnHandle handle)
Expand Down
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)),
new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), columnConverterProvider, session.getWarningCollector())),
typeManager,
hdfsEnvironment,
pageSorter,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,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());
MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector());
for (HiveColumnHandle partitionColumn : partitionColumns) {
Column key = new Column(
partitionColumn.getName(),
Expand Down Expand Up @@ -436,7 +436,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());
MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector());
Optional<Table> target = metastore.getTable(context, hiveTableHandle);
if (!target.isPresent()) {
throw new TableNotFoundException(hiveTableHandle.getSchemaTableName());
Expand All @@ -456,7 +456,7 @@ private List<String> getFilteredPartitionNames(ConnectorSession session, SemiTra
return ImmutableList.of();
}

MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider());
MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector());
// fetch the partition names
return metastore.getPartitionNamesByFilter(context, hiveTableHandle, partitionPredicates)
.orElseThrow(() -> new TableNotFoundException(hiveTableHandle.getSchemaTableName()));
Expand All @@ -468,7 +468,7 @@ private List<String> getAllPartitionNames(ConnectorSession session, SemiTransact
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());
MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector());
return metastore.getPartitionNames(context, hiveTableHandle)
.orElseThrow(() -> new TableNotFoundException(hiveTableHandle.getSchemaTableName()));
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,7 @@ public ConnectorSplitSource getSplits(
throw new PrestoException(HIVE_TRANSACTION_NOT_FOUND, format("Transaction not found: %s", transaction));
}
SemiTransactionalHiveMetastore metastore = metadata.getMetastore();
MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider());
MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector());
Table table = layout.getTable(metastore, metastoreContext);

if (!isOfflineDataDebugModeEnabled(session)) {
Expand Down Expand Up @@ -671,7 +671,7 @@ private Map<String, PartitionSplitInfo> getPartitionSplitInfo(
Map<String, HiveColumnHandle> predicateColumns,
Optional<Map<Subfield, Domain>> domains)
{
MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider());
MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), metastore.getColumnConverterProvider(), session.getWarningCollector());
Map<String, Optional<Partition>> partitions = metastore.getPartitionsByNames(
metastoreContext,
tableName.getSchemaName(),
Expand Down
Loading

0 comments on commit 496e98c

Please sign in to comment.