From d314a67b70d3f3712380f375e88f3b081fb530dc Mon Sep 17 00:00:00 2001
From: Nicolas Chaulet <nicolas.chaulet@elastic.co>
Date: Wed, 21 Sep 2022 15:21:55 -0400
Subject: [PATCH] [Fleet] Add experimental data stream features support to
 simplified package policy API

---
 .../simplified_package_policy_helper.test.ts  | 31 +++++++++++++++++++
 .../simplified_package_policy_helper.ts       | 10 +++++-
 .../server/routes/package_policy/handlers.ts  |  7 +++--
 .../server/types/models/package_policy.ts     | 21 +++++++------
 4 files changed, 56 insertions(+), 13 deletions(-)

diff --git a/x-pack/plugins/fleet/common/services/simplified_package_policy_helper.test.ts b/x-pack/plugins/fleet/common/services/simplified_package_policy_helper.test.ts
index eec59a647c39d..7d5a7966f91c3 100644
--- a/x-pack/plugins/fleet/common/services/simplified_package_policy_helper.test.ts
+++ b/x-pack/plugins/fleet/common/services/simplified_package_policy_helper.test.ts
@@ -115,5 +115,36 @@ describe('toPackagePolicy', () => {
         'nginx-nginx/metrics': ['nginx.stubstatus'],
       });
     });
+
+    it('should to pass experimental_data_stream_features', () => {
+      const res = simplifiedPackagePolicytoNewPackagePolicy(
+        {
+          name: 'nginx-1',
+          namespace: 'default',
+          policy_id: 'policy123',
+          description: 'Test description',
+        },
+        nginxPackageInfo as unknown as PackageInfo,
+        {
+          experimental_data_stream_features: [
+            {
+              data_stream: 'logs-nginx.access',
+              features: {
+                synthetic_source: true,
+              },
+            },
+          ],
+        }
+      );
+
+      expect(res.package?.experimental_data_stream_features).toEqual([
+        {
+          data_stream: 'logs-nginx.access',
+          features: {
+            synthetic_source: true,
+          },
+        },
+      ]);
+    });
   });
 });
diff --git a/x-pack/plugins/fleet/common/services/simplified_package_policy_helper.ts b/x-pack/plugins/fleet/common/services/simplified_package_policy_helper.ts
index c0d05d1eeab2a..a17796e4b8561 100644
--- a/x-pack/plugins/fleet/common/services/simplified_package_policy_helper.ts
+++ b/x-pack/plugins/fleet/common/services/simplified_package_policy_helper.ts
@@ -11,6 +11,7 @@ import type {
   PackagePolicyConfigRecord,
   NewPackagePolicy,
   PackageInfo,
+  ExperimentalDataStreamFeature,
 } from '../types';
 import { PackagePolicyValidationError } from '../errors';
 
@@ -66,7 +67,10 @@ type InputMap = Map<string, { input: NewPackagePolicyInput; streams: StreamsMap
 
 export function simplifiedPackagePolicytoNewPackagePolicy(
   data: SimplifiedPackagePolicy,
-  packageInfo: PackageInfo
+  packageInfo: PackageInfo,
+  options?: {
+    experimental_data_stream_features: ExperimentalDataStreamFeature[];
+  }
 ): NewPackagePolicy {
   const {
     policy_id: policyId,
@@ -77,6 +81,10 @@ export function simplifiedPackagePolicytoNewPackagePolicy(
     vars: packageLevelVars,
   } = data;
   const packagePolicy = packageToPackagePolicy(packageInfo, policyId, namespace, name, description);
+  if (packagePolicy.package && options?.experimental_data_stream_features) {
+    packagePolicy.package.experimental_data_stream_features =
+      options.experimental_data_stream_features;
+  }
 
   // Build a input and streams Map to easily find package policy stream
   const inputMap: InputMap = new Map();
diff --git a/x-pack/plugins/fleet/server/routes/package_policy/handlers.ts b/x-pack/plugins/fleet/server/routes/package_policy/handlers.ts
index 765c6c649db25..d96d574f77bed 100644
--- a/x-pack/plugins/fleet/server/routes/package_policy/handlers.ts
+++ b/x-pack/plugins/fleet/server/routes/package_policy/handlers.ts
@@ -205,7 +205,9 @@ export const createPackagePolicyHandler: FleetRequestHandler<
         pkgVersion: pkg.version,
         ignoreUnverified: force,
       });
-      newPackagePolicy = simplifiedPackagePolicytoNewPackagePolicy(newPolicy, pkgInfo);
+      newPackagePolicy = simplifiedPackagePolicytoNewPackagePolicy(newPolicy, pkgInfo, {
+        experimental_data_stream_features: pkg.experimental_data_stream_features,
+      });
     } else {
       newPackagePolicy = await packagePolicyService.enrichPolicyWithDefaultsFromPackage(soClient, {
         ...newPolicy,
@@ -293,7 +295,8 @@ export const updatePackagePolicyHandler: RequestHandler<
       });
       newData = simplifiedPackagePolicytoNewPackagePolicy(
         body as unknown as SimplifiedPackagePolicy,
-        pkgInfo
+        pkgInfo,
+        { experimental_data_stream_features: pkg.experimental_data_stream_features }
       );
     } else {
       // removed fields not recognized by schema
diff --git a/x-pack/plugins/fleet/server/types/models/package_policy.ts b/x-pack/plugins/fleet/server/types/models/package_policy.ts
index 9763d09e20555..82ff572cd1926 100644
--- a/x-pack/plugins/fleet/server/types/models/package_policy.ts
+++ b/x-pack/plugins/fleet/server/types/models/package_policy.ts
@@ -79,6 +79,15 @@ const PackagePolicyInputsSchema = {
   streams: schema.arrayOf(schema.object(PackagePolicyStreamsSchema)),
 };
 
+const ExperimentalDataStreamFeatures = schema.arrayOf(
+  schema.object({
+    data_stream: schema.string(),
+    features: schema.object({
+      synthetic_source: schema.boolean(),
+    }),
+  })
+);
+
 const PackagePolicyBaseSchema = {
   name: schema.string(),
   description: schema.maybe(schema.string()),
@@ -90,16 +99,7 @@ const PackagePolicyBaseSchema = {
       name: schema.string(),
       title: schema.string(),
       version: schema.string(),
-      experimental_data_stream_features: schema.maybe(
-        schema.arrayOf(
-          schema.object({
-            data_stream: schema.string(),
-            features: schema.object({
-              synthetic_source: schema.boolean(),
-            }),
-          })
-        )
-      ),
+      experimental_data_stream_features: schema.maybe(ExperimentalDataStreamFeatures),
     })
   ),
   // Deprecated TODO create remove issue
@@ -172,6 +172,7 @@ export const SimplifiedCreatePackagePolicyRequestBodySchema = schema.object({
   package: schema.object({
     name: schema.string(),
     version: schema.string(),
+    experimental_data_stream_features: schema.maybe(ExperimentalDataStreamFeatures),
   }),
   force: schema.maybe(schema.boolean()),
   vars: schema.maybe(SimplifiedVarsSchema),