CASSANDRA-20761 - support TRUNCATE KEYSPACE to truncate all tables in it

src/antlr/Parser.g

@@ -1312,6 +1312,7 @@ dropMaterializedViewStatement returns [DropViewStatement.Raw stmt]
   */
 truncateStatement returns [TruncateStatement stmt]
     : K_TRUNCATE (K_COLUMNFAMILY)? cf=columnFamilyName { $stmt = new TruncateStatement(cf); }
+    | K_TRUNCATE K_KEYSPACE ks=keyspaceName { $stmt = new TruncateStatement(new SemiQualifiedName(ks)); }
     ;
 
 /**

src/java/org/apache/cassandra/cql3/SemiQualifiedName.java

@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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
+ *
+ *     http://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.apache.cassandra.cql3;
+
+public class SemiQualifiedName extends QualifiedName
+{
+    private SemiQualifiedName() {}
+
+    public SemiQualifiedName(String keyspace)
+    {
+        super(keyspace, null);
+    }
+
+    @Override
+    public String toString()
+    {
+        return getKeyspace();
+    }
+
+    public String toCQLString()
+    {
+        return ColumnIdentifier.maybeQuote(getKeyspace());
+    }
+}

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

@@ -17,8 +17,12 @@
  */
 package org.apache.cassandra.cql3.statements;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.TimeoutException;
 
+import com.google.common.collect.Iterables;
+
 import org.apache.cassandra.audit.AuditLogContext;
 import org.apache.cassandra.audit.AuditLogEntryType;
 import org.apache.cassandra.auth.Permission;
@@ -27,6 +31,7 @@
 import org.apache.cassandra.db.Keyspace;
 import org.apache.cassandra.db.guardrails.Guardrails;
 import org.apache.cassandra.db.virtual.VirtualKeyspaceRegistry;
+import org.apache.cassandra.db.virtual.VirtualTable;
 import org.apache.cassandra.exceptions.*;
 import org.apache.cassandra.schema.Schema;
 import org.apache.cassandra.schema.TableId;
@@ -39,6 +44,8 @@
 import org.apache.commons.lang3.builder.ToStringBuilder;
 import org.apache.commons.lang3.builder.ToStringStyle;
 
