Polishing.

Introduce configuration option to disable schema metadata during schema actions. Tweak method naming. Apply schema metadata suspension also to destroy methods.

Related ticket: #990
Related ticket: #1253

Original pull request: #1255.

Author: mp911de

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/config/CqlSessionFactoryBean.java

@@ -23,6 +23,7 @@
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.CompletionStage;
 import java.util.function.IntFunction;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -57,6 +58,7 @@
 
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.core.CqlSessionBuilder;
+import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
 
 /**
  * Factory for creating and configuring a Cassandra {@link CqlSession}, which is a thread-safe singleton. As such, it is
@@ -108,6 +110,8 @@ public class CqlSessionFactoryBean
 
 	private SchemaAction schemaAction = SchemaAction.NONE;
 
+	private boolean suspendLifecycleSchemaRefresh = false;
+
 	private @Nullable SessionBuilderConfigurer sessionBuilderConfigurer;
 
 	private IntFunction<Collection<InetSocketAddress>> contactPoints = port -> createInetSocketAddresses(
@@ -348,8 +352,8 @@ protected CassandraMappingContext getMappingContext() {
 	 * Set the {@link SchemaAction}.
 	 *
 	 * @param schemaAction must not be {@literal null}.
-	 * @deprecated Use {@link CassandraSessionFactoryBean} with
-	 *             {@link CassandraSessionFactoryBean#setSchemaAction(SchemaAction)} instead.
+	 * @deprecated Use {@link SessionFactoryFactoryBean} with
+	 *             {@link SessionFactoryFactoryBean#setSchemaAction(SchemaAction)} instead.
 	 */
 	@Deprecated
 	public void setSchemaAction(SchemaAction schemaAction) {
@@ -366,6 +370,24 @@ public SchemaAction getSchemaAction() {
 		return this.schemaAction;
 	}
 
+	/**
+	 * Set whether to suspend schema refresh settings during {@link #afterPropertiesSet()} and {@link #destroy()}
+	 * lifecycle callbacks. Disabled by default to use schema metadata settings of the session configuration. When enabled
+	 * (set to {@code true}), then schema refresh during lifecycle methods is suspended until finishing schema actions to
+	 * avoid periodic schema refreshes for each DDL statement.
+	 * <p>
+	 * Suspending schema refresh can be useful to delay schema agreement until the entire schema is created. Note that
+	 * disabling schema refresh may interfere with schema actions. {@link SchemaAction#RECREATE_DROP_UNUSED} and
+	 * mapping-based schema creation rely on schema metadata.
+	 *
+	 * @param suspendLifecycleSchemaRefresh {@code true} to suspend the schema refresh during lifecycle callbacks;
+	 *          {@code false} otherwise to retain the session schema refresh configuration.
+	 * @since 2.7
+	 */
+	public void setSuspendLifecycleSchemaRefresh(boolean suspendLifecycleSchemaRefresh) {
+		this.suspendLifecycleSchemaRefresh = suspendLifecycleSchemaRefresh;
+	}
+
 	/**
 	 * Returns a reference to the connected Cassandra {@link CqlSession}.
 	 *
@@ -455,20 +477,30 @@ public void afterPropertiesSet() {
 
 		this.session = buildSession(sessionBuilder);
 
-		try {
-			SchemaRefreshUtils.withDisabledSchema(this.session, () -> {
-				executeCql(getStartupScripts().stream(), this.session);
-				performSchemaAction();
-			});
-		} catch (RuntimeException e) {
-			throw e;
-		} catch (Exception e) {
-			throw new IllegalStateException("Unexpected checked exception thrown", e);
+		initializeSchema(this.systemSession, this.session);
+	}
+
+	private void initializeSchema(CqlSession systemSession, CqlSession session) {
+
+		Runnable schemaActionRunnable = () -> {
+
+			executeCql(getStartupScripts().stream(), session);
+			performSchemaAction();
+		};
+
+		List<CompletionStage<?>> futures = new ArrayList<>(2);
+
+		if (this.suspendLifecycleSchemaRefresh) {
+			futures.add(SchemaUtils.withSuspendedAsyncSchemaRefresh(session, schemaActionRunnable));
+		} else {
+			futures.add(SchemaUtils.withAsyncSchemaRefresh(session, schemaActionRunnable));
 		}
 
-		if (this.systemSession.isSchemaMetadataEnabled()) {
-			this.systemSession.refreshSchema();
+		if (systemSession.isSchemaMetadataEnabled()) {
+			futures.add(systemSession.refreshSchemaAsync());
 		}
+
+		futures.forEach(CompletableFutures::getUninterruptibly);
 	}
 
 	protected CqlSessionBuilder buildBuilder() {
@@ -536,7 +568,15 @@ private void initializeCluster(CqlSession session) {
 		keyspaceStartupSpecifications.addAll(this.keyspaceCreations);
 		keyspaceStartupSpecifications.addAll(this.keyspaceAlterations);
 
-		executeSpecificationsAndScripts(keyspaceStartupSpecifications, this.keyspaceStartupScripts, session);
+		Runnable schemaActionRunnable = () -> {
+			executeSpecificationsAndScripts(keyspaceStartupSpecifications, this.keyspaceStartupScripts, session);
+		};
+
+		if (this.suspendLifecycleSchemaRefresh) {
+			SchemaUtils.withSuspendedAsyncSchemaRefresh(session, schemaActionRunnable);
+		} else {
+			schemaActionRunnable.run();
+		}
 	}
 
 	/**
@@ -564,7 +604,7 @@ private void generateSpecifications(Collection<KeyspaceActionSpecification> spec
 	}
 
 	/**
-	 * Perform the configure {@link SchemaAction} using {@link CassandraMappingContext} metadata.
+	 * Perform the configured {@link SchemaAction} using {@link CassandraMappingContext} metadata.
 	 */
 	protected void performSchemaAction() {
 
@@ -664,8 +704,23 @@ public DataAccessException translateExceptionIfPossible(RuntimeException e) {
 	public void destroy() {
 
 		if (this.session != null) {
-			executeCql(getShutdownScripts().stream(), this.session);
-			executeSpecificationsAndScripts(this.keyspaceDrops, this.keyspaceShutdownScripts, this.systemSession);
+
+			Runnable schemaActionRunnable = () -> {
+				executeCql(getShutdownScripts().stream(), this.session);
+			};
+
+			Runnable systemSchemaActionRunnable = () -> {
+				executeSpecificationsAndScripts(this.keyspaceDrops, this.keyspaceShutdownScripts, this.systemSession);
+			};
+
+			if (this.suspendLifecycleSchemaRefresh) {
+				SchemaUtils.withSuspendedAsyncSchemaRefresh(this.session, schemaActionRunnable);
+				SchemaUtils.withSuspendedAsyncSchemaRefresh(this.systemSession, systemSchemaActionRunnable);
+			} else {
+				schemaActionRunnable.run();
+				systemSchemaActionRunnable.run();
+			}
+
 			closeSession();
 			closeSystemSession();
 		}

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/config/SchemaRefreshUtils.java

@@ -0,0 +1,104 @@
+/*
+ * Copyright 2022 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      https://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.
+ */
+package org.springframework.data.cassandra.config;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+
+import com.datastax.oss.driver.api.core.session.Session;
+import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
+
+/**
+ * Utility methods for executing schema actions.
+ *
+ * @author Ammar Khaku
+ * @author Mark Paluch
+ * @since 2.7
+ */
+class SchemaUtils {
+
+	/**
+	 * Programmatically disables schema refresh on the session and runs the provided {@link Runnable}. Takes care to
+	 * restore the previous state of schema refresh on the provided session. Note that the session could have had schema
+	 * refreshes enabled/disabled either programmatically or via config.
+	 *
+	 * @param session the session to use.
+	 * @param schemaAction the runnable code block.
+	 */
+	static void withSuspendedSchemaRefresh(Session session, Runnable schemaAction) {
+		CompletableFutures.getUninterruptibly(withSuspendedAsyncSchemaRefresh(session, schemaAction));
+	}
+
+	/**
+	 * Programmatically disables schema refresh on the session and runs the provided {@link Runnable}. Takes care to
+	 * restore the previous state of schema refresh on the provided session. Note that the session could have had schema
+	 * refreshes enabled/disabled either programmatically or via config.
+	 *
+	 * @param session the session to use.
+	 * @param schemaAction the runnable code block.
+	 * @return a {@link CompletionStage} providing a handle to the schema refresh completion.
+	 */
+	static CompletionStage<?> withSuspendedAsyncSchemaRefresh(Session session, Runnable schemaAction) {
+
+		boolean schemaEnabledPreviously = session.isSchemaMetadataEnabled();
+		if (schemaEnabledPreviously) {
+			session.setSchemaMetadataEnabled(false);
+		}
+
+		CompletionStage<?> schemaRefresh;
+
+		try {
+			schemaAction.run();
+		} finally {
+			if (schemaEnabledPreviously) {
+				// user may have set it programmatically so set it back programmatically
+				schemaRefresh = session.setSchemaMetadataEnabled(null);
+			} else {
+				schemaRefresh = CompletableFuture.completedFuture(null);
+			}
+		}
+
+		return schemaRefresh;
+	}
+
+	/**
+	 * Run a {@link Runnable} and refresh the schema after finishing the runnable.
+	 *
+	 * @param session the session to use.
+	 * @param schemaAction the runnable code block.
+	 */
+	static void withSchemaRefresh(Session session, Runnable schemaAction) {
+		CompletableFutures.getUninterruptibly(withAsyncSchemaRefresh(session, schemaAction));
+	}
+
+	/**
+	 * Run a {@link Runnable} and refresh the schema after finishing the runnable.
+	 *
+	 * @param session the session to use.
+	 * @param schemaAction the runnable code block.
+	 * @return a {@link CompletionStage} providing a handle to the schema refresh completion.
+	 */
+	static CompletionStage<?> withAsyncSchemaRefresh(Session session, Runnable schemaAction) {
+
+		schemaAction.run();
+
+		if (session.isSchemaMetadataEnabled()) {
+			return session.refreshSchemaAsync();
+		}
+
+		return CompletableFuture.completedFuture(null);
+	}
+}