Merge branch 'ffi-sfa-test' into clp-s-ffi-sfa-archive-reader

Author: Bill-hbrhbr

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-mcp-server/clp_mcp_server/clp_connector.py

@@ -60,7 +60,7 @@ async def submit_query(
         job_config = msgpack.packb(
             {
                 "begin_timestamp": begin_ts,
-                "dataset": CLP_DEFAULT_DATASET_NAME,
+                "datasets": [CLP_DEFAULT_DATASET_NAME],
                 "end_timestamp": end_ts,
                 "ignore_case": True,
                 "max_num_results": SEARCH_MAX_NUM_RESULTS,

components/clp-package-utils/clp_package_utils/controller.py

@@ -191,6 +191,9 @@ def _set_up_env_for_database(self) -> EnvVarsDict:
         env_vars |= {
             "CLP_DB_NAME": self._clp_config.database.names[ClpDbNameType.CLP],
         }
+        if self._clp_config.compression_scheduler.type == OrchestrationType.SPIDER:
+            env_vars["SPIDER_DB_NAME"] = self._clp_config.database.names[ClpDbNameType.SPIDER]
+
         if BundledService.DATABASE not in self._clp_config.bundled:
             env_vars |= {
                 "CLP_DB_PORT": str(self._clp_config.database.port),
@@ -199,9 +202,8 @@ def _set_up_env_for_database(self) -> EnvVarsDict:
                     self._clp_config.database.host
                 ),
             }
-
-        if self._clp_config.compression_scheduler.type == OrchestrationType.SPIDER:
-            env_vars["SPIDER_DB_NAME"] = self._clp_config.database.names[ClpDbNameType.SPIDER]
+        else:
+            env_vars["CLP_DB_HOST"] = _get_ip_from_hostname(self._clp_config.database.host)
 
         # Credentials
         credentials = self._clp_config.database.credentials
@@ -275,6 +277,8 @@ def _set_up_env_for_queue(self) -> EnvVarsDict:
                 "CLP_EXTRA_HOST_QUEUE_NAME": QUEUE_COMPONENT_NAME,
                 "CLP_EXTRA_HOST_QUEUE_ADDR": _resolve_external_host(self._clp_config.queue.host),
             }
+        else:
+            env_vars["CLP_QUEUE_HOST"] = _get_ip_from_hostname(self._clp_config.queue.host)
 
         # Credentials
         env_vars |= {
@@ -359,6 +363,8 @@ def _set_up_env_for_redis(self) -> EnvVarsDict:
                 "CLP_EXTRA_HOST_REDIS_NAME": REDIS_COMPONENT_NAME,
                 "CLP_EXTRA_HOST_REDIS_ADDR": _resolve_external_host(self._clp_config.redis.host),
             }
+        else:
+            env_vars["CLP_REDIS_HOST"] = _get_ip_from_hostname(self._clp_config.redis.host)
 
         # Credentials
         env_vars |= {
@@ -463,6 +469,10 @@ def _set_up_env_for_results_cache(self) -> EnvVarsDict:
                     self._clp_config.results_cache.host
                 ),
             }
+        else:
+            env_vars["CLP_RESULTS_CACHE_HOST"] = _get_ip_from_hostname(
+                self._clp_config.results_cache.host
+            )
 
         return env_vars
 
@@ -701,6 +711,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/archive_manager.py

@@ -28,7 +28,7 @@
     END_TS_ARG,
     FIND_COMMAND,
 )
-from clp_package_utils.scripts.native.utils import validate_dataset_exists
+from clp_package_utils.scripts.native.utils import validate_datasets_exist
 
 logger: logging.Logger = logging.getLogger(__file__)
 
@@ -200,7 +200,7 @@ def main(argv: list[str]) -> int:
     dataset = parsed_args.dataset
     if dataset is not None:
         try:
-            validate_dataset_exists(database_config, dataset)
+            validate_datasets_exist(database_config, [dataset])
         except Exception as e:
             logger.error(e)
             return -1

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

@@ -12,6 +12,7 @@
     CLP_DB_PASS_ENV_VAR_NAME,
     CLP_DB_USER_ENV_VAR_NAME,
     CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
+    CLP_DEFAULT_DATASET_NAME,
     ClpConfig,
     ClpDbNameType,
     ClpDbUserType,
@@ -36,7 +37,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,
 )
 
@@ -144,11 +145,9 @@ def handle_extract_stream_cmd(
         )
     elif EXTRACT_JSON_CMD == command:
         dataset = parsed_args.dataset
-        if dataset is None:
-            logger.error(f"Dataset unspecified, but must be specified for command `{command}'.")
-            return -1
+        dataset = CLP_DEFAULT_DATASET_NAME if dataset is None else dataset
         try:
-            validate_dataset_exists(clp_config.database, dataset)
+            validate_datasets_exist(clp_config.database, [dataset])
         except Exception as e:
             logger.error(e)
             return -1

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"