+import static java.util.stream.Collectors.joining;
+
 public class TruncateStatement extends QualifiedStatement implements CQLStatement
 {
     public TruncateStatement(QualifiedName name)
@@ -53,45 +60,85 @@ public TruncateStatement prepare(ClientState state)
 
     public void authorize(ClientState state) throws InvalidRequestException, UnauthorizedException
     {
-        state.ensureTablePermission(keyspace(), name(), Permission.MODIFY);
+        if (name() == null)
+            state.ensureKeyspacePermission(keyspace(), Permission.DROP);
+        else
+            state.ensureTablePermission(keyspace(), name(), Permission.MODIFY);
     }
 
     public void validate(ClientState state) throws InvalidRequestException
     {
-        Schema.instance.validateTable(keyspace(), name());
+        if (name() == null)
+            Schema.instance.validateKeyspace(keyspace());
+        else
+            Schema.instance.validateTable(keyspace(), name());
+
         Guardrails.dropTruncateTableEnabled.ensureEnabled(state);
     }
 
     @Override
     public ResultMessage execute(QueryState state, QueryOptions options, Dispatcher.RequestTime requestTime) throws InvalidRequestException, TruncateException
+    {
+        if (name() == null)
+        {
+            Iterable<TableMetadata> tablesIterable = Iterables.filter(Schema.instance.getTablesAndViews(keyspace()), tmd -> !tmd.isView());
+            List<TableMetadata> tablesToTruncate = new ArrayList<>();
+            for (TableMetadata t : tablesIterable)
+                tablesToTruncate.add(t);
+
+            for (TableMetadata tmd : tablesIterable)
+            {
+                try
+                {
+                    truncateOne(tmd.keyspace, tmd.name);
+                    tablesToTruncate.remove(tmd);
+                }
+                catch (TruncateException ex)
+                {
+                    throw new TruncateException(ex, "Tables not truncated: " + tablesToTruncate.stream()
+                                                                                               .map(TableMetadata::toString)
+                                                                                               .collect(joining(",")));
+                }
+            }
+        }
+        else
+            truncateOne(keyspace(), name());
+
+        return null;
+    }
+
+    private void truncateOne(String keyspace, String table)
     {
         try
         {
-            TableMetadata metaData = Schema.instance.getTableMetadata(keyspace(), name());
+            TableMetadata metaData = Schema.instance.getTableMetadata(keyspace, table);
+
+            if (metaData == null)
+                throw new InvalidRequestException(String.format("Cannot TRUNCATE %s.%s as it does not exist.", keyspace(), name()));
+
             if (metaData.isView())
                 throw new InvalidRequestException("Cannot TRUNCATE materialized view directly; must truncate base table instead");
 
             if (metaData.isVirtual())
-            {
                 executeForVirtualTable(metaData.id);
-            }
             else
-            {
-                StorageProxy.truncateBlocking(keyspace(), name());
-            }
+                StorageProxy.truncateBlocking(keyspace, table);
         }
         catch (UnavailableException | TimeoutException e)
         {
             throw new TruncateException(e);
         }
-        return null;
     }
 
     public ResultMessage executeLocally(QueryState state, QueryOptions options)
     {
         try
         {
             TableMetadata metaData = Schema.instance.getTableMetadata(keyspace(), name());
+
+            if (metaData == null)
+                throw new InvalidRequestException(String.format("Cannot TRUNCATE %s.%s as it does not exist.", keyspace(), name()));
+
             if (metaData.isView())
                 throw new InvalidRequestException("Cannot TRUNCATE materialized view directly; must truncate base table instead");
 
@@ -114,7 +161,9 @@ public ResultMessage executeLocally(QueryState state, QueryOptions options)
 
     private void executeForVirtualTable(TableId id)
     {
-        VirtualKeyspaceRegistry.instance.getTableNullable(id).truncate();
+        VirtualTable maybeVTable = VirtualKeyspaceRegistry.instance.getTableNullable(id);
+        if (maybeVTable != null)
+            maybeVTable.truncate();
     }
 
     @Override
@@ -126,6 +175,9 @@ public String toString()
     @Override
     public AuditLogContext getAuditLogContext()
     {
-        return new AuditLogContext(AuditLogEntryType.TRUNCATE, keyspace(), name());
+        if (name() == null)
+            return new AuditLogContext(AuditLogEntryType.TRUNCATE, keyspace());
+        else
+            return new AuditLogContext(AuditLogEntryType.TRUNCATE, keyspace(), name());
     }
 }

src/java/org/apache/cassandra/exceptions/TruncateException.java

@@ -19,6 +19,11 @@
 
 public class TruncateException extends RequestExecutionException
 {
+    public TruncateException(Throwable e, String msg)
+    {
+        super(ExceptionCode.TRUNCATE_ERROR, msg, e);
+    }
+
     public TruncateException(Throwable e)
     {
         super(ExceptionCode.TRUNCATE_ERROR, "Error during truncate: " + e.getMessage(), e);

src/java/org/apache/cassandra/schema/SchemaProvider.java

@@ -104,6 +104,13 @@ default TableMetadata validateTable(String keyspaceName, String tableName)
         return metadata;
     }
 
+    default KeyspaceMetadata validateKeyspace(String keyspaceName)
+    {
+        return Optional.ofNullable(getKeyspaceInstance(keyspaceName))
+               .map(Keyspace::getMetadata)
+               .orElseThrow(() -> new InvalidRequestException(String.format("Keyspace %s does not exist", keyspaceName)));
+    }
+
     default ColumnFamilyStore getColumnFamilyStoreInstance(TableId id)
     {
         TableMetadata metadata = getTableMetadata(id);