/* * 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; } }