feat(shell): count_data return estimate count by default

src/shell/commands/data_operations.cpp

@@ -2082,7 +2082,8 @@ bool clear_data(command_executor *e, shell_context *sc, arguments args)
 
 bool count_data(command_executor *e, shell_context *sc, arguments args)
 {
-    static struct option long_options[] = {{"partition", required_argument, 0, 'p'},
+    static struct option long_options[] = {{"precise", no_argument, 0, 'c'},
+                                           {"partition", required_argument, 0, 'p'},
                                            {"max_batch_count", required_argument, 0, 'b'},
                                            {"timeout_ms", required_argument, 0, 't'},
                                            {"hash_key_filter_type", required_argument, 0, 'h'},
@@ -2097,6 +2098,11 @@ bool count_data(command_executor *e, shell_context *sc, arguments args)
                                            {"run_seconds", required_argument, 0, 'r'},
                                            {0, 0, 0, 0}};
 
+    // "count_data" usually need scan all online records to get precise result, which may affect
+    // cluster availability, so here define precise = false defaultly and it will return estimate
+    // count immediately.
+    bool precise = false;
+    bool need_scan = false;
     int32_t partition = -1;
     int max_batch_count = 500;
     int timeout_ms = sc->timeout_ms;
@@ -2119,10 +2125,14 @@ bool count_data(command_executor *e, shell_context *sc, arguments args)
         int option_index = 0;
         int c;
         c = getopt_long(
-            args.argc, args.argv, "p:b:t:h:x:s:y:v:z:dan:r:", long_options, &option_index);
+            args.argc, args.argv, "cp:b:t:h:x:s:y:v:z:dan:r:", long_options, &option_index);
         if (c == -1)
             break;
+        need_scan = true;
         switch (c) {
+        case 'c':
+            precise = true;
+            break;
         case 'p':
             if (!dsn::buf2int32(optarg, partition)) {
                 fprintf(stderr, "ERROR: parse %s as partition failed\n", optarg);
@@ -2201,6 +2211,44 @@ bool count_data(command_executor *e, shell_context *sc, arguments args)
         }
     }
 
+    // if input [-p|--partition] etc. means you want to get precise count by scanning, but no input
+    // [-c|--precise], it will return false
+    if (need_scan && !precise) {
+        fprintf(stderr,
+                "ERROR: you must input [-c|-pricise] flag when you expect to get pricise "
+                "result by scaning all record online\n");
+        return false;
+    }
+
+    if (!precise) {
+        // get estimate key number
+        std::vector<row_data> rows;
+        std::string app_name = sc->pg_client->get_app_name();
+        if (!get_app_stat(sc, app_name, rows)) {
+            fprintf(stderr, "ERROR: query app stat from server failed");
+            return true;
+        }
+
+        rows.resize(rows.size() + 1);
+        row_data &sum = rows.back();
+        sum.row_name = "(total:" + std::to_string(rows.size() - 1) + ")";
+        for (int i = 0; i < rows.size() - 1; ++i) {
+            row_data &row = rows[i];
+            sum.rdb_estimate_num_keys += row.rdb_estimate_num_keys;
+        }
+
+        ::dsn::utils::table_printer tp("count_data");
+        tp.add_title("pidx");
+        tp.add_column("estimate_count");
+        for (row_data &row : rows) {
+            tp.add_row(row.row_name);
+            tp.append_data(row.rdb_estimate_num_keys);
+        }
+
+        tp.output(std::cout, tp_output_format::kTabular);
+        return true;
+    }
+
     if (max_batch_count <= 1) {
         fprintf(stderr, "ERROR: max_batch_count should be greater than 1\n");
         return false;

src/shell/main.cpp

@@ -280,14 +280,15 @@ static command_executor commands[] = {
     {
         "count_data",
         "get app row count",
-        "[-p|--partition num] [-b|--max_batch_count num] [-t|--timeout_ms num] "
+        "[-c|--pricise][-p|--partition num] "
+        "[-b|--max_batch_count num][-t|--timeout_ms num] "
         "[-h|--hash_key_filter_type anywhere|prefix|postfix] "
         "[-x|--hash_key_filter_pattern str] "
         "[-s|--sort_key_filter_type anywhere|prefix|postfix|exact] "
         "[-y|--sort_key_filter_pattern str] "
         "[-v|--value_filter_type anywhere|prefix|postfix|exact] "
-        "[-z|--value_filter_pattern str] "
-        "[-d|--diff_hash_key] [-a|--stat_size] [-n|--top_count num] [-r|--run_seconds num]",
+        "[-z|--value_filter_pattern str][-d|--diff_hash_key] "
+        "[-a|--stat_size] [-n|--top_count num] [-r|--run_seconds num]",
         data_operations,
     },
     {