ported asyncio stream

examples_asyncio/pod_exec.py

@@ -0,0 +1,149 @@
+import asyncio
+
+from aiohttp.http import WSMsgType
+
+from kubernetes.aio import client, config, utils
+from kubernetes.aio.client.api_client import ApiClient
+from kubernetes.aio.stream import WsApiClient
+from kubernetes.aio.stream.ws_client import (
+    ERROR_CHANNEL, STDERR_CHANNEL, STDOUT_CHANNEL,
+)
+
+BUSYBOX_POD = "busybox-test"
+
+
+async def find_busybox_pod():
+    async with ApiClient() as api:
+        v1 = client.CoreV1Api(api)
+        ret = await v1.list_pod_for_all_namespaces()
+        for i in ret.items:
+            if i.metadata.namespace == 'default' and i.metadata.name == BUSYBOX_POD:
+                print(f"Found busybox pod: {i.metadata.name}")
+                return i.metadata.name
+    return None
+
+
+async def create_busybox_pod():
+    print(f"Pod {BUSYBOX_POD} does not exist. Creating it...")
+    manifest = {
+        'apiVersion': 'v1',
+        'kind': 'Pod',
+        'metadata': {
+            'name': BUSYBOX_POD,
+        },
+        'spec': {
+            'containers': [{
+                'image': 'busybox',
+                'name': 'sleep',
+                "args": [
+                    "/bin/sh",
+                    "-c",
+                    "while true; do date; sleep 5; done"
+                ]
+            }]
+        }
+    }
+    async with ApiClient() as api:
+        objects = await utils.create_from_dict(api, manifest, namespace="default")
+        pod = objects[0]
+        print(f"Created pod {pod.metadata.name}.")
+        return pod.metadata.name
+
+
+async def wait_busybox_pod_ready():
+    print(f"Waiting pod {BUSYBOX_POD} to be ready.")
+    async with ApiClient() as api:
+        v1 = client.CoreV1Api(api)
+        while True:
+            ret = await v1.read_namespaced_pod(name=BUSYBOX_POD, namespace="default")
+            if ret.status.phase != 'Pending':
+                break
+            await asyncio.sleep(1)
+
+
+async def main():
+    # Configs can be set in Configuration class directly or using helper
+    # utility. If no argument provided, the config will be loaded from
+    # default location.
+    await config.load_kube_config()
+
+    pod = await find_busybox_pod()
+    if not pod:
+        pod = await create_busybox_pod()
+    await wait_busybox_pod_ready()
+
+    # Execute a command in a pod non-interactively, and display its output
+    print("-------------")
+    async with WsApiClient() as ws_api:
+        v1_ws = client.CoreV1Api(api_client=ws_api)
+        exec_command = [
+            "/bin/sh",
+            "-c",
+            "echo This message goes to stderr >&2; echo This message goes to stdout",
+        ]
+        ret = await v1_ws.connect_get_namespaced_pod_exec(
+            pod,
+            "default",
+            command=exec_command,
+            stderr=True,
+            stdin=False,
+            stdout=True,
+            tty=False,
+        )
+        print(f"Response: {ret}")
+
+    # Execute a command interactively. If _preload_content=False is passed to
+    # connect_get_namespaced_pod_exec(), the returned object is an aiohttp ClientWebSocketResponse
+    # object, that can be manipulated directly.
+    print("-------------")
+    async with WsApiClient() as ws_api:
+        v1_ws = client.CoreV1Api(api_client=ws_api)
+        exec_command = ['/bin/sh']
+        websocket = await v1_ws.connect_get_namespaced_pod_exec(
+            BUSYBOX_POD,
+            "default",
+            command=exec_command,
+            stderr=True,
+            stdin=True,
+            stdout=True,
+            tty=False,
+            _preload_content=False,
+        )
+        commands = [
+            "echo 'This message goes to stdout'\n",
+            "echo 'This message goes to stderr' >&2\n",
+            "exit 1\n",
+        ]
+        error_data = ""
+        closed = False
+        async with websocket as ws:
+            while commands and not closed:
+                command = commands.pop(0)
+                stdin_channel_prefix = chr(0)
+                await ws.send_bytes((stdin_channel_prefix + command).encode("utf-8"))
+                while True:
+                    try:
+                        msg = await ws.receive(timeout=1)
+                    except asyncio.TimeoutError:
+                        break
+                    if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED):
+                        closed = True
+                        break
+                    channel = msg.data[0]
+                    data = msg.data[1:].decode("utf-8")
+                    if not data:
+                        continue
+                    if channel == STDOUT_CHANNEL:
+                        print(f"stdout: {data}")
+                    elif channel == STDERR_CHANNEL:
+                        print(f"stderr: {data}")
+                    elif channel == ERROR_CHANNEL:
+                        error_data += data
+            if error_data:
+                returncode = ws_api.parse_error_data(error_data)
+                print(f"Exit code: {returncode}")
+
+if __name__ == "__main__":
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(main())
+    loop.close()

