more strict quota for written bytes

src/Access/EnabledQuota.cpp

@@ -13,7 +13,7 @@ namespace DB
 {
 namespace ErrorCodes
 {
-    extern const int QUOTA_EXPIRED;
+    extern const int QUOTA_EXCEEDED;
 }
 
 
@@ -33,7 +33,7 @@ struct EnabledQuota::Impl
             "Quota for user " + backQuote(user_name) + " for " + to_string(duration) + " has been exceeded: "
                 + type_info.valueToStringWithName(used) + "/" + type_info.valueToString(max) + ". "
                 + "Interval will end at " + to_string(end_of_interval) + ". " + "Name of quota template: " + backQuote(quota_name),
-            ErrorCodes::QUOTA_EXPIRED);
+            ErrorCodes::QUOTA_EXCEEDED);
     }
 
 

src/Common/ErrorCodes.cpp

@@ -208,7 +208,7 @@
     M(198, DNS_ERROR) \
     M(199, UNKNOWN_QUOTA) \
     M(200, QUOTA_DOESNT_ALLOW_KEYS) \
-    M(201, QUOTA_EXPIRED) \
+    M(201, QUOTA_EXCEEDED) \
     M(202, TOO_MANY_SIMULTANEOUS_QUERIES) \
     M(203, NO_FREE_CONNECTION) \
     M(204, CANNOT_FSYNC) \

src/Interpreters/AsynchronousInsertQueue.cpp

@@ -198,13 +198,11 @@ void AsynchronousInsertQueue::push(ASTPtr query, ContextPtr query_context)
         copyData(*read_buf, write_buf);
     }
 
+    std::cerr << "bytes.size: " << bytes.size() << "\n";
+    std::cerr << "bytes: " << bytes << "\n";
+
     if (auto quota = query_context->getQuota())
-    {
+        quota->used(QuotaType::WRITTEN_BYTES, bytes.size());
-        /// Do not throw if quota exceded right now, because
-        /// bytes are not written now actually.
-        quota->checkExceeded(QuotaType::WRITTEN_BYTES);
-        quota->used(QuotaType::WRITTEN_BYTES, bytes.size(), /*check_exceeded=*/ false);
-    }
 
     auto entry = std::make_shared<InsertData::Entry>(std::move(bytes), query_context->getCurrentQueryId());
     InsertQuery key{query, settings};

src/Processors/Executors/ExecutionThreadContext.cpp

@@ -7,7 +7,7 @@ namespace DB
 namespace ErrorCodes
 {
     extern const int TOO_MANY_ROWS_OR_BYTES;
-    extern const int QUOTA_EXPIRED;
+    extern const int QUOTA_EXCEEDED;
     extern const int QUERY_WAS_CANCELLED;
 }
 
@@ -34,7 +34,7 @@ static bool checkCanAddAdditionalInfoToException(const DB::Exception & exception
 {
     /// Don't add additional info to limits and quota exceptions, and in case of kill query (to pass tests).
     return exception.code() != ErrorCodes::TOO_MANY_ROWS_OR_BYTES
-           && exception.code() != ErrorCodes::QUOTA_EXPIRED
+           && exception.code() != ErrorCodes::QUOTA_EXCEEDED
            && exception.code() != ErrorCodes::QUERY_WAS_CANCELLED;
 }
 

src/Processors/Transforms/CountingTransform.cpp

@@ -19,12 +19,7 @@ namespace DB
 void CountingTransform::onConsume(Chunk chunk)
 {
     if (quota)
-    {
+        quota->used(QuotaType::WRITTEN_BYTES, chunk.bytes());
-        /// Do not throw if quota exceded right now, because
-        /// bytes are not written now actually.
-        quota->checkExceeded(QuotaType::WRITTEN_BYTES);
-        quota->used(QuotaType::WRITTEN_BYTES, chunk.bytes(), /*check_exceeded=*/ false);
-    }
 
     Progress local_progress{WriteProgress(chunk.getNumRows(), chunk.bytes())};
     progress.incrementPiecewiseAtomically(local_progress);

src/Processors/Transforms/CountingTransform.h

@@ -52,6 +52,8 @@ protected:
     ProgressCallback progress_callback;
     QueryStatus * process_elem = nullptr;
     ThreadStatus * thread_status = nullptr;
+
+    /// Quota is used to limit amount of written bytes.
     std::shared_ptr<const EnabledQuota> quota;
     Chunk cur_chunk;
 };

tests/integration/test_quota/test.py

@@ -129,6 +129,7 @@ def test_quota_from_users_xml():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ]
         ]
     )
