async def _download_stream(
self,
bucket: str,
object_name: str,
*,
params: Optional[Dict[str, str]] = None,
headers: Optional[Dict[str, str]] = None,
timeout: int = 10,
session: Optional[Session] = None) -> StreamResponse:
# https://cloud.google.com/storage/docs/request-endpoints#encoding
encoded_object_name = quote(object_name, safe='')
url = f'{API_ROOT}/{bucket}/o/{encoded_object_name}'
headers = headers or {}
headers.update(await self._headers())
s = AioSession(session) if session else self.session
if BUILD_GCLOUD_REST:
# stream argument is only expected by requests.Session.
# pylint: disable=unexpected-keyword-arg
return StreamResponse(
s.get(url,
headers=headers,
params=params or {},
timeout=timeout,
stream=True))
return StreamResponse(await s.get(url,
headers=headers,
params=params or {},
timeout=timeout))
async def copy(self,
destination_project: str,
destination_dataset: str,
destination_table: str,
session: Optional[Session] = None,
timeout: int = 60) -> Dict[str, Any]:
"""
Copy BQ table to another table in BQ
https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/insert
"""
if not (destination_project and destination_dataset
and destination_table):
return {}
project = await self.project()
url = f'{API_ROOT}/projects/{project}/jobs'
body = self._make_copy_body(project, destination_project,
destination_dataset, destination_table)
payload = json.dumps(body).encode('utf-8')
headers = await self.headers()
headers.update({
'Content-Length': str(len(payload)),
'Content-Type': 'application/json'
})
s = AioSession(session) if session else self.session
resp = await s.post(url,
data=payload,
headers=headers,
params=None,
timeout=timeout)
return await resp.json()
async def load(self,
source_uris: List[str],
session: Optional[Session] = None,
timeout: int = 60) -> Dict[str, Any]:
"""
Loads entities from storage to big query.
https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/insert
"""
if not source_uris:
return {}
project = await self.project()
url = f'{API_ROOT}/projects/{project}/jobs'
body = self._make_load_body(source_uris, project)
payload = json.dumps(body).encode('utf-8')
headers = await self.headers()
headers.update({
'Content-Length': str(len(payload)),
'Content-Type': 'application/json'
})
s = AioSession(session) if session else self.session
resp = await s.post(url,
data=payload,
headers=headers,
params=None,
timeout=timeout)
return await resp.json()
async def _do_upload(self,
session_uri: str,
stream: IO[AnyStr],
headers: Dict[str, str],
*,
retries: int = 5,
session: Optional[Session] = None,
timeout: int = 30) -> Dict[str, Any]:
s = AioSession(session) if session else self.session
for tries in range(retries):
try:
resp = await s.put(session_uri,
headers=headers,
data=stream,
timeout=timeout)
except ResponseError:
headers.update({'Content-Range': '*/*'})
await sleep(2.**tries) # type: ignore[func-returns-value]
continue
break
data: Dict[str, Any] = await resp.json(content_type=None)
return data
async def runQuery(self, query: BaseQuery, transaction: str = None,
consistency: Consistency = Consistency.EVENTUAL,
session: Optional[Session] = None,
timeout: int = 10) -> QueryResultBatch:
project = await self.project()
url = f'{API_ROOT}/projects/{project}:runQuery'
if transaction:
options = {'transaction': transaction}
else:
options = {'readConsistency': consistency.value}
payload = json.dumps({
'partitionId': {
'projectId': project,
'namespaceId': self.namespace,
},
query.json_key: query.to_repr(),
'readOptions': options,
}).encode('utf-8')
headers = await self.headers()
headers.update({
'Content-Length': str(len(payload)),
'Content-Type': 'application/json',
})
s = AioSession(session) if session else self.session
resp = await s.post(url, data=payload, headers=headers, timeout=timeout)
data: dict = await resp.json()
return self.query_result_batch_kind.from_repr(data['batch'])
async def lookup(self, keys: List[Key], transaction: str = None,
consistency: Consistency = Consistency.STRONG,
session: Optional[Session] = None,
timeout: int = 10) -> Dict[str, Union[EntityResult, Key]]:
project = await self.project()
url = f'{API_ROOT}/projects/{project}:lookup'
if transaction:
options = {'transaction': transaction}
else:
options = {'readConsistency': consistency.value}
payload = json.dumps({
'keys': [k.to_repr() for k in keys],
'readOptions': options,
}).encode('utf-8')
headers = await self.headers()
headers.update({
'Content-Length': str(len(payload)),
'Content-Type': 'application/json',
})
s = AioSession(session) if session else self.session
resp = await s.post(url, data=payload, headers=headers, timeout=timeout)
data: dict = await resp.json()
return {
'found': [self.entity_result_kind.from_repr(e)
for e in data.get('found', [])],
'missing': [self.entity_result_kind.from_repr(e)
for e in data.get('missing', [])],
'deferred': [self.key_kind.from_repr(k)
for k in data.get('deferred', [])],
}
async def patch(self,
table: Dict[str, Any],
session: Optional[Session] = None,
timeout: int = 60) -> Dict[str, Any]:
"""Patch an existing table specified by tableId from the dataset."""
project = await self.project()
url = (f'{API_ROOT}/projects/{project}/datasets/'
f'{self.dataset_name}/tables/{self.table_name}')
table['tableReference'] = {
'projectId': project,
'datasetId': self.dataset_name,
'tableId': self.table_name
}
table_data = json.dumps(table).encode('utf-8')
headers = await self.headers()
s = AioSession(session) if session else self.session
resp = await s.patch(url,
data=table_data,
headers=headers,
timeout=timeout)
data: Dict[str, Any] = await resp.json()
return data
async def _upload_simple(self,
url: str,
object_name: str,
stream: io.IOBase,
params: dict,
headers: dict,
*,
session: Optional[Session] = None,
timeout: int = 30) -> dict:
# https://cloud.google.com/storage/docs/json_api/v1/how-tos/simple-upload
params['name'] = object_name
params['uploadType'] = 'media'
headers.update({
'Accept': 'application/json',
})
s = AioSession(session) if session else self.session
resp = await s.post(url,
data=stream,
headers=headers,
params=params,
timeout=timeout)
data: dict = await resp.json()
return data
async def export(self,
output_bucket_prefix: str,
kinds: Optional[List[str]] = None,
namespaces: Optional[List[str]] = None,
labels: Optional[Dict[str, str]] = None,
session: Optional[Session] = None,
timeout: int = 10) -> DatastoreOperation:
project = await self.project()
url = f'{API_ROOT}/projects/{project}:export'
payload = json.dumps({
'entityFilter': {
'kinds': kinds or [],
'namespaceIds': namespaces or [],
},
'labels': labels or {},
'outputUrlPrefix': f'gs://{output_bucket_prefix}',
}).encode('utf-8')
headers = await self.headers()
headers.update({
'Content-Length': str(len(payload)),
'Content-Type': 'application/json',
})
s = AioSession(session) if session else self.session
resp = await s.post(url,
data=payload,
headers=headers,
timeout=timeout)
data: Dict[str, Any] = await resp.json()
return self.datastore_operation_kind.from_repr(data)
async def publish(self,
topic: str,
messages: List[PubsubMessage],
session: Optional[Session] = None,
timeout: int = 10) -> Dict[str, Any]:
if not messages:
return {}
url = f'{API_ROOT}/{topic}:publish'
body = {'messages': [m.to_repr() for m in messages]}
payload = json.dumps(body).encode('utf-8')
headers = await self._headers()
headers.update({
'Content-Length': str(len(payload)),
'Content-Type': 'application/json',
})
s = AioSession(session) if session else self.session
resp = await s.post(url,
data=payload,
headers=headers,
timeout=timeout)
return await resp.json()
async def allocateIds(self,
keys: List[Key],
session: Optional[Session] = None,
timeout: int = 10) -> List[Key]:
project = await self.project()
url = f'{API_ROOT}/projects/{project}:allocateIds'
payload = json.dumps({
'keys': [k.to_repr() for k in keys],
}).encode('utf-8')
headers = await self.headers()
headers.update({
'Content-Length': str(len(payload)),
'Content-Type': 'application/json',
})
s = AioSession(session) if session else self.session
resp = await s.post(url,
data=payload,
headers=headers,
timeout=timeout)
data = await resp.json()
return [self.key_kind.from_repr(k) for k in data['keys']]
async def commit(self,
mutations: List[Dict[str, Any]],
transaction: Optional[str] = None,
mode: Mode = Mode.TRANSACTIONAL,
session: Optional[Session] = None,
timeout: int = 10) -> Dict[str, Any]:
project = await self.project()
url = f'{API_ROOT}/projects/{project}:commit'
body = self._make_commit_body(mutations,
transaction=transaction,
mode=mode)
payload = json.dumps(body).encode('utf-8')
headers = await self.headers()
headers.update({
'Content-Length': str(len(payload)),
'Content-Type': 'application/json',
})
s = AioSession(session) if session else self.session
resp = await s.post(url,
data=payload,
headers=headers,
timeout=timeout)
data: Dict[str, Any] = await resp.json()
return {
'mutationResults': [
self.mutation_result_kind.from_repr(r)
for r in data.get('mutationResults', [])
],
'indexUpdates':
data['indexUpdates'],
}
async def _do_upload(self,
session_uri: str,
stream: io.IOBase,
headers: dict,
*,
retries: int = 5,
session: Optional[Session] = None,
timeout: int = 30) -> dict:
s = AioSession(session) if session else self.session
for tries in range(retries):
try:
resp = await s.put(session_uri,
headers=headers,
data=stream,
timeout=timeout)
except ResponseError:
headers.update({'Content-Range': '*/*'})
await sleep(2.**tries)
continue
break
data: dict = await resp.json(content_type=None)
return data
async def patch_metadata(self,
bucket: str,
object_name: str,
metadata: Dict[str, Any],
*,
params: Optional[Dict[str, str]] = None,
headers: Optional[Dict[str, str]] = None,
session: Optional[Session] = None,
timeout: int = DEFAULT_TIMEOUT) -> Dict[str, Any]:
# https://cloud.google.com/storage/docs/json_api/v1/objects/patch
encoded_object_name = quote(object_name, safe='')
url = f'{API_ROOT}/{bucket}/o/{encoded_object_name}'
params = params or {}
headers = headers or {}
headers.update(await self._headers())
headers['Content-Type'] = 'application/json'
s = AioSession(session) if session else self.session
resp = await s.patch(url,
data=json.dumps(metadata).encode('utf-8'),
headers=headers,
params=params,
timeout=timeout)
data: Dict[str, Any] = await resp.json(content_type=None)
return data
async def delete(self,
session: Optional[Session] = None,
timeout: int = 60) -> Dict[str, Any]:
"""Deletes the table specified by tableId from the dataset."""
project = await self.project()
url = (f'{API_ROOT}/projects/{project}/datasets/'
f'{self.dataset_name}/tables/{self.table_name}')
headers = await self.headers()
s = AioSession(session) if session else self.session
resp = await s.session.delete(url,
headers=headers,
params=None,
timeout=timeout)
try:
data: Dict[str, Any] = await resp.json()
except Exception: # pylint: disable=broad-except
# For some reason, `gcloud-rest` seems to have intermittent issues
# parsing this response. In that case, fall back to returning the
# raw response body.
try:
data = {'response': await resp.text()}
except (AttributeError, TypeError):
data = {'response': resp.text}
return data
def __init__(self, *,
service_file: Optional[Union[str, IO[AnyStr]]] = None,
token: Optional[Token] = None,
session: Optional[Session] = None) -> None:
self.session = AioSession(session, verify_ssl=VERIFY_SSL)
self.token = token or Token(service_file=service_file, scopes=SCOPES,
session=self.session.session)
async def _initiate_upload(self,
url: str,
object_name: str,
params: dict,
headers: dict,
*,
metadata: dict = None,
session: Optional[Session] = None) -> str:
params['uploadType'] = 'resumable'
metadict = (metadata or {}).copy()
metadict.update({'name': object_name})
metadata = json.dumps(metadict)
post_headers = headers.copy()
post_headers.update({
'Content-Length':
str(len(metadata)),
'Content-Type':
'application/json; charset=UTF-8',
'X-Upload-Content-Type':
headers['Content-Type'],
'X-Upload-Content-Length':
headers['Content-Length']
})
s = AioSession(session) if session else self.session
resp = await s.post(url,
headers=post_headers,
params=params,
data=metadata,
timeout=10)
session_uri: str = resp.headers['Location']
return session_uri
async def _download(self,
bucket: str,
object_name: str,
*,
params: Optional[Dict[str, str]] = None,
headers: Optional[Dict[str, str]] = None,
timeout: int = 10,
session: Optional[Session] = None) -> bytes:
# https://cloud.google.com/storage/docs/request-endpoints#encoding
encoded_object_name = quote(object_name, safe='')
url = f'{API_ROOT}/{bucket}/o/{encoded_object_name}'
headers = headers or {}
headers.update(await self._headers())
s = AioSession(session) if session else self.session
response = await s.get(url,
headers=headers,
params=params or {},
timeout=timeout)
# N.B. the GCS API sometimes returns 'application/octet-stream' when a
# string was uploaded. To avoid potential weirdness, always return a
# bytes object.
try:
data: bytes = await response.read()
except (AttributeError, TypeError):
data = response.content
return data
async def pull(self,
subscription: str,
max_messages: int,
*,
session: Optional[Session] = None,
timeout: Optional[int] = 30) -> List[SubscriberMessage]:
"""
Pull messages from subscription
"""
url = f'{API_ROOT}/v1/{subscription}:pull'
headers = await self._headers()
payload = {
'maxMessages': max_messages,
}
encoded = json.dumps(payload).encode()
s = AioSession(session) if session else self.session
resp = await s.post(url,
data=encoded,
headers=headers,
timeout=timeout)
resp = await resp.json()
return [
SubscriberMessage.from_repr(m)
for m in resp.get('receivedMessages', [])
]
async def delete(self,
bucket: str,
object_name: str,
*,
timeout: int = 10,
params: Optional[Dict[str, str]] = None,
headers: Optional[Dict[str, str]] = None,
session: Optional[Session] = None) -> str:
# https://cloud.google.com/storage/docs/request-endpoints#encoding
encoded_object_name = quote(object_name, safe='')
url = f'{API_ROOT}/{bucket}/o/{encoded_object_name}'
headers = headers or {}
headers.update(await self._headers())
s = AioSession(session) if session else self.session
resp = await s.delete(url,
headers=headers,
params=params or {},
timeout=timeout)
try:
data: str = await resp.text()
except (AttributeError, TypeError):
data = str(resp.text)
return data
async def _download(self,
bucket: str,
object_name: str,
*,
params: dict = None,
timeout: int = 10,
session: Optional[Session] = None) -> bytes:
token = await self.token.get()
# https://cloud.google.com/storage/docs/json_api/#encoding
encoded_object_name = quote(object_name, safe='')
url = f'{API_ROOT}/{bucket}/o/{encoded_object_name}'
headers = {
'Authorization': f'Bearer {token}',
}
s = AioSession(session) if session else self.session
response = await s.get(url,
headers=headers,
params=params or {},
timeout=timeout)
# N.B. the GCS API sometimes returns 'application/octet-stream' when a
# string was uploaded. To avoid potential weirdness, always return a
# bytes object.
try:
data: bytes = await response.read()
except (AttributeError, TypeError):
data: bytes = response.content
return data
async def delete(self,
bucket: str,
object_name: str,
*,
params: dict = None,
timeout: int = 10,
session: Optional[Session] = None) -> str:
token = await self.token.get()
# https://cloud.google.com/storage/docs/json_api/#encoding
encoded_object_name = quote(object_name, safe='')
url = f'{API_ROOT}/{bucket}/o/{encoded_object_name}'
headers = {
'Authorization': f'Bearer {token}',
}
s = AioSession(session) if session else self.session
resp = await s.delete(url,
headers=headers,
params=params or {},
timeout=timeout)
try:
data: str = await resp.text()
except (AttributeError, TypeError):
data: str = str(resp.text)
return data
async def _upload_multipart(self,
url: str,
object_name: str,
stream: IO[AnyStr],
params: Dict[str, str],
headers: Dict[str, str],
metadata: Dict[str, Any],
*,
session: Optional[Session] = None,
timeout: int = 30) -> Dict[str, Any]:
# https://cloud.google.com/storage/docs/json_api/v1/how-tos/multipart-upload
params['uploadType'] = 'multipart'
metadata_headers = {'Content-Type': 'application/json; charset=UTF-8'}
metadata = {
self._format_metadata_key(k): v
for k, v in metadata.items()
}
if 'metadata' in metadata:
metadata['metadata'] = {
str(k): str(v) if v is not None else None
for k, v in metadata['metadata'].items()
}
metadata['name'] = object_name
raw_body: AnyStr = stream.read()
if isinstance(raw_body, str):
bytes_body: bytes = raw_body.encode('utf-8')
else:
bytes_body = raw_body
parts = [
(metadata_headers, json.dumps(metadata).encode('utf-8')),
({
'Content-Type': headers['Content-Type']
}, bytes_body),
]
boundary = choose_boundary()
body, content_type = encode_multipart_formdata(parts, boundary)
headers.update({
'Content-Type': content_type,
'Content-Length': str(len(body)),
'Accept': 'application/json'
})
s = AioSession(session) if session else self.session
if not BUILD_GCLOUD_REST:
# Wrap data in BytesIO to ensure aiohttp does not emit warning
# when payload size > 1MB
body = io.BytesIO(body) # type: ignore[assignment]
resp = await s.post(url,
data=body,
headers=headers,
params=params,
timeout=timeout)
data: Dict[str, Any] = await resp.json(content_type=None)
return data
async def copy(self, bucket: str, object_name: str,
destination_bucket: str, *, new_name: str = None,
headers: dict = None, params: dict = None,
timeout: int = 10,
session: Optional[Session] = None) -> bytes:
"""
When files are too large, multiple calls to `rewriteTo` are made. We
refer to the same copy job by using the `rewriteToken` from the
previous return payload in subsequent `rewriteTo` calls.
Using the `rewriteTo` GCS API is preferred in part because it is able
to make multiple calls to fully copy an object whereas the `copyTo` GCS
API only calls `rewriteTo` once under the hood, and thus may fail if
files are large.
In the rare case you need to resume a copy operation, include the
`rewriteToken` in the `params` dictionary. Once you begin a multi-part
copy operation, you then have 1 week to complete the copy job.
https://cloud.google.com/storage/docs/json_api/v1/objects/rewrite
"""
token = await self.token.get()
if not new_name:
new_name = object_name
url = (f"{API_ROOT}/{bucket}/o/{quote(object_name, safe='')}/rewriteTo"
f"/b/{destination_bucket}/o/{quote(new_name, safe='')}")
# We may optionally supply metadata* to apply to the rewritten
# object, which explains why `rewriteTo` is a POST endpoint; however,
# we don't expose that here so we have to send an empty body. Therefore
# the `Content-Length` and `Content-Type` indicate an empty body.
#
# * https://cloud.google.com/storage/docs/json_api/v1/objects#resource
headers = headers or {}
headers.update({
'Authorization': f'Bearer {token}',
'Content-Length': '0',
'Content-Type': '',
})
params = params or {}
s = AioSession(session) if session else self.session
resp = await s.post(url, headers=headers, params=params,
timeout=timeout)
data: dict = await resp.json()
while not data.get('done') and data.get('rewriteToken'):
params['rewriteToken'] = data['rewriteToken']
resp = await s.post(url, headers=headers, params=params,
timeout=timeout)
data = await resp.json()
return data
def __init__(self,
project: Optional[str] = None,
service_file: Optional[Union[str, io.IOBase]] = None,
session: Optional[Session] = None,
token: Optional[Token] = None) -> None:
self._project = project
self.session = AioSession(session)
self.token = token or Token(service_file=service_file, scopes=SCOPES,
session=self.session.session)
async def delete(self,
tname: str,
session: Optional[Session] = None) -> Any:
url = f'{self.base_api_root}/{tname}'
headers = await self.headers()
s = AioSession(session) if session else self.session
resp = await s.delete(url, headers=headers)
return await resp.json()
def __init__(self, project: str, taskqueue: str,
service_file: Optional[Union[str, io.IOBase]] = None,
location: str = LOCATION,
session: Optional[Session] = None,
token: Optional[Token] = None) -> None:
self.base_api_root = f'{API_ROOT}/v2beta3'
self.api_root = (f'{self.base_api_root}/projects/{project}/'
f'locations/{location}/queues/{taskqueue}')
self.session = AioSession(session)
self.token = token or Token(service_file=service_file, scopes=SCOPES,
session=self.session.session)
async def get_bucket_metadata(self, bucket: str, *, params: dict = None,
session: Optional[Session] = None,
timeout: int = 10) -> dict:
url = f'{API_ROOT}/{bucket}'
headers = await self._headers()
s = AioSession(session) if session else self.session
resp = await s.get(url, headers=headers, params=params or {},
timeout=timeout)
data: dict = await resp.json(content_type=None)
return data
async def _get_url(
self, url: str, session: Optional[Session],
timeout: int,
params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
headers = await self.headers()
s = AioSession(session) if session else self.session
resp = await s.get(url, headers=headers, timeout=timeout,
params=params or {})
data: Dict[str, Any] = await resp.json()
return data
async def delete_topic(self,
topic: str,
*,
session: Optional[Session] = None,
timeout: Optional[int] = 10) -> None:
"""
Delete topic.
"""
url = f'{API_ROOT}/{topic}'
headers = await self._headers()
s = AioSession(session) if session else self.session
await s.delete(url, headers=headers, timeout=timeout)