kubernetes/aio/__init__.py

@@ -19,7 +19,8 @@
 import kubernetes.aio.client as client
 import kubernetes.aio.config as config
 import kubernetes.aio.dynamic as dynamic
+import kubernetes.aio.stream as stream
 import kubernetes.aio.utils as utils
 import kubernetes.aio.watch as watch
 
-__all__ = ["client", "config", "dynamic", "utils", "watch"]
+__all__ = ["client", "config", "dynamic", "stream", "utils", "watch"]

kubernetes/aio/stream/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2017 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .ws_client import WsApiClient

kubernetes/aio/stream/ws_client.py

@@ -0,0 +1,111 @@
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+import json
+
+from six.moves.urllib.parse import urlencode, urlparse, urlunparse
+
+from kubernetes.aio.client import ApiClient
+from kubernetes.aio.client.rest import RESTResponse
+
+STDIN_CHANNEL = 0
+STDOUT_CHANNEL = 1
+STDERR_CHANNEL = 2
+ERROR_CHANNEL = 3
+RESIZE_CHANNEL = 4
+
+
+def get_websocket_url(url):
+    parsed_url = urlparse(url)
+    parts = list(parsed_url)
+    if parsed_url.scheme == 'http':
+        parts[0] = 'ws'
+    elif parsed_url.scheme == 'https':
+        parts[0] = 'wss'
+    return urlunparse(parts)
+
+
+class WsResponse(RESTResponse):
+
+    def __init__(self, status, data):
+        self.status = status
+        self.data = data
+        self.headers = {}
+        self.reason = None
+
+    def getheaders(self):
+        return self.headers
+
+    def getheader(self, name, default=None):
+        return self.headers.get(name, default)
+
+
+class WsApiClient(ApiClient):
+
+    def __init__(self, configuration=None, header_name=None, header_value=None,
+                 cookie=None, pool_threads=1, heartbeat=None):
+        super().__init__(configuration, header_name, header_value, cookie, pool_threads)
+        self.heartbeat = heartbeat
+
+    @classmethod
+    def parse_error_data(cls, error_data):
+        """
+        Parse data received on ERROR_CHANNEL and return the command exit code.
+        """
+        error_data_json = json.loads(error_data)
+        if error_data_json.get("status") == "Success":
+            return 0
+        return int(error_data_json["details"]["causes"][0]['message'])
+
+    async def request(self, method, url, query_params=None, headers=None,
+                      post_params=None, body=None, _preload_content=True,
+                      _request_timeout=None):
+
+        # Expand command parameter list to indivitual command params
+        if query_params:
+            new_query_params = []
+            for key, value in query_params:
+                if key == 'command' and isinstance(value, list):
+                    for command in value:
+                        new_query_params.append((key, command))
+                else:
+                    new_query_params.append((key, value))
+            query_params = new_query_params
+
+        if headers is None:
+            headers = {}
+        if 'sec-websocket-protocol' not in headers:
+            headers['sec-websocket-protocol'] = 'v4.channel.k8s.io'
+
+        if query_params:
+            url += '?' + urlencode(query_params)
+
+        url = get_websocket_url(url)
+
+        if _preload_content:
+
+            resp_all = ''
+            async with self.rest_client.pool_manager.ws_connect(url, headers=headers, heartbeat=self.heartbeat) as ws:
+                async for msg in ws:
+                    msg = msg.data.decode('utf-8')
+                    if len(msg) > 1:
+                        channel = ord(msg[0])
+                        data = msg[1:]
+                        if data:
+                            if channel in [STDOUT_CHANNEL, STDERR_CHANNEL]:
+                                resp_all += data
+
+            return WsResponse(200, resp_all.encode('utf-8'))
+
+        else:
+
+            return self.rest_client.pool_manager.ws_connect(url, headers=headers, heartbeat=self.heartbeat)