[Ingest] Support yaml variables in datasource (#64459)

x-pack/plugins/ingest_manager/server/services/datasource.ts

@@ -4,7 +4,6 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 import { SavedObjectsClientContract } from 'src/core/server';
-import { safeLoad } from 'js-yaml';
 import { AuthenticatedUser } from '../../../security/server';
 import {
   DeleteDatasourcesResponse,
@@ -239,20 +238,13 @@ async function _assignPackageStreamToStream(
     throw new Error(`Stream template not found for dataset ${dataset}`);
   }
 
-  // Populate template variables from input config and stream config
-  const data: { [k: string]: string | string[] } = {};
-  if (input.vars) {
-    for (const key of Object.keys(input.vars)) {
-      data[key] = input.vars[key].value;
-    }
-  }
-  if (stream.vars) {
-    for (const key of Object.keys(stream.vars)) {
-      data[key] = stream.vars[key].value;
-    }
-  }
-  const yaml = safeLoad(createStream(data, pkgStream.buffer.toString()));
+  const yaml = createStream(
+    // Populate template variables from input vars and stream vars
+    Object.assign({}, input.vars, stream.vars),
+    pkgStream.buffer.toString()
+  );
   stream.agent_stream = yaml;
+
   return { ...stream };
 }
 

x-pack/plugins/ingest_manager/server/services/epm/agent/agent.test.ts

@@ -6,29 +6,61 @@
 
 import { createStream } from './agent';
 
-test('Test creating a stream from template', () => {
-  const streamTemplate = `
-input: log
-paths:
-{{#each paths}}
-  - {{this}}
-{{/each}}
-exclude_files: [".gz$"]
-processors:
-  - add_locale: ~
-  `;
-  const vars = {
-    paths: ['/usr/local/var/log/nginx/access.log'],
-  };
+describe('createStream', () => {
+  it('should work', () => {
+    const streamTemplate = `
+    input: log
+    paths:
+    {{#each paths}}
+      - {{this}}
+    {{/each}}
+    exclude_files: [".gz$"]
+    processors:
+      - add_locale: ~
+      `;
+    const vars = {
+      paths: { value: ['/usr/local/var/log/nginx/access.log'] },
+    };
 
-  const output = createStream(vars, streamTemplate);
+    const output = createStream(vars, streamTemplate);
+    expect(output).toEqual({
+      input: 'log',
+      paths: ['/usr/local/var/log/nginx/access.log'],
+      exclude_files: ['.gz$'],
+      processors: [{ add_locale: null }],
+    });
+  });
 
-  expect(output).toBe(`
-input: log
-paths:
-  - /usr/local/var/log/nginx/access.log
-exclude_files: [".gz$"]
-processors:
-  - add_locale: ~
-  `);
+  it('should support yaml values', () => {
+    const streamTemplate = `
+    input: redis/metrics
+    metricsets: ["key"]
+    test: null
+    {{#if key.patterns}}
+    key.patterns: {{key.patterns}}
+    {{/if}}
+      `;
+    const vars = {
+      'key.patterns': {
+        type: 'yaml',
+        value: `
+        - limit: 20
+          pattern: '*'
+        `,
+      },
+    };
+
+    const output = createStream(vars, streamTemplate);
+    expect(output).toEqual({
+      input: 'redis/metrics',
+      metricsets: ['key'],
+      test: null,
+      'key.patterns': [
+        {
+          limit: 20,
+          pattern: '*',
+        },
+      ],
+    });
+  });
 });

x-pack/plugins/ingest_manager/server/services/epm/agent/agent.ts

@@ -5,12 +5,71 @@
  */
 
 import Handlebars from 'handlebars';
+import { safeLoad } from 'js-yaml';
+import { DatasourceConfigRecord } from '../../../../common';
 
-interface StreamVars {
-  [k: string]: string | string[];
+function isValidKey(key: string) {
+  return key !== '__proto__' && key !== 'constructor' && key !== 'prototype';
 }
 
-export function createStream(vars: StreamVars, streamTemplate: string) {
-  const template = Handlebars.compile(streamTemplate);
-  return template(vars);
+function replaceVariablesInYaml(yamlVariables: { [k: string]: any }, yaml: any) {
+  if (Object.keys(yamlVariables).length === 0 || !yaml) {
+    return yaml;
+  }
+
+  Object.entries(yaml).forEach(([key, value]: [string, any]) => {
+    if (typeof value === 'object') {
+      yaml[key] = replaceVariablesInYaml(yamlVariables, value);
+    }
+    if (typeof value === 'string' && value in yamlVariables) {
+      yaml[key] = yamlVariables[value];
+    }
+  });
+
+  return yaml;
+}
+
+function buildTemplateVariables(variables: DatasourceConfigRecord) {
+  const yamlValues: { [k: string]: any } = {};
+  const vars = Object.entries(variables).reduce((acc, [key, recordEntry]) => {
+    // support variables with . like key.patterns
+    const keyParts = key.split('.');
+    const lastKeyPart = keyParts.pop();
+
+    if (!lastKeyPart || !isValidKey(lastKeyPart)) {
+      throw new Error('Invalid key');
+    }
+
+    let varPart = acc;
+    for (const keyPart of keyParts) {
+      if (!isValidKey(keyPart)) {
+        throw new Error('Invalid key');
+      }
+      if (!varPart[keyPart]) {
+        varPart[keyPart] = {};
+      }
+      varPart = varPart[keyPart];
+    }
+
+    if (recordEntry.type && recordEntry.type === 'yaml') {
+      const yamlKeyPlaceholder = `##${key}##`;
+      varPart[lastKeyPart] = `"${yamlKeyPlaceholder}"`;
+      yamlValues[yamlKeyPlaceholder] = recordEntry.value ? safeLoad(recordEntry.value) : null;
+    } else {
+      varPart[lastKeyPart] = recordEntry.value;
+    }
+    return acc;
+  }, {} as { [k: string]: any });
+
+  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, {});
+
+  return replaceVariablesInYaml(yamlValues, yamlFromStream);
 }