SDK - Python support for arbitrary secret, similar to ".use_gcp_secret('user-gcp-sa')"

sdk/python/kfp/utils.py

@@ -0,0 +1,61 @@
+def use_secret(secret_name:str=None, volume_name:str=None, secret_volume_mount_path:str=None, env_variable:str=None, secret_file_path_in_volume:str=None):
+    """    
+       An operator that configures the container to use a secret.
+
+       This assumes that the secret is created and availabel in the k8s cluster.
+
+    Keyword Arguments:
+        secret_name {String} -- [Required] The k8s secret name (default: {None})
+        volume_name {String} -- The pod volume name (default: {None})
+        secret_volume_mount_path {String} -- [Required] The path to the secret that is mounted  (default: {None})
+        env_variable {String} -- Env variable pointing to the mounted secret file. Requires both the env_variable and secret_file_path_in_volume to be defined. 
+                                 The value is the path to the secret (default: {None})
+        secret_file_path_in_volume {String} -- The path to the secret in the volume. This will be the value of env_variable. 
+                                 Both env_variable and secret_file_path_in_volume needs to be set if any env variable should be created (default: {None})
+
+    Raises:
+        ValueError: If not the necessary variables (secret_name, volume_name", secret_volume_mount_path) are supplied.
+                    Or only one of  env_variable and secret_file_path_in_volume are supplied
+
+    Returns:
+        [ContainerOperator] -- Returns the container operator after it has been modified. 
+    """
+
+    if secret_name:
+        volume_name = volume_name or secret_name + '_volume'
+
+    params = [secret_name, volume_name, secret_volume_mount_path]
+    param_names = ["secret_name", "volume_name", "secret_volume_mount_path"]
+    for param, param_name in zip(params, param_names):
+        if param is None:
+            raise ValueError("'{}' needs to be specified, is: {}".format(param_name, param))
+
+    if bool(env_variable) != bool(secret_file_path_in_volume):
+        raise ValueError("Both {} and {} needs to be supplied together or not at all".format(env_variable, secret_file_path_in_volume))
+
+    def _use_secret(task):
+        import os 
+        from kubernetes import client as k8s_client
+        task = task.add_volume(
+            k8s_client.V1Volume(
+                name=volume_name,
+                secret=k8s_client.V1SecretVolumeSource(
+                    secret_name=secret_name
+                )
+            )
+        ).add_volume_mount(
+                k8s_client.V1VolumeMount(
+                    name=volume_name,
+                    mount_path=secret_volume_mount_path
+                )
+            )
+        if env_variable:
+            task.container.add_env_variable(
+                k8s_client.V1EnvVar(
+                    name=env_variable,
+                    value=os.path.join(secret_volume_mount_path, secret_file_path_in_volume),
+                )
+            )
+        return task
+
+    return _use_secret

sdk/python/tests/dsl/utils_tests.py

@@ -0,0 +1,66 @@
+from kfp.dsl import ContainerOp
+from kfp.utils import use_secret
+import unittest
+import inspect
+
+
+class TestAddSecrets(unittest.TestCase):
+    def test_default_use_secret(self):
+        spec = inspect.getfullargspec(use_secret)
+        assert len(spec.defaults) == 5
+        for spec in spec.defaults:
+            assert spec == None
+
+
+    def test_use_default_use_secret(self):
+        op1 = ContainerOp(name='op1', image='image')
+        secret_name = "my-secret"
+        secret_path = "/here/are/my/secret"
+        op1 = op1.apply(use_secret(secret_name=secret_name, 
+        secret_volume_mount_path=secret_path))
+        assert type(op1.container.env) == type(None)
+
+        container_dict = op1.container.to_dict()
+        volume_mounts = container_dict["volume_mounts"][0]
+        assert type(volume_mounts)==dict
+        assert volume_mounts["name"] == secret_name + '_volume'
+        assert volume_mounts["mount_path"] == secret_path
+
+
+    def test_use_set_volume_use_secret(self):
+        op1 = ContainerOp(name='op1', image='image')
+        secret_name = "my-secret"
+        secret_path = "/here/are/my/secret"
+        volume_name = "my_volume"
+        op1 = op1.apply(use_secret(secret_name=secret_name, 
+            secret_volume_mount_path=secret_path, 
+            volume_name = volume_name))
+        assert type(op1.container.env) == type(None)
+
+        container_dict = op1.container.to_dict()
+        volume_mounts = container_dict["volume_mounts"][0]
+        assert type(volume_mounts)==dict
+        assert volume_mounts["name"] == volume_name
+        assert volume_mounts["mount_path"] == secret_path
+
+
+    def test_use_set_env_ues_secret(self):
+        op1 = ContainerOp(name='op1', image='image')
+        secret_name = "my-secret"
+        secret_path = "/here/are/my/secret/"
+        env_variable = "MY_SECRET"
+        secret_file_path_in_volume = "secret.json"
+        op1 = op1.apply(use_secret(secret_name=secret_name, 
+            secret_volume_mount_path=secret_path,
+            env_variable=env_variable,
+            secret_file_path_in_volume=secret_file_path_in_volume))
+        assert len(op1.container.env) ==1
+
+        container_dict = op1.container.to_dict()
+        volume_mounts = container_dict["volume_mounts"][0]
+        assert type(volume_mounts)==dict
+        assert volume_mounts["name"] == secret_name + '_volume'
+        assert volume_mounts["mount_path"] == secret_path 
+        env_dict = op1.container.env[0].to_dict()
+        assert env_dict["name"] == env_variable
+        assert env_dict["value"] == secret_path + secret_file_path_in_volume