Skip to content

Commit 6a11fef

Browse files
hazefullyhazefully
committed
Make sure SQL user-defined types are preserved in RecordMetadata
1 parent eeb93c7 commit 6a11fef

21 files changed

Lines changed: 700 additions & 23 deletions

File tree

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaData.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,8 @@ public class RecordMetaData implements RecordMetaDataProvider {
8383
@Nonnull
8484
private final Map<Object, SyntheticRecordType<?>> recordTypeKeyToSyntheticTypeMap;
8585
@Nonnull
86+
private final Map<String, Descriptors.GenericDescriptor> auxiliaryTypeDescriptorsMap;
87+
@Nonnull
8688
private final Map<String, UserDefinedFunction> userDefinedFunctionMap;
8789
@Nonnull
8890
private final Map<String, View> viewMap;
@@ -118,6 +120,7 @@ protected RecordMetaData(@Nonnull RecordMetaData orig) {
118120
Collections.unmodifiableMap(orig.indexes),
119121
Collections.unmodifiableMap(orig.universalIndexes),
120122
Collections.unmodifiableList(orig.formerIndexes),
123+
Collections.unmodifiableMap(orig.auxiliaryTypeDescriptorsMap),
121124
Collections.unmodifiableMap(orig.userDefinedFunctionMap),
122125
Collections.unmodifiableMap(orig.viewMap),
123126
Collections.unmodifiableMap(orig.storedQueries),
@@ -140,6 +143,7 @@ protected RecordMetaData(@Nonnull Descriptors.FileDescriptor recordsDescriptor,
140143
@Nonnull Map<String, Index> indexes,
141144
@Nonnull Map<String, Index> universalIndexes,
142145
@Nonnull List<FormerIndex> formerIndexes,
146+
@Nonnull Map<String, Descriptors.GenericDescriptor> auxiliaryTypeDescriptorsMap,
143147
@Nonnull Map<String, UserDefinedFunction> userDefinedFunctionMap,
144148
@Nonnull Map<String, View> viewMap,
145149
@Nonnull Map<String, StoredQuery> storedQueries,
@@ -159,6 +163,7 @@ protected RecordMetaData(@Nonnull Descriptors.FileDescriptor recordsDescriptor,
159163
this.indexes = indexes;
160164
this.universalIndexes = universalIndexes;
161165
this.formerIndexes = formerIndexes;
166+
this.auxiliaryTypeDescriptorsMap = auxiliaryTypeDescriptorsMap;
162167
this.userDefinedFunctionMap = userDefinedFunctionMap;
163168
this.viewMap = viewMap;
164169
this.storedQueries = storedQueries;
@@ -748,6 +753,11 @@ public Map<String, StoredQuery> getStoredQueries() {
748753
return storedQueries;
749754
}
750755

756+
@Nonnull
757+
public Map<String, Descriptors.GenericDescriptor> getAuxiliaryTypeDescriptors() {
758+
return auxiliaryTypeDescriptorsMap;
759+
}
760+
751761
/**
752762
* A stored query is a SQL SELECT statement persisted in the metadata, optionally accompanied by an
753763
* list of DDL {@code CREATE [OR REPLACE]? TEMPORARY FUNCTION ...} that must be

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaDataBuilder.java‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ public class RecordMetaDataBuilder implements RecordMetaDataProvider {
9999

100100
private static final Descriptors.FileDescriptor[] emptyDependencyList = new Descriptors.FileDescriptor[0];
101101
public static final String DEFAULT_UNION_NAME = "RecordTypeUnion";
102+
public static final String DEFAULT_AUXILIARY_TYPE_UNION_NAME = "AuxiliaryTypeUnion";
102103

103104
@Nullable
104105
private Descriptors.FileDescriptor recordsDescriptor;
@@ -113,6 +114,8 @@ public class RecordMetaDataBuilder implements RecordMetaDataProvider {
113114
@Nonnull
114115
private final Map<String, SyntheticRecordTypeBuilder<?>> syntheticRecordTypes;
115116
@Nonnull
117+
private final Map<String, Descriptors.GenericDescriptor> auxiliaryTypesMap;
118+
@Nonnull
116119
private final Map<String, UserDefinedFunction> userDefinedFunctionMap;
117120
@Nonnull
118121
private final Map<String, View> viewMap;
@@ -156,6 +159,7 @@ public class RecordMetaDataBuilder implements RecordMetaDataProvider {
156159
userDefinedFunctionMap = new HashMap<>();
157160
viewMap = new HashMap<>();
158161
storedQueries = new HashMap<>();
162+
auxiliaryTypesMap = new HashMap<>();
159163
}
160164

161165
private void processSchemaOptions(boolean processExtensionOptions) {
@@ -316,6 +320,31 @@ private void initRecordTypesAndUnion(boolean processExtensionOptions) {
316320
unionDescriptor = fetchUnionDescriptor(recordsDescriptor);
317321
validateRecords(recordsDescriptor, unionDescriptor);
318322
fillUnionFields(processExtensionOptions);
323+
324+
fillAuxiliaryTypesMap(recordsDescriptor);
325+
}
326+
327+
private void fillAuxiliaryTypesMap(@Nonnull Descriptors.FileDescriptor fileDescriptor) {
328+
final var auxiliaryTypeUnionDescriptor = fileDescriptor
329+
.findMessageTypeByName(DEFAULT_AUXILIARY_TYPE_UNION_NAME);
330+
331+
if (auxiliaryTypeUnionDescriptor == null) {
332+
return;
333+
}
334+
335+
for (final var field : auxiliaryTypeUnionDescriptor.getFields()) {
336+
switch (field.getType()) {
337+
case MESSAGE:
338+
auxiliaryTypesMap.put(field.getName(), field.getMessageType());
339+
break;
340+
case ENUM:
341+
auxiliaryTypesMap.put(field.getName(), field.getEnumType());
342+
break;
343+
default:
344+
throw new MetaDataException("Unsupported user-defined auxiliary type",
345+
LogMessageKeys.FIELD_NAME, field.getName());
346+
}
347+
}
319348
}
320349

321350
@Nonnull
@@ -1475,7 +1504,7 @@ public RecordMetaData build(boolean validate) {
14751504
Map<Object, SyntheticRecordType<?>> recordTypeKeyToSyntheticRecordTypeMap = Maps.newHashMapWithExpectedSize(syntheticRecordTypes.size());
14761505
RecordMetaData metaData = new RecordMetaData(recordsDescriptor, getUnionDescriptor(), unionFields,
14771506
builtRecordTypes, builtSyntheticRecordTypes, recordTypeKeyToSyntheticRecordTypeMap,
1478-
indexes, universalIndexes, formerIndexes, userDefinedFunctionMap, viewMap, storedQueries,
1507+
indexes, universalIndexes, formerIndexes, auxiliaryTypesMap, userDefinedFunctionMap, viewMap, storedQueries,
14791508
splitLongRecords, storeRecordVersions, version, subspaceKeyCounter, usesSubspaceKeyCounter, recordCountKey, localFileDescriptor != null);
14801509
for (RecordTypeBuilder recordTypeBuilder : recordTypes.values()) {
14811510
KeyExpression primaryKey = recordTypeBuilder.getPrimaryKey();

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/plans/RecordQueryIndexPlan.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
import com.apple.foundationdb.record.provider.foundationdb.KeyValueCursorBase;
6161
import com.apple.foundationdb.record.provider.foundationdb.MultidimensionalIndexScanComparisons;
6262
import com.apple.foundationdb.record.provider.foundationdb.UnsupportedRemoteFetchIndexException;
63+
import com.apple.foundationdb.record.query.expressions.Comparisons;
6364
import com.apple.foundationdb.record.query.plan.AvailableFields;
6465
import com.apple.foundationdb.record.query.plan.QueryPlanConstraint;
6566
import com.apple.foundationdb.record.query.plan.ScanComparisons;
@@ -91,6 +92,7 @@
9192
import com.google.common.base.Verify;
9293
import com.google.common.collect.ImmutableList;
9394
import com.google.common.collect.ImmutableMap;
95+
import com.google.common.collect.ImmutableSet;
9496
import com.google.protobuf.ByteString;
9597
import com.google.protobuf.Message;
9698
import org.slf4j.Logger;
@@ -733,6 +735,21 @@ public QueryPlanConstraint getConstraint() {
733735
return constraint;
734736
}
735737

738+
@Nonnull
739+
@Override
740+
public Set<Type> getDynamicTypes() {
741+
final ImmutableSet.Builder<Type> resultBuilder = ImmutableSet.builder();
742+
resultBuilder.addAll(RecordQueryPlanWithNoChildren.super.getDynamicTypes());
743+
if (hasComparisons()) {
744+
getComparisons().stream()
745+
.map(Comparisons.Comparison::getValue)
746+
.filter(Objects::nonNull)
747+
.flatMap(value -> value.getDynamicTypes().stream())
748+
.forEach(resultBuilder::add);
749+
}
750+
return resultBuilder.build();
751+
}
752+
736753
@Nonnull
737754
@Override
738755
public Message toProto(@Nonnull final PlanSerializationContext serializationContext) {

‎fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/metadata/SchemaTemplate.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,16 @@ public interface SchemaTemplate extends Metadata {
140140
@Nonnull
141141
Map<String, StoredQuery> getStoredQueries() throws RelationalException;
142142

143+
/**
144+
* Retrieves a {@link DataType}, which can either be the underlying type of a {@link Table}
145+
* or a user-defined struct/enum type, by looking up its name.
146+
*
147+
* @param typeName The name of the type.
148+
* @return An {@link Optional} containing the {@link DataType} if it is found, otherwise {@code Empty}.
149+
*/
150+
@Nonnull
151+
Optional<DataType> findTypeByName(@Nonnull String typeName) throws RelationalException;
152+
143153
@Nonnull
144154
String getTransactionBoundMetadataAsString() throws RelationalException;
145155

‎fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/metadata/Visitor.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,4 +51,6 @@ default void visit(@Nonnull final Metadata metadata) {
5151
void visit(@Nonnull InvokedRoutine invokedRoutine);
5252

5353
void visit(@Nonnull View view);
54+
55+
void visit(@Nonnull DataType auxiliaryType);
5456
}

‎fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/NoOpSchemaTemplate.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424

2525
import com.apple.foundationdb.relational.api.exceptions.ErrorCode;
2626
import com.apple.foundationdb.relational.api.exceptions.RelationalException;
27+
import com.apple.foundationdb.relational.api.metadata.DataType;
2728
import com.apple.foundationdb.relational.api.metadata.InvokedRoutine;
2829
import com.apple.foundationdb.relational.api.metadata.Schema;
2930
import com.apple.foundationdb.relational.api.metadata.SchemaTemplate;
@@ -147,6 +148,12 @@ public Map<String, StoredQuery> getStoredQueries() throws RelationalException {
147148
throw new RelationalException("NoOpSchemaTemplate doesn't have stored queries!", ErrorCode.INVALID_PARAMETER);
148149
}
149150

151+
@Nonnull
152+
@Override
153+
public Optional<DataType> findTypeByName(@Nonnull final String typeName) throws RelationalException {
154+
throw new RelationalException("NoOpSchemaTemplate doesn't have types!", ErrorCode.INVALID_PARAMETER);
155+
}
156+
150157
@Nonnull
151158
@Override
152159
public String getTransactionBoundMetadataAsString() throws RelationalException {

‎fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/RecordLayerSchemaTemplate.java‎

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,9 @@ public final class RecordLayerSchemaTemplate implements SchemaTemplate {
7575
@Nonnull
7676
private final Set<RecordLayerTable> tables;
7777

78+
@Nonnull
79+
private final Map<String, DataType.Named> auxiliaryTypeMap;
80+
7881
@Nonnull
7982
private final Set<RecordLayerInvokedRoutine> invokedRoutines;
8083

@@ -109,6 +112,7 @@ public final class RecordLayerSchemaTemplate implements SchemaTemplate {
109112

110113
private RecordLayerSchemaTemplate(@Nonnull final String name,
111114
@Nonnull final Set<RecordLayerTable> tables,
115+
@Nonnull final Map<String, DataType.Named> auxiliaryTypeMap,
112116
@Nonnull final Set<RecordLayerInvokedRoutine> invokedRoutines,
113117
@Nonnull final Set<RecordLayerView> views,
114118
@Nonnull final Map<String, StoredQuery> storedQueries,
@@ -118,6 +122,7 @@ private RecordLayerSchemaTemplate(@Nonnull final String name,
118122
boolean intermingleTables) {
119123
this.name = name;
120124
this.tables = ImmutableSet.copyOf(tables);
125+
this.auxiliaryTypeMap = ImmutableMap.copyOf(auxiliaryTypeMap);
121126
this.invokedRoutines = ImmutableSet.copyOf(invokedRoutines);
122127
this.views = ImmutableSet.copyOf(views);
123128
this.storedQueries = ImmutableMap.copyOf(storedQueries);
@@ -134,6 +139,7 @@ private RecordLayerSchemaTemplate(@Nonnull final String name,
134139

135140
private RecordLayerSchemaTemplate(@Nonnull final String name,
136141
@Nonnull final Set<RecordLayerTable> tables,
142+
@Nonnull final Map<String, DataType.Named> auxiliaryTypeMap,
137143
@Nonnull final Set<RecordLayerInvokedRoutine> invokedRoutines,
138144
@Nonnull final Set<RecordLayerView> views,
139145
@Nonnull final Map<String, StoredQuery> storedQueries,
@@ -145,6 +151,7 @@ private RecordLayerSchemaTemplate(@Nonnull final String name,
145151
this.name = name;
146152
this.version = version;
147153
this.tables = ImmutableSet.copyOf(tables);
154+
this.auxiliaryTypeMap = ImmutableMap.copyOf(auxiliaryTypeMap);
148155
this.invokedRoutines = ImmutableSet.copyOf(invokedRoutines);
149156
this.views = ImmutableSet.copyOf(views);
150157
this.storedQueries = ImmutableMap.copyOf(storedQueries);
@@ -328,6 +335,11 @@ public BitSet getIndexEntriesAsBitset(@Nonnull final Optional<Set<String>> index
328335
return result;
329336
}
330337

338+
@Nonnull
339+
public Map<String, DataType.Named> getAuxiliaryTypes() {
340+
return auxiliaryTypeMap;
341+
}
342+
331343
@Nonnull
332344
@Override
333345
public Set<RecordLayerInvokedRoutine> getInvokedRoutines() {
@@ -352,6 +364,14 @@ public Map<String, StoredQuery> getStoredQueries() {
352364
return storedQueries;
353365
}
354366

367+
@Nonnull
368+
@Override
369+
public Optional<DataType> findTypeByName(@Nonnull final String typeName) {
370+
return findTableByName(typeName)
371+
.<DataType>map(Table::getDatatype)
372+
.or(() -> Optional.ofNullable((DataType)auxiliaryTypeMap.get(typeName)));
373+
}
374+
355375
@Nonnull
356376
@Override
357377
public Optional<? extends View> findViewByName(@Nonnull final String viewName) {
@@ -395,6 +415,9 @@ public void accept(@Nonnull final Visitor visitor) {
395415
for (final var table : getTables()) {
396416
table.accept(visitor);
397417
}
418+
for (final var auxiliaryType : auxiliaryTypeMap.values()) {
419+
visitor.visit((DataType)auxiliaryType);
420+
}
398421
for (final var invokedRoutine : getInvokedRoutines()) {
399422
invokedRoutine.accept(visitor);
400423
}
@@ -420,7 +443,7 @@ public static final class Builder {
420443
private final Map<String, RecordLayerTable> tables;
421444

422445
@Nonnull
423-
private final Map<String, DataType.Named> auxiliaryTypes; // for quick lookup
446+
private final Map<String, DataType.Named> auxiliaryTypes;
424447

425448
@Nonnull
426449
private final Map<String, RecordLayerInvokedRoutine> invokedRoutines;
@@ -662,10 +685,10 @@ public RecordLayerSchemaTemplate build() {
662685
}
663686

664687
if (cachedMetadata != null) {
665-
return new RecordLayerSchemaTemplate(name, new LinkedHashSet<>(tables.values()),
688+
return new RecordLayerSchemaTemplate(name, new LinkedHashSet<>(tables.values()), auxiliaryTypes,
666689
new LinkedHashSet<>(invokedRoutines.values()), new LinkedHashSet<>(views.values()), storedQueries, version, enableLongRows, storeRowVersions, intermingleTables, cachedMetadata);
667690
} else {
668-
return new RecordLayerSchemaTemplate(name, new LinkedHashSet<>(tables.values()),
691+
return new RecordLayerSchemaTemplate(name, new LinkedHashSet<>(tables.values()), auxiliaryTypes,
669692
new LinkedHashSet<>(invokedRoutines.values()), new LinkedHashSet<>(views.values()), storedQueries, version, enableLongRows, storeRowVersions, intermingleTables);
670693
}
671694
}
@@ -793,6 +816,7 @@ public Builder toBuilder() {
793816
.setEnableLongRows(enableLongRows)
794817
.setIntermingleTables(intermingleTables)
795818
.addTables(getTables())
819+
.addAuxiliaryTypes(auxiliaryTypeMap.values())
796820
.addInvokedRoutines(getInvokedRoutines())
797821
.addViews(getViews())
798822
.addStoredQueries(getStoredQueries());

‎fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/SkeletonVisitor.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.apple.foundationdb.annotation.API;
2424

2525
import com.apple.foundationdb.relational.api.metadata.Column;
26+
import com.apple.foundationdb.relational.api.metadata.DataType;
2627
import com.apple.foundationdb.relational.api.metadata.Index;
2728
import com.apple.foundationdb.relational.api.metadata.InvokedRoutine;
2829
import com.apple.foundationdb.relational.api.metadata.Schema;
@@ -83,4 +84,9 @@ public void visit(@Nonnull final InvokedRoutine invokedRoutine) {
8384
public void visit(@Nonnull final View view) {
8485
// no-op
8586
}
87+
88+
@Override
89+
public void visit(@Nonnull final DataType auxiliaryType) {
90+
// no-op
91+
}
8692
}

0 commit comments

Comments
 (0)