Merge branch 'main' into integration-tests/clp-ffi-js

Author: junhaoliao

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/controller.py

@@ -701,6 +701,7 @@ def _set_up_env_for_webui(self, container_clp_config: ClpConfig) -> EnvVarsDict:
             "ClpStorageEngine": self._clp_config.package.storage_engine,
             "ClpQueryEngine": self._clp_config.package.query_engine,
             "LogsInputType": self._clp_config.logs_input.type,
+            "MaxDatasetsPerQuery": self._clp_config.query_scheduler.max_datasets_per_query,
             "MongoDbSearchResultsMetadataCollectionName": (
                 self._clp_config.webui.results_metadata_collection_name
             ),

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,
@@ -235,9 +235,9 @@ def main(argv):
     args_parser.add_argument("wildcard_query", help="Wildcard query.")
     args_parser.add_argument(
         "--dataset",
-        type=str,
+        action="append",
         default=None,
-        help="The dataset that the archives belong to.",
+        help="A dataset to search. Can be specified multiple times.",
     )
     args_parser.add_argument(
         "--begin-time",
@@ -297,10 +297,18 @@ def main(argv):
         return -1
 
     database_config: Database = clp_config.database
-    dataset = parsed_args.dataset
-    if dataset is not None:
+    datasets = parsed_args.dataset
+    if datasets is not None:
+        max_datasets_per_query = clp_config.query_scheduler.max_datasets_per_query
+        if max_datasets_per_query is not None and len(datasets) > max_datasets_per_query:
+            logger.error(
+                "Number of datasets (%d) exceeds max_datasets_per_query=%s.",
+                len(datasets),
+                max_datasets_per_query,
+            )
+            return -1
         try:
-            validate_dataset_exists(database_config, dataset)
+            validate_datasets_exist(database_config, datasets)
         except Exception as e:
             logger.error(e)
             return -1
@@ -310,7 +318,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,11 @@ 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)
+        missing = [ds for ds in datasets if ds not in existing_datasets]
+        if len(missing) > 0:
+            err_msg = f"Dataset(s) {missing} don't exist."
+            raise ValueError(err_msg)
 
 
 def wait_for_query_job(sql_adapter: SqlAdapter, job_id: int) -> QueryJobStatus:

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

@@ -52,9 +52,9 @@ def main(argv):
     args_parser.add_argument("wildcard_query", help="Wildcard query.")
     args_parser.add_argument(
         "--dataset",
-        type=str,
+        action="append",
         default=None,
-        help="The dataset that the archives belong to.",
+        help="A dataset to search. Can be specified multiple times.",
     )
     args_parser.add_argument(
         "--begin-time",
@@ -113,16 +113,17 @@ def main(argv):
         )
         return -1
 
-    dataset = parsed_args.dataset
+    datasets = parsed_args.dataset
     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 +153,10 @@ 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:
+        for ds in datasets:
+            search_cmd.append("--dataset")
+            search_cmd.append(ds)
     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

@@ -433,6 +433,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;