@@ -349,6 +350,7 @@ def test_tracking_quota():
                 "\\N",
                 "\\N",
                 "\\N",
+                "\\N",
             ]
         ]
     )
@@ -454,7 +456,7 @@ def test_exceed_quota():
         ]
     )
     system_quota_limits(
-        [["myQuota", 31556952, 0, 1, 1, 1, 1, 1, "\\N", 1, "\\N", "\\N"]]
+        [["myQuota", 31556952, 0, 1, 1, 1, 1, 1, "\\N", 1, "\\N", "\\N", "\\N"]]
     )
     system_quota_usage(
         [
@@ -545,6 +547,7 @@ def test_exceed_quota():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ]
         ]
     )
@@ -634,6 +637,7 @@ def test_add_remove_interval():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ]
         ]
     )
@@ -695,6 +699,7 @@ def test_add_remove_interval():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ],
             [
                 "myQuota",
@@ -709,6 +714,7 @@ def test_add_remove_interval():
                 "\\N",
                 20000,
                 120,
+                "\\N",
             ],
         ]
     )
@@ -842,6 +848,7 @@ def test_add_remove_interval():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ]
         ]
     )
@@ -1003,6 +1010,7 @@ def test_add_remove_interval():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ]
         ]
     )
@@ -1064,6 +1072,7 @@ def test_add_remove_quota():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ]
         ]
     )
@@ -1136,6 +1145,7 @@ def test_add_remove_quota():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ],
             [
                 "myQuota2",
@@ -1150,6 +1160,7 @@ def test_add_remove_quota():
                 4000,
                 400000,
                 60,
+                "\\N",
             ],
             [
                 "myQuota2",
@@ -1164,6 +1175,7 @@ def test_add_remove_quota():
                 "\\N",
                 "\\N",
                 1800,
+                "\\N",
             ],
         ]
     )
@@ -1226,6 +1238,7 @@ def test_add_remove_quota():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ]
         ]
     )
@@ -1294,6 +1307,7 @@ def test_add_remove_quota():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ]
         ]
     )
@@ -1356,6 +1370,7 @@ def test_reload_users_xml_by_timer():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ]
         ]
     )
@@ -1382,7 +1397,7 @@ def test_reload_users_xml_by_timer():
     assert_eq_with_retry(
         instance,
         "SELECT * FROM system.quota_limits",
-        [["myQuota", 31556952, 0, 1, 1, 1, 1, 1, "\\N", 1, "\\N", "\\N"]],
+        [["myQuota", 31556952, 0, 1, 1, 1, 1, 1, "\\N", 1, "\\N", "\\N", "\\N"]],
     )
 
 
@@ -1481,15 +1496,15 @@ def test_dcl_management():
         == "CREATE QUOTA qA FOR INTERVAL 30 minute MAX execution_time = 0.5, FOR INTERVAL 5 quarter MAX queries = 321, errors = 10 TO default\n"
     )
     assert re.match(
-        "qA\\t\\t.*\\t1800\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t.*\\t0.5\n"
+        "qA\\t\\t.*\\t1800\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t.*\\t0.5\\t0\\t\\\\N\n"
-        "qA\\t\\t.*\\t39446190\\t1\\t321\\t1\\t\\\\N\\t0\\t\\\\N\\t0\\t10\\t50\\t\\\\N\\t200\\t\\\\N\\t50\\t\\\\N\\t200\\t\\\\N\\t.*\\t\\\\N\n",
+        "qA\\t\\t.*\\t39446190\\t1\\t321\\t1\\t\\\\N\\t0\\t\\\\N\\t0\\t10\\t50\\t\\\\N\\t200\\t\\\\N\\t50\\t\\\\N\\t200\\t\\\\N\\t.*\\t\\\\N\\t0\\t\\\\N\n",
         instance.query("SHOW QUOTA"),
     )
 
     instance.query("SELECT * from test_table")
     assert re.match(
-        "qA\\t\\t.*\\t1800\\t1\\t\\\\N\\t1\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t50\\t\\\\N\\t200\\t\\\\N\\t50\\t\\\\N\\t200\\t\\\\N\\t.*\\t0.5\n"
+        "qA\\t\\t.*\\t1800\\t1\\t\\\\N\\t1\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t50\\t\\\\N\\t200\\t\\\\N\\t50\\t\\\\N\\t200\\t\\\\N\\t.*\\t0.5\\t0\\t\\\\N\n"
-        "qA\\t\\t.*\\t39446190\\t2\\t321\\t2\\t\\\\N\\t0\\t\\\\N\\t0\\t10\\t100\\t\\\\N\\t400\\t\\\\N\\t100\\t\\\\N\\t400\\t\\\\N\\t.*\\t\\\\N\n",
+        "qA\\t\\t.*\\t39446190\\t2\\t321\\t2\\t\\\\N\\t0\\t\\\\N\\t0\\t10\\t100\\t\\\\N\\t400\\t\\\\N\\t100\\t\\\\N\\t400\\t\\\\N\\t.*\\t\\\\N\\t0\\t\\\\N\n",
         instance.query("SHOW QUOTA"),
     )
 
