From f349d0a398eb7944957b3fa45647be0b1fd62cb6 Mon Sep 17 00:00:00 2001
From: James Gowdy <jgowdy@elastic.co>
Date: Mon, 21 Dec 2020 16:56:35 +0000
Subject: [PATCH] [ML] Fix cloning of partition field in per-partition
 categorization jobs (#86635)

* [ML] Fix cloning of partition field in per-partition categorization jobs

* fixing cloned bucket span
---
 .../job_creator/categorization_job_creator.ts | 59 ++++++++++++-------
 .../new_job/common/job_creator/job_creator.ts |  2 +-
 2 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/x-pack/plugins/ml/public/application/jobs/new_job/common/job_creator/categorization_job_creator.ts b/x-pack/plugins/ml/public/application/jobs/new_job/common/job_creator/categorization_job_creator.ts
index a0f2832dc9f97a..aaf20ff1571eea 100644
--- a/x-pack/plugins/ml/public/application/jobs/new_job/common/job_creator/categorization_job_creator.ts
+++ b/x-pack/plugins/ml/public/application/jobs/new_job/common/job_creator/categorization_job_creator.ts
@@ -160,26 +160,6 @@ export class CategorizationJobCreator extends JobCreator {
     return this._categorizationAnalyzer;
   }
 
-  public cloneFromExistingJob(job: Job, datafeed: Datafeed) {
-    this._overrideConfigs(job, datafeed);
-    this.createdBy = CREATED_BY_LABEL.CATEGORIZATION;
-    const detectors = getRichDetectors(job, datafeed, this.additionalFields, false);
-
-    const dtr = detectors[0];
-    if (detectors.length && dtr.agg !== null && dtr.field !== null) {
-      this._detectorType =
-        dtr.agg.id === ML_JOB_AGGREGATION.COUNT
-          ? ML_JOB_AGGREGATION.COUNT
-          : ML_JOB_AGGREGATION.RARE;
-
-      const bs = job.analysis_config.bucket_span;
-      this.setDetectorType(this._detectorType);
-      // set the bucketspan back to the original value
-      // as setDetectorType applies a default
-      this.bucketSpan = bs;
-    }
-  }
-
   public get categorizationPerPartitionField() {
     return this._partitionFieldName;
   }
@@ -204,4 +184,43 @@ export class CategorizationJobCreator extends JobCreator {
       }
     }
   }
+
+  // override the setter and getter for the per-partition toggle
+  // so we can remove the partition field in the wizard when
+  // per-partition categorization is disabled.
+  public get perPartitionCategorization() {
+    return this._job_config.analysis_config.per_partition_categorization?.enabled === true;
+  }
+
+  public set perPartitionCategorization(enabled: boolean) {
+    this._initPerPartitionCategorization();
+    this._job_config.analysis_config.per_partition_categorization!.enabled = enabled;
+    if (enabled === false) {
+      this.categorizationPerPartitionField = null;
+    }
+  }
+
+  public cloneFromExistingJob(job: Job, datafeed: Datafeed) {
+    this._overrideConfigs(job, datafeed);
+    this.createdBy = CREATED_BY_LABEL.CATEGORIZATION;
+    const detectors = getRichDetectors(job, datafeed, this.additionalFields, false);
+
+    const dtr = detectors[0];
+    if (dtr !== undefined && dtr.agg !== null && dtr.field !== null) {
+      const detectorType =
+        dtr.agg.id === ML_JOB_AGGREGATION.COUNT
+          ? ML_JOB_AGGREGATION.COUNT
+          : ML_JOB_AGGREGATION.RARE;
+
+      const bs = job.analysis_config.bucket_span;
+      this.setDetectorType(detectorType);
+      if (dtr.partitionField !== null) {
+        this.categorizationPerPartitionField = dtr.partitionField.id;
+      }
+
+      // set the bucketspan back to the original value
+      // as setDetectorType applies a default
+      this.bucketSpan = bs;
+    }
+  }
 }
diff --git a/x-pack/plugins/ml/public/application/jobs/new_job/common/job_creator/job_creator.ts b/x-pack/plugins/ml/public/application/jobs/new_job/common/job_creator/job_creator.ts
index ba03e86724a150..c3b4471269bb86 100644
--- a/x-pack/plugins/ml/public/application/jobs/new_job/common/job_creator/job_creator.ts
+++ b/x-pack/plugins/ml/public/application/jobs/new_job/common/job_creator/job_creator.ts
@@ -631,7 +631,7 @@ export class JobCreator {
     return JSON.stringify(this._datafeed_config, null, 2);
   }
 
-  private _initPerPartitionCategorization() {
+  protected _initPerPartitionCategorization() {
     if (this._job_config.analysis_config.per_partition_categorization === undefined) {
       this._job_config.analysis_config.per_partition_categorization = {};
     }