Skip to content

Commit c31b8f6

Browse files
piax93piax93
committed
new remote-run cli
1 parent 960279f commit c31b8f6

File tree

3 files changed

+291
-10
lines changed

3 files changed

+291
-10
lines changed

‎paasta_tools/cli/cmds/remote_run.py

Lines changed: 174 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,21 +13,185 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515
import argparse
16+
import pty
17+
import shlex
18+
import time
1619

17-
from paasta_tools.utils import PaastaColors
20+
from paasta_tools.cli.utils import get_paasta_oapi_api_clustername
21+
from paasta_tools.cli.utils import get_paasta_oapi_client_with_auth
22+
from paasta_tools.cli.utils import lazy_choices_completer
23+
from paasta_tools.paastaapi.model.remote_run_start import RemoteRunStart
24+
from paasta_tools.paastaapi.model.remote_run_stop import RemoteRunStop
25+
from paasta_tools.utils import get_username
26+
from paasta_tools.utils import list_clusters
27+
from paasta_tools.utils import list_services
28+
from paasta_tools.utils import load_system_paasta_config
29+
from paasta_tools.utils import SystemPaastaConfig
1830

1931

20-
def add_subparser(subparsers: argparse._SubParsersAction):
21-
subparsers.add_parser(
32+
MAX_POLLING_ATTEMPTS = 20
33+
KUBECTL_CMD_TEMPLATE = (
34+
"kubectl-eks-{cluster} --token {token} exec -it -n {namespace} {pod} -- /bin/bash"
35+
)
36+
37+
38+
def paasta_remote_run_start(
39+
args: argparse.Namespace,
40+
system_paasta_config: SystemPaastaConfig,
41+
) -> int:
42+
client = get_paasta_oapi_client_with_auth(
43+
cluster=get_paasta_oapi_api_clustername(cluster=args.cluster, is_eks=True),
44+
system_paasta_config=system_paasta_config,
45+
)
46+
if not client:
47+
print("Cannot get a paasta-api client")
48+
return 1
49+
50+
user = get_username()
51+
start_response = client.remote_run.remote_run_start(
52+
args.service,
53+
args.instance,
54+
RemoteRunStart(
55+
user=user,
56+
interactive=args.interactive,
57+
recreate=args.recreate,
58+
max_duration=args.max_duration,
59+
),
60+
)
61+
if start_response.status >= 300:
62+
print(f"Error from PaaSTA APIs while starting job: {start_response.message}")
63+
return 1
64+
65+
for _ in range(MAX_POLLING_ATTEMPTS):
66+
poll_response = client.remote_run.remote_run_poll(
67+
args.service,
68+
args.instance,
69+
start_response.job_name,
70+
)
71+
if poll_response.status == 200:
72+
break
73+
time.sleep(3)
74+
else:
75+
print("Timed out while waiting for job to start")
76+
return 1
77+
78+
if not args.interactive:
79+
print("Successfully started remote-run job")
80+
return 0
81+
82+
token_response = client.remote_run.remote_run_token(
83+
args.service, args.instance, user
84+
)
85+
86+
exec_command = KUBECTL_CMD_TEMPLATE.format(
87+
cluster=args.cluster,
88+
namespace=poll_response.namespace,
89+
pod=poll_response.pod_name,
90+
token=token_response.token,
91+
)
92+
pty.spawn(shlex.split(exec_command))
93+
return 0
94+
95+
96+
def paasta_remote_run_stop(
97+
args: argparse.Namespace,
98+
system_paasta_config: SystemPaastaConfig,
99+
) -> int:
100+
client = get_paasta_oapi_client_with_auth(
101+
cluster=get_paasta_oapi_api_clustername(cluster=args.cluster, is_eks=True),
102+
system_paasta_config=system_paasta_config,
103+
)
104+
if not client:
105+
print("Cannot get a paasta-api client")
106+
return 1
107+
response = client.remote_run.remote_run_stop(
108+
args.service, args.instance, RemoteRunStop(user=get_username())
109+
)
110+
print(response.message)
111+
return 0 if response.status < 300 else 1
112+
113+
114+
def add_common_args_to_parser(parser: argparse.ArgumentParser):
115+
parser.add_argument(
116+
"-s",
117+
"--service",
118+
help="The name of the service you wish to inspect. Required.",
119+
required=True,
120+
).completer = lazy_choices_completer(
121+
list_services
122+
) # type: ignore
123+
parser.add_argument(
124+
"-i",
125+
"--instance",
126+
help=(
127+
"Simulate a docker run for a particular instance of the "
128+
"service, like 'main' or 'canary'. Required."
129+
),
130+
required=True,
131+
)
132+
parser.add_argument(
133+
"-c",
134+
"--cluster",
135+
help="The name of the cluster you wish to run your task on.",
136+
required=True,
137+
).completer = lazy_choices_completer(
138+
list_clusters
139+
) # type: ignore
140+
141+
142+
def add_subparser(subparsers: argparse._SubParsersAction) -> None:
143+
remote_run_parser = subparsers.add_parser(
22144
"remote-run",
23-
help="Schedule adhoc service sandbox on PaaSTA cluster",
24-
description=(
25-
"`paasta remote-run` is useful for running adhoc commands in "
26-
"context of a service's Docker image."
145+
help="Run stuff remotely.",
146+
description=("'paasta remote-run' runs stuff remotely "),
147+
)
148+
subparsers = remote_run_parser.add_subparsers(dest="remote_run_command")
149+
start_parser = subparsers.add_parser(
150+
"start",
151+
help="Start or connect to a remote-run job",
152+
description=("Starts or connects to a remote-run-job"),
153+
)
154+
start_parser.add_argument(
155+
"-I",
156+
"--interactive",
157+
help=(
158+
'Run container in interactive mode. If interactive is set the default command will be "bash" '
159+
'unless otherwise set by the "--cmd" flag'
27160
),
161+
action="store_true",
162+
default=False,
163+
)
164+
start_parser.add_argument(
165+
"-m",
166+
"--max-duration",
167+
help=(
168+
"Amount of time in seconds after which the job is "
169+
"automatically stopped (capped by the API backend)"
170+
),
171+
type=int,
172+
default=1800,
173+
)
174+
start_parser.add_argument(
175+
"-r",
176+
"--recreate",
177+
help="Recreate remote-run job if already existing",
178+
action="store_true",
179+
default=False,
180+
)
181+
stop_parser = subparsers.add_parser(
182+
"stop",
183+
help="Stop! In the name of Paasta",
184+
description="Stop your remote-run job if it exists",
28185
)
186+
add_common_args_to_parser(start_parser)
187+
add_common_args_to_parser(stop_parser)
188+
remote_run_parser.set_defaults(command=paasta_remote_run)
29189

30190

31-
def paasta_remote_run(args: argparse.Namespace):
32-
print(PaastaColors.red("Error: functionality under construction"))
33-
return 1
191+
def paasta_remote_run(args: argparse.Namespace) -> int:
192+
system_paasta_config = load_system_paasta_config()
193+
if args.remote_run_command == "start":
194+
return paasta_remote_run_start(args, system_paasta_config)
195+
elif args.remote_run_command == "stop":
196+
return paasta_remote_run_stop(args, system_paasta_config)
197+
raise ValueError(f"Unsupported subcommand: {args.remote_run_command}")

‎paasta_tools/cli/utils.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@
4343

4444
from paasta_tools import remote_git
4545
from paasta_tools.adhoc_tools import load_adhoc_job_config
46+
from paasta_tools.api.client import get_paasta_oapi_client
47+
from paasta_tools.api.client import PaastaOApiClient
4648
from paasta_tools.cassandracluster_tools import load_cassandracluster_instance_config
4749
from paasta_tools.eks_tools import EksDeploymentConfig
4850
from paasta_tools.eks_tools import load_eks_service_config
@@ -1135,3 +1137,16 @@ def get_sso_service_auth_token() -> str:
11351137
"""Generate an authentication token for the calling user from the Single Sign On provider"""
11361138
client_id = load_system_paasta_config().get_service_auth_sso_oidc_client_id()
11371139
return get_and_cache_jwt_default(client_id)
1140+
1141+
1142+
def get_paasta_oapi_client_with_auth(
1143+
cluster: str = None,
1144+
system_paasta_config: SystemPaastaConfig = None,
1145+
http_res: bool = False,
1146+
) -> Optional[PaastaOApiClient]:
1147+
return get_paasta_oapi_client(
1148+
cluster=cluster,
1149+
system_paasta_config=system_paasta_config,
1150+
http_res=http_res,
1151+
auth_token=get_sso_service_auth_token(),
1152+
)

‎tests/cli/test_cmds_remote_run.py

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
# Copyright 2015-2017 Yelp Inc.
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+
from unittest.mock import call
15+
from unittest.mock import MagicMock
16+
from unittest.mock import patch
17+
18+
from paasta_tools.cli.cmds.remote_run import paasta_remote_run_start
19+
from paasta_tools.cli.cmds.remote_run import paasta_remote_run_stop
20+
from paasta_tools.paastaapi.model.remote_run_start import RemoteRunStart
21+
from paasta_tools.paastaapi.model.remote_run_stop import RemoteRunStop
22+
23+
24+
@patch("paasta_tools.cli.cmds.remote_run.time", autospec=True)
25+
@patch(
26+
"paasta_tools.cli.cmds.remote_run.get_username", return_value="pippo", autospec=True
27+
)
28+
@patch("paasta_tools.cli.cmds.remote_run.pty", autospec=True)
29+
@patch(
30+
"paasta_tools.cli.cmds.remote_run.get_paasta_oapi_client_with_auth", autospec=True
31+
)
32+
def test_paasta_remote_run_start(mock_get_client, mock_pty, *_):
33+
mock_config = MagicMock()
34+
mock_args = MagicMock(
35+
service="foo",
36+
instance="bar",
37+
cluster="dev",
38+
interactive=True,
39+
recreate=False,
40+
max_duration=100,
41+
)
42+
mock_client = mock_get_client.return_value
43+
mock_client.remote_run.remote_run_start.return_value = MagicMock(
44+
status=200, message="started", job_name="foobar"
45+
)
46+
mock_client.remote_run.remote_run_poll.side_effect = [
47+
MagicMock(status=204, message="waiting"),
48+
MagicMock(status=204, message="waiting"),
49+
MagicMock(
50+
status=200, message="started", pod_name="foobar-123", namespace="svcfoo"
51+
),
52+
]
53+
mock_client.remote_run.remote_run_token.return_value = MagicMock(token="aaabbbccc")
54+
paasta_remote_run_start(mock_args, mock_config)
55+
mock_client.remote_run.remote_run_start.assert_called_once_with(
56+
"foo",
57+
"bar",
58+
RemoteRunStart(
59+
user="pippo", interactive=True, recreate=False, max_duration=100
60+
),
61+
)
62+
mock_client.remote_run.remote_run_poll.assert_has_calls(
63+
[call("foo", "bar", "foobar")] * 3
64+
)
65+
mock_client.remote_run.remote_run_token.assert_called_once_with(
66+
"foo", "bar", "pippo"
67+
)
68+
mock_pty.spawn.assert_called_once_with(
69+
[
70+
"kubectl-eks-dev",
71+
"--token",
72+
"aaabbbccc",
73+
"exec",
74+
"-it",
75+
"-n",
76+
"svcfoo",
77+
"foobar-123",
78+
"--",
79+
"/bin/bash",
80+
]
81+
)
82+
83+
84+
@patch(
85+
"paasta_tools.cli.cmds.remote_run.get_username", return_value="pippo", autospec=True
86+
)
87+
@patch(
88+
"paasta_tools.cli.cmds.remote_run.get_paasta_oapi_client_with_auth", autospec=True
89+
)
90+
def test_paasta_remote_run_stop(mock_get_client, _):
91+
mock_config = MagicMock()
92+
mock_args = MagicMock(service="foo", instance="bar")
93+
mock_client = mock_get_client.return_value
94+
mock_client.remote_run.remote_run_stop.return_value = MagicMock(
95+
status=200, message="stopped"
96+
)
97+
paasta_remote_run_stop(mock_args, mock_config)
98+
mock_client.remote_run.remote_run_stop.assert_called_once_with(
99+
"foo",
100+
"bar",
101+
RemoteRunStop(user="pippo"),
102+
)

0 commit comments

Comments
 (0)