from __future__ import annotations import re from contextlib import asynccontextmanager, contextmanager from http import HTTPStatus from typing import TYPE_CHECKING, Any from apify_client._docs import docs_group from apify_client._models import ( KeyValueStore, KeyValueStoreKey, KeyValueStoreResponse, ListOfKeys, ListOfKeysResponse, ) from apify_client._pagination import DEFAULT_CHUNK_SIZE, get_cursor_iterator, get_cursor_iterator_async from apify_client._resource_clients._resource_client import ResourceClient, ResourceClientAsync from apify_client._utils.crypto import create_hmac_signature, create_storage_content_signature from apify_client._utils.encoding import encode_key_value_store_record_value from apify_client._utils.errors import catch_not_found_or_throw from apify_client._utils.http import response_to_dict, to_path_segment from apify_client.errors import ApifyApiError, InvalidResponseBodyError if TYPE_CHECKING: from collections.abc import AsyncIterator, Iterator from datetime import timedelta from apify_client._literals import GeneralAccess from apify_client.http_clients import HttpResponse from apify_client.types import Timeout def _parse_get_record_response(response: HttpResponse) -> Any: """Parse an HTTP response based on its content type. Args: response: The HTTP response to parse. Returns: Parsed response data (JSON dict/list, text string, or raw bytes). Raises: InvalidResponseBodyError: If the response body cannot be parsed. """ if response.status_code == HTTPStatus.NO_CONTENT: return None content_type = '' if 'content-type' in response.headers: content_type = response.headers['content-type'].split(';')[0].strip() try: if re.search(r'^application/json', content_type, flags=re.IGNORECASE): response_data = response.json() elif re.search(r'^application/.*xml$', content_type, flags=re.IGNORECASE) or re.search( r'^text/', content_type, flags=re.IGNORECASE ): response_data = response.text else: response_data = response.content except ValueError as err: raise InvalidResponseBodyError(response) from err else: return response_data @docs_group('Resource clients') class KeyValueStoreClient(ResourceClient): """Sub-client for managing a specific key-value store. Provides methods to manage a specific key-value store, e.g. get it, update it, or manage its records. Obtain an instance via an appropriate method on the `ApifyClient` class. """ def __init__( self, *, resource_id: str | None = None, resource_path: str = 'key-value-stores', **kwargs: Any, ) -> None: super().__init__( resource_id=resource_id, resource_path=resource_path, **kwargs, ) def get(self, *, timeout: Timeout = 'short') -> KeyValueStore | None: """Retrieve the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/store-object/get-store Args: timeout: Timeout for the API HTTP request. Returns: The retrieved key-value store, or None if it does not exist. """ result = self._get(timeout=timeout) if result is None: return None return KeyValueStoreResponse.model_validate(result).data def update( self, *, name: str | None = None, general_access: GeneralAccess | None = None, timeout: Timeout = 'long', ) -> KeyValueStore: """Update the key-value store with specified fields. https://docs.apify.com/api/v2#/reference/key-value-stores/store-object/update-store Args: name: The new name for key-value store. general_access: Determines how others can access the key-value store. timeout: Timeout for the API HTTP request. Returns: The updated key-value store. """ result = self._update(timeout=timeout, name=name, generalAccess=general_access) return KeyValueStoreResponse.model_validate(result).data def delete(self, *, timeout: Timeout = 'short') -> None: """Delete the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/store-object/delete-store Args: timeout: Timeout for the API HTTP request. """ self._delete(timeout=timeout) def list_keys( self, *, limit: int | None = None, exclusive_start_key: str | None = None, collection: str | None = None, prefix: str | None = None, signature: str | None = None, timeout: Timeout = 'medium', ) -> ListOfKeys: """List the keys in the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/key-collection/get-list-of-keys Args: limit: Number of keys to be returned. Maximum value is 1000. exclusive_start_key: All keys up to this one (including) are skipped from the result. collection: The name of the collection in store schema to list keys from. prefix: The prefix of the keys to be listed. signature: Signature used to access the items. timeout: Timeout for the API HTTP request. Returns: The list of keys in the key-value store matching the given arguments. """ request_params = self._build_params( limit=limit, exclusiveStartKey=exclusive_start_key, collection=collection, prefix=prefix, signature=signature, ) response = self._http_client.call( url=self._build_url('keys'), method='GET', params=request_params, timeout=timeout, ) result = response_to_dict(response) return ListOfKeysResponse.model_validate(result).data def iterate_keys( self, *, limit: int | None = None, exclusive_start_key: str | None = None, collection: str | None = None, prefix: str | None = None, signature: str | None = None, chunk_size: int | None = None, timeout: Timeout = 'long', ) -> Iterator[KeyValueStoreKey]: """Iterate over the keys in the key-value store. Simple `list_keys` does only one API call, possibly not listing all items matching the criteria. This method returns an iterator that is capable of making multiple API calls to retrieve all items matching the criteria. https://docs.apify.com/api/v2#/reference/key-value-stores/key-collection/get-list-of-keys Args: limit: Maximum number of keys to return. By default there is no limit. exclusive_start_key: All keys up to this one (including) are skipped from the result. collection: The name of the collection in store schema to list keys from. prefix: The prefix of the keys to be listed. signature: Signature used to access the items. chunk_size: Maximum number of keys requested per API call when iterating across pages. timeout: Timeout for the API HTTP request. Yields: A key from the key-value store. """ def _callback(*, cursor: str | None = None, limit: int | None = None) -> ListOfKeys: return self.list_keys( limit=limit, exclusive_start_key=cursor, collection=collection, prefix=prefix, signature=signature, timeout=timeout, ) return get_cursor_iterator( _callback, cursor=exclusive_start_key, limit=limit, chunk_size=chunk_size or DEFAULT_CHUNK_SIZE, ) def get_record(self, key: str, *, signature: str | None = None, timeout: Timeout = 'long') -> dict | None: """Retrieve the given record from the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/record/get-record Args: key: Key of the record to retrieve. signature: Signature used to access the items. timeout: Timeout for the API HTTP request. Returns: The requested record, or None, if the record does not exist. """ try: response = self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='GET', params=self._build_params(signature=signature, attachment=True), timeout=timeout, ) return { 'key': key, 'value': _parse_get_record_response(response), 'content_type': response.headers['content-type'], } except ApifyApiError as exc: catch_not_found_or_throw(exc) return None def record_exists(self, key: str, *, timeout: Timeout = 'long') -> bool: """Check if given record is present in the key-value store. https://docs.apify.com/api/v2/key-value-store-record-head Args: key: Key of the record to check. timeout: Timeout for the API HTTP request. Returns: True if the record exists, False otherwise. """ try: response = self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='HEAD', params=self._build_params(), timeout=timeout, ) except ApifyApiError as exc: if exc.status_code == HTTPStatus.NOT_FOUND: return False raise return response.status_code == HTTPStatus.OK def get_record_as_bytes(self, key: str, *, signature: str | None = None, timeout: Timeout = 'long') -> dict | None: """Retrieve the given record from the key-value store, without parsing it. https://docs.apify.com/api/v2#/reference/key-value-stores/record/get-record Args: key: Key of the record to retrieve. signature: Signature used to access the items. timeout: Timeout for the API HTTP request. Returns: The requested record, or None, if the record does not exist. """ try: response = self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='GET', params=self._build_params(signature=signature, attachment=True), timeout=timeout, ) return { 'key': key, 'value': response.content, 'content_type': response.headers['content-type'], } except ApifyApiError as exc: catch_not_found_or_throw(exc) return None @contextmanager def stream_record( self, key: str, *, signature: str | None = None, timeout: Timeout = 'long' ) -> Iterator[dict | None]: """Retrieve the given record from the key-value store, as a stream. https://docs.apify.com/api/v2#/reference/key-value-stores/record/get-record Args: key: Key of the record to retrieve. signature: Signature used to access the items. timeout: Timeout for the API HTTP request. Returns: The requested record as a context-managed streaming Response, or None, if the record does not exist. """ response = None try: response = self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='GET', params=self._build_params(signature=signature, attachment=True), stream=True, timeout=timeout, ) yield { 'key': key, 'value': response, 'content_type': response.headers['content-type'], } except ApifyApiError as exc: catch_not_found_or_throw(exc) yield None finally: if response: response.close() def set_record( self, key: str, value: Any, *, content_type: str | None = None, content_encoding: str | None = None, timeout: Timeout = 'long', ) -> None: """Set a value to the given record in the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/record/put-record Args: key: The key of the record to save the value to. value: The value to save into the record. content_type: The content type of the saved value. content_encoding: The encoding already applied to `value`, sent as the `Content-Encoding` header. Pass it to upload a pre-compressed value - the client then forwards the bytes as they are instead of compressing them itself. The API accepts `gzip`, `br`, `deflate`, and `identity`, and stores the record exactly as uploaded, so this also becomes the encoding the record is served with. Only a bytes-like `value`, or a file-like one that reads into bytes, can carry a compression - anything else raises `TypeError` instead of being stored under a header that misdescribes it. timeout: Timeout for the API HTTP request. """ value, content_type = encode_key_value_store_record_value( value, content_type=content_type, content_encoding=content_encoding, ) headers = {'content-type': content_type} if content_encoding is not None: headers['content-encoding'] = content_encoding self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='PUT', params=self._build_params(), data=value, headers=headers, timeout=timeout, ) def delete_record(self, key: str, *, timeout: Timeout = 'short') -> None: """Delete the specified record from the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/record/delete-record Args: key: The key of the record which to delete. timeout: Timeout for the API HTTP request. """ self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='DELETE', params=self._build_params(), timeout=timeout, ) def get_record_public_url(self, key: str, *, timeout: Timeout = 'long') -> str: """Generate a URL that can be used to access key-value store record. If the client has permission to access the key-value store's URL signing key, the URL will include a signature to verify its authenticity. Args: key: The key for which the URL should be generated. timeout: Timeout for the API HTTP request. Returns: A public URL that can be used to access the value of the given key in the KVS. """ if self._resource_id is None: raise ValueError('resource_id cannot be None when generating a public URL') # Encode the key first, so a key that cannot address a record is refused before spending a request. record_path = f'records/{to_path_segment(key)}' metadata = self.get(timeout=timeout) request_params = self._build_params() # The signature covers the raw key, which is what the API sees once it decodes the path segment. if metadata and metadata.url_signing_secret_key: request_params['signature'] = create_hmac_signature(metadata.url_signing_secret_key, key) return self._build_public_url(record_path, request_params) def create_keys_public_url( self, *, limit: int | None = None, exclusive_start_key: str | None = None, collection: str | None = None, prefix: str | None = None, expires_in: timedelta | None = None, timeout: Timeout = 'long', ) -> str: """Generate a URL that can be used to access key-value store keys. If the client has permission to access the key-value store's URL signing key, the URL will include a signature to verify its authenticity. You can optionally control how long the signed URL should be valid using the `expires_in` option. This value sets the expiration duration from the time the URL is generated. If not provided, the URL will not expire. Any other options (like `limit` or `prefix`) will be included as query parameters in the URL. Args: limit: Number of keys to be returned. Maximum value is 1000. exclusive_start_key: All keys up to this one (including) are skipped from the result. collection: The name of the collection in store schema to list keys from. prefix: The prefix of the keys to be listed. expires_in: How long the signed URL should be valid from the time it is generated. timeout: Timeout for the API HTTP request. Returns: The public key-value store keys URL. """ metadata = self.get(timeout=timeout) request_params = self._build_params( limit=limit, exclusiveStartKey=exclusive_start_key, collection=collection, prefix=prefix, ) if metadata and metadata.url_signing_secret_key: signature = create_storage_content_signature( metadata.id, metadata.url_signing_secret_key, expires_in=expires_in, ) request_params['signature'] = signature return self._build_public_url('keys', request_params) @docs_group('Resource clients') class KeyValueStoreClientAsync(ResourceClientAsync): """Sub-client for managing a specific key-value store. Provides methods to manage a specific key-value store, e.g. get it, update it, or manage its records. Obtain an instance via an appropriate method on the `ApifyClientAsync` class. """ def __init__( self, *, resource_id: str | None = None, resource_path: str = 'key-value-stores', **kwargs: Any, ) -> None: super().__init__( resource_id=resource_id, resource_path=resource_path, **kwargs, ) async def get(self, *, timeout: Timeout = 'short') -> KeyValueStore | None: """Retrieve the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/store-object/get-store Args: timeout: Timeout for the API HTTP request. Returns: The retrieved key-value store, or None if it does not exist. """ result = await self._get(timeout=timeout) if result is None: return None return KeyValueStoreResponse.model_validate(result).data async def update( self, *, name: str | None = None, general_access: GeneralAccess | None = None, timeout: Timeout = 'long', ) -> KeyValueStore: """Update the key-value store with specified fields. https://docs.apify.com/api/v2#/reference/key-value-stores/store-object/update-store Args: name: The new name for key-value store. general_access: Determines how others can access the key-value store. timeout: Timeout for the API HTTP request. Returns: The updated key-value store. """ result = await self._update(timeout=timeout, name=name, generalAccess=general_access) return KeyValueStoreResponse.model_validate(result).data async def delete(self, *, timeout: Timeout = 'short') -> None: """Delete the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/store-object/delete-store Args: timeout: Timeout for the API HTTP request. """ await self._delete(timeout=timeout) async def list_keys( self, *, limit: int | None = None, exclusive_start_key: str | None = None, collection: str | None = None, prefix: str | None = None, signature: str | None = None, timeout: Timeout = 'medium', ) -> ListOfKeys: """List the keys in the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/key-collection/get-list-of-keys Args: limit: Number of keys to be returned. Maximum value is 1000. exclusive_start_key: All keys up to this one (including) are skipped from the result. collection: The name of the collection in store schema to list keys from. prefix: The prefix of the keys to be listed. signature: Signature used to access the items. timeout: Timeout for the API HTTP request. Returns: The list of keys in the key-value store matching the given arguments. """ request_params = self._build_params( limit=limit, exclusiveStartKey=exclusive_start_key, collection=collection, prefix=prefix, signature=signature, ) response = await self._http_client.call( url=self._build_url('keys'), method='GET', params=request_params, timeout=timeout, ) result = response_to_dict(response) return ListOfKeysResponse.model_validate(result).data def iterate_keys( self, *, limit: int | None = None, exclusive_start_key: str | None = None, collection: str | None = None, prefix: str | None = None, signature: str | None = None, chunk_size: int | None = None, timeout: Timeout = 'long', ) -> AsyncIterator[KeyValueStoreKey]: """Iterate over the keys in the key-value store. Simple `list_keys` does only one API call, possibly not listing all items matching the criteria. This method returns an iterator that is capable of making multiple API calls to retrieve all items matching the criteria. https://docs.apify.com/api/v2#/reference/key-value-stores/key-collection/get-list-of-keys Args: limit: Maximum number of keys to return. By default there is no limit. exclusive_start_key: All keys up to this one (including) are skipped from the result. collection: The name of the collection in store schema to list keys from. prefix: The prefix of the keys to be listed. signature: Signature used to access the items. chunk_size: Maximum number of keys requested per API call when iterating across pages. timeout: Timeout for the API HTTP request. Yields: A key from the key-value store. """ async def _callback(*, cursor: str | None = None, limit: int | None = None) -> ListOfKeys: return await self.list_keys( limit=limit, exclusive_start_key=cursor, collection=collection, prefix=prefix, signature=signature, timeout=timeout, ) return get_cursor_iterator_async( _callback, cursor=exclusive_start_key, limit=limit, chunk_size=chunk_size or DEFAULT_CHUNK_SIZE, ) async def get_record(self, key: str, *, signature: str | None = None, timeout: Timeout = 'long') -> dict | None: """Retrieve the given record from the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/record/get-record Args: key: Key of the record to retrieve. signature: Signature used to access the items. timeout: Timeout for the API HTTP request. Returns: The requested record, or None, if the record does not exist. """ try: response = await self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='GET', params=self._build_params(signature=signature, attachment=True), timeout=timeout, ) return { 'key': key, 'value': _parse_get_record_response(response), 'content_type': response.headers['content-type'], } except ApifyApiError as exc: catch_not_found_or_throw(exc) return None async def record_exists(self, key: str, *, timeout: Timeout = 'long') -> bool: """Check if given record is present in the key-value store. https://docs.apify.com/api/v2/key-value-store-record-head Args: key: Key of the record to check. timeout: Timeout for the API HTTP request. Returns: True if the record exists, False otherwise. """ try: response = await self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='HEAD', params=self._build_params(), timeout=timeout, ) except ApifyApiError as exc: if exc.status_code == HTTPStatus.NOT_FOUND: return False raise return response.status_code == HTTPStatus.OK async def get_record_as_bytes( self, key: str, *, signature: str | None = None, timeout: Timeout = 'long' ) -> dict | None: """Retrieve the given record from the key-value store, without parsing it. https://docs.apify.com/api/v2#/reference/key-value-stores/record/get-record Args: key: Key of the record to retrieve. signature: Signature used to access the items. timeout: Timeout for the API HTTP request. Returns: The requested record, or None, if the record does not exist. """ try: response = await self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='GET', params=self._build_params(signature=signature, attachment=True), timeout=timeout, ) return { 'key': key, 'value': response.content, 'content_type': response.headers['content-type'], } except ApifyApiError as exc: catch_not_found_or_throw(exc) return None @asynccontextmanager async def stream_record( self, key: str, *, signature: str | None = None, timeout: Timeout = 'long' ) -> AsyncIterator[dict | None]: """Retrieve the given record from the key-value store, as a stream. https://docs.apify.com/api/v2#/reference/key-value-stores/record/get-record Args: key: Key of the record to retrieve. signature: Signature used to access the items. timeout: Timeout for the API HTTP request. Returns: The requested record as a context-managed streaming Response, or None, if the record does not exist. """ response = None try: response = await self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='GET', params=self._build_params(signature=signature, attachment=True), stream=True, timeout=timeout, ) yield { 'key': key, 'value': response, 'content_type': response.headers['content-type'], } except ApifyApiError as exc: catch_not_found_or_throw(exc) yield None finally: if response: await response.aclose() async def set_record( self, key: str, value: Any, *, content_type: str | None = None, content_encoding: str | None = None, timeout: Timeout = 'long', ) -> None: """Set a value to the given record in the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/record/put-record Args: key: The key of the record to save the value to. value: The value to save into the record. content_type: The content type of the saved value. content_encoding: The encoding already applied to `value`, sent as the `Content-Encoding` header. Pass it to upload a pre-compressed value - the client then forwards the bytes as they are instead of compressing them itself. The API accepts `gzip`, `br`, `deflate`, and `identity`, and stores the record exactly as uploaded, so this also becomes the encoding the record is served with. Only a bytes-like `value`, or a file-like one that reads into bytes, can carry a compression - anything else raises `TypeError` instead of being stored under a header that misdescribes it. timeout: Timeout for the API HTTP request. """ value, content_type = encode_key_value_store_record_value( value, content_type=content_type, content_encoding=content_encoding, ) headers = {'content-type': content_type} if content_encoding is not None: headers['content-encoding'] = content_encoding await self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='PUT', params=self._build_params(), data=value, headers=headers, timeout=timeout, ) async def delete_record(self, key: str, *, timeout: Timeout = 'short') -> None: """Delete the specified record from the key-value store. https://docs.apify.com/api/v2#/reference/key-value-stores/record/delete-record Args: key: The key of the record which to delete. timeout: Timeout for the API HTTP request. """ await self._http_client.call( url=self._build_url(f'records/{to_path_segment(key)}'), method='DELETE', params=self._build_params(), timeout=timeout, ) async def get_record_public_url(self, key: str, *, timeout: Timeout = 'long') -> str: """Generate a URL that can be used to access key-value store record. If the client has permission to access the key-value store's URL signing key, the URL will include a signature to verify its authenticity. Args: key: The key for which the URL should be generated. timeout: Timeout for the API HTTP request. Returns: A public URL that can be used to access the value of the given key in the KVS. """ if self._resource_id is None: raise ValueError('resource_id cannot be None when generating a public URL') # Encode the key first, so a key that cannot address a record is refused before spending a request. record_path = f'records/{to_path_segment(key)}' metadata = await self.get(timeout=timeout) request_params = self._build_params() # The signature covers the raw key, which is what the API sees once it decodes the path segment. if metadata and metadata.url_signing_secret_key: request_params['signature'] = create_hmac_signature(metadata.url_signing_secret_key, key) return self._build_public_url(record_path, request_params) async def create_keys_public_url( self, *, limit: int | None = None, exclusive_start_key: str | None = None, collection: str | None = None, prefix: str | None = None, expires_in: timedelta | None = None, timeout: Timeout = 'long', ) -> str: """Generate a URL that can be used to access key-value store keys. If the client has permission to access the key-value store's URL signing key, the URL will include a signature to verify its authenticity. You can optionally control how long the signed URL should be valid using the `expires_in` option. This value sets the expiration duration from the time the URL is generated. If not provided, the URL will not expire. Any other options (like `limit` or `prefix`) will be included as query parameters in the URL. Args: limit: Number of keys to be returned. Maximum value is 1000. exclusive_start_key: All keys up to this one (including) are skipped from the result. collection: The name of the collection in store schema to list keys from. prefix: The prefix of the keys to be listed. expires_in: How long the signed URL should be valid from the time it is generated. timeout: Timeout for the API HTTP request. Returns: The public key-value store keys URL. """ metadata = await self.get(timeout=timeout) request_params = self._build_params( limit=limit, exclusiveStartKey=exclusive_start_key, collection=collection, prefix=prefix, ) if metadata and metadata.url_signing_secret_key: signature = create_storage_content_signature( metadata.id, metadata.url_signing_secret_key, expires_in=expires_in, ) request_params['signature'] = signature return self._build_public_url('keys', request_params)