[Ingest] Support root level yaml variables in agent stream template (#66120)

nchaulet · nchaulet · commit 64017f0ffeec · 2020-05-13T09:12:26.000-04:00
diff --git a/x-pack/plugins/ingest_manager/server/services/epm/agent/agent.test.ts b/x-pack/plugins/ingest_manager/server/services/epm/agent/agent.test.ts
@@ -9,18 +9,18 @@ import { createStream } from './agent';
 describe('createStream', () => {
   it('should work', () => {
     const streamTemplate = `
-    input: log
-    paths:
-    {{#each paths}}
-      - {{this}}
-    {{/each}}
-    exclude_files: [".gz$"]
-    processors:
-      - add_locale: ~
-    password: {{password}}
-    {{#if password}}
-    hidden_password: {{password}}
-    {{/if}}
+input: log
+paths:
+{{#each paths}}
+  - {{this}}
+{{/each}}
+exclude_files: [".gz$"]
+processors:
+  - add_locale: ~
+password: {{password}}
+{{#if password}}
+hidden_password: {{password}}
+{{/if}}
       `;
     const vars = {
       paths: { value: ['/usr/local/var/log/nginx/access.log'] },
@@ -39,13 +39,16 @@ describe('createStream', () => {
 
   it('should support yaml values', () => {
     const streamTemplate = `
-    input: redis/metrics
-    metricsets: ["key"]
-    test: null
-    password: {{password}}
-    {{#if key.patterns}}
-    key.patterns: {{key.patterns}}
-    {{/if}}
+input: redis/metrics
+metricsets: ["key"]
+test: null
+password: {{password}}
+{{custom}}
+custom: {{ custom }}
+{{#if key.patterns}}
+key.patterns: {{key.patterns}}
+{{/if}}
+{{ testEmpty }}
       `;
     const vars = {
       'key.patterns': {
@@ -55,6 +58,12 @@ describe('createStream', () => {
           pattern: '*'
         `,
       },
+      custom: {
+        type: 'yaml',
+        value: `
+foo: bar
+        `,
+      },
       password: { type: 'password', value: '' },
     };
 
@@ -70,6 +79,8 @@ describe('createStream', () => {
         },
       ],
       password: '',
+      foo: 'bar',
+      custom: { foo: 'bar' },
     });
   });
 });
diff --git a/x-pack/plugins/ingest_manager/server/services/epm/agent/agent.ts b/x-pack/plugins/ingest_manager/server/services/epm/agent/agent.ts
@@ -5,9 +5,32 @@
  */
 
 import Handlebars from 'handlebars';
-import { safeLoad } from 'js-yaml';
+import { safeLoad, safeDump } from 'js-yaml';
 import { DatasourceConfigRecord } from '../../../../common';
 
+export function createStream(variables: DatasourceConfigRecord, streamTemplate: string) {
+  const { vars, yamlValues } = buildTemplateVariables(variables, streamTemplate);
+
+  const template = Handlebars.compile(streamTemplate, { noEscape: true });
+  let stream = template(vars);
+  stream = replaceRootLevelYamlVariables(yamlValues, stream);
+
+  const yamlFromStream = safeLoad(stream, {});
+
+  // Hack to keep empty string ('') values around in the end yaml because
+  // `safeLoad` replaces empty strings with null
+  const patchedYamlFromStream = Object.entries(yamlFromStream).reduce((acc, [key, value]) => {
+    if (value === null && typeof vars[key] === 'string' && vars[key].trim() === '') {
+      acc[key] = '';
+    } else {
+      acc[key] = value;
+    }
+    return acc;
+  }, {} as { [k: string]: any });
+
+  return replaceVariablesInYaml(yamlValues, patchedYamlFromStream);
+}
+
 function isValidKey(key: string) {
   return key !== '__proto__' && key !== 'constructor' && key !== 'prototype';
 }
@@ -29,7 +52,7 @@ function replaceVariablesInYaml(yamlVariables: { [k: string]: any }, yaml: any)
   return yaml;
 }
 
-function buildTemplateVariables(variables: DatasourceConfigRecord) {
+function buildTemplateVariables(variables: DatasourceConfigRecord, streamTemplate: string) {
   const yamlValues: { [k: string]: any } = {};
   const vars = Object.entries(variables).reduce((acc, [key, recordEntry]) => {
     // support variables with . like key.patterns
@@ -64,23 +87,15 @@ function buildTemplateVariables(variables: DatasourceConfigRecord) {
   return { vars, yamlValues };
 }
 
-export function createStream(variables: DatasourceConfigRecord, streamTemplate: string) {
-  const { vars, yamlValues } = buildTemplateVariables(variables);
-
-  const template = Handlebars.compile(streamTemplate, { noEscape: true });
-  const stream = template(vars);
-  const yamlFromStream = safeLoad(stream, {});
+function replaceRootLevelYamlVariables(yamlVariables: { [k: string]: any }, yamlTemplate: string) {
+  if (Object.keys(yamlVariables).length === 0 || !yamlTemplate) {
+    return yamlTemplate;
+  }
 
-  // Hack to keep empty string ('') values around in the end yaml because
-  // `safeLoad` replaces empty strings with null
-  const patchedYamlFromStream = Object.entries(yamlFromStream).reduce((acc, [key, value]) => {
-    if (value === null && typeof vars[key] === 'string' && vars[key].trim() === '') {
-      acc[key] = '';
-    } else {
-      acc[key] = value;
-    }
-    return acc;
-  }, {} as { [k: string]: any });
+  let patchedTemplate = yamlTemplate;
+  Object.entries(yamlVariables).forEach(([key, val]) => {
+    patchedTemplate = patchedTemplate.replace(new RegExp(`^"${key}"`, 'gm'), safeDump(val));
+  });
 
-  return replaceVariablesInYaml(yamlValues, patchedYamlFromStream);
+  return patchedTemplate;
 }
