/*
* Licensed to CRATE Technology GmbH ("Crate") under one or more contributor
* license agreements. See the NOTICE file distributed with this work for
* additional information regarding copyright ownership. Crate licenses
* this file to you under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License. You may
* obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* However, if you have executed another commercial license agreement
* with Crate these terms will supersede the license and you may use the
* software solely pursuant to the terms of the relevant commercial agreement.
*/
package io.crate.analyze;
import com.google.common.base.Throwables;
import io.crate.metadata.ColumnIdent;
import io.crate.metadata.doc.DocSysColumns;
import org.apache.lucene.util.BytesRef;
import org.elasticsearch.common.UUIDs;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.io.stream.BytesStreamOutput;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.IOException;
import java.util.Base64;
import java.util.List;
import java.util.function.Function;
import static io.crate.collections.Lists2.getOnlyElement;
public class Id {
private final static Function<List<BytesRef>, String> RANDOM_ID = ignored -> UUIDs.base64UUID();
private final static Function<List<BytesRef>, String> ONLY_ITEM_NULL_VALIDATION = keyValues -> {
return ensureNonNull(getOnlyElement(keyValues)).utf8ToString();
};
private final static Function<List<BytesRef>, String> ONLY_ITEM = keyValues -> {
BytesRef element = getOnlyElement(keyValues);
if (element == null) {
return null;
}
return element.utf8ToString();
};
/**
* generates a function which can be used to generate an id and apply null validation.
* <p>
* This variant doesn't handle the pk = _id case.
*/
private static Function<List<BytesRef>, String> compileWithNullValidation(final int numPks,
final int clusteredByPosition) {
switch (numPks) {
case 0:
return RANDOM_ID;
case 1:
return ONLY_ITEM_NULL_VALIDATION;
default:
return keyValues -> {
if (keyValues.size() != numPks) {
throw new IllegalArgumentException("Missing primary key values");
}
return encode(keyValues, clusteredByPosition);
};
}
}
/**
* generates a function which can be used to generate an id.
* <p>
* This variant doesn't handle the pk = _id case.
*/
public static Function<List<BytesRef>, String> compile(final int numPks, final int clusteredByPosition) {
if (numPks == 1) {
return ONLY_ITEM;
}
return compileWithNullValidation(numPks, clusteredByPosition);
}
/**
* returns a function which can be used to generate an id with null validation.
*/
public static Function<List<BytesRef>, String> compileWithNullValidation(final List<ColumnIdent> pkColumns, final ColumnIdent clusteredBy) {
final int numPks = pkColumns.size();
if (numPks == 1 && getOnlyElement(pkColumns).equals(DocSysColumns.ID)) {
return RANDOM_ID;
}
return compileWithNullValidation(numPks, pkColumns.indexOf(clusteredBy));
}
@Nonnull
private static BytesRef ensureNonNull(@Nullable BytesRef pkValue) throws IllegalArgumentException {
if (pkValue == null) {
throw new IllegalArgumentException("A primary key value must not be NULL");
}
return pkValue;
}
private static String encode(List<BytesRef> values, int clusteredByPosition) {
try (BytesStreamOutput out = new BytesStreamOutput(estimateSize(values))) {
int size = values.size();
out.writeVInt(size);
if (clusteredByPosition >= 0) {
out.writeBytesRef(ensureNonNull(values.get(clusteredByPosition)));
}
for (int i = 0; i < size; i++) {
if (i != clusteredByPosition) {
out.writeBytesRef(ensureNonNull(values.get(i)));
}
}
return Base64.getEncoder().encodeToString(BytesReference.toBytes(out.bytes()));
} catch (IOException e) {
throw Throwables.propagate(e);
}
}
/**
* estimates the size the bytesRef values will take if written onto a StreamOutput using the String streamer
*/
private static int estimateSize(Iterable<BytesRef> values) {
int expectedEncodedSize = 0;
for (BytesRef value : values) {
// 5 bytes for the value of the length itself using vInt
expectedEncodedSize += 5 + (value != null ? value.length : 0);
}
return expectedEncodedSize;
}
}