kaqing 1.77.0__py3-none-any.whl → 2.0.171__py3-none-any.whl

{walker/k8s_utils → adam/utils_k8s}/jobs.py

@@ -1,7 +1,7 @@
 from kubernetes import client
 from time import sleep
 from .pods import Pods
-from walker.utils import log2
+from adam.utils import log2
 
 # utility collection on jobs; methods are all static
 class Jobs:
@@ -12,7 +12,7 @@ class Jobs:
             envs.append(client.V1EnvVar(name=k.upper(), value=str(v)))
         for k, v in env_from.items():
             envs.append(client.V1EnvVar(name=k.upper(), value_from=client.V1EnvVarSource(secret_key_ref=client.V1SecretKeySelector(key=k, name=v))))
-        template = Pods.create_pod_spec(job_name, image, image_pull_secret, envs, volume_name, pvc_name, mount_path, command)
+        template = Pods.create_pod_spec(job_name, image, image_pull_secret, envs, None, volume_name, pvc_name, mount_path, command)
         spec = client.V1JobSpec(template=client.V1PodTemplateSpec(spec=template), backoff_limit=1, ttl_seconds_after_finished=300)
         job = client.V1Job(
             api_version="batch/v1",
@@ -44,7 +44,7 @@ class Jobs:
         except Exception as e:
             log2("Exception when calling BatchV1Apii->delete_namespaced_job: %s\n" % e)
         return
-            
+
     def get_logs(job_name: str, namespace: str):
         v1 = client.CoreV1Api()
         try:

adam/utils_k8s/k8s.py

@@ -0,0 +1,87 @@
+from collections.abc import Callable
+import re
+import portforward
+
+from adam.commands.command import InvalidState
+from adam.repl_state import ReplState
+from adam.utils import log2
+from adam.utils_k8s.kube_context import KubeContext
+
+class PortForwardHandler:
+    connections: dict[str, int] = {}
+
+    def __init__(self, state: ReplState, local_port: int, svc_or_pod: Callable[[bool],str], target_port: int):
+        self.state = state
+        self.local_port = local_port
+        self.svc_or_pod = svc_or_pod
+        self.target_port = target_port
+        self.forward_connection = None
+        self.pod = None
+
+    def __enter__(self) -> tuple[str, str]:
+        state = self.state
+
+        if not self.svc_or_pod:
+            log2('No service or pod found.')
+
+            raise InvalidState(state)
+
+        if KubeContext.in_cluster():
+            svc_name = self.svc_or_pod(True)
+            if not svc_name:
+                log2('No service found.')
+
+                raise InvalidState(state)
+
+            # cs-a526330d23-cs-a526330d23-default-sts-0 ->
+            # curl http://cs-a526330d23-cs-a526330d23-reaper-service.stgawsscpsr.svc.cluster.local:8080
+            groups = re.match(r'^(.*?-.*?-.*?-.*?-).*', state.sts)
+            if groups:
+                svc = f'{groups[1]}{svc_name}.{state.namespace}.svc.cluster.local:{self.target_port}'
+                return (svc, svc)
+            else:
+                raise InvalidState(state)
+        else:
+            pod = self.svc_or_pod(False)
+            if not pod:
+                log2('No pod found.')
+
+                raise InvalidState(state)
+
+            self.pod = pod
+            self.forward_connection = portforward.forward(state.namespace, pod, self.local_port, self.target_port)
+            if self.inc_connection_cnt() == 1:
+                self.forward_connection.__enter__()
+
+            return (f'localhost:{self.local_port}', f'{pod}:{self.target_port}')
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.forward_connection:
+            if not self.dec_connection_cnt():
+                return self.forward_connection.__exit__(exc_type, exc_val, exc_tb)
+
+        return False
+
+    def inc_connection_cnt(self):
+        id = self.connection_id(self.pod)
+        if id not in PortForwardHandler.connections:
+            PortForwardHandler.connections[id] = 1
+        else:
+            PortForwardHandler.connections[id] += 1
+
+        return PortForwardHandler.connections[id]
+
+    def dec_connection_cnt(self):
+        id = self.connection_id(self.pod)
+        if id not in PortForwardHandler.connections:
+            PortForwardHandler.connections[id] = 0
+        elif PortForwardHandler.connections[id] > 0:
+            PortForwardHandler.connections[id] -= 1
+
+        return PortForwardHandler.connections[id]
+
+    def connection_id(self, pod: str):
+        return f'{self.local_port}:{pod}:{self.target_port}'
+
+def port_forwarding(state: ReplState, local_port: int, svc_or_pod: Callable[[bool],str], target_port: int):
+    return PortForwardHandler(state, local_port, svc_or_pod, target_port)

{walker/k8s_utils → adam/utils_k8s}/kube_context.py

@@ -2,8 +2,8 @@ import os
 import re
 from kubernetes import config as kconfig
 
-from walker.config import Config
-from walker.utils import lines_to_tabular, log2
+from adam.config import Config
+from adam.utils import idp_token_from_env, lines_to_tabular, log2
 
 class KubeContext:
     _in_cluster = False
@@ -44,7 +44,7 @@ class KubeContext:
             except kconfig.ConfigException:
                 pass
 
-        if msg and not is_user_entry:
+        if msg and not is_user_entry and not idp_token_from_env():
             log2(msg)
         if not loaded:
             exit(1)
@@ -105,4 +105,4 @@ class KubeContext:
         return s or Config().is_debug()
 
     def show_parallelism():
-        return Config().get('debug.show-parallelism', False)
+        return Config().get('debugs.show-parallelism', False)

adam/utils_k8s/pods.py

@@ -0,0 +1,290 @@
+from collections.abc import Callable
+from datetime import datetime
+import sys
+import time
+from typing import TypeVar
+from kubernetes import client
+from kubernetes.stream import stream
+from kubernetes.stream.ws_client import ERROR_CHANNEL, WSClient
+
+from adam.config import Config
+from adam.utils_k8s.volumes import ConfigMapMount
+from adam.pod_exec_result import PodExecResult
+from adam.utils import ParallelMapHandler, log2
+from .kube_context import KubeContext
+
+from websocket._core import WebSocket
+
+T = TypeVar('T')
+_TEST_POD_EXEC_OUTS: PodExecResult = None
+
+# utility collection on pods; methods are all static
+class Pods:
+    _TEST_POD_CLOSE_SOCKET: bool = False
+
+    def set_test_pod_exec_outs(outs: PodExecResult):
+        global _TEST_POD_EXEC_OUTS
+        _TEST_POD_EXEC_OUTS = outs
+
+        return _TEST_POD_EXEC_OUTS
+
+    def delete(pod_name: str, namespace: str, grace_period_seconds: int = None):
+        try:
+            v1 = client.CoreV1Api()
+            v1.delete_namespaced_pod(pod_name, namespace, grace_period_seconds=grace_period_seconds)
+        except Exception as e:
+            log2("Exception when calling CoreV1Api->delete_namespaced_pod: %s\n" % e)
+
+    def delete_with_selector(namespace: str, label_selector: str, grace_period_seconds: int = None):
+        v1 = client.CoreV1Api()
+
+        ret = v1.list_namespaced_pod(namespace=namespace, label_selector=label_selector)
+        for i in ret.items:
+            v1.delete_namespaced_pod(name=i.metadata.name, namespace=namespace, grace_period_seconds=grace_period_seconds)
+
+    def parallelize(collection: list, max_workers: int = 0, samples = sys.maxsize, msg: str = None, action: str = 'action'):
+        if not max_workers:
+            max_workers = Config().action_workers(action, 0)
+        if samples == sys.maxsize:
+            samples = Config().action_node_samples(action, sys.maxsize)
+
+        return ParallelMapHandler(collection, max_workers, samples = samples, msg = msg)
+
+    def exec(pod_name: str, container: str, namespace: str, command: str,
+             show_out = True, throw_err = False, shell = '/bin/sh',
+             background = False,
+             log_file = None,
+             interaction: Callable[[any, list[str]], any] = None,
+             env_prefix: str = None):
+        if _TEST_POD_EXEC_OUTS:
+            return _TEST_POD_EXEC_OUTS
+
+        show_out = KubeContext.show_out(show_out)
+
+        api = client.CoreV1Api()
+
+        tty = True
+        exec_command = [shell, '-c', command]
+        if env_prefix:
+            exec_command = [shell, '-c', f'{env_prefix} {command}']
+
+        if background or command.endswith(' &'):
+            # should be false for starting a background process
+            tty = False
+
+            if Config().get('repl.background-process.auto-nohup', True):
+                command = command.strip(' &')
+                cmd_name = ''
+                if command.startswith('nodetool '):
+                    cmd_name = f".{'_'.join(command.split(' ')[5:])}"
+
+                if not log_file:
+                    log_file = f'{log_prefix()}-{datetime.now().strftime("%d%H%M%S")}{cmd_name}.log'
+                command = f"nohup {command} > {log_file} 2>&1 &"
+                if env_prefix:
+                    command = f'{env_prefix} {command}'
+                exec_command = [shell, '-c', command]
+
+        k_command = f'kubectl exec {pod_name} -c {container} -n {namespace} -- {shell} -c "{command}"'
+        if Config().is_debug():
+            log2(k_command)
+
+        resp: WSClient = stream(
+            api.connect_get_namespaced_pod_exec,
+            pod_name,
+            namespace,
+            command=exec_command,
+            container=container,
+            stderr=True,
+            stdin=True,
+            stdout=True,
+            tty=tty,
+            _preload_content=False,
+        )
+
+        s: WebSocket = resp.sock
+        stdout = []
+        stderr = []
+        error_output = None
+        try:
+            while resp.is_open():
+                resp.update(timeout=1)
+                if resp.peek_stdout():
+                    frag = resp.read_stdout()
+                    stdout.append(frag)
+                    if show_out: print(frag, end="")
+
+                    if interaction:
+                        interaction(resp, stdout)
+                if resp.peek_stderr():
+                    frag = resp.read_stderr()
+                    stderr.append(frag)
+                    if show_out: print(frag, end="")
+
+            try:
+                # get the exit code from server
+                error_output = resp.read_channel(ERROR_CHANNEL)
+            except Exception as e:
+                pass
+        except Exception as e:
+            if throw_err:
+                raise e
+            else:
+                log2(e)
+        finally:
+            resp.close()
+            if s and s.sock and Pods._TEST_POD_CLOSE_SOCKET:
+                try:
+                    s.sock.close()
+                except:
+                    pass
+
+        return PodExecResult("".join(stdout), "".join(stderr), k_command, error_output, pod=pod_name, log_file=log_file)
+
+    def read_file(pod_name: str, container: str, namespace: str, file_path: str):
+        v1 = client.CoreV1Api()
+
+        resp = stream(
+            v1.connect_get_namespaced_pod_exec,
+            name=pod_name,
+            namespace=namespace,
+            container=container,
+            command=["cat", file_path],
+            stderr=True, stdin=False,
+            stdout=True, tty=False,
+            _preload_content=False, # Important for streaming
+        )
+
+        s: WebSocket = resp.sock
+        try:
+            while resp.is_open():
+                resp.update(timeout=1)
+                if resp.peek_stdout():
+                    yield resp.read_stdout()
+
+            try:
+                # get the exit code from server
+                error_output = resp.read_channel(ERROR_CHANNEL)
+            except Exception as e:
+                pass
+        except Exception as e:
+            raise e
+        finally:
+            resp.close()
+            if s and s.sock and Pods._TEST_POD_CLOSE_SOCKET:
+                try:
+                    s.sock.close()
+                except:
+                    pass
+    def get_container(namespace: str, pod_name: str, container_name: str):
+        pod = Pods.get(namespace, pod_name)
+        if not pod:
+            return None
+
+        for container in pod.spec.containers:
+            if container_name == container.name:
+                return container
+
+        return None
+
+    def get(namespace: str, pod_name: str):
+        v1 = client.CoreV1Api()
+        return v1.read_namespaced_pod(name=pod_name, namespace=namespace)
+
+    def get_with_selector(namespace: str, label_selector: str):
+        v1 = client.CoreV1Api()
+
+        ret = v1.list_namespaced_pod(namespace=namespace, label_selector=label_selector)
+        for i in ret.items:
+            return v1.read_namespaced_pod(name=i.metadata.name, namespace=namespace)
+
+    def create_pod_spec(name: str, image: str, image_pull_secret: str,
+                        envs: list, container_security_context: client.V1SecurityContext,
+                        volume_name: str, pvc_name:str, mount_path:str,
+                        command: list[str]=None, sa_name : str = None, config_map_mount: ConfigMapMount = None,
+                        restart_policy="Never"):
+        volume_mounts = []
+        if volume_name and pvc_name and mount_path:
+            volume_mounts=[client.V1VolumeMount(mount_path=mount_path, name=volume_name)]
+
+        if config_map_mount:
+            volume_mounts.append(client.V1VolumeMount(mount_path=config_map_mount.mount_path, sub_path=config_map_mount.sub_path, name=config_map_mount.name()))
+
+        container = client.V1Container(name=name, image=image, env=envs, security_context=container_security_context, command=command,
+                                    volume_mounts=volume_mounts)
+
+        volumes = []
+        if volume_name and pvc_name and mount_path:
+            volumes=[client.V1Volume(name=volume_name, persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource(claim_name=pvc_name))]
+
+        security_context = None
+        if not sa_name:
+            security_context=client.V1PodSecurityContext(run_as_user=1001, run_as_group=1001, fs_group=1001)
+
+        if config_map_mount:
+            volumes.append(client.V1Volume(name=config_map_mount.name(), config_map=client.V1ConfigMapVolumeSource(name=config_map_mount.config_map_name)))
+
+        return client.V1PodSpec(
+            restart_policy=restart_policy,
+            containers=[container],
+            image_pull_secrets=[client.V1LocalObjectReference(name=image_pull_secret)],
+            security_context=security_context,
+            service_account_name=sa_name,
+            volumes=volumes
+        )
+
+    def create(namespace: str, pod_name: str, image: str,
+               command: list[str] = None,
+               secret: str = None,
+               env: dict[str, any] = {},
+               container_security_context: client.V1SecurityContext = None,
+               labels: dict[str, str] = {},
+               volume_name: str = None,
+               pvc_name: str = None,
+               mount_path: str = None,
+               sa_name: str = None,
+               config_map_mount: ConfigMapMount = None):
+        v1 = client.CoreV1Api()
+        envs = []
+        for k, v in env.items():
+            envs.append(client.V1EnvVar(name=str(k), value=str(v)))
+        pod = Pods.create_pod_spec(pod_name, image, secret, envs, container_security_context, volume_name, pvc_name, mount_path, command=command,
+                                   sa_name=sa_name, config_map_mount=config_map_mount)
+        return v1.create_namespaced_pod(
+            namespace=namespace,
+            body=client.V1Pod(spec=pod, metadata=client.V1ObjectMeta(
+                name=pod_name,
+                labels=labels
+            ))
+        )
+
+    def wait_for_running(namespace: str, pod_name: str, msg: str = None, label_selector: str = None):
+        cnt = 2
+        while (cnt < 302 and Pods.get_with_selector(namespace, label_selector) if label_selector else Pods.get(namespace, pod_name)).status.phase != 'Running':
+            if not msg:
+                msg = f'Waiting for the {pod_name} pod to start up.'
+
+            max_len = len(msg) + 3
+            mod = cnt % 3
+            padded = ''
+            if mod == 0:
+                padded = f'\r{msg}'.ljust(max_len)
+            elif mod == 1:
+                padded = f'\r{msg}.'.ljust(max_len)
+            else:
+                padded = f'\r{msg}..'.ljust(max_len)
+            log2(padded, nl=False)
+            cnt += 1
+            time.sleep(1)
+
+        log2(f'\r{msg}..'.ljust(max_len), nl=False)
+        if cnt < 302:
+            log2(' OK')
+        else:
+            log2(' Timed Out')
+
+    def completed(namespace: str, pod_name: str):
+        return Pods.get(namespace, pod_name).status.phase in ['Succeeded', 'Failed']
+
+def log_prefix():
+    return Config().get('log-prefix', '/tmp/qing')

{walker/k8s_utils → adam/utils_k8s}/secrets.py

@@ -1,15 +1,19 @@
 import base64
+import functools
 import re
 from typing import cast
 from kubernetes import client
 from kubernetes.client import V1Secret
 
-from walker.config import Config
-from walker.utils import log2
+from adam.config import Config
+from adam.utils import log2, wait_log
 
 # utility collection on secrets; methods are all static
 class Secrets:
+    @functools.lru_cache()
     def list_secrets(namespace: str = None, name_pattern: str = None):
+        wait_log('Inspecting Cassandra Instances...')
+
         secrets_names = []
 
         v1 = client.CoreV1Api()
@@ -35,14 +39,14 @@ class Secrets:
 
         return secrets_names
 
-    def get_user_pass(ss_name: str, namespace: str, secret_path: str = 'cql.secret'):
+    def get_user_pass(sts_or_pod_name: str, namespace: str, secret_path: str = 'cql.secret'):
         # cs-d0767a536f-cs-d0767a536f-default-sts ->
         # cs-d0767a536f-superuser
         # cs-d0767a536f-reaper-ui
         user = 'superuser'
         if secret_path == 'reaper.secret':
             user = 'reaper-ui'
-        groups = re.match(Config().get(f'{secret_path}.cluster-regex', r'(.*?-.*?)-.*'), ss_name)
+        groups = re.match(Config().get(f'{secret_path}.cluster-regex', r'(.*?-.*?)-.*'), sts_or_pod_name)
         secret_name = Config().get(f'{secret_path}.name', '{cluster}-' + user).replace('{cluster}', groups[1], 1)
 
         secret = Secrets.get_data(namespace, secret_name)

adam/utils_k8s/service_accounts.py

@@ -0,0 +1,170 @@
+from kubernetes import client, config
+
+from adam.config import Config
+from adam.utils import debug
+
+# utility collection on service accounts; methods are all static
+class ServiceAccounts:
+    def delete(namespace: str, label_selector: str):
+        ServiceAccounts.delete_cluster_role_bindings(label_selector)
+        ServiceAccounts.delete_role_bindings(namespace, label_selector)
+        ServiceAccounts.delete_service_account(namespace, label_selector)
+
+    def replicate(to_sa: str, namespace: str, from_sa: str, labels: dict[str, str] = {}, add_cluster_roles: list[str] = []):
+        ServiceAccounts.create_service_account(to_sa, namespace, labels=labels)
+        for b in ServiceAccounts.get_role_bindings(from_sa, namespace):
+            n = f'{to_sa}-{b.role_ref.name}'
+            ServiceAccounts.create_role_binding(n, namespace, to_sa, b.role_ref.name, labels=labels)
+
+        for b in ServiceAccounts.get_cluster_role_bindings(from_sa):
+            n = f'{to_sa}-{b.role_ref.name}'
+            ServiceAccounts.create_cluster_role_binding(n, namespace, to_sa, b.role_ref.name, labels=labels)
+
+        for cr in add_cluster_roles:
+            n = f'{to_sa}-{cr}'
+            ServiceAccounts.create_cluster_role_binding(n, namespace, to_sa, cr, labels=labels)
+
+    def create_service_account(name: str, namespace: str, labels: dict[str, str] = {}):
+        config.load_kube_config()
+
+        v1 = client.CoreV1Api()
+
+        service_account = client.V1ServiceAccount(
+            metadata=client.V1ObjectMeta(
+                name=name,
+                labels=labels)
+        )
+        api_response = v1.create_namespaced_service_account(
+            namespace=namespace,
+            body=service_account
+        )
+        debug(f"Service Account '{api_response.metadata.name}' created in namespace '{namespace}'.")
+
+    def delete_service_account(namespace: str, label_selector: str) -> list:
+        refs = []
+
+        v1 = client.CoreV1Api()
+        sas = v1.list_namespaced_service_account(namespace=namespace, label_selector=label_selector).items
+        for sa in sas:
+            debug(f'delete {sa.metadata.name}')
+            v1.delete_namespaced_service_account(name=sa.metadata.name, namespace=namespace)
+            refs.append(sa)
+
+        return refs
+
+    def create_role_binding(name: str, namespace: str, sa_name: str, role_name: str, labels: dict[str, str] = {}):
+        api = client.RbacAuthorizationV1Api()
+
+        metadata = client.V1ObjectMeta(
+            name=name,
+            namespace=namespace,
+            labels=labels
+        )
+        role_ref = client.V1RoleRef(
+            api_group="rbac.authorization.k8s.io",
+            kind="Role",
+            name=role_name
+        )
+
+        subjects = [
+            client.RbacV1Subject(
+                kind="ServiceAccount",
+                name=sa_name, # Name of the service account
+                namespace=namespace # Namespace of the service account
+            )
+        ]
+
+        role_binding = client.V1RoleBinding(
+            api_version="rbac.authorization.k8s.io/v1",
+            kind="RoleBinding",
+            metadata=metadata,
+            role_ref=role_ref,
+            subjects=subjects
+        )
+
+        api.create_namespaced_role_binding(namespace=namespace, body=role_binding)
+
+    def get_role_bindings(service_account_name: str, namespace: str) -> list:
+        refs = []
+
+        rbac_api = client.RbacAuthorizationV1Api()
+        role_bindings = rbac_api.list_namespaced_role_binding(namespace=namespace)
+        for binding in role_bindings.items:
+            if binding.subjects:
+                for subject in binding.subjects:
+                    if subject.kind == "ServiceAccount" and subject.name == service_account_name:
+                        refs.append(binding)
+
+        return refs
+
+    def delete_role_bindings(namespace: str, label_selector: str) -> list:
+        refs = []
+
+        v1_rbac = client.RbacAuthorizationV1Api()
+        cluster_role_bindings = v1_rbac.list_namespaced_role_binding(namespace=namespace, label_selector=label_selector).items
+        for binding in cluster_role_bindings:
+            debug(f'delete {binding.metadata.name}')
+            v1_rbac.delete_namespaced_role_binding(name=binding.metadata.name, namespace=namespace)
+            refs.append(binding)
+
+        return refs
+
+    def create_cluster_role_binding(name: str, namespace: str, sa_name: str, role_name: str, labels: dict[str, str] = {}):
+        api = client.RbacAuthorizationV1Api()
+
+        metadata = client.V1ObjectMeta(
+            name=name,
+            namespace=namespace,
+            labels=labels
+        )
+        role_ref = client.V1RoleRef(
+            api_group="rbac.authorization.k8s.io",
+            kind="ClusterRole",
+            name=role_name
+        )
+
+        subjects = [
+            client.RbacV1Subject(
+                kind="ServiceAccount",
+                name=sa_name,
+                namespace=namespace
+            )
+        ]
+
+        role_binding = client.V1ClusterRoleBinding(
+            api_version="rbac.authorization.k8s.io/v1",
+            metadata=metadata,
+            role_ref=role_ref,
+            subjects=subjects
+        )
+
+        try:
+            api.create_cluster_role_binding(body=role_binding)
+        except client.ApiException as e:
+            print(f"Error creating ClusterRoleBinding: {e}")
+
+    def get_cluster_role_bindings(service_account_name: str) -> list:
+        refs = []
+
+        v1_rbac = client.RbacAuthorizationV1Api()
+        cluster_role_bindings = v1_rbac.list_cluster_role_binding().items
+        for binding in cluster_role_bindings:
+            if binding.subjects:
+                for subject in binding.subjects:
+                    if subject.kind == "ServiceAccount" and subject.name == service_account_name:
+                        refs.append(binding)
+
+        return refs
+
+
+    def delete_cluster_role_bindings(label_selector: str) -> list:
+        refs = []
+
+        v1_rbac = client.RbacAuthorizationV1Api()
+        cluster_role_bindings = v1_rbac.list_cluster_role_binding(label_selector=label_selector).items
+        for binding in cluster_role_bindings:
+            debug(f'delete {binding.metadata.name}')
+            v1_rbac.delete_cluster_role_binding(binding.metadata.name)
+            refs.append(binding)
+
+        return refs

{walker/k8s_utils → adam/utils_k8s}/services.py

@@ -1,8 +1,8 @@
 from typing import List
 from kubernetes import client
 
-from walker.config import Config
-from walker.utils import log2
+from adam.config import Config
+from adam.utils import debug, log2
 
 from .kube_context import KubeContext
 
@@ -71,8 +71,7 @@ class Services:
                     namespace=namespace,
                     body=delete_options
                 )
-                if Config().is_debug():
-                    log2(f"200 Service '{name}' in namespace '{namespace}' deleted successfully.")
+                debug(f"200 Service '{name}' in namespace '{namespace}' deleted successfully.")
         except client.ApiException as e:
             log2(f"Error deleting Service '{name}': {e}")
         except Exception as e: