Merge branch 'y-scope:main' into 2024-anonymous-usage-telemetry-framework-for-clp

Author: Nathan903

Cargo.lock

@@ -38,7 +38,5 @@ thiserror = "2.0.18"
 tokio = { version = "1.49.0", features = ["full"] }
 tower-http = { version = "0.6.8", features = ["cors"] }
 tracing = "0.1.44"
-tracing-appender = "0.2.4"
-tracing-subscriber = { version = "0.3.22", features = ["json", "env-filter", "fmt", "std"] }
 utoipa = { version = "5.4.0", features = ["axum_extras"] }
 utoipa-axum = "0.2.0"

components/api-server/Cargo.toml

@@ -1,11 +1,6 @@
 use anyhow::Context;
 use clap::Parser;
 use clp_rust_utils::{clp_config::package, serde::yaml};
-use tracing_appender::{
-    non_blocking::WorkerGuard,
-    rolling::{RollingFileAppender, Rotation},
-};
-use tracing_subscriber::{self, fmt::writer::MakeWriterExt};
 
 #[derive(Parser)]
 #[command(version, about = "API Server for CLP.")]
@@ -40,29 +35,6 @@ fn read_config_and_credentials(
     Ok((config, credentials))
 }
 
-fn set_up_logging() -> anyhow::Result<WorkerGuard> {
-    let logs_directory =
-        std::env::var("CLP_LOGS_DIR").context("Expect `CLP_LOGS_DIR` environment variable.")?;
-    let logs_directory = std::path::Path::new(logs_directory.as_str());
-    let file_appender =
-        RollingFileAppender::new(Rotation::HOURLY, logs_directory, "api_server.log");
-    let (non_blocking_writer, guard) = tracing_appender::non_blocking(file_appender);
-    tracing_subscriber::fmt()
-        .event_format(
-            tracing_subscriber::fmt::format()
-                .with_level(true)
-                .with_target(false)
-                .with_file(true)
-                .with_line_number(true)
-                .json(),
-        )
-        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
-        .with_ansi(false)
-        .with_writer(std::io::stdout.and(non_blocking_writer))
-        .init();
-    Ok(guard)
-}
-
 async fn shutdown_signal() {
     let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
         .expect("failed to listen for SIGTERM");
@@ -79,7 +51,7 @@ async fn main() -> anyhow::Result<()> {
     let args = Args::parse();
 
     let (config, credentials) = read_config_and_credentials(&args)?;
-    let _guard = set_up_logging()?;
+    let _guard = clp_rust_utils::logging::set_up_logging("api_server.log");
 
     let api_server_config = config
         .api_server

components/api-server/src/bin/api_server.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/client.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/api-server/src/routes.rs

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

@@ -1,15 +1,13 @@
 """CLP MCP Server entry point."""
 
 import ipaddress
-import logging
-import os
 import socket
 import sys
 from pathlib import Path
 
 import click
 from clp_py_utils.clp_config import ClpConfig, MCP_SERVER_COMPONENT_NAME
-from clp_py_utils.clp_logging import get_logger, get_logging_formatter, set_logging_level
+from clp_py_utils.clp_logging import configure_logging, get_logger
 from clp_py_utils.core import read_yaml_config_file
 from pydantic import ValidationError
 
@@ -38,12 +36,8 @@ def main(host: str, port: int, config_path: Path) -> int:
     :param config_path: The path to server's configuration file.
     :return: Exit code (0 for success, non-zero for failure).
     """
-    # Setup logging to file
-    log_file_path = Path(os.getenv("CLP_LOGS_DIR")) / "mcp_server.log"
-    logging_file_handler = logging.FileHandler(filename=log_file_path, encoding="utf-8")
-    logging_file_handler.setFormatter(get_logging_formatter())
-    logger.addHandler(logging_file_handler)
-    set_logging_level(logger, os.getenv("CLP_LOGGING_LEVEL"))
+    # Setup optional file logging and logging level.
+    configure_logging(logger, "mcp_server")
 
     exit_code = 0
 

components/clp-mcp-server/clp_mcp_server/clp_mcp_server.py

@@ -1,3 +1,4 @@
+import ipaddress
 import json
 import logging
 import multiprocessing
@@ -188,13 +189,25 @@ def _set_up_env_for_database(self) -> EnvVarsDict:
 
         # Connection config
         env_vars |= {
-            "CLP_DB_HOST": _get_ip_from_hostname(self._clp_config.database.host),
             "CLP_DB_NAME": self._clp_config.database.names[ClpDbNameType.CLP],
-            "CLP_DB_PORT": str(self._clp_config.database.port),
         }
         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_CONNECT_PORT": str(self._clp_config.database.port),
+                "CLP_EXTRA_HOST_DATABASE_NAME": DB_COMPONENT_NAME,
+                "CLP_EXTRA_HOST_DATABASE_ADDR": _resolve_external_host(
+                    self._clp_config.database.host
+                ),
+            }
+        else:
+            env_vars |= {
+                "CLP_DB_HOST": _get_ip_from_hostname(self._clp_config.database.host),
+                "CLP_DB_PORT": str(self._clp_config.database.port),
+            }
+
         # Credentials
         credentials = self._clp_config.database.credentials
         env_vars |= {
@@ -261,10 +274,17 @@ def _set_up_env_for_queue(self) -> EnvVarsDict:
         env_vars = EnvVarsDict()
 
         # Connection config
-        env_vars |= {
-            "CLP_QUEUE_HOST": _get_ip_from_hostname(self._clp_config.queue.host),
-            "CLP_QUEUE_PORT": str(self._clp_config.queue.port),
-        }
+        if BundledService.QUEUE not in self._clp_config.bundled:
+            env_vars |= {
+                "CLP_QUEUE_CONNECT_PORT": str(self._clp_config.queue.port),
+                "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),
+                "CLP_QUEUE_PORT": str(self._clp_config.queue.port),
+            }
 
         # Credentials
         env_vars |= {
@@ -343,10 +363,17 @@ def _set_up_env_for_redis(self) -> EnvVarsDict:
         env_vars = EnvVarsDict()
 
         # Connection config
-        env_vars |= {
-            "CLP_REDIS_HOST": _get_ip_from_hostname(self._clp_config.redis.host),
-            "CLP_REDIS_PORT": str(self._clp_config.redis.port),
-        }
+        if BundledService.REDIS not in self._clp_config.bundled:
+            env_vars |= {
+                "CLP_REDIS_CONNECT_PORT": str(self._clp_config.redis.port),
+                "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),
+                "CLP_REDIS_PORT": str(self._clp_config.redis.port),
+            }
 
         # Credentials
         env_vars |= {
@@ -442,9 +469,22 @@ def _set_up_env_for_results_cache(self) -> EnvVarsDict:
         # Connection config
         env_vars |= {
             "CLP_RESULTS_CACHE_DB_NAME": self._clp_config.results_cache.db_name,
-            "CLP_RESULTS_CACHE_HOST": _get_ip_from_hostname(self._clp_config.results_cache.host),
-            "CLP_RESULTS_CACHE_PORT": str(self._clp_config.results_cache.port),
         }
+        if BundledService.RESULTS_CACHE not in self._clp_config.bundled:
+            env_vars |= {
+                "CLP_RESULTS_CACHE_CONNECT_PORT": str(self._clp_config.results_cache.port),
+                "CLP_EXTRA_HOST_RESULTS_CACHE_NAME": RESULTS_CACHE_COMPONENT_NAME,
+                "CLP_EXTRA_HOST_RESULTS_CACHE_ADDR": _resolve_external_host(
+                    self._clp_config.results_cache.host
+                ),
+            }
+        else:
+            env_vars |= {
+                "CLP_RESULTS_CACHE_HOST": _get_ip_from_hostname(
+                    self._clp_config.results_cache.host
+                ),
+                "CLP_RESULTS_CACHE_PORT": str(self._clp_config.results_cache.port),
+            }
 
         return env_vars
 
@@ -683,6 +723,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
             ),
@@ -1160,3 +1201,20 @@ def _get_ip_from_hostname(hostname: str) -> str:
     :return: The resolved IP address.
     """
     return socket.gethostbyname(hostname)
+
+
+def _resolve_external_host(hostname: str) -> str:
+    """
+    Resolves a hostname to an address suitable for Docker's ``extra_hosts``.
+
+    When the hostname resolves to a loopback address, returns Docker's ``host-gateway`` token so
+    that containers can reach services running on the Docker host. For any other hostname, falls
+    back to standard DNS resolution.
+
+    :param hostname:
+    :return: The resolved address.
+    """
+    resolved_ip = _get_ip_from_hostname(hostname)
+    if ipaddress.ip_address(resolved_ip).is_loopback:
+        return "host-gateway"
+    return resolved_ip

components/clp-package-utils/clp_package_utils/controller.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/archive_manager.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