STAR-1575 port cndb-1496 and 1145b3e61f4 from DSE

* port cndb-1496

Enables overriding keypace with a prefixed keyspace.

* port 1145b3e61f4 from DSE

"Add KS override for schema altering statements (#19372)"

(cherry picked from commit 2eaf23d)
(cherry picked from commit 0ccdffa)
(cherry picked from commit a3962ad)

src/java/org/apache/cassandra/cql3/statements/BatchStatement.java

@@ -28,6 +28,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.function.UnaryOperator;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -645,6 +646,12 @@ private void setKeyspace(ClientState state) throws InvalidRequestException
                 statement.setKeyspace(state);
         }
 
+        public void setKeyspace(Function<ModificationStatement.Parsed, String> convertKeyspace) throws InvalidRequestException
+        {
+            for (ModificationStatement.Parsed statement : parsedStatements)
+                statement.setKeyspace(convertKeyspace.apply(statement));
+        }
+
         @Override
         public BatchStatement prepare(ClientState state, UnaryOperator<String> keyspaceMapper)
         {

src/java/org/apache/cassandra/cql3/statements/schema/AlterSchemaStatement.java

@@ -18,6 +18,7 @@
 package org.apache.cassandra.cql3.statements.schema;
 
 import java.util.Set;
+import java.util.function.Function;
 
 import com.google.common.collect.ImmutableSet;
 
@@ -41,7 +42,7 @@
 public abstract class AlterSchemaStatement implements CQLStatement.SingleKeyspaceCqlStatement, SchemaTransformation
 {
     private final String rawCQLStatement;
-    protected final String keyspaceName; // name of the keyspace affected by the statement
+    protected String keyspaceName; // name of the keyspace affected by the statement
 
     protected AlterSchemaStatement(String queryString, String keyspaceName)
     {
@@ -72,6 +73,11 @@ public String keyspace()
         return keyspaceName;
     }
 
+    public void overrideKeyspace(Function<String, String> overrideKeyspace)
+    {
+        this.keyspaceName = overrideKeyspace.apply(keyspaceName);
+    }
+
     public ResultMessage executeLocally(QueryState state, QueryOptions options)
     {
         return execute(state, true);