furyhawk/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-07-21 10:15:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: validate external Arrow types before field decode (#49400)

Browse Source 
## Summary

External table loading previously trusted schemaless Arrow arrays after
minimal
layout normalization. Several paths then passed those arrays to
FieldData or
vector conversion code that expects Milvus-native physical types.
Mismatches
could either crash during load or silently reinterpret buffers as a
different
Milvus type.

This PR adds semantic validation before external Arrow arrays reach
those decode
paths:

- reject scalar Arrow types that do not match the Milvus field type
- keep the existing safe integer narrowing for narrow int fields
- reject invalid ARRAY element types before protobuf serialization
- reject null ARRAY elements while preserving nullable ARRAY rows
- validate dense vector byte width, element type and dimension
- reject unsupported list input for binary and bfloat16 vectors
- reject null VECTOR_ARRAY rows and invalid inner vector layouts
- normalize mixed vector chunks independently instead of dispatching the
whole
  batch from the first chunk type

It also preserves external spec extension properties while redacting
sensitive
values and keeps the generated external virtual primary key description
implementation-agnostic.

## Why

This prevents external table data from being accepted when its physical
Arrow
type does not match the collection schema. In particular, it prevents
vector
dimension corruption, scalar buffer reinterpretation, invalid ARRAY
decoding,
and load-time crashes for mismatched external columns.

Preserving extension properties avoids dropping external table metadata
when
redacted specs are propagated through proxy paths.

## Issues

- https://github.com/milvus-io/milvus/issues/49388
- https://github.com/milvus-io/milvus/issues/49392

## Test plan

- `ninja -C cmake_build -j1 unittest/all_tests`
- `DYLD_LIBRARY_PATH=cmake_build/src ./unittest/all_tests
--gtest_filter='NormalizeVectorArrays.*:NormalizeVectorArraysToFixedSizeBinary.*:NormalizeExternalArrow.*'`
- `go test ./util/externalspec` from `pkg` module was attempted locally,
but the local Go toolchain/cache is inconsistent: cached packages were
built with `go1.24.12` while the active tool is `go1.26.1`.

---------

Signed-off-by: Wei Liu <wei.liu@zilliz.com>
This commit is contained in:
wei liu
2026-04-28 13:39:49 +08:00
committed by GitHub
parent 9fc3a1e998
commit 54c25d00a1
8 changed files with 1039 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/core/src/storage/Util.cpp
+285 -36
View File
@@ -1693,6 +1693,103 @@ GetFieldIDList(FieldId column_group_id,
return field_id_list;
}
void
ValidateFixedSizeBinaryVectorWidth(const std::shared_ptr<arrow::Array>& array,
DataType data_type,
int dim) {
auto fsb_array =
std::static_pointer_cast<arrow::FixedSizeBinaryArray>(array);
int byte_width = GetDataTypeSize(data_type, dim);
AssertInfo(fsb_array->byte_width() == byte_width,
"vector byte width mismatch, expected {} bytes for "
"dim {}, actual {} bytes",
byte_width,
dim,
fsb_array->byte_width());
}
void
ValidateBinaryVectorWidth(const std::shared_ptr<arrow::Array>& array,
DataType data_type,
int dim) {
auto binary_array = std::static_pointer_cast<arrow::BinaryArray>(array);
int byte_width = GetDataTypeSize(data_type, dim);
for (int64_t i = 0; i < binary_array->length(); ++i) {
if (binary_array->IsNull(i)) {
continue;
}
auto actual_width = binary_array->value_length(i);
AssertInfo(actual_width == byte_width,
"vector byte width mismatch, expected {} bytes for "
"dim {}, actual {} bytes at row {}",
byte_width,
dim,
actual_width,
i);
}
}
void
ValidateNoNullValuesInRange(const std::shared_ptr<arrow::Array>& values,
int64_t begin,
int64_t end,
const char* context) {
if (values->null_count() == 0) {
return;
}
for (int64_t i = begin; i < end; ++i) {
AssertInfo(values->IsValid(i),
"{} contains null element at child offset {}",
context,
i);
}
}
arrow::Type::type
ExpectedVectorListElementArrowType(DataType data_type) {
switch (data_type) {
case DataType::VECTOR_FLOAT:
return arrow::Type::FLOAT;
case DataType::VECTOR_INT8:
return arrow::Type::INT8;
case DataType::VECTOR_FLOAT16:
return arrow::Type::HALF_FLOAT;
case DataType::VECTOR_BINARY:
case DataType::VECTOR_BFLOAT16:
ThrowInfo(ErrorCode::Unsupported,
"vector list input is not supported for {}",
data_type);
default:
ThrowInfo(ErrorCode::Unsupported,
"unsupported vector list input for {}",
data_type);
}
}
const char*
ArrowTypeName(arrow::Type::type type) {
switch (type) {
case arrow::Type::FLOAT:
return "float";
case arrow::Type::INT8:
return "int8";
case arrow::Type::HALF_FLOAT:
return "halffloat";
default:
return "unsupported";
}
}
void
ValidateVectorListElementType(
const std::shared_ptr<arrow::DataType>& actual_type, DataType data_type) {
auto expected_type = ExpectedVectorListElementArrowType(data_type);
AssertInfo(actual_type->id() == expected_type,
"vector element type mismatch, expected {}, actual {}",
ArrowTypeName(expected_type),
actual_type->ToString());
}
arrow::ArrayVector
NormalizeVectorArraysToFixedSizeBinary(const arrow::ArrayVector& arrays,
DataType data_type,
@@ -1706,6 +1803,7 @@ NormalizeVectorArraysToFixedSizeBinary(const arrow::ArrayVector& arrays,
for (const auto& array : arrays) {
auto type_id = array->type_id();
if (type_id == arrow::Type::FIXED_SIZE_BINARY) {
ValidateFixedSizeBinaryVectorWidth(array, data_type, dim);
result.push_back(array);
continue;
}
@@ -1724,6 +1822,7 @@ NormalizeVectorArraysToFixedSizeBinary(const arrow::ArrayVector& arrays,
if (type_id == arrow::Type::LIST) {
auto list_array = std::static_pointer_cast<arrow::ListArray>(array);
auto values = list_array->values();
ValidateVectorListElementType(values->type(), data_type);
int elem_bit_width = values->type()->bit_width();
AssertInfo(elem_bit_width > 0 && elem_bit_width % 8 == 0,
"Vector list element must be fixed-width byte-aligned "
@@ -1735,6 +1834,15 @@ NormalizeVectorArraysToFixedSizeBinary(const arrow::ArrayVector& arrays,
for (int64_t i = 0; i < num_rows; i++) {
if (array->IsValid(i)) {
auto offset = list_array->value_offset(i);
auto actual_dim = list_array->value_offset(i + 1) - offset;
AssertInfo(actual_dim == dim,
"vector dimension mismatch, expected {}, "
"actual {} at row {}",
dim,
actual_dim,
i);
ValidateNoNullValuesInRange(
values, offset, offset + actual_dim, "vector list");
memcpy(dst + i * byte_width,
raw + offset * elem_byte_size,
byte_width);
@@ -1744,17 +1852,24 @@ NormalizeVectorArraysToFixedSizeBinary(const arrow::ArrayVector& arrays,
auto fsl_array =
std::static_pointer_cast<arrow::FixedSizeListArray>(array);
auto values = fsl_array->values();
ValidateVectorListElementType(values->type(), data_type);
int elem_bit_width = values->type()->bit_width();
AssertInfo(elem_bit_width > 0 && elem_bit_width % 8 == 0,
"Vector list element must be fixed-width byte-aligned "
"type, got bit_width={}",
elem_bit_width);
int elem_byte_size = elem_bit_width / 8;
AssertInfo(fsl_array->value_length() == dim,
"vector dimension mismatch, expected {}, actual {}",
dim,
fsl_array->value_length());
auto raw = reinterpret_cast<const uint8_t*>(
values->data()->buffers[1]->data());
for (int64_t i = 0; i < num_rows; i++) {
if (array->IsValid(i)) {
auto offset = fsl_array->value_offset(i);
ValidateNoNullValuesInRange(
values, offset, offset + dim, "vector list");
memcpy(dst + i * byte_width,
raw + offset * elem_byte_size,
byte_width);
@@ -2377,8 +2492,45 @@ ConvertTimestampToInt64(const arrow::ArrayVector& arrays) {
return result;
}
DataType
ArrowListElementTypeToMilvus(const std::shared_ptr<arrow::Array>& values) {
switch (values->type_id()) {
case arrow::Type::BOOL:
return DataType::BOOL;
case arrow::Type::INT8:
return DataType::INT8;
case arrow::Type::INT16:
return DataType::INT16;
case arrow::Type::INT32:
return DataType::INT32;
case arrow::Type::INT64:
return DataType::INT64;
case arrow::Type::FLOAT:
return DataType::FLOAT;
case arrow::Type::DOUBLE:
return DataType::DOUBLE;
case arrow::Type::STRING:
case arrow::Type::LARGE_STRING:
case arrow::Type::STRING_VIEW:
return DataType::STRING;
default:
ThrowInfo(ErrorCode::Unsupported,
"unsupported array element arrow type: {}",
values->type()->ToString());
}
}
bool
IsCompatibleArrayElementType(DataType actual_type, DataType expected_type) {
if (actual_type == expected_type) {
return true;
}
return actual_type == DataType::STRING && IsStringDataType(expected_type);
}
arrow::ArrayVector
ConvertListToProtobufBinary(const arrow::ArrayVector& arrays) {
ConvertListToProtobufBinary(const arrow::ArrayVector& arrays,
DataType element_type) {
arrow::ArrayVector result;
result.reserve(arrays.size());
for (const auto& arr : arrays) {
@@ -2387,6 +2539,13 @@ ConvertListToProtobufBinary(const arrow::ArrayVector& arrays) {
continue;
}
auto list_arr = std::static_pointer_cast<arrow::ListArray>(arr);
auto actual_element_type =
ArrowListElementTypeToMilvus(list_arr->values());
AssertInfo(
IsCompatibleArrayElementType(actual_element_type, element_type),
"array element type mismatch, expected {}, actual {}",
element_type,
actual_element_type);
arrow::BinaryBuilder builder;
auto status = builder.Reserve(list_arr->length());
AssertInfo(status.ok(), "BinaryBuilder reserve failed");
@@ -2394,6 +2553,10 @@ ConvertListToProtobufBinary(const arrow::ArrayVector& arrays) {
if (list_arr->IsNull(i)) {
status = builder.AppendNull();
} else {
auto start = list_arr->value_offset(i);
auto end = list_arr->value_offset(i + 1);
ValidateNoNullValuesInRange(
list_arr->values(), start, end, "array list");
auto proto = ArrowListToScalarFieldProto(list_arr, i);
std::string serialized;
proto.SerializeToString(&serialized);
@@ -2417,30 +2580,35 @@ NormalizeVectorArrays(const arrow::ArrayVector& arrays,
if (arrays.empty()) {
return arrays;
}
auto type_id = arrays[0]->type_id();
// Already the target type
if (type_id == arrow::Type::FIXED_SIZE_BINARY && !nullable) {
return arrays;
}
if (type_id == arrow::Type::BINARY && nullable) {
return arrays;
}
// Step 1: List/FixedSizeList → FixedSizeBinary
arrow::ArrayVector fsb;
if (type_id == arrow::Type::FIXED_SIZE_BINARY) {
fsb =
arrays; // already FixedSizeBinary (but needs Binary for nullable)
} else {
fsb = NormalizeVectorArraysToFixedSizeBinary(
arrays, data_type, static_cast<int>(dim));
}
arrow::ArrayVector result;
result.reserve(arrays.size());
for (const auto& array : arrays) {
auto type_id = array->type_id();
if (type_id == arrow::Type::BINARY && nullable) {
ValidateBinaryVectorWidth(array, data_type, static_cast<int>(dim));
result.push_back(array);
continue;
}
// Step 2: nullable → FixedSizeBinary → BinaryArray
if (nullable) {
return ConvertFixedSizeBinaryToBinary(fsb);
arrow::ArrayVector fsb;
if (type_id == arrow::Type::FIXED_SIZE_BINARY) {
ValidateFixedSizeBinaryVectorWidth(
array, data_type, static_cast<int>(dim));
fsb.push_back(array);
} else {
fsb = NormalizeVectorArraysToFixedSizeBinary(
{array}, data_type, static_cast<int>(dim));
}
if (nullable) {
auto binary = ConvertFixedSizeBinaryToBinary(fsb);
result.push_back(binary[0]);
} else {
result.push_back(fsb[0]);
}
}
return fsb;
return result;
}
arrow::ArrayVector
@@ -2455,7 +2623,11 @@ NormalizeVectorArrayInner(const arrow::ArrayVector& arrays,
continue;
}
auto list_arr = std::static_pointer_cast<arrow::ListArray>(arr);
AssertInfo(list_arr->null_count() == 0,
"VECTOR_ARRAY does not support null rows");
if (list_arr->values()->type_id() == arrow::Type::FIXED_SIZE_BINARY) {
ValidateFixedSizeBinaryVectorWidth(
list_arr->values(), element_type, static_cast<int>(dim));
result.push_back(arr);
continue;
}
@@ -2577,6 +2749,52 @@ MaybeNarrowInt(DataType data_type, const std::shared_ptr<arrow::Array>& array) {
return nullptr;
}
std::shared_ptr<arrow::DataType>
ExpectedScalarArrowType(DataType data_type) {
switch (data_type) {
case DataType::BOOL:
return arrow::boolean();
case DataType::INT8:
return arrow::int8();
case DataType::INT16:
return arrow::int16();
case DataType::INT32:
return arrow::int32();
case DataType::INT64:
return arrow::int64();
case DataType::FLOAT:
return arrow::float32();
case DataType::DOUBLE:
return arrow::float64();
default:
return nullptr;
}
}
void
ValidateScalarArrowType(DataType data_type,
const std::shared_ptr<arrow::Array>& array) {
auto expected_type = ExpectedScalarArrowType(data_type);
if (expected_type == nullptr) {
return;
}
AssertInfo(array->type()->Equals(*expected_type),
"field type mismatch, expected Arrow {}, actual Arrow {}",
expected_type->ToString(),
array->type()->ToString());
}
void
AssertExternalArrowType(DataType data_type,
const std::shared_ptr<arrow::Array>& array,
const std::string& expected) {
ThrowInfo(ErrorCode::Unsupported,
"field type mismatch for {}, expected Arrow {}, actual Arrow {}",
data_type,
expected,
array->type()->ToString());
}
// Overload for callers without FieldMeta.
std::shared_ptr<arrow::Array>
NormalizeExternalArrow(const std::shared_ptr<arrow::Array>& array_in,
@@ -2598,24 +2816,51 @@ NormalizeExternalArrow(const std::shared_ptr<arrow::Array>& array_in,
}
auto type_id = array->type_id();
if (data_type == DataType::TIMESTAMPTZ) {
if (type_id == arrow::Type::TIMESTAMP) {
auto result = ConvertTimestampToInt64({array});
return result[0];
}
if (type_id == arrow::Type::INT64) {
return array;
}
AssertExternalArrowType(data_type, array, "timestamp or int64");
}
ValidateScalarArrowType(data_type, array);
if (IsSparseFloatVectorDataType(data_type)) {
if (type_id == arrow::Type::BINARY) {
return array;
}
AssertExternalArrowType(data_type, array, "binary");
}
// Dense vectors
if (IsVectorDataType(data_type) &&
!IsSparseFloatVectorDataType(data_type) &&
!IsVectorArrayDataType(data_type)) {
if (type_id == arrow::Type::FIXED_SIZE_BINARY && !nullable) {
ValidateFixedSizeBinaryVectorWidth(
array, data_type, static_cast<int>(dim));
return array;
}
if (type_id == arrow::Type::BINARY && nullable) {
ValidateBinaryVectorWidth(array, data_type, static_cast<int>(dim));
return array;
}
auto result = NormalizeVectorArrays({array}, data_type, dim, nullable);
return result[0];
}
// VectorArray inner: List<List<scalar>> → List<FixedSizeBinary>
if (IsVectorArrayDataType(data_type) && type_id == arrow::Type::LIST &&
element_type != DataType::NONE) {
auto result = NormalizeVectorArrayInner({array}, element_type, dim);
return result[0];
if (IsVectorArrayDataType(data_type)) {
AssertInfo(element_type != DataType::NONE,
"element_type must be specified for VECTOR_ARRAY");
if (type_id == arrow::Type::LIST) {
auto result = NormalizeVectorArrayInner({array}, element_type, dim);
return result[0];
}
AssertExternalArrowType(data_type, array, "list");
}
// Geometry STRING input -> WKT, convert to WKB.
// (View/large-string variants already canonicalized to STRING above.)
@@ -2627,31 +2872,35 @@ NormalizeExternalArrow(const std::shared_ptr<arrow::Array>& array_in,
if (data_type == DataType::GEOMETRY && type_id == arrow::Type::BINARY) {
return array;
}
if (data_type == DataType::GEOMETRY) {
AssertExternalArrowType(data_type, array, "string or binary");
}
// JSON/VARCHAR/STRING/TEXT: String -> Binary (canonical internal repr).
// (View/large variants already canonicalized to STRING above; issue #49352.)
if ((data_type == DataType::VARCHAR || data_type == DataType::STRING ||
data_type == DataType::TEXT || data_type == DataType::JSON) &&
if ((IsStringDataType(data_type) || data_type == DataType::JSON) &&
type_id == arrow::Type::STRING) {
auto result = ConvertStringArrayToBinary({array});
return result[0];
}
// JSON/VARCHAR/STRING/TEXT: BINARY input is already canonical.
if ((data_type == DataType::VARCHAR || data_type == DataType::STRING ||
data_type == DataType::TEXT || data_type == DataType::JSON) &&
if ((IsStringDataType(data_type) || data_type == DataType::JSON) &&
type_id == arrow::Type::BINARY) {
return array;
}
// Timestamptz: Timestamp → Int64
if (data_type == DataType::TIMESTAMPTZ &&
type_id == arrow::Type::TIMESTAMP) {
auto result = ConvertTimestampToInt64({array});
return result[0];
if (IsStringDataType(data_type) || data_type == DataType::JSON) {
AssertExternalArrowType(data_type, array, "string or binary");
}
// Array: List → Protobuf Binary
if (data_type == DataType::ARRAY && type_id == arrow::Type::LIST) {
auto result = ConvertListToProtobufBinary({array});
auto result = ConvertListToProtobufBinary({array}, element_type);
return result[0];
}
if (data_type == DataType::ARRAY && type_id == arrow::Type::BINARY) {
return array;
}
if (data_type == DataType::ARRAY) {
AssertExternalArrowType(data_type, array, "list or binary");
}
return array;
}
internal/core/src/storage/Util.h
+2 -1
View File
@@ -483,7 +483,8 @@ ConvertTimestampToInt64(const arrow::ArrayVector& arrays);
// Convert ListArray<Scalar> → BinaryArray (protobuf-serialized ScalarFieldProto).
arrow::ArrayVector
ConvertListToProtobufBinary(const arrow::ArrayVector& arrays);
ConvertListToProtobufBinary(const arrow::ArrayVector& arrays,
DataType element_type);
// Unified vector normalization: List/FSList → final format.
// Non-nullable → FixedSizeBinaryArray; nullable → BinaryArray.
internal/core/unittest/test_arrow_canonicalize.cpp
+523
View File
@@ -291,6 +291,38 @@ MakeInt32Array(const std::vector<int32_t>& vals,
EXPECT_TRUE(b.Finish(&out).ok());
return out;
}
std::shared_ptr<arrow::Array>
MakeInt64Array(const std::vector<int64_t>& vals,
const std::vector<bool>& valid) {
arrow::Int64Builder b;
for (size_t i = 0; i < vals.size(); ++i) {
if (!valid[i]) {
EXPECT_TRUE(b.AppendNull().ok());
} else {
EXPECT_TRUE(b.Append(vals[i]).ok());
}
}
std::shared_ptr<arrow::Array> out;
EXPECT_TRUE(b.Finish(&out).ok());
return out;
}
std::shared_ptr<arrow::Array>
MakeDoubleArray(const std::vector<double>& vals,
const std::vector<bool>& valid) {
arrow::DoubleBuilder b;
for (size_t i = 0; i < vals.size(); ++i) {
if (!valid[i]) {
EXPECT_TRUE(b.AppendNull().ok());
} else {
EXPECT_TRUE(b.Append(vals[i]).ok());
}
}
std::shared_ptr<arrow::Array> out;
EXPECT_TRUE(b.Finish(&out).ok());
return out;
}
} // namespace
TEST(IntegerNarrowing, Int32ToInt8) {
@@ -345,6 +377,497 @@ TEST(IntegerNarrowing, NoNarrowOnExactMatch) {
EXPECT_EQ(out.get(), in.get());
}
TEST(NormalizeExternalArrow, Int64RejectsString) {
arrow::StringBuilder builder;
ASSERT_TRUE(builder.AppendValues({"1", "2"}).ok());
std::shared_ptr<arrow::Array> input;
ASSERT_TRUE(builder.Finish(&input).ok());
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(
input, milvus::DataType::INT64, 0, false, milvus::DataType::NONE),
std::exception);
}
TEST(NormalizeExternalArrow, Issue49392ScalarMismatchesReject) {
arrow::StringBuilder string_builder;
ASSERT_TRUE(string_builder.AppendValues({"abcd", "efgh"}).ok());
std::shared_ptr<arrow::Array> string_input;
ASSERT_TRUE(string_builder.Finish(&string_input).ok());
for (auto data_type : {milvus::DataType::BOOL,
milvus::DataType::INT8,
milvus::DataType::INT16,
milvus::DataType::INT32,
milvus::DataType::INT64,
milvus::DataType::FLOAT,
milvus::DataType::DOUBLE,
milvus::DataType::TIMESTAMPTZ}) {
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(
string_input, data_type, 0, false, milvus::DataType::NONE),
std::exception);
}
auto int64_input = MakeInt64Array({1, 2}, {true, true});
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(int64_input,
milvus::DataType::FLOAT,
0,
false,
milvus::DataType::NONE),
std::exception);
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(int64_input,
milvus::DataType::DOUBLE,
0,
false,
milvus::DataType::NONE),
std::exception);
auto double_input = MakeDoubleArray({1.0, 2.0}, {true, true});
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(double_input,
milvus::DataType::INT64,
0,
false,
milvus::DataType::NONE),
std::exception);
}
TEST(NormalizeExternalArrow, TimestamptzAcceptsTimestampAndInt64) {
arrow::TimestampBuilder builder(arrow::timestamp(arrow::TimeUnit::MICRO),
arrow::default_memory_pool());
ASSERT_TRUE(builder.AppendValues({1000, 2000}).ok());
std::shared_ptr<arrow::Array> timestamp_input;
ASSERT_TRUE(builder.Finish(&timestamp_input).ok());
auto timestamp_out =
milvus::storage::NormalizeExternalArrow(timestamp_input,
milvus::DataType::TIMESTAMPTZ,
0,
false,
milvus::DataType::NONE);
ASSERT_EQ(timestamp_out->type_id(), arrow::Type::INT64);
auto int64_input = MakeInt64Array({1000, 2000}, {true, true});
auto int64_out =
milvus::storage::NormalizeExternalArrow(int64_input,
milvus::DataType::TIMESTAMPTZ,
0,
false,
milvus::DataType::NONE);
EXPECT_EQ(int64_out.get(), int64_input.get());
}
TEST(NormalizeExternalArrow, StringLikeFieldsRejectInt64) {
auto int64_input = MakeInt64Array({1, 2}, {true, true});
for (auto data_type : {milvus::DataType::VARCHAR,
milvus::DataType::STRING,
milvus::DataType::TEXT,
milvus::DataType::JSON,
milvus::DataType::GEOMETRY}) {
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(
int64_input, data_type, 0, false, milvus::DataType::NONE),
std::exception);
}
}
TEST(NormalizeVectorArraysToFixedSizeBinary, ListDimMismatchAsserts) {
arrow::FloatBuilder values_builder;
ASSERT_TRUE(values_builder.AppendValues({1.0f, 2.0f, 3.0f}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
arrow::Int32Builder offsets_builder;
ASSERT_TRUE(offsets_builder.AppendValues({0, 3}).ok());
std::shared_ptr<arrow::Array> offsets;
ASSERT_TRUE(offsets_builder.Finish(&offsets).ok());
auto input = *arrow::ListArray::FromArrays(*offsets, *values);
EXPECT_THROW(milvus::storage::NormalizeVectorArraysToFixedSizeBinary(
{input}, milvus::DataType::VECTOR_FLOAT, 2),
std::exception);
}
TEST(NormalizeVectorArraysToFixedSizeBinary, ListElementTypeMismatchAsserts) {
arrow::Int64Builder values_builder;
ASSERT_TRUE(values_builder.AppendValues({1, 2}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
arrow::Int32Builder offsets_builder;
ASSERT_TRUE(offsets_builder.AppendValues({0, 2}).ok());
std::shared_ptr<arrow::Array> offsets;
ASSERT_TRUE(offsets_builder.Finish(&offsets).ok());
auto input = *arrow::ListArray::FromArrays(*offsets, *values);
EXPECT_THROW(milvus::storage::NormalizeVectorArraysToFixedSizeBinary(
{input}, milvus::DataType::VECTOR_FLOAT, 2),
std::exception);
}
TEST(NormalizeVectorArraysToFixedSizeBinary, ListNullElementAsserts) {
arrow::FloatBuilder values_builder;
ASSERT_TRUE(values_builder.Append(1.0f).ok());
ASSERT_TRUE(values_builder.AppendNull().ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
arrow::Int32Builder offsets_builder;
ASSERT_TRUE(offsets_builder.AppendValues({0, 2}).ok());
std::shared_ptr<arrow::Array> offsets;
ASSERT_TRUE(offsets_builder.Finish(&offsets).ok());
auto input = *arrow::ListArray::FromArrays(*offsets, *values);
EXPECT_THROW(milvus::storage::NormalizeVectorArraysToFixedSizeBinary(
{input}, milvus::DataType::VECTOR_FLOAT, 2),
std::exception);
}
TEST(NormalizeVectorArraysToFixedSizeBinary, FixedSizeListDimMismatchAsserts) {
arrow::FloatBuilder values_builder;
ASSERT_TRUE(values_builder.AppendValues({1.0f, 2.0f, 3.0f}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
auto input = std::make_shared<arrow::FixedSizeListArray>(
arrow::fixed_size_list(arrow::float32(), 3), 1, values);
EXPECT_THROW(milvus::storage::NormalizeVectorArraysToFixedSizeBinary(
{input}, milvus::DataType::VECTOR_FLOAT, 2),
std::exception);
}
TEST(NormalizeVectorArraysToFixedSizeBinary,
FixedSizeListElementTypeMismatchAsserts) {
arrow::Int64Builder values_builder;
ASSERT_TRUE(values_builder.AppendValues({1, 2}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
auto input = std::make_shared<arrow::FixedSizeListArray>(
arrow::fixed_size_list(arrow::int64(), 2), 1, values);
EXPECT_THROW(milvus::storage::NormalizeVectorArraysToFixedSizeBinary(
{input}, milvus::DataType::VECTOR_FLOAT, 2),
std::exception);
}
TEST(NormalizeVectorArraysToFixedSizeBinary, BinaryVectorRejectsFixedSizeList) {
arrow::Int8Builder values_builder;
ASSERT_TRUE(values_builder.AppendValues({1, 0, 1, 0, 1, 0, 1, 0}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
auto input = std::make_shared<arrow::FixedSizeListArray>(
arrow::fixed_size_list(arrow::int8(), 8), 1, values);
EXPECT_THROW(milvus::storage::NormalizeVectorArraysToFixedSizeBinary(
{input}, milvus::DataType::VECTOR_BINARY, 8),
std::exception);
}
TEST(NormalizeVectorArraysToFixedSizeBinary,
BFloat16VectorRejectsFixedSizeList) {
arrow::Int16Builder values_builder;
ASSERT_TRUE(values_builder.AppendValues({1, 2}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
auto input = std::make_shared<arrow::FixedSizeListArray>(
arrow::fixed_size_list(arrow::int16(), 2), 1, values);
EXPECT_THROW(milvus::storage::NormalizeVectorArraysToFixedSizeBinary(
{input}, milvus::DataType::VECTOR_BFLOAT16, 2),
std::exception);
}
TEST(NormalizeVectorArraysToFixedSizeBinary,
FixedSizeBinaryWidthMismatchAsserts) {
auto fsb_type = arrow::fixed_size_binary(3 * sizeof(float));
arrow::FixedSizeBinaryBuilder builder(fsb_type);
float values[3] = {1.0f, 2.0f, 3.0f};
ASSERT_TRUE(builder.Append(reinterpret_cast<const uint8_t*>(values)).ok());
std::shared_ptr<arrow::Array> input;
ASSERT_TRUE(builder.Finish(&input).ok());
EXPECT_THROW(milvus::storage::NormalizeVectorArraysToFixedSizeBinary(
{input}, milvus::DataType::VECTOR_FLOAT, 2),
std::exception);
}
TEST(NormalizeVectorArrays, MixedChunkTypesNormalizeIndependently) {
auto fsb_type = arrow::fixed_size_binary(2 * sizeof(float));
arrow::FixedSizeBinaryBuilder fsb_builder(fsb_type);
float fsb_values[2] = {1.0f, 2.0f};
ASSERT_TRUE(
fsb_builder.Append(reinterpret_cast<const uint8_t*>(fsb_values)).ok());
std::shared_ptr<arrow::Array> fsb_input;
ASSERT_TRUE(fsb_builder.Finish(&fsb_input).ok());
arrow::FloatBuilder values_builder;
ASSERT_TRUE(values_builder.AppendValues({3.0f, 4.0f}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
arrow::Int32Builder offsets_builder;
ASSERT_TRUE(offsets_builder.AppendValues({0, 2}).ok());
std::shared_ptr<arrow::Array> offsets;
ASSERT_TRUE(offsets_builder.Finish(&offsets).ok());
auto list_input = *arrow::ListArray::FromArrays(*offsets, *values);
auto out = milvus::storage::NormalizeVectorArrays(
{fsb_input, list_input}, milvus::DataType::VECTOR_FLOAT, 2, false);
ASSERT_EQ(out.size(), 2);
EXPECT_EQ(out[0]->type_id(), arrow::Type::FIXED_SIZE_BINARY);
EXPECT_EQ(out[1]->type_id(), arrow::Type::FIXED_SIZE_BINARY);
EXPECT_EQ(std::static_pointer_cast<arrow::FixedSizeBinaryArray>(out[0])
->byte_width(),
2 * sizeof(float));
EXPECT_EQ(std::static_pointer_cast<arrow::FixedSizeBinaryArray>(out[1])
->byte_width(),
2 * sizeof(float));
}
TEST(NormalizeExternalArrow, NullableFloatVectorRejectsBinaryWidthMismatch) {
arrow::BinaryBuilder builder;
float values[3] = {1.0f, 2.0f, 3.0f};
ASSERT_TRUE(
builder.Append(reinterpret_cast<const uint8_t*>(values), sizeof(values))
.ok());
std::shared_ptr<arrow::Array> input;
ASSERT_TRUE(builder.Finish(&input).ok());
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(input,
milvus::DataType::VECTOR_FLOAT,
2,
true,
milvus::DataType::NONE),
std::exception);
}
TEST(NormalizeExternalArrow, FloatVectorRejectsFixedSizeBinaryWidthMismatch) {
auto fsb_type = arrow::fixed_size_binary(3 * sizeof(float));
arrow::FixedSizeBinaryBuilder builder(fsb_type);
float values[3] = {1.0f, 2.0f, 3.0f};
ASSERT_TRUE(builder.Append(reinterpret_cast<const uint8_t*>(values)).ok());
std::shared_ptr<arrow::Array> input;
ASSERT_TRUE(builder.Finish(&input).ok());
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(input,
milvus::DataType::VECTOR_FLOAT,
2,
false,
milvus::DataType::NONE),
std::exception);
}
TEST(NormalizeExternalArrow,
NullableFloatVectorRejectsFixedSizeBinaryWidthMismatch) {
auto fsb_type = arrow::fixed_size_binary(3 * sizeof(float));
arrow::FixedSizeBinaryBuilder builder(fsb_type);
float values[3] = {1.0f, 2.0f, 3.0f};
ASSERT_TRUE(builder.Append(reinterpret_cast<const uint8_t*>(values)).ok());
std::shared_ptr<arrow::Array> input;
ASSERT_TRUE(builder.Finish(&input).ok());
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(input,
milvus::DataType::VECTOR_FLOAT,
2,
true,
milvus::DataType::NONE),
std::exception);
}
TEST(NormalizeExternalArrow, FloatVectorRejectsFixedSizeListInt64) {
arrow::Int64Builder values_builder;
ASSERT_TRUE(values_builder.AppendValues({1, 2}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
auto input = std::make_shared<arrow::FixedSizeListArray>(
arrow::fixed_size_list(arrow::int64(), 2), 1, values);
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(input,
milvus::DataType::VECTOR_FLOAT,
2,
false,
milvus::DataType::NONE),
std::exception);
}
TEST(NormalizeExternalArrow, VectorArrayRejectsFixedSizeBinaryWidthMismatch) {
auto fsb_type = arrow::fixed_size_binary(3 * sizeof(float));
arrow::FixedSizeBinaryBuilder values_builder(fsb_type);
float values[3] = {1.0f, 2.0f, 3.0f};
ASSERT_TRUE(
values_builder.Append(reinterpret_cast<const uint8_t*>(values)).ok());
std::shared_ptr<arrow::Array> values_array;
ASSERT_TRUE(values_builder.Finish(&values_array).ok());
arrow::Int32Builder offsets_builder;
ASSERT_TRUE(offsets_builder.AppendValues({0, 1}).ok());
std::shared_ptr<arrow::Array> offsets;
ASSERT_TRUE(offsets_builder.Finish(&offsets).ok());
auto input = *arrow::ListArray::FromArrays(*offsets, *values_array);
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(input,
milvus::DataType::VECTOR_ARRAY,
2,
false,
milvus::DataType::VECTOR_FLOAT),
std::exception);
}
TEST(NormalizeExternalArrow, VectorArrayRejectsNonListInput) {
auto input = MakeInt64Array({1, 2}, {true, true});
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(input,
milvus::DataType::VECTOR_ARRAY,
2,
false,
milvus::DataType::VECTOR_FLOAT),
std::exception);
}
TEST(NormalizeExternalArrow, VectorArrayRejectsOuterNullRow) {
arrow::FloatBuilder values_builder;
ASSERT_TRUE(values_builder.AppendValues({1.0f, 2.0f}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
arrow::Int32Builder inner_offsets_builder;
ASSERT_TRUE(inner_offsets_builder.AppendValues({0, 2}).ok());
std::shared_ptr<arrow::Array> inner_offsets;
ASSERT_TRUE(inner_offsets_builder.Finish(&inner_offsets).ok());
auto inner_list = *arrow::ListArray::FromArrays(*inner_offsets, *values);
arrow::Int32Builder outer_offsets_builder;
ASSERT_TRUE(outer_offsets_builder.AppendValues({0, 1, 1}).ok());
std::shared_ptr<arrow::Array> outer_offsets;
ASSERT_TRUE(outer_offsets_builder.Finish(&outer_offsets).ok());
arrow::TypedBufferBuilder<bool> null_bitmap_builder;
ASSERT_TRUE(null_bitmap_builder.Reserve(2).ok());
null_bitmap_builder.UnsafeAppend(true);
null_bitmap_builder.UnsafeAppend(false);
std::shared_ptr<arrow::Buffer> null_bitmap;
ASSERT_TRUE(null_bitmap_builder.Finish(&null_bitmap).ok());
auto outer_offsets_values =
std::static_pointer_cast<arrow::Int32Array>(outer_offsets)->values();
auto input =
std::make_shared<arrow::ListArray>(arrow::list(inner_list->type()),
2,
outer_offsets_values,
inner_list,
null_bitmap,
1);
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(input,
milvus::DataType::VECTOR_ARRAY,
2,
false,
milvus::DataType::VECTOR_FLOAT),
std::exception);
}
TEST(NormalizeExternalArrow, SparseVectorRejectsString) {
arrow::StringBuilder builder;
ASSERT_TRUE(builder.AppendValues({"not_sparse"}).ok());
std::shared_ptr<arrow::Array> input;
ASSERT_TRUE(builder.Finish(&input).ok());
EXPECT_THROW(milvus::storage::NormalizeExternalArrow(
input,
milvus::DataType::VECTOR_SPARSE_U32_F32,
0,
false,
milvus::DataType::NONE),
std::exception);
}
TEST(NormalizeExternalArrow, ArrayInt64RejectsStringList) {
arrow::StringBuilder values_builder;
ASSERT_TRUE(values_builder.AppendValues({"1", "2"}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
arrow::Int32Builder offsets_builder;
ASSERT_TRUE(offsets_builder.AppendValues({0, 2}).ok());
std::shared_ptr<arrow::Array> offsets;
ASSERT_TRUE(offsets_builder.Finish(&offsets).ok());
auto input = *arrow::ListArray::FromArrays(*offsets, *values);
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(
input, milvus::DataType::ARRAY, 0, false, milvus::DataType::INT64),
std::exception);
}
TEST(NormalizeExternalArrow, ArrayInt64RejectsNullElement) {
arrow::Int64Builder values_builder;
ASSERT_TRUE(values_builder.Append(1).ok());
ASSERT_TRUE(values_builder.AppendNull().ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
arrow::Int32Builder offsets_builder;
ASSERT_TRUE(offsets_builder.AppendValues({0, 2}).ok());
std::shared_ptr<arrow::Array> offsets;
ASSERT_TRUE(offsets_builder.Finish(&offsets).ok());
auto input = *arrow::ListArray::FromArrays(*offsets, *values);
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(
input, milvus::DataType::ARRAY, 0, false, milvus::DataType::INT64),
std::exception);
}
TEST(NormalizeExternalArrow, ArrayVarcharRejectsInt64List) {
arrow::Int64Builder values_builder;
ASSERT_TRUE(values_builder.AppendValues({1, 2}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
arrow::Int32Builder offsets_builder;
ASSERT_TRUE(offsets_builder.AppendValues({0, 2}).ok());
std::shared_ptr<arrow::Array> offsets;
ASSERT_TRUE(offsets_builder.Finish(&offsets).ok());
auto input = *arrow::ListArray::FromArrays(*offsets, *values);
EXPECT_THROW(
milvus::storage::NormalizeExternalArrow(input,
milvus::DataType::ARRAY,
0,
false,
milvus::DataType::VARCHAR),
std::exception);
}
TEST(NormalizeExternalArrow, ArrayVarcharAcceptsStringList) {
arrow::StringBuilder values_builder;
ASSERT_TRUE(values_builder.AppendValues({"a", "b"}).ok());
std::shared_ptr<arrow::Array> values;
ASSERT_TRUE(values_builder.Finish(&values).ok());
arrow::Int32Builder offsets_builder;
ASSERT_TRUE(offsets_builder.AppendValues({0, 2}).ok());
std::shared_ptr<arrow::Array> offsets;
ASSERT_TRUE(offsets_builder.Finish(&offsets).ok());
auto input = *arrow::ListArray::FromArrays(*offsets, *values);
auto out = milvus::storage::NormalizeExternalArrow(
input, milvus::DataType::ARRAY, 0, false, milvus::DataType::VARCHAR);
ASSERT_EQ(out->type_id(), arrow::Type::BINARY);
ASSERT_EQ(out->length(), 1);
EXPECT_FALSE(out->IsNull(0));
}
// ===== CoerceToList =====
TEST(CoerceToList, LargeListToList) {
internal/metastore/kv/rootcoord/kv_catalog.go
+2
View File
@@ -749,6 +749,8 @@ func (kc *Catalog) alterModifyCollection(ctx context.Context, oldColl *model.Col
oldCollClone.EnableDynamicField = newColl.EnableDynamicField
oldCollClone.SchemaVersion = newColl.SchemaVersion
oldCollClone.ShardInfos = newColl.ShardInfos
oldCollClone.ExternalSource = newColl.ExternalSource
oldCollClone.ExternalSpec = newColl.ExternalSpec
newKey := BuildCollectionKey(newColl.DBID, oldColl.CollectionID)
value, err := proto.Marshal(model.MarshalCollectionModel(oldCollClone))
internal/metastore/kv/rootcoord/kv_catalog_test.go
+16 -2
View File
@@ -1056,8 +1056,20 @@ func TestCatalog_AlterCollection(t *testing.T) {
kc := NewCatalog(snapshot).(*Catalog)
ctx := context.Background()
var collectionID int64 = 1
oldC := &model.Collection{CollectionID: collectionID, State: pb.CollectionState_CollectionCreated}
newC := &model.Collection{CollectionID: collectionID, State: pb.CollectionState_CollectionCreated, UpdateTimestamp: rand.Uint64()}
oldC := &model.Collection{
CollectionID: collectionID,
State: pb.CollectionState_CollectionCreated,
ExternalSource: "s3://bucket/old/",
ExternalSpec: `{"format":"parquet"}`,
}
newSpec := `{"format":"parquet","my_extra":{"k":"v"}}`
newC := &model.Collection{
CollectionID: collectionID,
State: pb.CollectionState_CollectionCreated,
UpdateTimestamp: rand.Uint64(),
ExternalSource: "s3://bucket/new/",
ExternalSpec: newSpec,
}
err := kc.AlterCollection(ctx, oldC, newC, metastore.MODIFY, 0, true)
assert.NoError(t, err)
key := BuildCollectionKey(0, collectionID)
@@ -1069,6 +1081,8 @@ func TestCatalog_AlterCollection(t *testing.T) {
got := model.UnmarshalCollectionModel(&collPb)
assert.Equal(t, pb.CollectionState_CollectionCreated, got.State)
assert.Equal(t, newC.UpdateTimestamp, got.UpdateTimestamp)
assert.Equal(t, newC.ExternalSource, got.ExternalSource)
assert.Equal(t, newSpec, got.ExternalSpec)
})
t.Run("modify EnableDynamicField and SchemaVersion", func(t *testing.T) {
internal/proxy/util.go
+1 -1
View File
@@ -817,7 +817,7 @@ func injectVirtualPKForExternalCollection(schema *schemapb.CollectionSchema) err
// will assign the actual field ID during collection creation.
virtualPKField := &schemapb.FieldSchema{
Name: common.VirtualPKFieldName,
Description: "Virtual primary key for external collection: (segmentID << 32) | offset",
Description: "auto-generated primary key for external collection",
DataType: schemapb.DataType_Int64,
IsPrimaryKey: true,
AutoID: true, // Virtual PKs are auto-generated
pkg/util/externalspec/external_spec.go
+67 -15
View File
@@ -511,28 +511,24 @@ var awsFamilyScheme = map[string]bool{
// RedactExternalSpec returns a log-safe representation of an external spec
// JSON string. Secret extfs values (see secretExtfsKeys) are replaced with
// "***" so that AK/SK/PEM material never reaches log sinks. On parse failure
// it returns "<invalid spec>" rather than the raw input — the input itself
// may already contain a partially-recognized credential blob, so we never
// echo it back. Empty input returns empty string for log readability.
// "***" so that AK/SK/PEM material never reaches log sinks. Unknown fields
// are preserved so API callers can still observe extension metadata. On parse
// failure it returns "<invalid spec>" rather than the raw input — the input
// itself may already contain a partially-recognized credential blob, so we
// never echo it back. Empty input returns empty string for log readability.
func RedactExternalSpec(specStr string) string {
if specStr == "" {
return ""
}
var spec ExternalSpec
var spec map[string]json.RawMessage
if err := json.Unmarshal([]byte(specStr), &spec); err != nil {
return "<invalid spec>"
}
if len(spec.Extfs) > 0 {
redacted := make(map[string]string, len(spec.Extfs))
for k, v := range spec.Extfs {
if secretExtfsKeys[k] && v != "" {
redacted[k] = "***"
} else {
redacted[k] = v
}
}
spec.Extfs = redacted
if err := normalizeSnapshotID(spec); err != nil {
return "<invalid spec>"
}
if err := redactExtfsSecrets(spec); err != nil {
return "<invalid spec>"
}
out, err := json.Marshal(spec)
if err != nil {
@@ -540,3 +536,59 @@ func RedactExternalSpec(specStr string) string {
}
return string(out)
}
func normalizeSnapshotID(spec map[string]json.RawMessage) error {
snapshotRaw, ok := spec["snapshot_id"]
if !ok || len(snapshotRaw) == 0 || string(snapshotRaw) == "null" {
return nil
}
var n int64
if err := json.Unmarshal(snapshotRaw, &n); err == nil {
spec["snapshot_id"] = quotedInt64JSON(n)
return nil
}
var str string
if err := json.Unmarshal(snapshotRaw, &str); err != nil {
return err
}
parsed, err := strconv.ParseInt(str, 10, 64)
if err != nil {
return err
}
spec["snapshot_id"] = quotedInt64JSON(parsed)
return nil
}
func quotedInt64JSON(v int64) json.RawMessage {
return json.RawMessage(strconv.Quote(strconv.FormatInt(v, 10)))
}
func redactExtfsSecrets(spec map[string]json.RawMessage) error {
extfsRaw, ok := spec["extfs"]
if !ok || len(extfsRaw) == 0 || string(extfsRaw) == "null" {
return nil
}
var extfs map[string]json.RawMessage
if err := json.Unmarshal(extfsRaw, &extfs); err != nil {
return err
}
for k, v := range extfs {
if !secretExtfsKeys[k] {
continue
}
var str string
if err := json.Unmarshal(v, &str); err == nil && str == "" {
continue
}
extfs[k] = json.RawMessage(`"***"`)
}
redactedExtfs, err := json.Marshal(extfs)
if err != nil {
return err
}
spec["extfs"] = redactedExtfs
return nil
}
pkg/util/externalspec/external_spec_test.go
+143
View File
@@ -17,6 +17,7 @@
package externalspec
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/assert"
@@ -301,6 +302,10 @@ func TestRedactExternalSpec(t *testing.T) {
assert.Equal(t, "<invalid spec>", RedactExternalSpec("{not json"))
})
t.Run("invalid_extfs_returns_placeholder", func(t *testing.T) {
assert.Equal(t, "<invalid spec>", RedactExternalSpec(`{"format":"parquet","extfs":"not-an-object"}`))
})
t.Run("secrets_masked", func(t *testing.T) {
out := RedactExternalSpec(`{"format":"parquet","extfs":{"access_key_id":"AKIA","access_key_value":"SEC","region":"us"}}`)
assert.Contains(t, out, `"access_key_id":"***"`)
@@ -314,6 +319,144 @@ func TestRedactExternalSpec(t *testing.T) {
out := RedactExternalSpec(`{"format":"parquet","extfs":{"access_key_id":""}}`)
assert.Contains(t, out, `"access_key_id":""`)
})
t.Run("non_string_secret_values_masked", func(t *testing.T) {
out := RedactExternalSpec(`{"format":"parquet","extfs":{"access_key_id":123,"ssl_ca_cert":true}}`)
assert.Contains(t, out, `"access_key_id":"***"`)
assert.Contains(t, out, `"ssl_ca_cert":"***"`)
})
t.Run("snapshot_id_string_preserved_as_string", func(t *testing.T) {
out := RedactExternalSpec(`{"format":"iceberg-table","snapshot_id":"5320540205222981137","extfs":null}`)
assert.Contains(t, out, `"snapshot_id":"5320540205222981137"`)
})
t.Run("invalid_snapshot_id_returns_placeholder", func(t *testing.T) {
assert.Equal(t, "<invalid spec>", RedactExternalSpec(`{"format":"iceberg-table","snapshot_id":"abc"}`))
assert.Equal(t, "<invalid spec>", RedactExternalSpec(`{"format":"iceberg-table","snapshot_id":{}}`))
})
t.Run("preserves_unknown_top_level_fields", func(t *testing.T) {
out := RedactExternalSpec(`{
"format":"parquet",
"cloud_extra":{
"volume_uri":"volume://tk-stagexxx2222222/test/external-collection/",
"volume_id":"volume-ifi5qkwp89z4ljtkvali",
"integration_id":"integ-lir5xfbcgrkla6fjc39w15qjk",
"path":"test/external-collection/"
},
"extfs":{
"cloud_provider":"aws",
"region":"us-west-2",
"use_iam":"true",
"role_arn":"arn:aws:iam::306787409409:role/lentitude-bucket-role",
"external_id":"zilliz-external-sO1cjGS2Vgpyan"
}
}`)
var got map[string]any
require.NoError(t, json.Unmarshal([]byte(out), &got))
require.Contains(t, got, "cloud_extra")
cloudExtra, ok := got["cloud_extra"].(map[string]any)
require.True(t, ok)
assert.Equal(t, "volume://tk-stagexxx2222222/test/external-collection/", cloudExtra["volume_uri"])
assert.Equal(t, "volume-ifi5qkwp89z4ljtkvali", cloudExtra["volume_id"])
assert.Equal(t, "integ-lir5xfbcgrkla6fjc39w15qjk", cloudExtra["integration_id"])
assert.Equal(t, "test/external-collection/", cloudExtra["path"])
extfs, ok := got["extfs"].(map[string]any)
require.True(t, ok)
assert.Equal(t, "aws", extfs["cloud_provider"])
assert.Equal(t, "arn:aws:iam::306787409409:role/lentitude-bucket-role", extfs["role_arn"])
assert.Equal(t, "***", extfs["external_id"])
assert.NotContains(t, out, "zilliz-external-sO1cjGS2Vgpyan")
})
}
func TestExternalSpecMarshalJSON(t *testing.T) {
t.Run("snapshot_id_marshaled_as_string", func(t *testing.T) {
snapshotID := int64(5320540205222981137)
out, err := json.Marshal(ExternalSpec{
Format: FormatIcebergTable,
Columns: []string{"id", "vec"},
Extfs: map[string]string{"cloud_provider": "aws"},
SnapshotID: &snapshotID,
})
require.NoError(t, err)
assert.Contains(t, string(out), `"snapshot_id":"5320540205222981137"`)
var got map[string]any
require.NoError(t, json.Unmarshal(out, &got))
assert.Equal(t, FormatIcebergTable, got["format"])
assert.Equal(t, "5320540205222981137", got["snapshot_id"])
})
t.Run("nil_snapshot_id_omitted", func(t *testing.T) {
out, err := json.Marshal(ExternalSpec{Format: FormatParquet})
require.NoError(t, err)
assert.NotContains(t, string(out), "snapshot_id")
assert.Contains(t, string(out), `"format":"parquet"`)
})
}
func TestIsCloudEndpointHost(t *testing.T) {
for _, host := range []string{
"s3.us-west-2.amazonaws.com",
"s3.cn-north-1.amazonaws.com.cn",
"storage.googleapis.com",
"oss-cn-hangzhou.aliyuncs.com",
"cos.ap-shanghai.myqcloud.com",
"obs.cn-north-4.myhuaweicloud.com",
"acct.blob.core.windows.net",
"acct.blob.core.chinacloudapi.cn",
"acct.blob.core.usgovcloudapi.net",
"acct.blob.core.cloudapi.de",
"S3.US-WEST-2.AMAZONAWS.COM",
} {
assert.True(t, IsCloudEndpointHost(host), "host %s should be recognized", host)
}
for _, host := range []string{
"localhost:9000",
"bucket",
"amazonaws.com",
} {
assert.False(t, IsCloudEndpointHost(host), "host %s should not be recognized", host)
}
}
func TestDeriveEndpoint(t *testing.T) {
tests := []struct {
name string
cloudProvider string
region string
expected string
}{
{"aws_requires_region", CloudProviderAWS, "", ""},
{"aws_standard", CloudProviderAWS, "us-west-2", "https://s3.us-west-2.amazonaws.com"},
{"aws_china", CloudProviderAWS, "cn-north-1", "https://s3.cn-north-1.amazonaws.com.cn"},
{"gcp_global", CloudProviderGCP, "", "https://storage.googleapis.com"},
{"aliyun_requires_region", CloudProviderAliyun, "", ""},
{"aliyun_region", CloudProviderAliyun, "cn-hangzhou", "https://oss-cn-hangzhou.aliyuncs.com"},
{"tencent_requires_region", CloudProviderTencent, "", ""},
{"tencent_region", CloudProviderTencent, "ap-shanghai", "https://cos.ap-shanghai.myqcloud.com"},
{"huawei_requires_region", CloudProviderHuawei, "", ""},
{"huawei_region", CloudProviderHuawei, "cn-north-4", "https://obs.cn-north-4.myhuaweicloud.com"},
{"azure_requires_region", CloudProviderAzure, "", ""},
{"azure_china", CloudProviderAzure, "china", "core.chinacloudapi.cn"},
{"azure_usgov", CloudProviderAzure, "usgov", "core.usgovcloudapi.net"},
{"azure_usdod", CloudProviderAzure, "usdod", "core.usgovcloudapi.net"},
{"azure_germany", CloudProviderAzure, "germany", "core.cloudapi.de"},
{"azure_public", CloudProviderAzure, "public", "core.windows.net"},
{"case_insensitive_provider", "AWS", "us-east-1", "https://s3.us-east-1.amazonaws.com"},
{"unknown_provider", "minio", "us-east-1", ""},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert.Equal(t, tt.expected, DeriveEndpoint(tt.cloudProvider, tt.region))
})
}
}
func TestParseExternalSpec_ArnKeys(t *testing.T) {