Add `DBSIZE` query

actions.jsonc

@@ -111,5 +111,13 @@
         "args": "SUPDATE <key1> <value1> <key2> <value2> ...",
         "desc": "Update all keys if all of the keys exist",
         "return": "(Code: 0) if all keys were updated, otherwise (Code: 1)"
+    },
+    {
+        "name": "DBSIZE",
+        "since": "0.4.3",
+        "complexity": "O(1)",
+        "args": "DBSIZE",
+        "desc": "Returns the number of key/value pairs stored in the database",
+        "return": "Number of keys that exist in the database as an unsigned int"
     }
 ]

server/src/kvengine/dbsize.rs

@@ -0,0 +1,41 @@
+/*
+ * Created on Thu Sep 24 2020
+ *
+ * This file is a part of TerrabaseDB
+ * Copyright (c) 2020, Sayan Nandan <ohsayan at outlook dot com>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program. If not, see <https://www.gnu.org/licenses/>.
+ *
+*/
+
+use crate::coredb::{self, CoreDB};
+use crate::protocol::{responses, ActionGroup, Connection};
+use crate::resp::GroupBegin;
+use libtdb::TResult;
+
+/// Get the number of keys in the database
+pub async fn dbsize(handle: &CoreDB, con: &mut Connection, act: ActionGroup) -> TResult<()> {
+    if act.howmany() != 0 {
+        return con
+            .write_response(responses::fresp::R_ACTION_ERR.to_owned())
+            .await;
+    }
+    let mut len = 0;
+    {
+        len = handle.acquire_read().get_ref().len();
+    }
+    con.write_response(GroupBegin(1)).await?;
+    con.write_response(len).await?;
+    Ok(())
+}

server/src/kvengine/mod.rs

@@ -23,6 +23,7 @@
 //! This is TerrabaseDB's K/V engine. It contains utilities to interface with
 //! TDB's K/V store
 
+pub mod dbsize;
 pub mod del;
 pub mod exists;
 pub mod get;
@@ -31,8 +32,8 @@ pub mod mget;
 pub mod mset;
 pub mod mupdate;
 pub mod set;
-pub mod update;
 pub mod strong;
+pub mod update;
 pub mod heya {
     //! Respond to `HEYA` queries
     use crate::protocol;

server/src/queryengine/mod.rs

@@ -53,6 +53,8 @@ mod tags {
     pub const TAG_SDEL: &'static str = "SDEL";
     /// `SUPDATE` action tag
     pub const TAG_SUPDATE: &'static str = "SUPDATE";
+    /// `DBSIZE` action tag
+    pub const TAG_DBSIZE: &'static str = "DBSIZE";
 }
 
 /// Execute a simple(*) query
@@ -78,6 +80,7 @@ pub async fn execute_simple(db: &CoreDB, con: &mut Connection, buf: ActionGroup)
         tags::TAG_SSET => kvengine::strong::sset(db, con, buf).await?,
         tags::TAG_SDEL => kvengine::strong::sdel(db, con, buf).await?,
         tags::TAG_SUPDATE => kvengine::strong::supdate(db, con, buf).await?,
+        tags::TAG_DBSIZE => kvengine::dbsize::dbsize(db, con, buf).await?,
         _ => {
             con.write_response(responses::fresp::R_UNKNOWN_ACTION.to_owned())
                 .await?

server/src/tests/kvengine_tests.rs

@@ -70,6 +70,8 @@ async fn test_queries() {
     queries.add(test_sdel_multiple_nil).await;
     queries.add(test_sdel_multiple_okay).await;
     queries.add(test_sdel_syntax_error).await;
+    queries.add(test_dbsize_mixed).await;
+    queries.add(test_dbsize_syntax_error).await;
     queries.run_queries_and_close_sockets();
 
     // Clean up everything else
@@ -748,3 +750,29 @@ async fn test_sdel_syntax_error(mut stream: TcpStream) -> TcpStream {
     assert_eq!(response, fresp::R_ACTION_ERR.to_owned(), "{}", __func__!());
     stream
 }
+
+/// Test a `DBSIZE` query
+async fn test_dbsize_mixed(stream: TcpStream) -> TcpStream {
+    let mut stream = set_values(
+        "x ex y why z zed a firstalphabet b secondalphabet",
+        5,
+        stream,
+    )
+    .await;
+    let query = terrapipe::proc_query("DBSIZE");
+    stream.write_all(&query).await.unwrap();
+    let res_should_be = "#2\n*1\n#2\n&1\n:1\n5\n".to_owned().into_bytes();
+    let mut response = vec![0; res_should_be.len()];
+    stream.read_exact(&mut response).await.unwrap();
+    assert_eq!(response, res_should_be, "{}", __func__!());
+    stream
+}
+
+async fn test_dbsize_syntax_error(mut stream: TcpStream) -> TcpStream {
+    let query = terrapipe::proc_query("DBSIZE x y z");
+    stream.write_all(&query).await.unwrap();
+    let mut response = vec![0; fresp::R_ACTION_ERR.len()];
+    stream.read_exact(&mut response).await.unwrap();
+    assert_eq!(response, fresp::R_ACTION_ERR.to_owned(), "{}", __func__!());
+    stream
+}