From ac1a0d67e9ea4ba7d92977207f2e1c7ebd74dd6d Mon Sep 17 00:00:00 2001
From: mck <mck@apache.org>
Date: Mon, 29 Apr 2024 15:40:33 +0200
Subject: [PATCH] =?UTF-8?q?Avoid=20calling=20`CREATE=20KEYSPACE|TABLE=20IF?=
 =?UTF-8?q?=20NOT=20EXISTS=20=E2=80=A6`=20unnecessarily?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This also is a fix for AstraDB, which throws an exception if you try to `CREATE KEYSPACE IF NOT EXISTS …`
And use more compatible `USING 'StorageAttachedIndex'` index creation syntax.
---
 .../CassandraVectorStoreConfig.java           | 77 ++++++++++---------
 1 file changed, 40 insertions(+), 37 deletions(-)

diff --git a/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java
index 66dff35db48..528da6c053b 100644
--- a/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java
+++ b/vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java
@@ -441,7 +441,7 @@ private void ensureIndexesExists() {
 		{
 			SimpleStatement indexStmt = SchemaBuilder.createIndex(this.schema.index)
 				.ifNotExists()
-				.custom("SAI")
+				.custom("StorageAttachedIndex")
 				.onTable(this.schema.keyspace, this.schema.table)
 				.andColumn(this.schema.embedding)
 				.build();
@@ -457,7 +457,7 @@ private void ensureIndexesExists() {
 
 				SimpleStatement indexStmt = SchemaBuilder.createIndex(String.format("%s_idx", metadata.name()))
 					.ifNotExists()
-					.custom("SAI")
+					.custom("StorageAttachedIndex")
 					.onTable(this.schema.keyspace, this.schema.table)
 					.andColumn(metadata.name())
 					.build();
@@ -468,39 +468,41 @@ private void ensureIndexesExists() {
 	}
 
 	private void ensureTableExists(int vectorDimension) {
+		if (this.session.getMetadata().getKeyspace(this.schema.keyspace).get().getTable(this.schema.table).isEmpty()) {
 
-		CreateTable createTable = null;
+			CreateTable createTable = null;
 
-		CreateTableStart createTableStart = SchemaBuilder.createTable(this.schema.keyspace, this.schema.table)
-			.ifNotExists();
+			CreateTableStart createTableStart = SchemaBuilder.createTable(this.schema.keyspace, this.schema.table)
+				.ifNotExists();
 
-		for (SchemaColumn partitionKey : this.schema.partitionKeys) {
-			createTable = (null != createTable ? createTable : createTableStart).withPartitionKey(partitionKey.name,
-					partitionKey.type);
-		}
-		for (SchemaColumn clusteringKey : this.schema.clusteringKeys) {
-			createTable = createTable.withClusteringColumn(clusteringKey.name, clusteringKey.type);
-		}
+			for (SchemaColumn partitionKey : this.schema.partitionKeys) {
+				createTable = (null != createTable ? createTable : createTableStart).withPartitionKey(partitionKey.name,
+						partitionKey.type);
+			}
+			for (SchemaColumn clusteringKey : this.schema.clusteringKeys) {
+				createTable = createTable.withClusteringColumn(clusteringKey.name, clusteringKey.type);
+			}
 
-		createTable = createTable.withColumn(this.schema.content, DataTypes.TEXT);
+			createTable = createTable.withColumn(this.schema.content, DataTypes.TEXT);
 
-		for (SchemaColumn metadata : this.schema.metadataColumns) {
-			createTable = createTable.withColumn(metadata.name(), metadata.type());
-		}
-
-		// https://datastax-oss.atlassian.net/browse/JAVA-3118
-		// .withColumn(config.embedding, new DefaultVectorType(DataTypes.FLOAT,
-		// vectorDimension));
-
-		StringBuilder tableStmt = new StringBuilder(createTable.asCql());
-		tableStmt.setLength(tableStmt.length() - 1);
-		tableStmt.append(',')
-			.append(this.schema.embedding)
-			.append(" vector<float,")
-			.append(vectorDimension)
-			.append(">)");
-		logger.debug("Executing {}", tableStmt.toString());
-		this.session.execute(tableStmt.toString());
+			for (SchemaColumn metadata : this.schema.metadataColumns) {
+				createTable = createTable.withColumn(metadata.name(), metadata.type());
+			}
+
+			// https://datastax-oss.atlassian.net/browse/JAVA-3118
+			// .withColumn(config.embedding, new DefaultVectorType(DataTypes.FLOAT,
+			// vectorDimension));
+
+			StringBuilder tableStmt = new StringBuilder(createTable.asCql());
+			tableStmt.setLength(tableStmt.length() - 1);
+			tableStmt.append(',')
+				.append(this.schema.embedding)
+				.append(" vector<float,")
+				.append(vectorDimension)
+				.append(">)");
+			logger.debug("Executing {}", tableStmt.toString());
+			this.session.execute(tableStmt.toString());
+		}
 	}
 
 	private void ensureTableColumnsExist(int vectorDimension) {
@@ -563,14 +565,15 @@ private void ensureTableColumnsExist(int vectorDimension) {
 	}
 
 	private void ensureKeyspaceExists() {
+		if (this.session.getMetadata().getKeyspace(this.schema.keyspace).isEmpty()) {
+			SimpleStatement keyspaceStmt = SchemaBuilder.createKeyspace(this.schema.keyspace)
+				.ifNotExists()
+				.withSimpleStrategy(1)
+				.build();
 
-		SimpleStatement keyspaceStmt = SchemaBuilder.createKeyspace(this.schema.keyspace)
-			.ifNotExists()
-			.withSimpleStrategy(1)
-			.build();
-
-		logger.debug("Executing {}", keyspaceStmt.getQuery());
-		this.session.execute(keyspaceStmt);
+			logger.debug("Executing {}", keyspaceStmt.getQuery());
+			this.session.execute(keyspaceStmt);
+		}
 	}
 
 }