Skip to content

Commit e8e4ccb

Browse files
Merge pull request #2635 from yliaog/stream
ported asyncio stream
2 parents e21c2e8 + 1140a22 commit e8e4ccb

5 files changed

Lines changed: 401 additions & 1 deletion

File tree

examples_asyncio/pod_exec.py

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
import asyncio
2+
3+
from aiohttp.http import WSMsgType
4+
5+
from kubernetes.aio import client, config, utils
6+
from kubernetes.aio.client.api_client import ApiClient
7+
from kubernetes.aio.stream import WsApiClient
8+
from kubernetes.aio.stream.ws_client import (
9+
ERROR_CHANNEL, STDERR_CHANNEL, STDOUT_CHANNEL,
10+
)
11+
12+
BUSYBOX_POD = "busybox-test"
13+
14+
15+
async def find_busybox_pod():
16+
async with ApiClient() as api:
17+
v1 = client.CoreV1Api(api)
18+
ret = await v1.list_pod_for_all_namespaces()
19+
for i in ret.items:
20+
if i.metadata.namespace == 'default' and i.metadata.name == BUSYBOX_POD:
21+
print(f"Found busybox pod: {i.metadata.name}")
22+
return i.metadata.name
23+
return None
24+
25+
26+
async def create_busybox_pod():
27+
print(f"Pod {BUSYBOX_POD} does not exist. Creating it...")
28+
manifest = {
29+
'apiVersion': 'v1',
30+
'kind': 'Pod',
31+
'metadata': {
32+
'name': BUSYBOX_POD,
33+
},
34+
'spec': {
35+
'containers': [{
36+
'image': 'busybox',
37+
'name': 'sleep',
38+
"args": [
39+
"/bin/sh",
40+
"-c",
41+
"while true; do date; sleep 5; done"
42+
]
43+
}]
44+
}
45+
}
46+
async with ApiClient() as api:
47+
objects = await utils.create_from_dict(api, manifest, namespace="default")
48+
pod = objects[0]
49+
print(f"Created pod {pod.metadata.name}.")
50+
return pod.metadata.name
51+
52+
53+
async def wait_busybox_pod_ready():
54+
print(f"Waiting pod {BUSYBOX_POD} to be ready.")
55+
async with ApiClient() as api:
56+
v1 = client.CoreV1Api(api)
57+
while True:
58+
ret = await v1.read_namespaced_pod(name=BUSYBOX_POD, namespace="default")
59+
if ret.status.phase != 'Pending':
60+
break
61+
await asyncio.sleep(1)
62+
63+
64+
async def main():
65+
# Configs can be set in Configuration class directly or using helper
66+
# utility. If no argument provided, the config will be loaded from
67+
# default location.
68+
await config.load_kube_config()
69+
70+
pod = await find_busybox_pod()
71+
if not pod:
72+
pod = await create_busybox_pod()
73+
await wait_busybox_pod_ready()
74+
75+
# Execute a command in a pod non-interactively, and display its output
76+
print("-------------")
77+
async with WsApiClient() as ws_api:
78+
v1_ws = client.CoreV1Api(api_client=ws_api)
79+
exec_command = [
80+
"/bin/sh",
81+
"-c",
82+
"echo This message goes to stderr >&2; echo This message goes to stdout",
83+
]
84+
ret = await v1_ws.connect_get_namespaced_pod_exec(
85+
pod,
86+
"default",
87+
command=exec_command,
88+
stderr=True,
89+
stdin=False,
90+
stdout=True,
91+
tty=False,
92+
)
93+
print(f"Response: {ret}")
94+
95+
# Execute a command interactively. If _preload_content=False is passed to
96+
# connect_get_namespaced_pod_exec(), the returned object is an aiohttp ClientWebSocketResponse
97+
# object, that can be manipulated directly.
98+
print("-------------")
99+
async with WsApiClient() as ws_api:
100+
v1_ws = client.CoreV1Api(api_client=ws_api)
101+
exec_command = ['/bin/sh']
102+
websocket = await v1_ws.connect_get_namespaced_pod_exec(
103+
BUSYBOX_POD,
104+
"default",
105+
command=exec_command,
106+
stderr=True,
107+
stdin=True,
108+
stdout=True,
109+
tty=False,
110+
_preload_content=False,
111+
)
112+
commands = [
113+
"echo 'This message goes to stdout'\n",
114+
"echo 'This message goes to stderr' >&2\n",
115+
"exit 1\n",
116+
]
117+
error_data = ""
118+
closed = False
119+
async with websocket as ws:
120+
while commands and not closed:
121+
command = commands.pop(0)
122+
stdin_channel_prefix = chr(0)
123+
await ws.send_bytes((stdin_channel_prefix + command).encode("utf-8"))
124+
while True:
125+
try:
126+
msg = await ws.receive(timeout=1)
127+
except asyncio.TimeoutError:
128+
break
129+
if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED):
130+
closed = True
131+
break
132+
channel = msg.data[0]
133+
data = msg.data[1:].decode("utf-8")
134+
if not data:
135+
continue
136+
if channel == STDOUT_CHANNEL:
137+
print(f"stdout: {data}")
138+
elif channel == STDERR_CHANNEL:
139+
print(f"stderr: {data}")
140+
elif channel == ERROR_CHANNEL:
141+
error_data += data
142+
if error_data:
143+
returncode = ws_api.parse_error_data(error_data)
144+
print(f"Exit code: {returncode}")
145+
146+
if __name__ == "__main__":
147+
loop = asyncio.get_event_loop()
148+
loop.run_until_complete(main())
149+
loop.close()

