cql3, storage_proxy: add support for TRUNCATE USING TIMEOUT

Extend the cql3 truncate statement to accept attributes,
similar to modification statements.

To achieve that we define cql3::statements::raw::truncate_statement
derived from raw::cf_statement, and implement its pure virtual
prepare() method to make a prepared truncate_statement.

The latter, statements::truncate_statement, is no longer derived
from raw::cf_statement, and just stores a schema_ptr to get to the
keyspace and column_family names.

`test_truncate_using_timeout` cql-pytest was added to test
the new USING TIMEOUT feature.

Fixes #11408

Signed-off-by: Benny Halevy <bhalevy@scylladb.com>

cql3/Cql.g

@@ -44,7 +44,7 @@ options {
 #include "cql3/statements/drop_aggregate_statement.hh"
 #include "cql3/statements/drop_service_level_statement.hh"
 #include "cql3/statements/detach_service_level_statement.hh"
-#include "cql3/statements/truncate_statement.hh"
+#include "cql3/statements/raw/truncate_statement.hh"
 #include "cql3/statements/raw/update_statement.hh"
 #include "cql3/statements/raw/insert_statement.hh"
 #include "cql3/statements/raw/delete_statement.hh"
@@ -1072,10 +1072,18 @@ dropIndexStatement returns [std::unique_ptr<drop_index_statement> expr]
     ;
 
 /**
-  * TRUNCATE <CF>;
+  * TRUNCATE [TABLE] <CF>
+  * [USING TIMEOUT <duration>];
   */
-truncateStatement returns [std::unique_ptr<truncate_statement> stmt]
-    : K_TRUNCATE (K_COLUMNFAMILY)? cf=columnFamilyName { $stmt = std::make_unique<truncate_statement>(cf); }
+truncateStatement returns [std::unique_ptr<raw::truncate_statement> stmt]
+    @init {
+        auto attrs = std::make_unique<cql3::attributes::raw>();
+    }
+    : K_TRUNCATE (K_COLUMNFAMILY)? cf=columnFamilyName
+      ( usingTimeoutClause[attrs] )?
+      {
+        $stmt = std::make_unique<raw::truncate_statement>(std::move(cf), std::move(attrs));
+      }
     ;
 
 /**

cql3/statements/raw/modification_statement.hh

@@ -39,7 +39,7 @@ private:
     const bool _if_not_exists;
     const bool _if_exists;
 protected:
-    modification_statement(cf_name name, std::unique_ptr<attributes::raw> attrs, conditions_vector conditions, bool if_not_exists, bool if_exists);
+    modification_statement(cf_name name, std::unique_ptr<attributes::raw> attrs, conditions_vector conditions = {}, bool if_not_exists = false, bool if_exists = false);
 
 public:
     virtual std::unique_ptr<prepared_statement> prepare(data_dictionary::database db, cql_stats& stats) override;

cql3/statements/raw/truncate_statement.hh

@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2022-present ScyllaDB
+ *
+ * Modified by ScyllaDB
+ */
+
+/*
+ * SPDX-License-Identifier: (AGPL-3.0-or-later and Apache-2.0)
+ */
+
+#pragma once
+
+#include "cql3/statements/raw/cf_statement.hh"
+#include "cql3/attributes.hh"
+
+namespace cql3 {
+
+namespace statements {
+
+namespace raw {
+
+class truncate_statement : public raw::cf_statement {
+private:
+    std::unique_ptr<attributes::raw> _attrs;
+public:
+    /**
+     * Creates a new truncate_statement from a column family name, and attributes.
+     *
+     * @param name column family being operated on
+     * @param attrs additional attributes for statement (timeout)
+     */
+    truncate_statement(cf_name name, std::unique_ptr<attributes::raw> attrs);
+
+    virtual std::unique_ptr<prepared_statement> prepare(data_dictionary::database db, cql_stats& stats) override;
+};
+
+} // namespace raw
+
+} // namespace statements
+
+} // namespace cql3

cql3/statements/truncate_statement.cc

@@ -8,6 +8,7 @@
  * SPDX-License-Identifier: (AGPL-3.0-or-later and Apache-2.0)
  */
 
+#include "cql3/statements/raw/truncate_statement.hh"
 #include "cql3/statements/truncate_statement.hh"
 #include "cql3/statements/prepared_statement.hh"
 #include "cql3/cql_statement.hh"
@@ -15,15 +16,54 @@
 #include "cql3/query_processor.hh"
 #include "service/storage_proxy.hh"
 #include <optional>
+#include "validation.hh"
 
 namespace cql3 {
 
 namespace statements {
 
-truncate_statement::truncate_statement(cf_name name)
-    : cf_statement{std::move(name)}
-    , cql_statement_no_metadata(&timeout_config::truncate_timeout)
+namespace raw {
+
+truncate_statement::truncate_statement(cf_name name, std::unique_ptr<attributes::raw> attrs)
+        : cf_statement(std::move(name))
+        , _attrs(std::move(attrs))
 {
+    // Validate the attributes.
+    // Currently, TRUNCATE supports only USING TIMEOUT
+    assert(!_attrs->timestamp.has_value());
+    assert(!_attrs->time_to_live.has_value());
+}
+
+std::unique_ptr<prepared_statement> truncate_statement::prepare(data_dictionary::database db, cql_stats& stats) {
+    schema_ptr schema = validation::validate_column_family(db, keyspace(), column_family());
+    auto prepared_attributes = _attrs->prepare(db, keyspace(), column_family());
+    auto ctx = get_prepare_context();
+    prepared_attributes->fill_prepare_context(ctx);
+    auto stmt = ::make_shared<cql3::statements::truncate_statement>(std::move(schema), std::move(prepared_attributes));
+    return std::make_unique<prepared_statement>(std::move(stmt));
+}
+
+} // namespace raw
+
+truncate_statement::truncate_statement(schema_ptr schema, std::unique_ptr<attributes> prepared_attrs)
+    : cql_statement_no_metadata(&timeout_config::truncate_timeout)
+    , _schema{std::move(schema)}
+    , _attrs(std::move(prepared_attrs))
+{
+}
+
+truncate_statement::truncate_statement(const truncate_statement& ts)
+    : cql_statement_no_metadata(ts)
+    , _schema(ts._schema)
+    , _attrs(std::make_unique<attributes>(*ts._attrs))
+{ }
+
+const sstring& truncate_statement::keyspace() const {
+    return _schema->ks_name();
+}
+
+const sstring& truncate_statement::column_family() const {
+    return _schema->cf_name();
 }
 
 uint32_t truncate_statement::get_bound_terms() const
@@ -31,11 +71,6 @@ uint32_t truncate_statement::get_bound_terms() const
     return 0;
 }
 
-std::unique_ptr<prepared_statement> truncate_statement::prepare(data_dictionary::database db,cql_stats& stats)
-{
-    return std::make_unique<prepared_statement>(::make_shared<truncate_statement>(*this));
-}
-
 bool truncate_statement::depends_on(std::string_view ks_name, std::optional<std::string_view> cf_name) const
 {
     return false;
@@ -60,13 +95,18 @@ truncate_statement::execute(query_processor& qp, service::query_state& state, co
     if (qp.db().find_schema(keyspace(), column_family())->is_view()) {
         throw exceptions::invalid_request_exception("Cannot TRUNCATE materialized view directly; must truncate base table instead");
     }
-    return qp.proxy().truncate_blocking(keyspace(), column_family()).handle_exception([](auto ep) {
+    auto timeout_in_ms = std::chrono::duration_cast<std::chrono::milliseconds>(get_timeout(state.get_client_state(), options));
+    return qp.proxy().truncate_blocking(keyspace(), column_family(), timeout_in_ms).handle_exception([](auto ep) {
         throw exceptions::truncate_exception(ep);
     }).then([] {
         return ::shared_ptr<cql_transport::messages::result_message>{};
     });
 }
 
+db::timeout_clock::duration truncate_statement::get_timeout(const service::client_state& state, const query_options& options) const {
+    return _attrs->is_timeout_set() ? _attrs->get_timeout(options) : state.get_timeout_config().truncate_timeout;
+}
+
 }
 
 }

cql3/statements/truncate_statement.hh

@@ -12,6 +12,7 @@
 
 #include "cql3/statements/raw/cf_statement.hh"
 #include "cql3/cql_statement.hh"
+#include "cql3/attributes.hh"
 
 namespace cql3 {
 
@@ -19,14 +20,20 @@ class query_processor;
 
 namespace statements {
 
-class truncate_statement : public raw::cf_statement, public cql_statement_no_metadata {
+class truncate_statement : public cql_statement_no_metadata {
+    schema_ptr _schema;
+    const std::unique_ptr<attributes> _attrs;
 public:
-    truncate_statement(cf_name name);
+    truncate_statement(schema_ptr schema, std::unique_ptr<attributes> prepared_attrs);
+    truncate_statement(const truncate_statement&);
+    truncate_statement(truncate_statement&&) = default;
+
+    const sstring& keyspace() const;
+
+    const sstring& column_family() const;
 
     virtual uint32_t get_bound_terms() const override;
 
-    virtual std::unique_ptr<prepared_statement> prepare(data_dictionary::database db, cql_stats& stats) override;
-
     virtual bool depends_on(std::string_view ks_name, std::optional<std::string_view> cf_name) const override;
 
     virtual future<> check_access(query_processor& qp, const service::client_state& state) const override;
@@ -35,6 +42,8 @@ public:
 
     virtual future<::shared_ptr<cql_transport::messages::result_message>>
     execute(query_processor& qp, service::query_state& state, const query_options& options) const override;
+private:
+    db::timeout_clock::duration get_timeout(const service::client_state& state, const query_options& options) const;
 };
 
 }

service/storage_proxy.cc

@@ -322,7 +322,7 @@ public:
         return ser::storage_proxy_rpc_verbs::send_paxos_prune(&_ms, addr, timeout, schema_id, key, ballot, tracing::make_trace_info(tr_state));
     }
 
-    future<> send_truncate_blocking(sstring keyspace, sstring cfname) {
+    future<> send_truncate_blocking(sstring keyspace, sstring cfname, std::optional<std::chrono::milliseconds> timeout_in_ms) {
         slogger.debug("Starting a blocking truncate operation on keyspace {}, CF {}", keyspace, cfname);
 
         if (!_gossiper.get_unreachable_token_owners().empty()) {
@@ -338,7 +338,7 @@ public:
         }
 
         auto all_endpoints = _gossiper.get_live_token_owners();
-        auto timeout = clock_type::now() + std::chrono::milliseconds(_sp._db.local().get_config().truncate_request_timeout_in_ms());
+        auto timeout = clock_type::now() + timeout_in_ms.value_or(std::chrono::milliseconds(_sp._db.local().get_config().truncate_request_timeout_in_ms()));
 
         slogger.trace("Enqueuing truncate messages to hosts {}", all_endpoints);
 
@@ -5896,8 +5896,8 @@ db::hints::manager& storage_proxy::hints_manager_for(db::write_type type) {
     return type == db::write_type::VIEW ? _hints_for_views_manager : _hints_manager;
 }
 
-future<> storage_proxy::truncate_blocking(sstring keyspace, sstring cfname) {
-    return remote().send_truncate_blocking(std::move(keyspace), std::move(cfname));
+future<> storage_proxy::truncate_blocking(sstring keyspace, sstring cfname, std::optional<std::chrono::milliseconds> timeout_in_ms) {
+    return remote().send_truncate_blocking(std::move(keyspace), std::move(cfname), timeout_in_ms);
 }
 
 void storage_proxy::init_messaging_service(migration_manager* mm) {

service/storage_proxy.hh

@@ -585,8 +585,9 @@ public:
      * the column family cfname
      * @param keyspace
      * @param cfname
+     * @param timeout (default: use truncate_request_timeout_in_ms config)
      */
-    future<> truncate_blocking(sstring keyspace, sstring cfname);
+    future<> truncate_blocking(sstring keyspace, sstring cfname, std::optional<std::chrono::milliseconds> timeout_in_ms = std::nullopt);
 
     /*
      * Executes data query on the whole cluster.

test/cql-pytest/test_using_timeout.py

@@ -7,6 +7,7 @@
 from util import new_test_keyspace, unique_name, unique_key_int
 import pytest
 from cassandra.protocol import InvalidRequest, ReadTimeout, WriteTimeout, SyntaxException
+from cassandra.cluster import NoHostAvailable
 from cassandra.util import Duration
 
 def r(regex):
@@ -157,3 +158,21 @@ def test_invalid_timeout(scylla_only, cql, table1):
     invalid_syntax(f"SELECT * FROM {table} USING TIMEOUT 60s AND TTL 10000")
     invalid_syntax(f"SELECT * FROM {table} USING TIMEOUT 60s AND TTL 123 AND TIMESTAMP 911")
     invalid_syntax(f"DELETE FROM {table} USING TIMEOUT 60s AND TTL 42 WHERE p = 42")
+
+def test_truncate_using_timeout(scylla_only, cql, table1):
+    table = table1
+    key = unique_key_int()
+    cql.execute(f"INSERT INTO {table} (p,c,v) VALUES ({key},1,1)")
+    res = list(cql.execute(f"SELECT * FROM {table} WHERE p = {key} and c = 1"))
+    assert len(res) == 1
+    cql.execute(f"TRUNCATE TABLE {table} USING TIMEOUT 1000s")
+    res = list(cql.execute(f"SELECT * FROM {table} WHERE p = {key} and c = 1"))
+    assert len(res) == 0
+    with pytest.raises(NoHostAvailable):
+        cql.execute(f"TRUNCATE TABLE {table} USING TIMEOUT 0s")
+    with pytest.raises(SyntaxException):
+        cql.execute(f"TRUNCATE TABLE {table} USING TTL 1")
+    with pytest.raises(SyntaxException):
+        cql.execute(f"TRUNCATE TABLE {table} USING TIMESTAMP 123456789")
+    with pytest.raises(SyntaxException):
+        cql.execute(f"TRUNCATE TABLE {table} USING TIMEOUT 1h AND TTL 42")