Followup to "Fixes k8s restore pvc"

MANIFEST.in

@@ -2,7 +2,7 @@ include MANIFEST.in README.rst TODO.rst LICENSE.txt .dockerignore
 include .style.yapf alembic
 recursive-include docs *.rst *.html Makefile *.gif *.png *.cast *.js *.css
 recursive-include src *.py *.ini *.yaml
-recursive-include images Dockerfile *.sh
+recursive-include images Dockerfile ceph.repo bashrc *.sh *.in *.py
 recursive-include scripts *.sh
 recursive-include charts *.yaml *.tpl .helmignore
 recursive-include tests Makefile *.yaml

images/benji-k8s/Dockerfile

@@ -18,8 +18,10 @@ LABEL org.label-schema.schema-version="1.0" \
 RUN curl -o /usr/bin/kubectl -sSL https://storage.googleapis.com/kubernetes-release/release/$(curl -s https://storage.googleapis.com/kubernetes-release/release/stable.txt)/bin/linux/amd64/kubectl && \
 	chmod a+x /usr/bin/kubectl
 
-COPY images/benji-k8s/bin/ $VENV_DIR/bin/
-RUN chmod a+x $VENV_DIR/bin/*
+COPY images/benji-k8s/k8s-tools /k8s-tools-source
+RUN . $VENV_DIR/bin/activate && \
+    pip install /k8s-tools-source && \
+    rm -rf /k8s-tools-source
 
 ENTRYPOINT ["/bin/bash"]
 CMD ["-c", "sleep 3650d"]

images/benji-k8s/k8s-tools/MANIFEST.in

@@ -0,0 +1,2 @@
+include MANIFEST.in
+recursive-include src *.py

images/benji-k8s/k8s-tools/setup.py

@@ -0,0 +1,22 @@
+from setuptools import setup, find_packages
+
+setup(name='benji-k8s-tools',
+      version='0.1',
+      description='Small tools for using Benji with Kubernetes',
+      url='https://github.com/elemental-lf/benji',
+      author='Lars Fenneberg',
+      author_email='lf@elemental.net',
+      license='LGPG-3',
+      python_requires='~=3.6',
+      packages=find_packages('src'),
+      package_dir={
+          '': 'src',
+      },
+      install_requires=['benji', 'kubernetes>=10.0.0,<11'],
+      entry_points="""
+        [console_scripts]
+            benji-backup-pvc = benji.k8s_tools.scripts.backup_pvc:main
+            benji-command = benji.k8s_tools.scripts.command:main
+            benji-restore-pvc = benji.k8s_tools.scripts.restore_pvc:main
+            benji-versions-status = benji.k8s_tools.scripts.versions_status:main
+    """)

images/benji-k8s/k8s-tools/src/benji/__init__.py

@@ -0,0 +1 @@
+__path__ = __import__('pkgutil').extend_path(__path__, __name__)

src/benji/helpers/kubernetes.py → images/benji-k8s/k8s-tools/src/benji/k8s_tools/kubernetes.py

@@ -5,7 +5,7 @@
 import time
 import uuid
 from subprocess import TimeoutExpired, CalledProcessError
-from typing import List, Union, Tuple
+from typing import List, Union, Tuple, Optional
 
 import kubernetes
 from kubernetes.stream import stream
@@ -49,7 +49,11 @@ def service_account_namespace() -> str:
 # connect_post_namespaced_pod_exec. The examples from the kubernetes client use connect_get_namespaced_pod_exec instead.
 # There shouldn't be any differences in functionality but the settings in the RBAC role are different (create vs. get)
 # which is why we follow the kubectl implementation here.
-def pod_exec(args: List[str], *, name: str, namespace: str, container: str = None,
+def pod_exec(args: List[str],
+             *,
+             name: str,
+             namespace: str,
+             container: str = None,
              timeout: float = float("inf")) -> Tuple[str, str]:
     core_v1_api = kubernetes.client.CoreV1Api()
     logger.debug('Running command in pod {}/{}: {}.'.format(namespace, name, ' '.join(args)))
@@ -171,8 +175,8 @@ def create_pvc_event(*, type: str, reason: str, message: str, pvc_namespace: str
     return core_v1_api.create_namespaced_event(namespace=pvc_namespace, body=event)
 
 
-def create_pvc(pvc_name: str, pvc_namespace: int,
-               pvc_size: str) -> kubernetes.client.models.v1_persistent_volume_claim.V1PersistentVolumeClaim:
+def create_pvc(pvc_name: str, pvc_namespace: int, pvc_size: str,
+               pvc_storage_class: str) -> kubernetes.client.models.v1_persistent_volume_claim.V1PersistentVolumeClaim:
     pvc = {
         'kind': 'PersistentVolumeClaim',
         'apiVersion': 'v1',
@@ -181,7 +185,7 @@ def create_pvc(pvc_name: str, pvc_namespace: int,
             'name': pvc_name,
         },
         'spec': {
-            'storageClassName': 'rbd',
+            'storageClassName': pvc_storage_class,
             'accessModes': ['ReadWriteOnce'],
             'resources': {
                 'requests': {
@@ -195,6 +199,34 @@ def create_pvc(pvc_name: str, pvc_namespace: int,
     return core_v1_api.create_namespaced_persistent_volume_claim(namespace=pvc_namespace, body=pvc)
 
 
+def determine_rbd_image_from_pv(
+        pv: kubernetes.client.models.v1_persistent_volume.V1PersistentVolume) -> Tuple[Optional[str], Optional[str]]:
+    pool, image = None, None
+
+    if hasattr(pv.spec, 'rbd') and hasattr(pv.spec.rbd, 'pool') and hasattr(pv.spec.rbd, 'image'):
+        # Native Kubernetes RBD PV
+        pool, image = pv.spec.rbd.pool, pv.spec.rbd.image
+    elif hasattr(pv.spec, 'flex_volume') and hasattr(pv.spec.flex_volume, 'options') and hasattr(
+            pv.spec.flex_volume, 'driver'):
+        # Rook Ceph PV
+        options = pv.spec.flex_volume.options
+        driver = pv.spec.flex_volume.driver
+        if driver.startswith('ceph.rook.io/') and options.get('pool') and options.get('image'):
+            pool, image = options['pool'], options['image']
+    elif (hasattr(pv.spec, 'csi') and hasattr(pv.spec.csi, 'driver') and
+          pv.spec.csi.driver in ('rook-ceph.rbd.csi.ceph.com', 'rbd.csi.ceph.com') and
+          hasattr(pv.spec.csi, 'volume_handle') and pv.spec.csi.volume_handle and
+          hasattr(pv.spec.csi, 'volume_attributes') and pv.spec.csi.volume_attributes.get('pool')):
+        attributes = pv.spec.csi.volume_attributes
+        volume_handle_parts = pv.spec.csi.volume_handle.split('-')
+        if len(volume_handle_parts) >= 9:
+            image_prefix = attributes.get('volumeNamePrefix', 'csi-vol-')
+            image_suffix = '-'.join(volume_handle_parts[len(volume_handle_parts) - 5:])
+            pool, image = attributes['pool'], image_prefix + image_suffix
+
+    return pool, image
+
+
 # This is taken from https://github.com/kubernetes-client/python/pull/855 with minimal changes.
 #
 # Copyright 2019 The Kubernetes Authors.