Avoid calling CREATE KEYSPACE|TABLE IF NOT EXISTS … unnecessarily

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.
spring-projects · May 4, 2024 · ac1a0d6 · ac1a0d6
1 parent 3c40268
commit ac1a0d6
Showing 1 changed file with 40 additions and 37 deletions.
diff --git a/...assandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java b/...assandra/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);
+		}
 	}
 
 }