From 8a085957ef8a0c329503000cfa1730e57d3ff6ee Mon Sep 17 00:00:00 2001
From: Aseem Bansal <asmbansal2@gmail.com>
Date: Tue, 23 Aug 2022 19:18:31 +0530
Subject: [PATCH] feat(upgrade): allow setting batch args

---
 charts/datahub/Chart.yaml                                     | 2 +-
 .../datahub-upgrade/datahub-restore-indices-job-template.yml  | 4 ++++
 .../datahub/templates/datahub-upgrade/datahub-upgrade-job.yml | 4 ++--
 charts/datahub/values.yaml                                    | 2 ++
 4 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/charts/datahub/Chart.yaml b/charts/datahub/Chart.yaml
index eb260475e..863a3cbb3 100644
--- a/charts/datahub/Chart.yaml
+++ b/charts/datahub/Chart.yaml
@@ -4,7 +4,7 @@ description: A Helm chart for LinkedIn DataHub
 type: application
 # This is the chart version. This version number should be incremented each time you make changes
 # to the chart and its templates, including the app version.
-version: 0.2.89
+version: 0.2.90
 # This is the version number of the application being deployed. This version number should be
 # incremented each time you make changes to the application.
 appVersion: 0.8.43
diff --git a/charts/datahub/templates/datahub-upgrade/datahub-restore-indices-job-template.yml b/charts/datahub/templates/datahub-upgrade/datahub-restore-indices-job-template.yml
index d47dcab7a..4dfcb2f08 100644
--- a/charts/datahub/templates/datahub-upgrade/datahub-restore-indices-job-template.yml
+++ b/charts/datahub/templates/datahub-upgrade/datahub-restore-indices-job-template.yml
@@ -67,6 +67,10 @@ spec:
               args:
                 - "-u"
                 - "RestoreIndices"
+                - "-a"
+                - "batchSize={{ .Values.datahubUpgrade.batchSize }}"
+                - "-a"
+                - "batchDelayMs={{ .Values.datahubUpgrade.batchDelayMs }}"
               env:
                 {{- include "datahub.upgrade.env" . | nindent 16}}
               {{- with .Values.datahubUpgrade.extraEnvs }}
diff --git a/charts/datahub/templates/datahub-upgrade/datahub-upgrade-job.yml b/charts/datahub/templates/datahub-upgrade/datahub-upgrade-job.yml
index 51550f81b..f22cd43a4 100644
--- a/charts/datahub/templates/datahub-upgrade/datahub-upgrade-job.yml
+++ b/charts/datahub/templates/datahub-upgrade/datahub-upgrade-job.yml
@@ -66,9 +66,9 @@ spec:
             - "-u"
             - "NoCodeDataMigration"
             - "-a"
-            - "batchSize=1000"
+            - "batchSize={{ .Values.datahubUpgrade.batchSize }}"
             - "-a"
-            - "batchDelayMs=100"
+            - "batchDelayMs={{ .Values.datahubUpgrade.batchDelayMs }}"
             - "-a"
             - "dbType={{ .Values.datahubUpgrade.noCodeDataMigration.sqlDbType }}"
           env:
diff --git a/charts/datahub/values.yaml b/charts/datahub/values.yaml
index 0e38c0771..9a897be83 100644
--- a/charts/datahub/values.yaml
+++ b/charts/datahub/values.yaml
@@ -92,6 +92,8 @@ datahubUpgrade:
   image:
     repository: acryldata/datahub-upgrade
     tag: "v0.8.43"
+  batchSize: 1000
+  batchDelayMs: 100
   noCodeDataMigration:
     sqlDbType: "MYSQL"
   podSecurityContext: {}