From c6178f86367d4825aa98f63590aae67e630b235e Mon Sep 17 00:00:00 2001
From: Sidharth Menon <menon.sid.k@gmail.com>
Date: Thu, 17 Jun 2021 15:19:02 -0400
Subject: [PATCH] [helm] Add k8s fields to pydantic schema

Summary: Adding more k8s fields to the pydantic schema, including `InitContainers` & `VolumeMounts`.

Test Plan: `helm lint` + bk

Reviewers: rexledesma, jordansanders

Reviewed By: rexledesma

Subscribers: johann

Differential Revision: https://dagster.phacility.com/D8439
---
 .../dagster-buildkite/dagster_buildkite/steps/helm.py  |  2 +-
 helm/dagster/schema/schema/charts/utils/kubernetes.py  | 10 ++++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/.buildkite/dagster-buildkite/dagster_buildkite/steps/helm.py b/.buildkite/dagster-buildkite/dagster_buildkite/steps/helm.py
index 301c25de168eb..2886a8585b53e 100644
--- a/.buildkite/dagster-buildkite/dagster_buildkite/steps/helm.py
+++ b/.buildkite/dagster-buildkite/dagster_buildkite/steps/helm.py
@@ -48,7 +48,7 @@ def helm_steps() -> List[dict]:
         supported_pythons=[SupportedPython.V3_8],
         buildkite_label="dagster-helm-schema",
         upload_coverage=False,
-        retries=2,
+        retries=4,
     ).get_tox_build_steps()
 
     return tests
diff --git a/helm/dagster/schema/schema/charts/utils/kubernetes.py b/helm/dagster/schema/schema/charts/utils/kubernetes.py
index eb0779259c9c9..edd178f8f7f96 100644
--- a/helm/dagster/schema/schema/charts/utils/kubernetes.py
+++ b/helm/dagster/schema/schema/charts/utils/kubernetes.py
@@ -95,6 +95,11 @@ class Config:
         schema_extra = {"$ref": create_definition_ref("io.k8s.api.core.v1.SecurityContext")}
 
 
+class InitContainer(BaseModel):
+    class Config:
+        schema_extra = {"$ref": create_definition_ref("io.k8s.api.core.v1.Container")}
+
+
 class Resources(BaseModel):
     __root__: Dict[str, Any]
 
@@ -132,3 +137,8 @@ class Config:
 class ConfigMapEnvSource(BaseModel):
     class Config:
         schema_extra = {"$ref": create_definition_ref("io.k8s.api.core.v1.ConfigMapEnvSource")}
+
+
+class VolumeMount(BaseModel):
+    class Config:
+        schema_extra = {"$ref": create_definition_ref("io.k8s.api.core.v1.VolumeMount")}