forked from ydb-platform/ydb
-
Notifications
You must be signed in to change notification settings - Fork 0
/
ya
executable file
·383 lines (296 loc) · 10.3 KB
/
ya
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
#!/usr/bin/env sh
# Please, keep this script in sync with arcadia/ya
# Shell commands follow
# Next line is bilingual: it starts a comment in Python, but do nothing in shell
""":"
# Find a suitable python interpreter
for cmd in python3 python; do
command -v > /dev/null $cmd && exec "$(command -v $cmd)" "$0" "$@"
done
echo "Python interpreter is not found in this system, please, install python" >&2
exit 2
":"""
# Previous line is bilingual: it ends a comment in Python, but do nothing in shell
# Shell commands end here
# Python script follows
import os
import sys
import platform
import time
def add_stage_start_to_environ(stage_name):
stages = os.environ.get('YA_STAGES', '')
os.environ['YA_STAGES'] = stages + (':' if stages else '') + '{}@{}'.format(stage_name, time.time())
RETRIES = 5
HASH_PREFIX = 10
REGISTRY_ENDPOINT = os.environ.get("YA_REGISTRY_ENDPOINT", "https://devtools-registry.s3.yandex.net")
# Please do not change this dict, it is updated automatically
# Start of mapping
PLATFORM_MAP = {
"data": {
"darwin": {
"md5": "2dc168a3aaf4ec96f956951525e36a17",
"urls": [
f"{REGISTRY_ENDPOINT}/6892379666"
]
},
"darwin-arm64": {
"md5": "d76d597fcf8bf09bff05b21c9d5e6e0f",
"urls": [
f"{REGISTRY_ENDPOINT}/6892378821"
]
},
"linux-aarch64": {
"md5": "7c0d5054bf02b83e10893dc5b754e038",
"urls": [
f"{REGISTRY_ENDPOINT}/6892378087"
]
},
"win32-clang-cl": {
"md5": "0af14f489d5da14d1becb98e5ddf9408",
"urls": [
f"{REGISTRY_ENDPOINT}/6892380563"
]
},
"linux": {
"md5": "124ceda651b91c93275b8c589301fd1e",
"urls": [
f"{REGISTRY_ENDPOINT}/6892381380"
]
}
}
} # End of mapping
add_stage_start_to_environ('ya-script-initialization')
def create_dirs(path):
try:
os.makedirs(path)
except OSError as e:
import errno
if e.errno != errno.EEXIST:
raise
return path
def home_dir():
# Do not trust $HOME, as it is unreliable in certain environments
# Temporarily delete os.environ["HOME"] to force reading current home directory from /etc/passwd
home_from_env = os.environ.pop("HOME", None)
try:
home_from_passwd = os.path.expanduser("~")
if os.path.isabs(home_from_passwd):
# This home dir is valid, prefer it over $HOME
return home_from_passwd
else:
# When python is built with musl (this is quire weird though),
# only users from /etc/passwd will be properly resolved,
# as musl does not have nss module for LDAP integration.
return home_from_env
finally:
if home_from_env is not None:
os.environ["HOME"] = home_from_env
def misc_root():
return create_dirs(os.getenv('YA_CACHE_DIR') or os.path.join(home_dir(), '.ya'))
def tool_root():
return create_dirs(os.getenv('YA_CACHE_DIR_TOOLS') or os.path.join(misc_root(), 'tools'))
# TODO: remove when switched to S3, won't be needed in OSS
def ya_token():
def get_token_from_file():
try:
with open(os.environ.get('YA_TOKEN_PATH', os.path.join(home_dir(), '.ya_token')), 'r') as f:
return f.read().strip()
except:
pass
return os.getenv('YA_TOKEN') or get_token_from_file()
TOOLS_DIR = tool_root()
def uniq(size=6):
import string
import random
return ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(size))
_ssl_is_tuned = False
def _tune_ssl():
global _ssl_is_tuned
if _ssl_is_tuned:
return
try:
import ssl
ssl._create_default_https_context = ssl._create_unverified_context
except AttributeError:
pass
try:
import urllib3
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
except (AttributeError, ImportError):
pass
_ssl_is_tuned = True
def _fetch(url, into):
import hashlib
_tune_ssl()
add_stage_start_to_environ('ya-script-download')
from urllib.request import urlopen
from urllib.request import Request
from urllib.parse import urlparse
request = Request(str(url))
# TODO: Remove when switched to S3 distribution
request.add_header('User-Agent', 'ya-bootstrap')
token = ya_token()
if token:
request.add_header('Authorization', 'OAuth {}'.format(token))
md5 = hashlib.md5()
sys.stderr.write('Downloading %s ' % url)
sys.stderr.flush()
conn = urlopen(request, timeout=10)
sys.stderr.write('[')
sys.stderr.flush()
try:
with open(into, 'wb') as f:
while True:
block = conn.read(1024 * 1024)
sys.stderr.write('.')
sys.stderr.flush()
if block:
md5.update(block)
f.write(block)
else:
break
return md5.hexdigest()
finally:
sys.stderr.write('] ')
sys.stderr.flush()
def _atomic_fetch(url, into, md5):
tmp_dest = into + '.' + uniq()
try:
real_md5 = _fetch(url, tmp_dest)
if real_md5 != md5:
raise Exception('MD5 mismatched: %s differs from %s' % (real_md5, md5))
os.rename(tmp_dest, into)
sys.stderr.write('OK\n')
except Exception as e:
sys.stderr.write('ERROR: ' + str(e) + '\n')
raise
finally:
try:
os.remove(tmp_dest)
except OSError:
pass
def _extract(path, into):
import tarfile
tar = tarfile.open(path, errorlevel=2)
# tar.extractall() will try to set file ownership according to the attributes stored in the archive
# by calling TarFile.chown() method.
# As this information is hardly relevant to the point of deployment / extraction,
# it will just fail (python2) if ya is executed with root euid, or silently set non-existent numeric owner (python3)
# to the files being extracted.
# mock it with noop to retain current user ownership.
tar.chown = lambda *args, **kwargs: None
if sys.version_info >= (3, 12):
tar.extractall(path=into, filter='data')
else:
tar.extractall(path=into)
tar.close()
def _get(urls, md5):
dest_path = os.path.join(TOOLS_DIR, md5[:HASH_PREFIX])
if not os.path.exists(dest_path):
for iter in range(RETRIES):
try:
_atomic_fetch(urls[iter % len(urls)], dest_path, md5)
break
except Exception:
if iter + 1 == RETRIES:
raise
else:
time.sleep(iter)
return dest_path
def _get_dir(urls, md5, ya_name):
dest_dir = os.path.join(TOOLS_DIR, md5[:HASH_PREFIX] + '_d')
if os.path.isfile(os.path.join(dest_dir, ya_name)):
return dest_dir
try:
packed_path = _get(urls, md5)
except Exception:
if os.path.isfile(os.path.join(dest_dir, ya_name)):
return dest_dir
raise
add_stage_start_to_environ('ya-script-extract')
tmp_dir = dest_dir + '.' + uniq()
try:
try:
_extract(packed_path, tmp_dir)
except Exception:
if os.path.isfile(os.path.join(dest_dir, ya_name)):
return dest_dir
raise
try:
os.rename(tmp_dir, dest_dir)
except OSError as e:
import errno
if e.errno != errno.ENOTEMPTY:
raise
return dest_dir
finally:
import shutil
shutil.rmtree(tmp_dir, ignore_errors=True)
try:
os.remove(packed_path)
except Exception:
pass
def _mine_repo_root():
# We think that this script is located in the root of the repo.
return os.path.dirname(os.path.realpath(__file__))
def main():
if not os.path.exists(TOOLS_DIR):
os.makedirs(TOOLS_DIR)
result_args = sys.argv[1:]
meta = PLATFORM_MAP['data']
my_platform = platform.system().lower()
my_machine = platform.machine().lower()
if my_platform == 'linux':
if 'ppc64le' in platform.platform():
my_platform = 'linux-ppc64le'
elif 'aarch64' in platform.platform():
my_platform = 'linux-aarch64'
else:
my_platform = 'linux_musl'
if my_platform == 'darwin' and my_machine == 'arm64':
my_platform = 'darwin-arm64'
def _platform_key(target_platform):
"""match by max prefix length, prefer shortest"""
def _key_for_platform(platform):
return len(os.path.commonprefix([target_platform, platform])), -len(platform)
return _key_for_platform
best_key = max(meta.keys(), key=_platform_key(my_platform))
value = meta[best_key]
ya_name = {'win32': 'ya-bin.exe', 'win32-clang-cl': 'ya-bin.exe'}.get(best_key, 'ya-bin') # XXX
ya_dir = _get_dir(value['urls'], value['md5'], ya_name)
add_stage_start_to_environ('ya-script-prepare')
# Popen `args` must have `str` type
ya_path = str(os.path.join(ya_dir, ya_name))
env = os.environ.copy()
if 'YA_SOURCE_ROOT' not in env:
src_root = _mine_repo_root()
if src_root is not None:
env['YA_SOURCE_ROOT'] = src_root
# Disable respawn for opensource/ya
if __file__.endswith('ya/opensource/ya'):
env['YA_NO_RESPAWN'] = os.environ.get('YA_NO_RESPAWN', '1')
for env_name in [
'LD_PRELOAD',
'Y_PYTHON_SOURCE_ROOT',
]:
if env_name in os.environ:
sys.stderr.write(
"Warn: {}='{}' is specified and may affect the correct operation of the ya\n".format(
env_name, env[env_name]
)
)
if os.name == 'nt':
import subprocess
p = subprocess.Popen([ya_path] + result_args, env=env)
p.wait()
sys.exit(p.returncode)
else:
os.execve(ya_path, [ya_path] + result_args, env)
if __name__ == '__main__':
try:
main()
except Exception as e:
sys.stderr.write('ERROR: ' + str(e) + '\n')
from traceback import format_exc
sys.stderr.write(format_exc() + "\n")
sys.exit(1)