From 82b85beb596bd7afa6ad074d57e63a014abbc7ba Mon Sep 17 00:00:00 2001 From: Eric Hare Date: Tue, 16 Jun 2026 17:33:23 -0700 Subject: [PATCH] fix: remove vector index profile from owning table on dropIndex createVectorIndex persists a per-index profile record (name + applied SAI options) under the owning table's VECTOR_INDEX_PROFILES extension. dropIndex previously left that entry behind, so profile metadata outlived the dropped index (metadata bloat that could later mislead a profile-echo on listIndexes). dropIndex now resolves the owning table from the keyspace metadata and, when a stored profile exists for the index, runs a sequential ALTER TABLE ... WITH extensions after the drop to remove just that entry. The rewrite reuses the clobber-safe TableExtensions.createCustomProperties so vectorize config and the other indexes' profiles are preserved. A keyspace-typed DropVectorIndexProfileDBTask is added so it can share a TaskGroup with the keyspace-typed DropIndexDBTask (a TableSchemaObject-typed AlterTableDBTask could not). KeyspaceSchemaObject exposes keyspaceMetadata() to support owning-table resolution. The happy path is not integration-testable on the default dse-server:6.9.21 backend (creating a profiled index needs custom SAI params it rejects), consistent with the parent PR; the resolution + entry-removal logic is unit tested instead. --- .../cqldriver/executor/TableExtensions.java | 58 ++++++++ .../tables/DropVectorIndexProfileDBTask.java | 60 +++++++++ .../DropVectorIndexProfileDBTaskBuilder.java | 56 ++++++++ .../resolver/DropIndexCommandResolver.java | 108 ++++++++++----- .../service/schema/KeyspaceSchemaObject.java | 13 ++ .../executor/TableExtensionsTest.java | 125 ++++++++++++++++++ .../DropVectorIndexProfileDBTaskTest.java | 53 ++++++++ 7 files changed, 443 insertions(+), 30 deletions(-) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTask.java create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTaskBuilder.java create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTaskTest.java diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/TableExtensions.java b/src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/TableExtensions.java index ca8166dfcf..f713e7d892 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/TableExtensions.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/TableExtensions.java @@ -4,6 +4,7 @@ import com.datastax.oss.driver.api.core.CqlIdentifier; import com.datastax.oss.driver.api.core.data.ByteUtils; +import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata; import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; @@ -13,6 +14,7 @@ import java.util.HashMap; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.stream.Collectors; import org.slf4j.Logger; @@ -114,6 +116,62 @@ public static Map createCustomProperties( return customProperties; } + /** + * Computes the extensions payload that drops {@code indexName}'s vector-index profile from the + * table that owns it. Used to keep the {@link + * SchemaConstants.MetadataFieldsNames#VECTOR_INDEX_PROFILES} extension in sync when an index is + * dropped, so a profile record does not outlive its index. + * + *

The owning table is found by scanning {@code keyspaceMetadata} for the table whose indexes + * contain {@code indexName}. Returns empty when there is nothing to do — no table owns the index, + * or the owning table has no stored profile for it — so the caller can skip the extra DDL. + * + *

When a rewrite is needed, the existing vectorize config and the other indexes' profiles are + * read back and included so the clobbering extension write does not lose them (the same approach + * as the create side, see {@link #createCustomProperties(Map, Map, ObjectMapper)}). + */ + public static Optional removeIndexProfile( + KeyspaceMetadata keyspaceMetadata, CqlIdentifier indexName, ObjectMapper objectMapper) { + Objects.requireNonNull(keyspaceMetadata, "keyspaceMetadata must not be null"); + Objects.requireNonNull(indexName, "indexName must not be null"); + Objects.requireNonNull(objectMapper, "objectMapper must not be null"); + + var owningTable = + keyspaceMetadata.getTables().values().stream() + .filter(table -> table.getIndexes().containsKey(indexName)) + .findFirst(); + if (owningTable.isEmpty()) { + return Optional.empty(); + } + + var tableMetadata = owningTable.get(); + var profiles = VectorIndexProfileDefinition.from(tableMetadata, objectMapper); + // null def => remove; false return => no entry existed, so there is nothing to rewrite. + if (!VectorIndexProfileDefinition.putOrRemove( + profiles, cqlIdentifierToJsonKey(indexName), null)) { + return Optional.empty(); + } + + // Read the vectorize config back so the full-replace extension write preserves it. The stored + // keys are the column identifiers' internal form, so reconstruct the CqlIdentifier keys that + // createCustomProperties expects. + var vectorDefs = + VectorizeDefinition.from(tableMetadata, objectMapper).entrySet().stream() + .collect( + Collectors.toMap( + entry -> CqlIdentifier.fromInternal(entry.getKey()), Map.Entry::getValue)); + + var customProperties = createCustomProperties(vectorDefs, profiles, objectMapper); + return Optional.of(new IndexProfileRemoval(tableMetadata.getName(), customProperties)); + } + + /** + * The result of {@link #removeIndexProfile}: the table to alter and the complete extensions + * payload to write (with the dropped index's profile removed and everything else preserved). + */ + public record IndexProfileRemoval( + CqlIdentifier tableName, Map customProperties) {} + private static String writeJson(Object value, ObjectMapper objectMapper) { try { return objectMapper.writeValueAsString(value); diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTask.java b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTask.java new file mode 100644 index 0000000000..c14079814c --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTask.java @@ -0,0 +1,60 @@ +package io.stargate.sgv2.jsonapi.service.operation.tables; + +import static com.datastax.oss.driver.api.querybuilder.SchemaBuilder.alterTable; + +import com.datastax.oss.driver.api.core.CqlIdentifier; +import com.datastax.oss.driver.api.core.cql.SimpleStatement; +import io.stargate.sgv2.jsonapi.service.cqldriver.executor.DefaultDriverExceptionHandler; +import io.stargate.sgv2.jsonapi.service.cqldriver.executor.TableExtensions; +import io.stargate.sgv2.jsonapi.service.operation.SchemaDBTask; +import io.stargate.sgv2.jsonapi.service.schema.KeyspaceSchemaObject; +import java.util.Map; +import java.util.Objects; + +/** + * Removes a dropped index's entry from its owning table's vector-index-profiles extension, so a + * profile record does not outlive the index it described. + * + *

This runs as a keyspace-scoped sibling to {@link DropIndexDBTask} so the two can share one + * {@link io.stargate.sgv2.jsonapi.service.operation.tasks.TaskGroup}; a {@link + * io.stargate.sgv2.jsonapi.service.schema.tables.TableSchemaObject}-typed {@link AlterTableDBTask} + * (used by the create side) could not, because a TaskGroup has a single schema-object type. The + * owning table and the rewritten extensions payload are resolved at command-resolve time via {@link + * TableExtensions#removeIndexProfile}; this task only issues the {@code ALTER TABLE ... WITH + * extensions = {...}}. + */ +public class DropVectorIndexProfileDBTask extends SchemaDBTask { + + private final CqlIdentifier tableName; + private final Map customProperties; + + public DropVectorIndexProfileDBTask( + int position, + KeyspaceSchemaObject schemaObject, + SchemaDBTask.SchemaRetryPolicy schemaRetryPolicy, + DefaultDriverExceptionHandler.Factory exceptionHandlerFactory, + CqlIdentifier tableName, + Map customProperties) { + super(position, schemaObject, schemaRetryPolicy, exceptionHandlerFactory); + + this.tableName = Objects.requireNonNull(tableName, "tableName must not be null"); + this.customProperties = + Objects.requireNonNull(customProperties, "customProperties must not be null"); + setStatus(TaskStatus.READY); + } + + public static DropVectorIndexProfileDBTaskBuilder builder(KeyspaceSchemaObject schemaObject) { + return new DropVectorIndexProfileDBTaskBuilder(schemaObject); + } + + @Override + protected SimpleStatement buildStatement() { + + // The owning table lives in this keyspace; take the keyspace from the schema object identifier, + // mirroring DropIndexDBTask which builds its statement the same way. + var extensions = TableExtensions.toExtensions(customProperties); + return alterTable(schemaObject.identifier().keyspace(), tableName) + .withOption(TableExtensions.TABLE_OPTIONS_EXTENSION_KEY.asInternal(), extensions) + .build(); + } +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTaskBuilder.java b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTaskBuilder.java new file mode 100644 index 0000000000..1f60b57fe8 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTaskBuilder.java @@ -0,0 +1,56 @@ +package io.stargate.sgv2.jsonapi.service.operation.tables; + +import com.datastax.oss.driver.api.core.CqlIdentifier; +import io.stargate.sgv2.jsonapi.service.operation.SchemaDBTask; +import io.stargate.sgv2.jsonapi.service.operation.tasks.TaskBuilder; +import io.stargate.sgv2.jsonapi.service.schema.KeyspaceSchemaObject; +import java.util.Map; +import java.util.Objects; + +/** Builds a {@link DropVectorIndexProfileDBTask}. */ +public class DropVectorIndexProfileDBTaskBuilder + extends TaskBuilder< + DropVectorIndexProfileDBTask, KeyspaceSchemaObject, DropVectorIndexProfileDBTaskBuilder> { + + private CqlIdentifier tableName; + private Map customProperties; + private SchemaDBTask.SchemaRetryPolicy schemaRetryPolicy; + + protected DropVectorIndexProfileDBTaskBuilder(KeyspaceSchemaObject schemaObject) { + super(schemaObject); + } + + public DropVectorIndexProfileDBTaskBuilder withTableName(CqlIdentifier tableName) { + this.tableName = Objects.requireNonNull(tableName, "tableName must not be null"); + return this; + } + + public DropVectorIndexProfileDBTaskBuilder withCustomProperties( + Map customProperties) { + this.customProperties = + Objects.requireNonNull(customProperties, "customProperties must not be null"); + return this; + } + + public DropVectorIndexProfileDBTaskBuilder withSchemaRetryPolicy( + SchemaDBTask.SchemaRetryPolicy schemaRetryPolicy) { + this.schemaRetryPolicy = + Objects.requireNonNull(schemaRetryPolicy, "schemaRetryPolicy cannot be null"); + return this; + } + + public DropVectorIndexProfileDBTask build() { + + Objects.requireNonNull(tableName, "tableName must not be null"); + Objects.requireNonNull(customProperties, "customProperties must not be null"); + Objects.requireNonNull(schemaRetryPolicy, "schemaRetryPolicy cannot be null"); + + return new DropVectorIndexProfileDBTask( + nextPosition(), + schemaObject, + schemaRetryPolicy, + getExceptionHandlerFactory(), + tableName, + customProperties); + } +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/DropIndexCommandResolver.java b/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/DropIndexCommandResolver.java index cf6ff8e511..c324f34910 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/DropIndexCommandResolver.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/DropIndexCommandResolver.java @@ -2,20 +2,26 @@ import static io.stargate.sgv2.jsonapi.util.CqlIdentifierUtil.cqlIdentifierFromUserInput; +import com.datastax.oss.driver.api.core.CqlIdentifier; +import com.fasterxml.jackson.databind.ObjectMapper; import io.stargate.sgv2.jsonapi.api.model.command.CommandContext; import io.stargate.sgv2.jsonapi.api.model.command.impl.DropIndexCommand; import io.stargate.sgv2.jsonapi.config.OperationsConfig; import io.stargate.sgv2.jsonapi.service.cqldriver.executor.DefaultDriverExceptionHandler; +import io.stargate.sgv2.jsonapi.service.cqldriver.executor.TableExtensions; import io.stargate.sgv2.jsonapi.service.operation.Operation; import io.stargate.sgv2.jsonapi.service.operation.SchemaDBTask; import io.stargate.sgv2.jsonapi.service.operation.SchemaDBTaskPage; +import io.stargate.sgv2.jsonapi.service.operation.keyspaces.KeyspaceDriverExceptionHandler; import io.stargate.sgv2.jsonapi.service.operation.tables.DropIndexDBTask; import io.stargate.sgv2.jsonapi.service.operation.tables.DropIndexExceptionHandler; +import io.stargate.sgv2.jsonapi.service.operation.tables.DropVectorIndexProfileDBTask; import io.stargate.sgv2.jsonapi.service.operation.tasks.TaskGroup; import io.stargate.sgv2.jsonapi.service.operation.tasks.TaskOperation; import io.stargate.sgv2.jsonapi.service.schema.KeyspaceSchemaObject; import io.stargate.sgv2.jsonapi.util.ApiOptionUtils; import jakarta.enterprise.context.ApplicationScoped; +import jakarta.inject.Inject; import java.time.Duration; /** Resolver for the {@link DropIndexCommand}. */ @@ -24,6 +30,8 @@ public class DropIndexCommandResolver implements CommandResolver getCommandClass() { return DropIndexCommand.class; @@ -33,39 +41,79 @@ public Class getCommandClass() { public Operation resolveKeyspaceCommand( CommandContext commandContext, DropIndexCommand command) { + var schemaObject = commandContext.schemaObject(); var indexName = cqlIdentifierFromUserInput(command.name()); // Check if the index exists, we check if columns exist before trying to drop them so do for // indexes as well - var taskBuilder = - DropIndexDBTask.builder(commandContext.schemaObject()) - .withSchemaRetryPolicy( - new SchemaDBTask.SchemaRetryPolicy( - commandContext - .config() - .get(OperationsConfig.class) - .databaseConfig() - .ddlRetries(), - Duration.ofMillis( - commandContext - .config() - .get(OperationsConfig.class) - .databaseConfig() - .ddlRetryDelayMillis()))); - - taskBuilder.withExceptionHandlerFactory( - DefaultDriverExceptionHandler.Factory.withIdentifier( - DropIndexExceptionHandler::new, indexName)); - - taskBuilder - .withIndexName(indexName) - .withIfExists( - ApiOptionUtils.getOrDefault( - command.options(), DropIndexCommand.Options::ifExists, IF_EXISTS_DEFAULT)); - - var taskGroup = new TaskGroup<>(taskBuilder.build()); - - return new TaskOperation<>( - taskGroup, SchemaDBTaskPage.accumulator(DropIndexDBTask.class, commandContext)); + var schemaRetryPolicy = + new SchemaDBTask.SchemaRetryPolicy( + commandContext.config().get(OperationsConfig.class).databaseConfig().ddlRetries(), + Duration.ofMillis( + commandContext + .config() + .get(OperationsConfig.class) + .databaseConfig() + .ddlRetryDelayMillis())); + + var dropIndexTask = + DropIndexDBTask.builder(schemaObject) + .withSchemaRetryPolicy(schemaRetryPolicy) + .withExceptionHandlerFactory( + DefaultDriverExceptionHandler.Factory.withIdentifier( + DropIndexExceptionHandler::new, indexName)) + .withIndexName(indexName) + .withIfExists( + ApiOptionUtils.getOrDefault( + command.options(), DropIndexCommand.Options::ifExists, IF_EXISTS_DEFAULT)) + .build(); + + // Also drop the index's vector-index profile (if any) from the owning table's extensions, so + // the profile record does not outlive the index. Null when the keyspace metadata is unknown or + // the owning table has no stored profile for this index, in which case only the drop runs. + var profileCleanupTask = buildProfileCleanupTask(schemaObject, indexName, schemaRetryPolicy); + + if (profileCleanupTask == null) { + return new TaskOperation<>( + new TaskGroup<>(dropIndexTask), + SchemaDBTaskPage.accumulator(DropIndexDBTask.class, commandContext)); + } + + // Sequential so the extension cleanup only runs if the index drop succeeded. + TaskGroup, KeyspaceSchemaObject> taskGroup = + new TaskGroup<>(true); + taskGroup.add(dropIndexTask); + taskGroup.add(profileCleanupTask); + + @SuppressWarnings("unchecked") + Class> taskClass = + (Class>) (Class) SchemaDBTask.class; + return new TaskOperation<>(taskGroup, SchemaDBTaskPage.accumulator(taskClass, commandContext)); + } + + /** + * Builds the cleanup task that removes the dropped index's profile from its owning table's + * extensions, or null when there is nothing to clean up (keyspace metadata unknown, no owning + * table, or no stored profile for this index). + */ + private DropVectorIndexProfileDBTask buildProfileCleanupTask( + KeyspaceSchemaObject schemaObject, + CqlIdentifier indexName, + SchemaDBTask.SchemaRetryPolicy schemaRetryPolicy) { + + return schemaObject + .keyspaceMetadata() + .flatMap( + keyspaceMetadata -> + TableExtensions.removeIndexProfile(keyspaceMetadata, indexName, objectMapper)) + .map( + removal -> + DropVectorIndexProfileDBTask.builder(schemaObject) + .withSchemaRetryPolicy(schemaRetryPolicy) + .withExceptionHandlerFactory(KeyspaceDriverExceptionHandler::new) + .withTableName(removal.tableName()) + .withCustomProperties(removal.customProperties()) + .build()) + .orElse(null); } } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/schema/KeyspaceSchemaObject.java b/src/main/java/io/stargate/sgv2/jsonapi/service/schema/KeyspaceSchemaObject.java index 11bdef3f59..7f86f5d3ad 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/schema/KeyspaceSchemaObject.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/schema/KeyspaceSchemaObject.java @@ -6,6 +6,7 @@ import io.stargate.sgv2.jsonapi.service.cqldriver.executor.IndexUsage; import io.stargate.sgv2.jsonapi.service.cqldriver.executor.VectorConfig; import java.util.Objects; +import java.util.Optional; /** * A Keyspace in the API. @@ -39,6 +40,18 @@ public KeyspaceSchemaObject(Tenant tenant, KeyspaceMetadata keyspaceMetadata) { Objects.requireNonNull(keyspaceMetadata, "keyspaceMetadata must not be null"); } + /** + * The Cassandra metadata for this keyspace, when known. + * + *

Empty when the object was built via the {@link + * #KeyspaceSchemaObject(SchemaObjectIdentifier)} test constructor, which carries no metadata. + * Present for objects built from live schema, where it lets callers reach the keyspace's tables + * and their indexes (e.g. to find the table that owns a named index). + */ + public Optional keyspaceMetadata() { + return Optional.ofNullable(keyspaceMetadata); + } + @Override public VectorConfig vectorConfig() { return VectorConfig.NOT_ENABLED_CONFIG; diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/TableExtensionsTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/TableExtensionsTest.java index 29a476b5d8..4ac19cd453 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/TableExtensionsTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/TableExtensionsTest.java @@ -1,11 +1,21 @@ package io.stargate.sgv2.jsonapi.service.cqldriver.executor; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import com.datastax.oss.driver.api.core.CqlIdentifier; +import com.datastax.oss.driver.api.core.metadata.schema.IndexMetadata; +import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata; +import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata; import com.fasterxml.jackson.databind.ObjectMapper; import io.stargate.sgv2.jsonapi.config.constants.SchemaConstants; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; import java.util.Map; +import java.util.Set; +import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; class TableExtensionsTest { @@ -71,4 +81,119 @@ void twoArgOverloadOmitsProfiles() { assertThat(props).doesNotContainKey(SchemaConstants.MetadataFieldsNames.VECTOR_INDEX_PROFILES); } + + @Nested + class RemoveIndexProfile { + + private static final CqlIdentifier MY_IDX = CqlIdentifier.fromInternal("my_idx"); + + @Test + void emptyWhenNoTableOwnsTheIndex() { + // the only table in the keyspace carries a different index + var keyspace = + keyspace( + table( + "other_table", + Set.of(CqlIdentifier.fromInternal("some_other_idx")), + Map.of( + SchemaConstants.MetadataFieldsNames.VECTOR_INDEX_PROFILES, + profilesJson("some_other_idx")))); + + assertThat(TableExtensions.removeIndexProfile(keyspace, MY_IDX, MAPPER)).isEmpty(); + } + + @Test + void emptyWhenOwningTableHasNoProfileForTheIndex() { + // the owning table has a profiles blob, but not for the index being dropped + var keyspace = + keyspace( + table( + "my_table", + Set.of(MY_IDX), + Map.of( + SchemaConstants.MetadataFieldsNames.VECTOR_INDEX_PROFILES, + profilesJson("unrelated_idx")))); + + assertThat(TableExtensions.removeIndexProfile(keyspace, MY_IDX, MAPPER)).isEmpty(); + } + + @Test + void removesProfileAndPreservesOtherProfilesAndVectorize() { + var keyspace = + keyspace( + table( + "my_table", + Set.of(MY_IDX), + Map.of( + SchemaConstants.MetadataFieldsNames.VECTOR_INDEX_PROFILES, + profilesJson("my_idx", "kept_idx"), + SchemaConstants.MetadataFieldsNames.VECTORIZE_CONFIG, + "{\"v\":{\"provider\":\"openai\",\"modelName\":\"text-embedding-3-small\"}}"))); + + var removal = TableExtensions.removeIndexProfile(keyspace, MY_IDX, MAPPER); + + assertThat(removal).isPresent(); + assertThat(removal.get().tableName()).isEqualTo(CqlIdentifier.fromInternal("my_table")); + + var customProperties = removal.get().customProperties(); + // schema type/version always written + assertThat(customProperties) + .containsKey(SchemaConstants.MetadataFieldsNames.SCHEMA_TYPE) + .containsKey(SchemaConstants.MetadataFieldsNames.SCHEMA_VERSION) + // vectorize config is read back and preserved + .containsKey(SchemaConstants.MetadataFieldsNames.VECTORIZE_CONFIG); + + // the dropped index's profile is gone, the other index's profile is kept + var profiles = + VectorIndexProfileDefinition.fromJson( + customProperties.get(SchemaConstants.MetadataFieldsNames.VECTOR_INDEX_PROFILES), + MAPPER); + assertThat(profiles).containsOnlyKeys("kept_idx"); + } + + /** Builds a {@code {index: {profile, options}}} blob for the given index keys. */ + private static String profilesJson(String... indexKeys) { + var profiles = new HashMap(); + for (var key : indexKeys) { + profiles.put(key, new VectorIndexProfileDefinition("small-high-recall", Map.of())); + } + try { + return MAPPER.writeValueAsString(profiles); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static KeyspaceMetadata keyspace(TableMetadata... tables) { + var keyspaceMetadata = mock(KeyspaceMetadata.class); + Map tableMap = new HashMap<>(); + for (var table : tables) { + tableMap.put(table.getName(), table); + } + when(keyspaceMetadata.getTables()).thenReturn(tableMap); + return keyspaceMetadata; + } + + private static TableMetadata table( + String name, Set indexNames, Map extensions) { + var tableMetadata = mock(TableMetadata.class); + when(tableMetadata.getName()).thenReturn(CqlIdentifier.fromInternal(name)); + + Map indexes = new HashMap<>(); + for (var indexName : indexNames) { + indexes.put(indexName, mock(IndexMetadata.class)); + } + when(tableMetadata.getIndexes()).thenReturn(indexes); + + Map extensionBuffers = new HashMap<>(); + extensions.forEach( + (key, value) -> + extensionBuffers.put(key, ByteBuffer.wrap(value.getBytes(StandardCharsets.UTF_8)))); + Map options = new HashMap<>(); + options.put(TableExtensions.TABLE_OPTIONS_EXTENSION_KEY, extensionBuffers); + when(tableMetadata.getOptions()).thenReturn(options); + + return tableMetadata; + } + } } diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTaskTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTaskTest.java new file mode 100644 index 0000000000..fdb75f7040 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/operation/tables/DropVectorIndexProfileDBTaskTest.java @@ -0,0 +1,53 @@ +package io.stargate.sgv2.jsonapi.service.operation.tables; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.datastax.oss.driver.api.core.CqlIdentifier; +import com.fasterxml.jackson.databind.ObjectMapper; +import io.stargate.sgv2.jsonapi.service.cqldriver.executor.TableExtensions; +import io.stargate.sgv2.jsonapi.service.cqldriver.executor.VectorIndexProfileDefinition; +import io.stargate.sgv2.jsonapi.service.operation.SchemaDBTask; +import io.stargate.sgv2.jsonapi.service.operation.keyspaces.KeyspaceDriverExceptionHandler; +import io.stargate.sgv2.jsonapi.service.schema.KeyspaceSchemaObject; +import io.stargate.sgv2.jsonapi.service.schema.SchemaObjectIdentifier; +import java.time.Duration; +import java.util.Map; +import org.junit.jupiter.api.Test; + +class DropVectorIndexProfileDBTaskTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + @Test + void buildsAlterTableExtensionsStatementForOwningTable() { + var identifier = mock(SchemaObjectIdentifier.class); + when(identifier.keyspace()).thenReturn(CqlIdentifier.fromInternal("my_ks")); + var schemaObject = mock(KeyspaceSchemaObject.class); + when(schemaObject.identifier()).thenReturn(identifier); + + var customProperties = + TableExtensions.createCustomProperties( + Map.of(), + Map.of("kept_idx", new VectorIndexProfileDefinition("small-high-recall", Map.of())), + MAPPER); + + var task = + DropVectorIndexProfileDBTask.builder(schemaObject) + .withSchemaRetryPolicy(new SchemaDBTask.SchemaRetryPolicy(1, Duration.ofMillis(1))) + .withExceptionHandlerFactory(KeyspaceDriverExceptionHandler::new) + .withTableName(CqlIdentifier.fromInternal("my_table")) + .withCustomProperties(customProperties) + .build(); + + var query = task.buildStatement().getQuery(); + + // ALTER TABLE on the owning table in the schema object's keyspace, updating the extensions map. + assertThat(query) + .contains("ALTER TABLE") + .contains("my_ks") + .contains("my_table") + .contains("extensions"); + } +}