[Ingest] Support yaml variables in datasource

elastic · Apr 24, 2020 · 41ba84b · 41ba84b
1 parent 9703d85
commit 41ba84b
Show file tree

Hide file tree

Showing 3 changed files with 111 additions and 41 deletions.
diff --git a/x-pack/plugins/ingest_manager/server/services/datasource.ts b/x-pack/plugins/ingest_manager/server/services/datasource.ts
@@ -239,19 +239,11 @@ 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.config) {
-    for (const key of Object.keys(input.config)) {
-      data[key] = input.config[key].value;
-    }
-  }
-  if (stream.config) {
-    for (const key of Object.keys(stream.config)) {
-      data[key] = stream.config[key].value;
-    }
-  }
-  const yaml = safeLoad(createStream(data, pkgStream.buffer.toString()));
+  const yaml = createStream(
+    // Populate template variables from input config and stream config
+    Object.assign({}, input.config, stream.config),
+    pkgStream.buffer.toString()
+  );
   stream.pkg_stream = yaml;
   return { ...stream };
 }

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
@@ -6,29 +6,59 @@
 
 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"]
+    {{#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'],
+      'key.patterns': [
+        {
+          limit: 20,
+          pattern: '*',
+        },
+      ],
+    });
+  });
 });
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,12 +5,60 @@
  */
 
 import Handlebars from 'handlebars';
+import { safeLoad } from 'js-yaml';
+import { DatasourceConfigRecord } from '../../../../common';
 
-interface StreamVars {
-  [k: string]: string | string[];
+function replaceVariablesInYaml(yamlVariables: { [k: string]: any }, yaml: any) {
+  const yamlKeys = Object.keys(yamlVariables);
+  if (yamlKeys.length === 0) {
+    return yaml;
+  }
+
+  Object.entries(yaml).forEach(([key, value]: [string, any]) => {
+    if (typeof value === 'object') {
+      yaml[key] = replaceVariablesInYaml(yamlVariables, value);
+    }
+
+    if (typeof value === 'string' && yamlKeys.indexOf(value) >= 0) {
+      yaml[key] = yamlVariables[value];
+    }
+  });
+
+  return yaml;
 }
 
-export function createStream(vars: StreamVars, streamTemplate: string) {
-  const template = Handlebars.compile(streamTemplate);
-  return template(vars);
+export function createStream(variables: DatasourceConfigRecord, streamTemplate: string) {
+  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) {
+      throw new Error('Invalid key');
+    }
+
+    let varPart = acc;
+    for (const keyPart of keyParts) {
+      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 });
+
+  const template = Handlebars.compile(streamTemplate, { noEscape: true });
+  const stream = template(vars);
+  const yamlFromStream = safeLoad(stream, {});
+
+  return replaceVariablesInYaml(yamlValues, yamlFromStream);
 }