feat(upgrade-job): add support for yq yaml obj update

Signed-off-by: Niladri Halder <[email protected]>

k8s/upgrade/src/bin/upgrade-job/common/constants.rs

@@ -44,3 +44,6 @@ pub(crate) const TWO_DOT_FOUR: &str = "2.4.0";
 
 /// Version value for the earliest possible 2.5 release.
 pub(crate) const TWO_DOT_FIVE: &str = "2.5.0";
+
+/// Version value for the earliest possible 2.6 release.
+pub(crate) const TWO_DOT_SIX: &str = "2.6.0";

k8s/upgrade/src/bin/upgrade-job/common/error.rs

@@ -601,6 +601,19 @@ pub(crate) enum Error {
         std_err: String,
     },
 
+    /// Error for when the yq command to update a yaml object returns an error.
+    #[snafu(display(
+        "`yq` set-object-command returned an error,\ncommand: {},\nargs: {:?},\nstd_err: {}",
+        command,
+        args,
+        std_err,
+    ))]
+    YqSetObjCommand {
+        command: String,
+        args: Vec<String>,
+        std_err: String,
+    },
+
     /// Error for when we fail to read the entries of a directory.
     #[snafu(display("Failed to read the contents of directory {}: {}", path.display(), source))]
     ReadingDirectoryContents {

k8s/upgrade/src/bin/upgrade-job/helm/values.rs

@@ -1,6 +1,8 @@
 use crate::{
     common::{
-        constants::{TWO_DOT_FIVE, TWO_DOT_FOUR, TWO_DOT_ONE, TWO_DOT_O_RC_ONE, TWO_DOT_THREE},
+        constants::{
+            TWO_DOT_FIVE, TWO_DOT_FOUR, TWO_DOT_ONE, TWO_DOT_O_RC_ONE, TWO_DOT_SIX, TWO_DOT_THREE,
+        },
         error::{Result, SemverParse},
         file::write_to_tempfile,
     },
@@ -57,7 +59,8 @@ where
     // Resultant values yaml for helm upgrade command.
     // Merge the source values with the target values.
     let yq = YqV4::new()?;
-    let upgrade_values_yaml = yq.merge_files(source_values_file.path(), target_values_filepath)?;
+    let upgrade_values_yaml =
+        yq.merge_files(source_values_file.path(), target_values_filepath.as_ref())?;
     let upgrade_values_file: TempFile =
         write_to_tempfile(Some(workdir), upgrade_values_yaml.as_slice())?;
 
@@ -162,6 +165,23 @@ where
         }
     }
 
+    // Special-case values for 2.6.x.
+    let two_dot_six = Version::parse(TWO_DOT_SIX).context(SemverParse {
+        version_string: TWO_DOT_SIX.to_string(),
+    })?;
+    if source_version.ge(&two_dot_o_rc_zero) && source_version.lt(&two_dot_six) {
+        yq.set_obj(
+            YamlKey::try_from(".loki-stack.loki")?,
+            target_values_filepath.as_ref(),
+            upgrade_values_file.path(),
+        )?;
+        yq.set_obj(
+            YamlKey::try_from(".loki-stack.promtail")?,
+            target_values_filepath.as_ref(),
+            upgrade_values_file.path(),
+        )?;
+    }
+
     // Default options.
     // Image tag is set because the high_priority file is the user's source options file.
     // The target's image tag needs to be set for PRODUCT upgrade.

k8s/upgrade/src/bin/upgrade-job/helm/yaml/yq.rs

@@ -1,7 +1,7 @@
 use crate::{
     common::error::{
         NotAValidYamlKeyForStringValue, NotYqV4, RegexCompile, Result, U8VectorToString,
-        YqCommandExec, YqMergeCommand, YqSetCommand, YqVersionCommand,
+        YqCommandExec, YqMergeCommand, YqSetCommand, YqSetObjCommand, YqVersionCommand,
     },
     vec_to_strings,
 };
@@ -109,16 +109,24 @@ impl YqV4 {
     /// preferred over those of the other file's. In case there are values absent in the latter one
     /// which exist in the other file, the values of the other file are taken. The 'latter' file in
     /// this function is the one called 'high_priority' and the other file is the 'low_priority'
-    /// one. The output of the command is stripped of the 'COMPUTED VALUES:' suffix.
+    /// one.
     /// E.g:
-    ///       high_priority file:                         low_priority file:
-    ///       ===================                         ==================
-    ///       foo:                                        foo:
-    ///         bar: "foobar"                               bar: "foobaz"
-    ///         baz:                                        baz:
-    ///           - "alpha"                                   - "gamma"
-    ///           - "beta"                                    - "delta" friend: "ferris"
+    ///       high_priority file:
+    ///       ===================
+    ///       foo:
+    ///         bar: "foobar"
+    ///         baz:
+    ///           - "alpha"
+    ///           - "beta"
     ///
+    ///       low_priority file:
+    ///       ==================
+    ///       foo:
+    ///         bar: "foobaz"
+    ///         baz:
+    ///           - "gamma"
+    ///           - "delta"
+    ///       friend: "ferris"
     ///
     ///       result:
     ///       =======
@@ -225,6 +233,41 @@ impl YqV4 {
         Ok(())
     }
 
+    /// This sets yaml objects to a file from the same yaml object in another file.
+    pub(crate) fn set_obj(&self, key: YamlKey, obj_source: &Path, filepath: &Path) -> Result<()> {
+        let yq_set_obj_args = vec_to_strings![
+            "-i",
+            format!(
+                r#"{} = load("{}"){}"#,
+                key.as_str(),
+                obj_source.to_string_lossy(),
+                key.as_str()
+            ),
+            filepath.to_string_lossy()
+        ];
+        let yq_set_obj_output = self
+            .command()
+            .args(yq_set_obj_args.clone())
+            .output()
+            .context(YqCommandExec {
+                command: self.command_as_str().to_string(),
+                args: yq_set_obj_args.clone(),
+            })?;
+
+        ensure!(
+            yq_set_obj_output.status.success(),
+            YqSetObjCommand {
+                command: self.command_as_str().to_string(),
+                args: yq_set_obj_args,
+                std_err: str::from_utf8(yq_set_obj_output.stderr.as_slice())
+                    .context(U8VectorToString)?
+                    .to_string()
+            }
+        );
+
+        Ok(())
+    }
+
     /// Returns an std::process::Command using the command_as_str member's value.
     fn command(&self) -> Command {
         Command::new(self.command_name.clone())