From f6fc2bcce4a0cbf64819d44fb09a7ad2770fd52e Mon Sep 17 00:00:00 2001
From: Matt Mundell <matt@mundell.me>
Date: Mon, 19 Jun 2023 22:44:37 +0200
Subject: [PATCH] Add --vt-sev-insert-size

---
 src/gvmd.c            |  8 ++++++++
 src/manage_sql_nvts.c | 35 ++++++++++++++++++++---------------
 src/manage_sql_nvts.h | 10 ++++++++++
 3 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/src/gvmd.c b/src/gvmd.c
index c5d73a51b..a75b391da 100644
--- a/src/gvmd.c
+++ b/src/gvmd.c
@@ -1887,6 +1887,7 @@ gvmd (int argc, char** argv, char *env[])
   static gchar *feed_lock_path = NULL;
   static int feed_lock_timeout = 0;
   static int vt_ref_insert_size = VT_REF_INSERT_SIZE_DEFAULT;
+  static int vt_sev_insert_size = VT_SEV_INSERT_SIZE_DEFAULT;
   static gchar *vt_verification_collation = NULL;
 
   GString *full_disable_commands = g_string_new ("");
@@ -2216,6 +2217,11 @@ gvmd (int argc, char** argv, char *env[])
           "Max number of VT refs to insert per statement during VT update,"
           " 0 for unlimited, default: "
           G_STRINGIFY (VT_REF_INSERT_SIZE_DEFAULT), "<number>" },
+        { "vt-sev-insert-size", '\0', 0, G_OPTION_ARG_INT,
+          &vt_sev_insert_size,
+          "Max number of VT severities to insert per statement during VT update,"
+          " 0 for unlimited, default: "
+          G_STRINGIFY (VT_SEV_INSERT_SIZE_DEFAULT), "<number>" },
         { "vt-verification-collation", '\0', 0, G_OPTION_ARG_STRING,
           &vt_verification_collation,
           "Set collation for VT verification to <collation>, omit or leave"
@@ -2305,6 +2311,8 @@ gvmd (int argc, char** argv, char *env[])
 
   set_vt_ref_insert_size (vt_ref_insert_size);
 
+  set_vt_sev_insert_size (vt_sev_insert_size);
+
   /* Set VT verification collation override */
   set_vt_verification_collation (vt_verification_collation);
 
diff --git a/src/manage_sql_nvts.c b/src/manage_sql_nvts.c
index f90314476..97815969d 100644
--- a/src/manage_sql_nvts.c
+++ b/src/manage_sql_nvts.c
@@ -50,20 +50,6 @@
  */
 #define G_LOG_DOMAIN "md manage"
 
-/**
- * @brief Rows per statement when inserting VT refs for update/rebuild.
- *
- * There are about 500k vt_refs.
- */
-#define VT_REFS_BATCH_SIZE 50000
-
-/**
- * @brief Rows per statement when inserting VT severities for update/rebuild.
- *
- * There are about 80k vt_severities.
- */
-#define VT_SEVS_BATCH_SIZE 100000
-
 
 /* Headers from backend specific manage_xxx.c file. */
 
@@ -78,6 +64,11 @@ create_tables_nvt (const gchar *);
  */
 static int vt_ref_insert_size = VT_REF_INSERT_SIZE_DEFAULT;
 
+/**
+ * @brief Max number of rows inserted per statement.
+ */
+static int vt_sev_insert_size = VT_SEV_INSERT_SIZE_DEFAULT;
+
 /**
  * @brief File socket for OSP NVT update.
  */
@@ -155,6 +146,20 @@ set_vt_ref_insert_size (int new_size)
     vt_ref_insert_size = new_size;
 }
 
+/**
+ * @brief Set the VT severity insert size.
+ *
+ * @param new_size  New size.
+ */
+void
+set_vt_sev_insert_size (int new_size)
+{
+  if (new_size < 0)
+    vt_sev_insert_size = 0;
+  else
+    vt_sev_insert_size = new_size;
+}
+
 /**
  * @brief Ensures the sanity of nvts cache in DB.
  */
@@ -1766,7 +1771,7 @@ update_nvts_from_vts (element_t *get_vts_response,
     sql ("TRUNCATE nvt_preferences;");
 
   vt_refs_batch = batch_start (vt_ref_insert_size);
-  vt_sevs_batch = batch_start (VT_SEVS_BATCH_SIZE);
+  vt_sevs_batch = batch_start (vt_sev_insert_size);
   vt = element_first_child (vts);
   while (vt)
     {
diff --git a/src/manage_sql_nvts.h b/src/manage_sql_nvts.h
index 117b2a779..5e04edf83 100644
--- a/src/manage_sql_nvts.h
+++ b/src/manage_sql_nvts.h
@@ -105,6 +105,16 @@
 void
 set_vt_ref_insert_size (int);
 
+/**
+ * @brief Default for vt_sev_insert_size.
+ *
+ * There are about 80k vt_severities.
+ */
+#define VT_SEV_INSERT_SIZE_DEFAULT 100000
+
+void
+set_vt_sev_insert_size (int);
+
 const char *
 get_osp_vt_update_socket ();