elastic: automatically try to update field limit on startup

jasonish · Nov 29, 2024 · 057edc3 · 057edc3
1 parent 0a3c754
commit 057edc3
Show file tree

Hide file tree

Showing 6 changed files with 109 additions and 50 deletions.
diff --git a/src/cli/elastic/main.rs b/src/cli/elastic/main.rs
@@ -3,7 +3,7 @@
 
 use anyhow::Result;
 use clap::{Command, CommandFactory, FromArgMatches, Parser, Subcommand};
-use tracing::info;
+use tracing::{info, warn};
 
 use super::set_field_limit;
 
@@ -79,14 +79,18 @@ async fn get_field_limit(args: &Args) -> Result<()> {
         client = client.with_password(password);
     }
     let client = client.build();
+    let template_name = &args.options.template;
 
-    let template = client.get_template(&args.options.template).await?;
-    let template_field_limit =
-        &template[&args.options.template]["settings"]["index"]["mapping"]["total_fields"]["limit"];
-    info!(
-        "Template {}: {:?}",
-        &args.options.template, template_field_limit
-    );
+    match client.get_template(template_name).await {
+        Ok(template) => {
+            let template_field_limit =
+                &template["settings"]["index"]["mapping"]["total_fields"]["limit"];
+            info!("Template: {}: {:?}", template_name, template_field_limit);
+        }
+        Err(err) => {
+            warn!("Failed to fetch template {}: {:?}", template_name, err);
+        }
+    }
 
     for index in client.get_indices_pattern("*").await? {
         // Only look at indices that match the name-YYYY.MM.DD
@@ -102,7 +106,7 @@ async fn get_field_limit(args: &Args) -> Result<()> {
         if let Some(map) = settings.as_object() {
             for (index, settings) in map {
                 let limit = &settings["settings"]["index"]["mapping"]["total_fields"]["limit"];
-                info!("{}: {:?}", index, limit);
+                info!("Index: {}: {:?}", index, limit);
             }
         }
     }

diff --git a/src/cli/elastic/set_field_limit.rs b/src/cli/elastic/set_field_limit.rs
@@ -6,7 +6,7 @@ use clap::Parser;
 use serde_json::json;
 use tracing::info;
 
-use crate::elastic::Client;
+use crate::elastic::{util::update_template_field_limit, Client};
 
 #[derive(Debug, Clone, Parser)]
 pub(crate) struct Args {
@@ -42,7 +42,7 @@ pub(crate) async fn main(args: Args) -> Result<()> {
     }
 
     info!("Updating template for pattern {}*", args.index);
-    update_template(&client, &args.index, args.limit).await?;
+    update_template_field_limit(&client, &args.index, args.limit).await?;
 
     Ok(())
 }
@@ -62,32 +62,3 @@ async fn update_index(client: &Client, index: &str, limit: usize) -> Result<()>
     info!("Index {}: status: {}, body: {}", index, status, body);
     Ok(())
 }
-
-async fn update_template(client: &crate::elastic::Client, index: &str, limit: usize) -> Result<()> {
-    #[rustfmt::skip]
-    let request = json!({
-	"index_patterns": [
-            format!("{}*", index),
-	],
-	"settings": {
-            "index": {
-		"mapping": {
-		    "total_fields": {
-			"limit": limit,
-		    }
-		}
-            }
-	}
-    });
-
-    let response = client
-        .put(&format!("_template/{index}"))?
-        .json(&request)
-        .send()
-        .await?;
-    let status = response.status();
-    let body = response.text().await?;
-    info!("Template {}: status: {}, body: {}", index, status, body);
-
-    Ok(())
-}
diff --git a/src/elastic/client.rs b/src/elastic/client.rs
@@ -196,9 +196,10 @@ impl Client {
 
     pub(crate) async fn get_template(&self, name: &str) -> anyhow::Result<serde_json::Value> {
         let response = self.get(&format!("_template/{}", name))?.send().await?;
+        let response = response.error_for_status()?;
         let text = response.text().await?;
-        let response = serde_json::from_str(&text)?;
-        Ok(response)
+        let mut response: serde_json::Value = serde_json::from_str(&text)?;
+        Ok(response[name].take())
     }
 }
 

