def open(self,
path: str,
mode: str = 'r',
buffer_size: int = 8192) -> io.IOBase:
del buffer_size
if mode not in ('r', 'rb', 'w', 'wb'):
raise ValueError(f'Unsupported mode: {repr(mode)}')
strm: io.IOBase
if mode[0] == 'r':
strm = SyncReadableStream(async_to_blocking(self.afs.open(path)),
path)
else:
assert mode[0] == 'w'
strm = SyncWritableStream(async_to_blocking(self.afs.create(path)),
path)
if path[-3:] == '.gz' or path[-4:] == '.bgz':
strm = gzip.GzipFile(fileobj=strm, mode=mode)
if 'b' not in mode:
strm = io.TextIOWrapper(
strm, encoding='utf-8'
) # type: ignore # typeshed is wrong, this *is* an IOBase
return strm
def __init__(self,
gcs_project=None,
fs=None,
deploy_config=None,
session=None,
headers=None,
_token=None):
self._client = MemoryClient(gcs_project, fs, deploy_config, session,
headers, _token)
async_to_blocking(self._client.async_init())
def exists(self, path: str) -> bool:
async def _exists():
dir_path = path
if dir_path[-1] != '/':
dir_path = dir_path + '/'
return any(await asyncio.gather(self.afs.isfile(path),
self.afs.isdir(dir_path)))
return async_to_blocking(_exists())
def create_user(username: str,
email: Optional[str],
is_developer: bool = False,
is_service_account: bool = False):
return async_to_blocking(
async_create_user(username,
email,
is_developer=is_developer,
is_service_account=is_service_account))
def test_small_write_around(self):
async def read(url):
async with await self.fs.open(url) as f:
return await f.read()
cases = [('empty_file', b''), ('null', b'\0'),
('small', b'hello world')]
for file, data in cases:
handle = async_to_blocking(
self.add_temp_file_from_string(file, data))
expected = async_to_blocking(read(handle))
self.assertEqual(expected, data)
i = 0
cached = self.client._get_file_if_exists(handle)
while cached is None and i < 10:
cached = self.client._get_file_if_exists(handle)
i += 1
self.assertEqual(cached, expected)
def copy(self, src: str, dest: str, *, max_simultaneous_transfers=75):
transfer = Transfer(src, dest)
async def _copy():
sema = asyncio.Semaphore(max_simultaneous_transfers)
async with sema:
await Copier.copy(self.afs, asyncio.Semaphore, transfer)
return async_to_blocking(_copy())
def create_user(username: str,
login_id: str,
is_developer: bool,
is_service_account: bool,
namespace: Optional[str] = None):
return async_to_blocking(
async_create_user(username,
login_id,
is_developer,
is_service_account,
namespace=namespace))
def ls(self,
path: str,
_max_simultaneous_files: int = 50) -> List[StatResult]:
async def _ls():
return await bounded_gather(*[
functools.partial(self._fle_to_dict, fle)
async for fle in await self.afs.listfiles(path)
],
parallelism=_max_simultaneous_files)
return async_to_blocking(_ls())
def stat(self, path: str) -> StatResult:
async def size_bytes_or_none():
try:
return await (await self.afs.statfile(path)).size()
except FileNotFoundError:
return None
size_bytes, is_dir = async_to_blocking(
asyncio.gather(size_bytes_or_none(), self._async_is_dir(path)))
if size_bytes is None:
if not is_dir:
raise FileNotFoundError(path)
return _stat_result(True, 0, path)
return _stat_result(is_dir, size_bytes, path)
def _run(
self,
batch: 'batch.Batch',
dry_run: bool,
verbose: bool,
delete_scratch_on_exit: bool,
wait: bool = True,
open: bool = False,
disable_progress_bar: bool = False,
callback: Optional[str] = None,
token: Optional[str] = None,
**backend_kwargs
) -> bc.Batch: # pylint: disable-msg=too-many-statements
"""Execute a batch.
Warning
-------
This method should not be called directly. Instead, use :meth:`.batch.Batch.run`
and pass :class:`.ServiceBackend` specific arguments as key-word arguments.
Parameters
----------
batch:
Batch to execute.
dry_run:
If `True`, don't execute code.
verbose:
If `True`, print debugging output.
delete_scratch_on_exit:
If `True`, delete temporary directories with intermediate files.
wait:
If `True`, wait for the batch to finish executing before returning.
open:
If `True`, open the UI page for the batch.
disable_progress_bar:
If `True`, disable the progress bar.
callback:
If not `None`, a URL that will receive at most one POST request
after the entire batch completes.
token:
If not `None`, a string used for idempotency of batch submission.
"""
return async_to_blocking(
self._async_run(batch, dry_run, verbose, delete_scratch_on_exit,
wait, open, disable_progress_bar, callback, token,
**backend_kwargs))
def get_userinfo(deploy_config=None):
return async_to_blocking(async_get_userinfo(deploy_config))
def get_userinfo(deploy_config=None, session_id=None, client_session=None):
return async_to_blocking(
async_get_userinfo(deploy_config=deploy_config,
session_id=session_id,
client_session=client_session))
def copy_paste_login(copy_paste_token, namespace=None):
return async_to_blocking(
async_copy_paste_login(copy_paste_token, namespace))
def readinto(self, b):
b[:] = async_to_blocking(self.ars.readexactly(len(b)))
def __init__(self, cm: AsyncContextManager[WritableStream], name: str):
super().__init__()
self.cm = cm
self.aws = async_to_blocking(self.cm.__aenter__())
self._mode = 'wb'
self._name = name
def is_file(self, path: str) -> bool:
return async_to_blocking(self.afs.isfile(path))
def readall(self) -> bytes:
return async_to_blocking(self.ars.read(-1))
def rmtree(self, path: str):
return async_to_blocking(self.afs.rmtree(None, path))
def close(self):
self.ars.close()
async_to_blocking(self.ars.wait_closed())
def delete_user(username: str):
return async_to_blocking(async_delete_user(username))
def remove(self, path: str):
return async_to_blocking(self.afs.remove(path))
return None
return row['id']
return await db.execute_insertone(
'''
INSERT INTO users (state, username, email, is_developer, is_service_account)
VALUES (%s, %s, %s, %s, %s);
''',
('creating', username, email, 1, 0),
)
return await insert(db)
async def main():
parser = argparse.ArgumentParser(
description='Create initial Hail as a service account.')
parser.add_argument('username', help='The username of the initial user.')
parser.add_argument('email', help='The email of the initial user.')
args = parser.parse_args()
db = Database()
await db.async_init(maxsize=50)
await insert_user_if_not_exists(db, args.username, args.email)
async_to_blocking(main())
def close(self):
self.aws.close()
async_to_blocking(self.cm.__aexit__())
def is_dir(self, path: str) -> bool:
return async_to_blocking(self._async_is_dir(path))
def close(self):
if self._session is not None:
async_to_blocking(self._session.close())
self._session = None
def read(self, size=-1) -> bytes:
return async_to_blocking(self.ars.read(size))
def request(self, endpoint, **data):
return async_to_blocking(retry_transient_errors(self.async_request, endpoint, **data))
def write(self, b):
return async_to_blocking(self.aws.write(b))
def __exit__(self, exc_type, exc_value, traceback):
async_to_blocking(self.real_session.close())
def mkdir(self, path: str):
return async_to_blocking(self.afs.mkdir(path))