kubernetes/aio/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@
1919
import kubernetes.aio.client as client
2020
import kubernetes.aio.config as config
2121
import kubernetes.aio.dynamic as dynamic
22+
import kubernetes.aio.stream as stream
2223
import kubernetes.aio.utils as utils
2324
import kubernetes.aio.watch as watch
2425

25-
__all__ = ["client", "config", "dynamic", "utils", "watch"]
26+
__all__ = ["client", "config", "dynamic", "stream", "utils", "watch"]

kubernetes/aio/stream/__init__.py

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

kubernetes/aio/stream/ws_client.py

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
# Licensed under the Apache License, Version 2.0 (the "License"); you may
2+
# not use this file except in compliance with the License. You may obtain
3+
# a copy of the License at
4+
#
5+
# http://www.apache.org/licenses/LICENSE-2.0
6+
#
7+
# Unless required by applicable law or agreed to in writing, software
8+
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
9+
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
10+
# License for the specific language governing permissions and limitations
11+
# under the License.
12+
13+
import json
14+
15+
from six.moves.urllib.parse import urlencode, urlparse, urlunparse
16+
17+
from kubernetes.aio.client import ApiClient
18+
from kubernetes.aio.client.rest import RESTResponse
19+
20+
STDIN_CHANNEL = 0
21+
STDOUT_CHANNEL = 1
22+
STDERR_CHANNEL = 2
23+
ERROR_CHANNEL = 3
24+
RESIZE_CHANNEL = 4
25+
26+
27+
def get_websocket_url(url):
28+
parsed_url = urlparse(url)
29+
parts = list(parsed_url)
30+
if parsed_url.scheme == 'http':
31+
parts[0] = 'ws'
32+
elif parsed_url.scheme == 'https':
33+
parts[0] = 'wss'
34+
return urlunparse(parts)
35+
36+
37+
class WsResponse(RESTResponse):
38+
39+
def __init__(self, status, data):
40+
self.status = status
41+
self.data = data
42+
self.headers = {}
43+
self.reason = None
44+
45+
def getheaders(self):
46+
return self.headers
47+
48+
def getheader(self, name, default=None):
49+
return self.headers.get(name, default)
50+
51+
52+
class WsApiClient(ApiClient):
53+
54+
def __init__(self, configuration=None, header_name=None, header_value=None,
55+
cookie=None, pool_threads=1, heartbeat=None):
56+
super().__init__(configuration, header_name, header_value, cookie, pool_threads)
57+
self.heartbeat = heartbeat
58+
59+
@classmethod
60+
def parse_error_data(cls, error_data):
61+
"""
62+
Parse data received on ERROR_CHANNEL and return the command exit code.
63+
"""
64+
error_data_json = json.loads(error_data)
65+
if error_data_json.get("status") == "Success":
66+
return 0
67+
return int(error_data_json["details"]["causes"][0]['message'])
68+
69+
async def request(self, method, url, query_params=None, headers=None,
70+
post_params=None, body=None, _preload_content=True,
71+
_request_timeout=None):
72+
73+
# Expand command parameter list to indivitual command params
74+
if query_params:
75+
new_query_params = []
76+
for key, value in query_params:
77+
if key == 'command' and isinstance(value, list):
78+
for command in value:
79+
new_query_params.append((key, command))
80+
else:
81+
new_query_params.append((key, value))
82+
query_params = new_query_params
83+
84+
if headers is None:
85+
headers = {}
86+
if 'sec-websocket-protocol' not in headers:
87+
headers['sec-websocket-protocol'] = 'v4.channel.k8s.io'
88+
89+
if query_params:
90+
url += '?' + urlencode(query_params)
91+
92+
url = get_websocket_url(url)
93+
94+
if _preload_content:
95+
96+
resp_all = ''
97+
async with self.rest_client.pool_manager.ws_connect(url, headers=headers, heartbeat=self.heartbeat) as ws:
98+
async for msg in ws:
99+
msg = msg.data.decode('utf-8')
100+
if len(msg) > 1:
101+
channel = ord(msg[0])
102+
data = msg[1:]
103+
if data:
104+
if channel in [STDOUT_CHANNEL, STDERR_CHANNEL]:
105+
resp_all += data
106+
107+
return WsResponse(200, resp_all.encode('utf-8'))
108+
109+
else:
110+
111+
return self.rest_client.pool_manager.ws_connect(url, headers=headers, heartbeat=self.heartbeat)

0 commit comments

Comments
 (0)