diff --git a/src/elastic/mod.rs b/src/elastic/mod.rs
@@ -1,20 +1,23 @@
 // SPDX-FileCopyrightText: (C) 2020 Jason Ish <[email protected]>
 // SPDX-License-Identifier: MIT
 
+use serde::{Deserialize, Serialize};
+use serde_json::json;
+use thiserror::Error;
+
 use crate::datetime::DateTime;
 use crate::eventrepo::DatastoreError;
+
 pub(crate) use client::Version;
 pub(crate) use client::{Client, ClientBuilder};
 pub(crate) use eventrepo::ElasticEventRepo;
 pub(crate) use importer::ElasticEventSink;
-use serde::{Deserialize, Serialize};
-use serde_json::json;
-use thiserror::Error;
 
-pub mod client;
-pub mod eventrepo;
-pub mod importer;
-pub mod request;
+pub(crate) mod client;
+pub(crate) mod eventrepo;
+pub(crate) mod importer;
+pub(crate) mod request;
+pub(crate) mod util;
 
 pub(crate) const TAG_ESCALATED: &str = "evebox.escalated";
 pub(crate) const TAGS_ESCALATED: [&str; 1] = [TAG_ESCALATED];

diff --git a/src/elastic/util.rs b/src/elastic/util.rs
@@ -0,0 +1,77 @@
+// SPDX-FileCopyrightText: (C) 2024 Jason Ish <[email protected]>
+// SPDX-License-Identifier: MIT
+
+use tracing::{error, info};
+
+use super::Client;
+
+pub(crate) async fn check_and_set_field_limit(client: &Client, template_name: &str) {
+    match client.get_template(template_name).await {
+        Ok(template) => {
+            let field_limit = &template["settings"]["index"]["mapping"]["total_fields"]["limit"];
+            let limit: Option<i64> = match field_limit {
+                serde_json::Value::Number(n) => n.as_i64(),
+                serde_json::Value::String(s) => s.parse::<i64>().ok(),
+                _ => None,
+            };
+            if let Some(limit) = limit {
+                if limit >= 5000 {
+                    info!("Field limit of {} OK, will not increase", limit);
+                    return;
+                }
+            }
+        }
+        Err(err) => {
+            info!(
+                "Failed to find template for index {}: {:?}",
+                template_name, err
+            );
+        }
+    }
+
+    info!("Attempting to increase Elasticsearch field limit to 5000");
+    match update_template_field_limit(client, template_name, 5000).await {
+        Ok(_ok) => {
+            info!("Successfully updated Elasticsearch template field limit");
+        }
+        Err(err) => {
+            error!(
+                "Failed to update Elasticsearch template field limit: {:?}",
+                err
+            );
+        }
+    }
+}
+
+pub(crate) async fn update_template_field_limit(
+    client: &Client,
+    index: &str,
+    limit: usize,
+) -> anyhow::Result<()> {
+    #[rustfmt::skip]
+    let request = json!({
+	"index_patterns": [
+            format!("{}*", index),
+	],
+	"settings": {
+            "index": {
+		"mapping": {
+		    "total_fields": {
+			"limit": limit,
+		    }
+		}
+            }
+	}
+    });
+
+    let response = client
+        .put(&format!("_template/{index}"))?
+        .json(&request)
+        .send()
+        .await?;
+    let status = response.status();
+    let body = response.text().await?;
+    info!("Template {}: status: {}, body: {}", index, status, body);
+
+    Ok(())
+}
diff --git a/src/server/main.rs b/src/server/main.rs
@@ -488,7 +488,7 @@ async fn configure_datastore(config: Config, server_config: &ServerConfig) -> Re
             let eventstore = elastic::ElasticEventRepo {
                 base_index: server_config.elastic_index.clone(),
                 index_pattern,
-                client,
+                client: client.clone(),
                 ecs: server_config.elastic_ecs,
             };
             debug!("Elasticsearch base index: {}", &eventstore.base_index);
@@ -497,6 +497,9 @@ async fn configure_datastore(config: Config, server_config: &ServerConfig) -> Re
                 &eventstore.index_pattern
             );
             debug!("Elasticsearch ECS mode: {}", eventstore.ecs);
+
+            elastic::util::check_and_set_field_limit(&client, &eventstore.base_index).await;
+
             Ok(EventRepo::Elastic(eventstore))
         }
         "sqlite" => {