feat!: Add multi-dataset query support (resolves #1933).

components/api-server/src/client.rs

@@ -28,9 +28,9 @@ pub struct QueryConfig {
     /// The search query as a KQL string.
     pub query_string: String,
 
-    /// The dataset to search within. If not provided, only `default` dataset will be searched.
+    /// The datasets to search within. If not provided, only `default` dataset will be searched.
     #[serde(default)]
-    pub dataset: Option<String>,
+    pub datasets: Option<Vec<String>>,
 
     /// The maximum number of results to return. Set to `0` for no limit.
     #[serde(default)]
@@ -58,7 +58,7 @@ pub struct QueryConfig {
 impl From<QueryConfig> for SearchJobConfig {
     fn from(value: QueryConfig) -> Self {
         Self {
-            dataset: value.dataset,
+            datasets: value.datasets,
             query_string: value.query_string,
             max_num_results: value.max_num_results,
             begin_timestamp: value.time_range_begin_millisecs,
@@ -128,10 +128,10 @@ impl Client {
     /// * Forwards [`sqlx::query::Query::execute`]'s return values on failure.
     pub async fn submit_query(&self, query_config: QueryConfig) -> Result<u64, ClientError> {
         let mut search_job_config: SearchJobConfig = query_config.into();
-        if search_job_config.dataset.is_none() {
-            search_job_config.dataset = match self.config.package.storage_engine {
+        if search_job_config.datasets.is_none() {
+            search_job_config.datasets = match self.config.package.storage_engine {
                 StorageEngine::Clp => None,
-                StorageEngine::ClpS => Some("default".to_owned()),
+                StorageEngine::ClpS => Some(vec!["default".to_owned()]),
             }
         }
         if search_job_config.max_num_results == 0 {

components/api-server/src/routes.rs

@@ -85,7 +85,7 @@ async fn health() -> String {
         content= QueryConfig,
         example = json!({
             "query_string": "*",
-            "dataset": "default",
+            "datasets": ["default"],
             "time_range_begin_millisecs": 0,
             "time_range_end_millisecs": 17_356_896,
             "ignore_case": true,

components/clp-package-utils/clp_package_utils/scripts/native/search.py

@@ -27,7 +27,7 @@
 from clp_package_utils.scripts.native.utils import (
     run_function_in_process,
     submit_query_job,
-    validate_dataset_exists,
+    validate_datasets_exist,
     wait_for_query_job,
 )
 
@@ -37,7 +37,7 @@
 def create_and_monitor_job_in_db(
     db_config: Database,
     results_cache: ResultsCache,
-    dataset: str | None,
+    datasets: list[str] | None,
     wildcard_query: str,
     begin_timestamp: int | None,
     end_timestamp: int | None,
@@ -48,7 +48,7 @@ def create_and_monitor_job_in_db(
     count_by_time_bucket_size: int | None,
 ):
     search_config = SearchJobConfig(
-        dataset=dataset,
+        datasets=datasets,
         query_string=wildcard_query,
         begin_timestamp=begin_timestamp,
         end_timestamp=end_timestamp,
@@ -115,7 +115,7 @@ async def worker_connection_handler(reader: asyncio.StreamReader, writer: asynci
 async def do_search_without_aggregation(
     db_config: Database,
     results_cache: ResultsCache,
-    dataset: str | None,
+    datasets: list[str] | None,
     wildcard_query: str,
     begin_timestamp: int | None,
     end_timestamp: int | None,
@@ -144,7 +144,7 @@ async def do_search_without_aggregation(
             create_and_monitor_job_in_db,
             db_config,
             results_cache,
-            dataset,
+            datasets,
             wildcard_query,
             begin_timestamp,
             end_timestamp,
@@ -181,7 +181,7 @@ async def do_search_without_aggregation(
 async def do_search(
     db_config: Database,
     results_cache: ResultsCache,
-    dataset: str | None,
+    datasets: list[str] | None,
     wildcard_query: str,
     begin_timestamp: int | None,
     end_timestamp: int | None,
@@ -195,7 +195,7 @@ async def do_search(
         await do_search_without_aggregation(
             db_config,
             results_cache,
-            dataset,
+            datasets,
             wildcard_query,
             begin_timestamp,
             end_timestamp,
@@ -208,7 +208,7 @@ async def do_search(
             create_and_monitor_job_in_db,
             db_config,
             results_cache,
-            dataset,
+            datasets,
             wildcard_query,
             begin_timestamp,
             end_timestamp,
@@ -234,10 +234,11 @@ def main(argv):
     )
     args_parser.add_argument("wildcard_query", help="Wildcard query.")
     args_parser.add_argument(
-        "--dataset",
+        "--datasets",
         type=str,
+        nargs="+",
         default=None,
-        help="The dataset that the archives belong to.",
+        help="The datasets that the archives belong to.",
     )
     args_parser.add_argument(
         "--begin-time",
@@ -297,10 +298,10 @@ def main(argv):
         return -1
 
     database_config: Database = clp_config.database
-    dataset = parsed_args.dataset
-    if dataset is not None:
+    datasets = parsed_args.datasets
+    if datasets is not None:
         try:
-            validate_dataset_exists(database_config, dataset)
+            validate_datasets_exist(database_config, datasets)
         except Exception as e:
             logger.error(e)
             return -1
@@ -310,7 +311,7 @@ def main(argv):
             do_search(
                 database_config,
                 clp_config.results_cache,
-                dataset,
+                datasets,
                 parsed_args.wildcard_query,
                 parsed_args.begin_time,
                 parsed_args.end_time,

components/clp-package-utils/clp_package_utils/scripts/native/utils.py

@@ -72,13 +72,13 @@ def submit_query_job(
         return db_cursor.lastrowid
 
 
-def validate_dataset_exists(db_config: Database, dataset: str) -> None:
+def validate_datasets_exist(db_config: Database, datasets: list[str]) -> None:
     """
-    Validates that `dataset` exists in the metadata database.
+    Validates that all datasets in `datasets` exist in the metadata database.
 
     :param db_config:
-    :param dataset:
-    :raise: ValueError if the dataset doesn't exist.
+    :param datasets:
+    :raise: ValueError if any dataset doesn't exist.
     """
     sql_adapter = SqlAdapter(db_config)
     clp_db_connection_params = db_config.get_clp_connection_params_and_type(True)
@@ -87,8 +87,10 @@ def validate_dataset_exists(db_config: Database, dataset: str) -> None:
         closing(sql_adapter.create_connection(True)) as db_conn,
         closing(db_conn.cursor(dictionary=True)) as db_cursor,
     ):
-        if dataset not in fetch_existing_datasets(db_cursor, table_prefix):
-            raise ValueError(f"Dataset `{dataset}` doesn't exist.")
+        existing_datasets = fetch_existing_datasets(db_cursor, table_prefix)
+        for dataset in datasets:
+            if dataset not in existing_datasets:
+                raise ValueError(f"Dataset `{dataset}` doesn't exist.")
 
 
 def wait_for_query_job(sql_adapter: SqlAdapter, job_id: int) -> QueryJobStatus:

components/clp-package-utils/clp_package_utils/scripts/search.py

@@ -51,10 +51,11 @@ def main(argv):
     )
     args_parser.add_argument("wildcard_query", help="Wildcard query.")
     args_parser.add_argument(
-        "--dataset",
+        "--datasets",
         type=str,
+        nargs="+",
         default=None,
-        help="The dataset that the archives belong to.",
+        help="The datasets that the archives belong to.",
     )
     args_parser.add_argument(
         "--begin-time",
@@ -113,16 +114,17 @@ def main(argv):
         )
         return -1
 
-    dataset = parsed_args.dataset
+    datasets = parsed_args.datasets
     if StorageEngine.CLP_S == storage_engine:
-        dataset = CLP_DEFAULT_DATASET_NAME if dataset is None else dataset
+        datasets = [CLP_DEFAULT_DATASET_NAME] if datasets is None else datasets
         try:
             clp_db_connection_params = clp_config.database.get_clp_connection_params_and_type(True)
-            validate_dataset_name(clp_db_connection_params["table_prefix"], dataset)
+            for ds in datasets:
+                validate_dataset_name(clp_db_connection_params["table_prefix"], ds)
         except Exception as e:
             logger.error(e)
             return -1
-    elif dataset is not None:
+    elif datasets is not None:
         logger.error(f"Dataset selection is not supported for storage engine: {storage_engine}.")
         return -1
 
@@ -152,9 +154,9 @@ def main(argv):
     # fmt: on
     if parsed_args.verbose:
         search_cmd.append("--verbose")
-    if dataset is not None:
-        search_cmd.append("--dataset")
-        search_cmd.append(dataset)
+    if datasets is not None:
+        search_cmd.append("--datasets")
+        search_cmd.extend(datasets)
     if parsed_args.begin_time is not None:
         search_cmd.append("--begin-time")
         search_cmd.append(str(parsed_args.begin_time))

components/clp-py-utils/clp_py_utils/clp_config.py

@@ -432,6 +432,7 @@ class QueryScheduler(BaseModel):
     host: DomainStr = "localhost"
     port: Port = DEFAULT_PORT
     jobs_poll_delay: PositiveFloat = 0.1  # seconds
+    max_datasets_per_query: PositiveInt | None = 10
     num_archives_to_search_per_sub_job: PositiveInt = 16
     logging_level: LoggingLevel = "INFO"
 

components/clp-rust-utils/src/job_config/search.rs

@@ -11,7 +11,7 @@ pub const QUERY_JOBS_TABLE_NAME: &str = "query_jobs";
 #[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
 #[serde(default)]
 pub struct SearchJobConfig {
-    pub dataset: Option<String>,
+    pub datasets: Option<Vec<String>>,
     pub query_string: String,
     pub max_num_results: u32,
     pub begin_timestamp: Option<i64>,

components/core/src/clp_s/CommandLineArguments.cpp

@@ -605,6 +605,10 @@ CommandLineArguments::parse_arguments(int argc, char const** argv) {
                     po::value<uint64_t>(&m_max_num_results)->value_name("MAX")->
                             default_value(m_max_num_results),
                     "The maximum number of results to output"
+            )(
+                    "dataset",
+                    po::value<std::string>(&m_dataset)->value_name("DATASET"),
+                    "The dataset name to include in each result document"
             );
 
             po::options_description file_output_handler_options("File Output Handler Options");

components/core/src/clp_s/CommandLineArguments.hpp

@@ -123,6 +123,8 @@ class CommandLineArguments {
 
     bool get_record_log_order() const { return false == m_disable_log_order; }
 
+    std::string const& get_dataset() const { return m_dataset; }
+
 private:
     // Methods
     /**
@@ -227,6 +229,7 @@ class CommandLineArguments {
     std::optional<epochtime_t> m_search_end_ts;
     bool m_ignore_case{false};
     std::vector<std::string> m_projection_columns;
+    std::string m_dataset;
 
     // Search aggregation variables
     std::string m_reducer_host;

components/core/src/clp_s/OutputHandlerImpl.cpp

@@ -70,11 +70,13 @@ ResultsCacheOutputHandler::ResultsCacheOutputHandler(
         string const& collection,
         uint64_t batch_size,
         uint64_t max_num_results,
+        string dataset,
         bool should_output_timestamp
 )
         : ::clp_s::search::OutputHandler(should_output_timestamp, true),
           m_batch_size(batch_size),
-          m_max_num_results(max_num_results) {
+          m_max_num_results(max_num_results),
+          m_dataset(std::move(dataset)) {
     try {
         auto mongo_uri = mongocxx::uri(uri);
         m_client = mongocxx::client(mongo_uri);
@@ -114,6 +116,10 @@ ErrorCode ResultsCacheOutputHandler::flush() {
                                     bsoncxx::builder::basic::kvp(
                                             constants::results_cache::search::cLogEventIx,
                                             result.log_event_idx
+                                    ),
+                                    bsoncxx::builder::basic::kvp(
+                                            constants::results_cache::search::cDataset,
+                                            std::move(result.dataset)
                                     )
                             )
                     )
@@ -154,7 +160,8 @@ void ResultsCacheOutputHandler::write(
                         message,
                         timestamp,
                         archive_id,
-                        log_event_idx
+                        log_event_idx,
+                        m_dataset
                 )
         );
     } else if (m_latest_results.top()->timestamp < timestamp) {
@@ -165,7 +172,8 @@ void ResultsCacheOutputHandler::write(
                         message,
                         timestamp,
                         archive_id,
-                        log_event_idx
+                        log_event_idx,
+                        m_dataset
                 )
         );
     }

components/core/src/clp_s/OutputHandlerImpl.hpp

@@ -128,19 +128,22 @@ class ResultsCacheOutputHandler : public ::clp_s::search::OutputHandler {
                 std::string_view message,
                 epochtime_t timestamp,
                 std::string_view archive_id,
-                int64_t log_event_idx
+                int64_t log_event_idx,
+                std::string_view dataset
         )
                 : original_path(original_path),
                   message(message),
                   timestamp(timestamp),
                   archive_id(archive_id),
-                  log_event_idx(log_event_idx) {}
+                  log_event_idx(log_event_idx),
+                  dataset(dataset) {}
 
         std::string original_path;
         std::string message;
         epochtime_t timestamp;
         std::string archive_id;
         int64_t log_event_idx;
+        std::string dataset;
     };
 
     struct QueryResultGreaterTimestampComparator {
@@ -165,6 +168,7 @@ class ResultsCacheOutputHandler : public ::clp_s::search::OutputHandler {
             std::string const& collection,
             uint64_t batch_size,
             uint64_t max_num_results,
+            std::string dataset,
             bool should_output_metadata = true
     );
 
@@ -191,6 +195,7 @@ class ResultsCacheOutputHandler : public ::clp_s::search::OutputHandler {
     std::vector<bsoncxx::document::value> m_results;
     uint64_t m_batch_size;
     uint64_t m_max_num_results;
+    std::string m_dataset;
     std::priority_queue<
             std::unique_ptr<QueryResult>,
             std::vector<std::unique_ptr<QueryResult>>,

components/core/src/clp_s/archive_constants.hpp

@@ -58,6 +58,7 @@ constexpr char cLogEventIx[]{"log_event_ix"};
 constexpr char cTimestamp[]{"timestamp"};
 constexpr char cMessage[]{"message"};
 constexpr char cArchiveId[]{"archive_id"};
+constexpr char cDataset[]{"dataset"};
 }  // namespace results_cache::search
 }  // namespace clp_s::constants
 #endif  // CLP_S_ARCHIVE_CONSTANTS_HPP

components/core/src/clp_s/clp-s.cpp

@@ -265,7 +265,8 @@ bool search_archive(
                         command_line_arguments.get_mongodb_uri(),
                         command_line_arguments.get_mongodb_collection(),
                         command_line_arguments.get_batch_size(),
-                        command_line_arguments.get_max_num_results()
+                        command_line_arguments.get_max_num_results(),
+                        command_line_arguments.get_dataset()
                 );
                 break;
             case CommandLineArguments::OutputHandlerType::Stdout: