apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: {{ template "airflow.fullname" . }}-worker
  labels:
    app.kubernetes.io/name: {{ include "airflow.name" . }}
    helm.sh/chart: {{ include "airflow.chart" . }}
    app.kubernetes.io/instance: {{ .Release.Name }}
    app.kubernetes.io/managed-by: {{ .Release.Service }}
    app.kubernetes.io/component: worker
spec:
  {{- if .Values.airflow.worker.podManagementPolicy }}
  podManagementPolicy: {{ .Values.airflow.worker.podManagementPolicy | quote }}
  {{- end }}
  selector:
    matchLabels:
      app.kubernetes.io/name: {{ include "airflow.name" . }}
      app.kubernetes.io/instance: {{ .Release.Name }}
      app.kubernetes.io/component: worker
  serviceName: {{ template "airflow.fullname" . }}-headless
  replicas: {{ .Values.airflow.worker.replicas }}
  updateStrategy:
    type: {{ .Values.updateStrategy | quote }}
    {{- if (eq "Recreate" .Values.updateStrategy) }}
    rollingUpdate: null
    {{- else if .Values.rollingUpdatePartition }}
    rollingUpdate:
      partition: {{ .Values.rollingUpdatePartition }}
    {{- end }}
  template:
    metadata:
      annotations: {{- include "airflow.tplValue" (dict "value" .Values.airflow.worker.podAnnotations "context" $) | nindent 8 }}
      labels:
        app.kubernetes.io/name: {{ include "airflow.name" . }}
        app.kubernetes.io/instance: {{ .Release.Name }}
        app.kubernetes.io/component: worker
    spec:
      {{- include "airflow.imagePullSecrets" . | nindent 6 }}
      {{- if .Values.securityContext.enabled }}
      securityContext:
        fsGroup: {{ .Values.securityContext.fsGroup }}
        runAsUser: {{ .Values.securityContext.runAsUser }}
      {{- end }}
      {{- if .Values.affinity }}
      affinity: {{- include "airflow.tplValue" (dict "value" .Values.affinity "context" $) | nindent 8 }}
      {{- else }}
      affinity:
        podAffinity: {{- include "common.affinities.pods" (dict "type" .Values.podAffinityPreset "component" "worker" "context" $) | nindent 10 }}
        podAntiAffinity: {{- include "common.affinities.pods" (dict "type" .Values.podAntiAffinityPreset "component" "worker" "context" $) | nindent 10 }}
        nodeAffinity: {{- include "common.affinities.nodes" (dict "type" .Values.nodeAffinityPreset.type "key" .Values.nodeAffinityPreset.key "values" .Values.nodeAffinityPreset.values) | nindent 10 }}
      {{- end }}
      {{- if .Values.nodeSelector }}
      nodeSelector: {{- include "airflow.tplValue" (dict "value" .Values.nodeSelector "context" $) | nindent 8 }}
      {{- end }}
      {{- if .Values.tolerations }}
      tolerations: {{- include "airflow.tplValue" (dict "value" .Values.tolerations "context" $) | nindent 8 }}
      {{- end }}
      {{- if or .Values.airflow.cloneDagFilesFromGit.enabled .Values.airflow.clonePluginsFromGit.enabled }}
      initContainers:
        - name: git-clone-repository
          image: {{ include "git.image" . }}
          imagePullPolicy: {{ .Values.git.pullPolicy | quote }}
          {{- if .Values.airflow.gitCloneContainer.command }}
          command: {{- include "airflow.tplValue" (dict "value" .Values.airflow.gitCloneContainer.command "context" $) | nindent 12 }}
          {{- else }}
          command:
            - /bin/bash
            - -ec
            - |
              {{- if .Values.airflow.cloneDagFilesFromGit.enabled }}
              {{- if and .Values.airflow.cloneDagFilesFromGit.repository .Values.airflow.cloneDagFilesFromGit.branch }}
                git clone {{ .Values.airflow.cloneDagFilesFromGit.repository }} --branch {{ .Values.airflow.cloneDagFilesFromGit.branch }} /dags
              {{- end }}
              {{- range .Values.airflow.cloneDagFilesFromGit.repositories }}
                git clone {{ .repository }} --branch {{ .branch }} /dags-{{ .name | kebabcase }}
              {{- end }}
              {{- end }}
              {{- if .Values.airflow.clonePluginsFromGit.enabled }}
                git clone {{ .Values.airflow.clonePluginsFromGit.repository }} --branch {{ .Values.airflow.clonePluginsFromGit.branch }} /plugins
              {{- end }}
          {{- end }}
          {{- if .Values.airflow.gitCloneContainer.args }}
          args: {{- include "airflow.tplValue" (dict "value" .Values.airflow.gitCloneContainer.args "context" $) | nindent 12 }}
          {{- end }}
          volumeMounts:
            {{- if .Values.airflow.cloneDagFilesFromGit.enabled }}
            {{- if and .Values.airflow.cloneDagFilesFromGit.repository .Values.airflow.cloneDagFilesFromGit.branch }}
            - name: git-cloned-dag-files
              mountPath: /dags
            {{- end }}
            {{- range .Values.airflow.cloneDagFilesFromGit.repositories }}
            - name: git-cloned-dag-files-{{ .name | kebabcase }}
              mountPath: /dags-{{ .name | kebabcase }}
            {{- end }}
            {{- end }}
            {{- if .Values.airflow.clonePluginsFromGit.enabled }}
            - name: git-cloned-plugins
              mountPath: /plugins
            {{- end }}
            {{- if .Values.airflow.gitCloneContainer.extraVolumeMounts }}
            {{- include "airflow.tplValue" (dict "value" .Values.airflow.gitCloneContainer.extraVolumeMounts "context" $) | nindent 12 }}
            {{- end }}
          {{- if .Values.airflow.gitCloneContainer.extraEnvVars }}
          env: {{- include "airflow.tplValue" (dict "value" .Values.airflow.gitCloneContainer.extraEnvVars "context" $) | nindent 12 }}
          {{- end }}
          {{- if or .Values.airflow.gitCloneContainer.extraEnvVarsCM .Values.airflow.gitCloneContainer.extraEnvVarsSecret }}
          envFrom:
            {{- if .Values.airflow.gitCloneContainer.extraEnvVarsCM }}
            - configMapRef:
                name: {{ .Values.airflow.gitCloneContainer.extraEnvVarsCM }}
            {{- end }}
            {{- if .Values.airflow.gitCloneContainer.extraEnvVarsSecret }}
            - secretRef:
                name: {{ .Values.airflow.gitCloneContainer.extraEnvVarsSecret }}
            {{- end }}
          {{- end }}
      containers:
        - name: git-repo-syncer
          image: "{{ template "git.image" . }}"
          imagePullPolicy: {{ .Values.git.pullPolicy | quote }}
          {{- if .Values.airflow.gitSyncContainer.command }}
          command: {{- include "airflow.tplValue" (dict "value" .Values.airflow.gitSyncContainer.command "context" $) | nindent 12 }}
          {{- else }}
          command:
            - /bin/bash
            - -ec
            - |
              while true; do
                  {{- if .Values.airflow.cloneDagFilesFromGit.enabled }}
                  {{- if and .Values.airflow.cloneDagFilesFromGit.repository .Values.airflow.cloneDagFilesFromGit.branch }}
                  cd /dags && git pull origin {{ .Values.airflow.cloneDagFilesFromGit.branch }}
                  {{- end }}
                  {{- range .Values.airflow.cloneDagFilesFromGit.repositories }}
                  cd /dags-{{ .name | kebabcase }} && git pull origin {{ .branch }}
                  {{- end }}
                  {{- end }}
                  {{- if .Values.airflow.clonePluginsFromGit.enabled }}
                  cd /plugins && git pull origin {{ .Values.airflow.clonePluginsFromGit.branch }}
                  {{- end }}
                  sleep {{ default "60" .Values.airflow.gitSyncInterval }}
              done
          {{- end }}
          {{- if .Values.airflow.gitSyncContainer.args }}
          args: {{- include "airflow.tplValue" (dict "value" .Values.airflow.gitSyncContainer.args "context" $) | nindent 12 }}
          {{- end }}
          volumeMounts:
            {{- if .Values.airflow.cloneDagFilesFromGit.enabled }}
            {{- if and .Values.airflow.cloneDagFilesFromGit.repository .Values.airflow.cloneDagFilesFromGit.branch }}
            - name: git-cloned-dag-files
              mountPath: /dags
            {{- end }}
            {{- range .Values.airflow.cloneDagFilesFromGit.repositories }}
            - name: git-cloned-dag-files-{{ .name | kebabcase }}
              mountPath: /dags-{{ .name | kebabcase }}
            {{- end }}
            {{- end }}
            {{- if .Values.airflow.clonePluginsFromGit.enabled }}
            - name: git-cloned-plugins
              mountPath: /plugins
            {{- end }}
            {{- if .Values.airflow.gitSyncContainer.extraVolumeMounts }}
            {{- include "airflow.tplValue" (dict "value" .Values.airflow.gitSyncContainer.extraVolumeMounts "context" $) | nindent 12 }}
            {{- end }}
          {{- if .Values.airflow.gitSyncContainer.extraEnvVars }}
          env: {{- include "airflow.tplValue" (dict "value" .Values.airflow.gitSyncContainer.extraEnvVars "context" $) | nindent 12 }}
          {{- end }}
          {{- if or .Values.airflow.gitSyncContainer.extraEnvVarsCM .Values.airflow.gitSyncContainer.extraEnvVarsSecret }}
          envFrom:
            {{- if .Values.airflow.gitSyncContainer.extraEnvVarsCM }}
            - configMapRef:
                name: {{ .Values.airflow.gitSyncContainer.extraEnvVarsCM }}
            {{- end }}
            {{- if .Values.airflow.gitSyncContainer.extraEnvVarsSecret }}
            - secretRef:
                name: {{ .Values.airflow.gitSyncContainer.extraEnvVarsSecret }}
            {{- end }}
          {{- end }}
      {{- else }}
      containers:
      {{- end }}
        - name: airflow-worker
          image: {{ include "airflow.workerImage" . }}
          imagePullPolicy: {{ .Values.workerImage.pullPolicy | quote }}
          {{- if .Values.airflow.worker.command }}
          command: {{- include "airflow.tplValue" (dict "value" .Values.airflow.worker.command "context" $) | nindent 12 }}
          {{- end }}
          {{- if .Values.airflow.worker.args }}
          args: {{- include "airflow.tplValue" (dict "value" .Values.airflow.worker.args "context" $) | nindent 12 }}
          {{- end }}
          {{- if .Values.airflow.worker.resources }}
          resources: {{- toYaml .Values.airflow.worker.resources | nindent 12 }}
          {{- end }}
          env:
            {{- if .Values.workerImage.debug }}
            - name: BASH_DEBUG
              value: "1"
            - name: NAMI_DEBUG
              value: "1"
            - name: NAMI_LOG_LEVEL
              value: "trace8"
            {{- end }}
            - name: AIRFLOW_DATABASE_NAME
              {{- if .Values.postgresql.enabled }}
              value: {{ .Values.postgresql.postgresqlDatabase }}
              {{- else }}
              value: {{ .Values.externalDatabase.database | quote }}
              {{- end }}
            - name: AIRFLOW_DATABASE_USERNAME
              {{- if .Values.postgresql.enabled }}
              value: {{ .Values.postgresql.postgresqlUsername }}
              {{- else }}
              value: {{ .Values.externalDatabase.user | quote }}
              {{- end }}
            - name: AIRFLOW_DATABASE_PASSWORD
              valueFrom:
                secretKeyRef:
                  name: {{ template "airflow.postgresql.secretName" . }}
                  key: postgresql-password
            - name: AIRFLOW_DATABASE_HOST
              {{- if .Values.postgresql.enabled }}
              value: {{ template "airflow.postgresql.fullname" . }}
              {{- else }}
              value: {{ .Values.externalDatabase.host | quote }}
              {{- end }}
            - name: AIRFLOW_DATABASE_PORT_NUMBER
              {{- if .Values.postgresql.enabled }}
              value: "5432"
              {{- else }}
              value: {{ .Values.externalDatabase.port | quote }}
              {{- end }}
            - name: REDIS_HOST
              {{- if .Values.redis.enabled }}
              value: {{ template "airflow.redis.fullname" . }}
              {{- else }}
              value: {{ .Values.externalRedis.host | quote }}
              {{- end }}
            - name: REDIS_PORT_NUMBER
              {{- if .Values.redis.enabled }}
              value: "6379"
              {{- else }}
              value: {{ .Values.externalRedis.port | quote }}
              {{- end }}
            {{- if and (not .Values.redis.enabled) .Values.externalRedis.username }}
            - name: REDIS_USER
              value: {{ .Values.externalRedis.username | quote }}
            {{- end }}
            - name: REDIS_PASSWORD
              valueFrom:
                secretKeyRef:
                  name: {{ template "airflow.redis.secretName" . }}
                  key: redis-password
            - name: AIRFLOW_EXECUTOR
              value: "CeleryExecutor"
            - name: AIRFLOW_FERNET_KEY
              valueFrom:
                secretKeyRef:
                  name: {{ if .Values.airflow.auth.existingSecret }}{{ .Values.airflow.auth.existingSecret }}{{ else }}{{ template "airflow.fullname" . }}{{ end }}
                  key: airflow-fernetKey
            - name: AIRFLOW_WEBSERVER_HOST
              value: {{ template "airflow.fullname" . }}
            - name: AIRFLOW_WEBSERVER_PORT_NUMBER
              value: {{ .Values.service.port | quote }}
            {{- if .Values.airflow.extraEnvVars }}
              {{ toYaml .Values.airflow.extraEnvVars | nindent 12 }}
            {{- end }}
            {{- if .Values.airflow.worker.extraEnvVars }}
              {{- include "airflow.tplValue" (dict "value" .Values.airflow.worker.extraEnvVars "context" $) | nindent 12 }}
            {{- end }}
          envFrom:
            {{- if .Values.airflow.extraEnvVarsCM }}
            - configMapRef:
                name: {{ .Values.airflow.extraEnvVarsCM }}
            {{- end }}
            {{- if .Values.airflow.extraEnvVarsSecret }}
            - secretRef:
                name: {{ .Values.airflow.extraEnvVarsSecret }}
            {{- end }}
            {{- if .Values.airflow.worker.extraEnvVarsCM }}
            - configMapRef:
                name: {{ .Values.airflow.worker.extraEnvVarsCM }}
            {{- end }}
            {{- if .Values.airflow.worker.extraEnvVarsSecret }}
            - secretRef:
                name: {{ .Values.airflow.worker.extraEnvVarsSecret }}
            {{- end }}
          ports:
            - name: worker
              containerPort: {{ .Values.airflow.worker.port }}
          {{- if .Values.livenessProbe.enabled }}
          livenessProbe:
            tcpSocket:
              port: worker
            initialDelaySeconds: {{ .Values.livenessProbe.initialDelaySeconds }}
            periodSeconds: {{ .Values.livenessProbe.periodSeconds }}
            timeoutSeconds: {{ .Values.livenessProbe.timeoutSeconds }}
            successThreshold: {{ .Values.livenessProbe.successThreshold }}
            failureThreshold: {{ .Values.livenessProbe.failureThreshold }}
          {{- end }}
          {{- if .Values.readinessProbe.enabled }}
          readinessProbe:
            tcpSocket:
              port: worker
            initialDelaySeconds: {{ .Values.readinessProbe.initialDelaySeconds }}
            periodSeconds: {{ .Values.readinessProbe.periodSeconds }}
            timeoutSeconds: {{ .Values.readinessProbe.timeoutSeconds }}
            successThreshold: {{ .Values.readinessProbe.successThreshold }}
            failureThreshold: {{ .Values.readinessProbe.failureThreshold }}
          {{- end }}
          volumeMounts:
            {{- if .Files.Glob "files/dags/*.py" }}
            - name: local-dag-files
              mountPath: /opt/bitnami/airflow/dags/local
            {{- end }}
            {{- if .Values.airflow.dagsConfigMap }}
            - name: external-dag-files
              mountPath: /opt/bitnami/airflow/dags/external
            {{- end }}
            {{- if .Values.airflow.cloneDagFilesFromGit.enabled }}
            {{- if and .Values.airflow.cloneDagFilesFromGit.repository .Values.airflow.cloneDagFilesFromGit.branch }}
            - name: git-cloned-dag-files
              mountPath: /opt/bitnami/airflow/dags/git
              {{- if .Values.airflow.cloneDagFilesFromGit.path }}
              subPath: {{ .Values.airflow.cloneDagFilesFromGit.path }}
              {{- end }}
            {{- end }}
            {{- range .Values.airflow.cloneDagFilesFromGit.repositories }}
            - name: git-cloned-dag-files-{{ .name | kebabcase }}
              mountPath: /opt/bitnami/airflow/dags/git-{{ .name | kebabcase }}
              {{- if .path }}
              subPath: {{ .path }}
              {{- end }}
            {{- end }}
            {{- end }}
            {{- if .Values.airflow.clonePluginsFromGit.enabled }}
            - name: git-cloned-plugins
              mountPath: /opt/bitnami/airflow/plugins
              {{- if .Values.airflow.clonePluginsFromGit.path }}
              subPath: {{ .Values.airflow.clonePluginsFromGit.path }}
              {{- end }}
            {{- end }}
            {{- if .Values.airflow.configurationConfigMap }}
            - name: custom-configuration-file
              mountPath: /opt/bitnami/airflow/airflow.cfg
              subPath: airflow.cfg
            {{- end }}
          {{- if .Values.airflow.worker.extraVolumeMounts }}
            {{- include "airflow.tplValue" (dict "value" .Values.airflow.worker.extraVolumeMounts "context" $) | nindent 12 }}
          {{- end }}
      volumes:
      {{- if .Files.Glob "files/dags/*.py" }}
      - name: local-dag-files
        configMap:
          name: {{ template "airflow.fullname" . }}-dag-files
      {{- end }}
      {{- if .Values.airflow.dagsConfigMap }}
      - name: external-dag-files
        configMap:
          name: {{ .Values.airflow.dagsConfigMap }}
      {{- end }}
      {{- if .Values.airflow.cloneDagFilesFromGit.enabled }}
      {{- if and .Values.airflow.cloneDagFilesFromGit.repository .Values.airflow.cloneDagFilesFromGit.branch }}
      - name: git-cloned-dag-files
        emptyDir: {}
      {{- end }}
      {{- range .Values.airflow.cloneDagFilesFromGit.repositories }}
      - name: git-cloned-dag-files-{{ .name | kebabcase }}
        emptyDir: {}
      {{- end }}
      {{- end }}
      {{- if .Values.airflow.clonePluginsFromGit.enabled }}
      - name: git-cloned-plugins
        emptyDir: {}
      {{- end }}
      {{- if .Values.airflow.configurationConfigMap }}
      - name: custom-configuration-file
        configMap:
          name: {{ .Values.airflow.configurationConfigMap }}
      {{- end }}
      {{- if .Values.airflow.worker.extraVolumes }}
      {{- include "airflow.tplValue" (dict "value" .Values.airflow.worker.extraVolumes "context" $) | nindent 6 }}
      {{- end }}