@@ -1503,7 +1518,7 @@ def test_dcl_management():
 
     instance.query("SELECT * from test_table")
     assert re.match(
-        "qA\\t\\t.*\\t42075936\\t1\\t\\\\N\\t1\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t50\\t\\\\N\\t200\\t\\\\N\\t50\\t\\\\N\\t200\\t\\\\N\\t.*\\t\\\\N\n",
+        "qA\\t\\t.*\\t42075936\\t1\\t\\\\N\\t1\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t50\\t\\\\N\\t200\\t\\\\N\\t50\\t\\\\N\\t200\\t\\\\N\\t.*\\t\\\\N\\t0\\t\\\\N\n",
         instance.query("SHOW QUOTA"),
     )
 
@@ -1519,7 +1534,7 @@ def test_dcl_management():
 
     instance.query("SELECT * from test_table")
     assert re.match(
-        "qB\\t\\t.*\\t42075936\\t2\\t\\\\N\\t2\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t100\\t\\\\N\\t400\\t\\\\N\\t100\\t\\\\N\\t400\\t\\\\N\\t.*\\t\\\\N\n",
+        "qB\\t\\t.*\\t42075936\\t2\\t\\\\N\\t2\\t\\\\N\\t0\\t\\\\N\\t0\\t\\\\N\\t100\\t\\\\N\\t400\\t\\\\N\\t100\\t\\\\N\\t400\\t\\\\N\\t.*\\t\\\\N\\t0\\t\\\\N\n",
         instance.query("SHOW QUOTA"),
     )
 
@@ -1563,6 +1578,7 @@ def test_query_inserts():
                 1000,
                 "\\N",
                 "\\N",
+                "\\N",
             ]
         ]
     )

tests/queries/0_stateless/02246_async_insert_quota.reference

@@ -1,2 +1,2 @@
-QUOTA_EXPIRED
+QUOTA_EXCEEDED
 2

tests/queries/0_stateless/02246_async_insert_quota.sh

@@ -20,7 +20,7 @@ ${CLICKHOUSE_CLIENT} -q "CREATE QUOTA q02246 FOR INTERVAL 1 HOUR MAX QUERY INSER
 
 ${CLICKHOUSE_CLIENT} --user u02246 --async_insert 1 -q "INSERT INTO async_inserts_02246 VALUES (1, 'a')"
 ${CLICKHOUSE_CLIENT} --user u02246 --async_insert 1 -q "INSERT INTO async_inserts_02246 VALUES (2, 'b')"
-${CLICKHOUSE_CLIENT} --user u02246 --async_insert 1 -q "INSERT INTO async_inserts_02246 VALUES (3, 'c')" 2>&1 | grep -m1 -o QUOTA_EXPIRED
+${CLICKHOUSE_CLIENT} --user u02246 --async_insert 1 -q "INSERT INTO async_inserts_02246 VALUES (3, 'c')" 2>&1 | grep -m1 -o QUOTA_EXCEEDED
 
 sleep 1.0
 

tests/queries/0_stateless/02247_written_bytes_quota.reference

@@ -1,9 +1,7 @@
-QUOTA_EXPIRED
+QUOTA_EXCEEDED
-QUOTA_EXPIRED
+QUOTA_EXCEEDED
 1
 2
-QUOTA_EXPIRED
+QUOTA_EXCEEDED
-QUOTA_EXPIRED
-QUOTA_EXPIRED
 1
 50

tests/queries/0_stateless/02247_written_bytes_quota.sh

@@ -16,25 +16,22 @@ ${CLICKHOUSE_CLIENT} -q "CREATE ROLE r02247"
 ${CLICKHOUSE_CLIENT} -q "CREATE USER u02247"
 ${CLICKHOUSE_CLIENT} -q "GRANT ALL ON *.* TO r02247"
 ${CLICKHOUSE_CLIENT} -q "GRANT r02247 to u02247"
-${CLICKHOUSE_CLIENT} -q "CREATE QUOTA q02247 FOR INTERVAL 100 YEAR MAX WRITTEN BYTES = 10 TO r02247"
+${CLICKHOUSE_CLIENT} -q "CREATE QUOTA q02247 FOR INTERVAL 100 YEAR MAX WRITTEN BYTES = 25 TO r02247"
 
-${CLICKHOUSE_CLIENT} --user u02247 --async_insert 1 -q "INSERT INTO written_bytes_02247 VALUES ('qwqwqw')"
+${CLICKHOUSE_CLIENT} --user u02247 --async_insert 1 -q "INSERT INTO written_bytes_02247 VALUES ('qwqw')"
-${CLICKHOUSE_CLIENT} --user u02247 --async_insert 0 -q "INSERT INTO written_bytes_02247 VALUES ('qwqwqw')"
+${CLICKHOUSE_CLIENT} --user u02247 --async_insert 0 -q "INSERT INTO written_bytes_02247 VALUES ('qwqw')"
-${CLICKHOUSE_CLIENT} --user u02247 --async_insert 1 -q "INSERT INTO written_bytes_02247 VALUES ('qwqwqw')" 2>&1 | grep -m1 -o QUOTA_EXPIRED
+${CLICKHOUSE_CLIENT} --user u02247 --async_insert 1 -q "INSERT INTO written_bytes_02247 VALUES ('qwqw')" 2>&1 | grep -m1 -o QUOTA_EXCEEDED
-${CLICKHOUSE_CLIENT} --user u02247 --async_insert 0 -q "INSERT INTO written_bytes_02247 VALUES ('qwqwqw')" 2>&1 | grep -m1 -o QUOTA_EXPIRED
+${CLICKHOUSE_CLIENT} --user u02247 --async_insert 0 -q "INSERT INTO written_bytes_02247 VALUES ('qwqw')" 2>&1 | grep -m1 -o QUOTA_EXCEEDED
 
 ${CLICKHOUSE_CLIENT} -q "SELECT written_bytes > 10 FROM system.quotas_usage WHERE quota_name = 'q02247'"
 ${CLICKHOUSE_CLIENT} -q "SELECT count() FROM written_bytes_02247"
 
 ${CLICKHOUSE_CLIENT} -q "DROP QUOTA q02247"
-${CLICKHOUSE_CLIENT} -q "CREATE QUOTA q02247 FOR INTERVAL 100 YEAR MAX WRITTEN BYTES = 100 TO r02247"
+${CLICKHOUSE_CLIENT} -q "CREATE QUOTA q02247 FOR INTERVAL 100 YEAR MAX WRITTEN BYTES = 1000 TO r02247"
 ${CLICKHOUSE_CLIENT} -q "TRUNCATE TABLE written_bytes_02247"
 
 ${CLICKHOUSE_CLIENT} --user u02247 -q "INSERT INTO written_bytes_02247 SELECT toString(number) FROM numbers(50)"
-
+${CLICKHOUSE_CLIENT} --user u02247 -q "INSERT INTO written_bytes_02247 SELECT toString(number) FROM numbers(100)" 2>&1 | grep -m1 -o QUOTA_EXCEEDED
-${CLICKHOUSE_CLIENT} --user u02247 -q "INSERT INTO written_bytes_02247 SELECT toString(number) FROM numbers(1)" 2>&1 | grep -m1 -o QUOTA_EXPIRED
-${CLICKHOUSE_CLIENT} --user u02247 --async_insert 1 -q "INSERT INTO written_bytes_02247 VALUES ('qwqwqw')" 2>&1 | grep -m1 -o QUOTA_EXPIRED
-${CLICKHOUSE_CLIENT} --user u02247 --async_insert 0 -q "INSERT INTO written_bytes_02247 VALUES ('qwqwqw')" 2>&1 | grep -m1 -o QUOTA_EXPIRED
 
 ${CLICKHOUSE_CLIENT} -q "SELECT written_bytes > 100 FROM system.quotas_usage WHERE quota_name = 'q02247'"
 ${CLICKHOUSE_CLIENT} -q "SELECT count() FROM written_bytes_02247"