# Copyright 2021-2022 The NATS Authors # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from __future__ import annotations import asyncio import datetime import logging import re from dataclasses import dataclass from typing import TYPE_CHECKING, List, Optional import nats.errors import nats.js.errors from nats.js import api if TYPE_CHECKING: from nats.js import JetStreamContext KV_OP = "KV-Operation" KV_DEL = "DEL" KV_PURGE = "PURGE" MSG_ROLLUP_SUBJECT = "sub" logger = logging.getLogger(__name__) VALID_KEY_RE = re.compile(r"^[-/_=\.a-zA-Z0-9]+$") def _is_key_valid(key: str) -> bool: if len(key) == 0 or key[0] == "." or key[-1] == ".": return False return bool(VALID_KEY_RE.match(key)) class StopIterSentinel: """A sentinel class used to indicate that iteration should stop.""" pass class KeyValue: """ KeyValue uses the JetStream KeyValue functionality. :: import asyncio import nats async def main(): nc = await nats.connect() js = nc.jetstream() # Create a KV kv = await js.create_key_value(bucket='MY_KV') # Set and retrieve a value await kv.put('hello', b'world') entry = await kv.get('hello') print(f'KeyValue.Entry: key={entry.key}, value={entry.value}') # KeyValue.Entry: key=hello, value=world await nc.close() if __name__ == '__main__': asyncio.run(main()) """ @dataclass class Entry: """ An entry from a KeyValue store in JetStream. """ bucket: str key: str value: Optional[bytes] revision: Optional[int] delta: Optional[int] created: Optional[int] operation: Optional[str] @dataclass(frozen=True) class BucketStatus: """ BucketStatus is the status of a KeyValue bucket. """ stream_info: api.StreamInfo bucket: str @property def values(self) -> int: """ values returns the number of stored messages in the stream. """ return self.stream_info.state.messages @property def history(self) -> int: """ history returns the max msgs per subject. """ return self.stream_info.config.max_msgs_per_subject @property def ttl(self) -> Optional[float]: """ ttl returns the max age in seconds. """ if self.stream_info.config.max_age is None: return None return self.stream_info.config.max_age def __init__( self, name: str, stream: str, pre: str, js: JetStreamContext, direct: bool, ) -> None: self._name = name self._stream = stream self._pre = pre self._js = js self._direct = direct async def get(self, key: str, revision: Optional[int] = None, validate_keys: bool = True) -> Entry: """ get returns the latest value for the key. """ if validate_keys and not _is_key_valid(key): raise nats.js.errors.InvalidKeyError entry = None try: entry = await self._get(key, revision) except nats.js.errors.KeyDeletedError as err: raise nats.js.errors.KeyNotFoundError(err.entry, err.op) return entry async def _get(self, key: str, revision: Optional[int] = None) -> Entry: msg = None subject = f"{self._pre}{key}" try: if revision: msg = await self._js.get_msg( self._stream, seq=revision, direct=self._direct, ) else: msg = await self._js.get_msg( self._stream, subject=subject, seq=revision, direct=self._direct, ) except nats.js.errors.NotFoundError: raise nats.js.errors.KeyNotFoundError # Check whether the revision from the stream does not match the key. if subject != msg.subject: raise nats.js.errors.KeyNotFoundError(message=f"expected '{subject}', but got '{msg.subject}'") entry = KeyValue.Entry( bucket=self._name, key=key, value=msg.data, revision=msg.seq, delta=None, created=None, operation=None, ) # Check headers to see if deleted or purged. if msg.headers: op = msg.headers.get(KV_OP, None) if op == KV_DEL or op == KV_PURGE: raise nats.js.errors.KeyDeletedError(entry, op) return entry async def put(self, key: str, value: bytes, validate_keys: bool = True) -> int: """ put will place the new value for the key into the store and return the revision number. Note: This method does not support TTL. Use create() if you need TTL support. :param key: The key to put :param value: The value to store :param validate_keys: Whether to validate the key format """ if validate_keys and not _is_key_valid(key): raise nats.js.errors.InvalidKeyError(key) pa = await self._js.publish(f"{self._pre}{key}", value) return pa.seq async def create(self, key: str, value: bytes, validate_keys: bool = True, msg_ttl: Optional[float] = None) -> int: """ create will add the key/value pair iff it does not exist. :param key: The key to create :param value: The value to store :param validate_keys: Whether to validate the key format :param msg_ttl: Optional TTL (time-to-live) in seconds for this specific message """ if validate_keys and not _is_key_valid(key): raise nats.js.errors.InvalidKeyError(key) pa = None try: pa = await self.update(key, value, last=0, validate_keys=validate_keys, msg_ttl=msg_ttl) except nats.js.errors.KeyWrongLastSequenceError as err: # In case of attempting to recreate an already deleted key, # the client would get a KeyWrongLastSequenceError. When this happens, # it is needed to fetch latest revision number and attempt to update. try: # NOTE: This reimplements the following behavior from Go client. # # Since we have tombstones for DEL ops for watchers, this could be from that # so we need to double check. # # Get latest revision to update in case it was deleted but if it was not await self._get(key) # No exception so not a deleted key, so reraise the original KeyWrongLastSequenceError. # If it was deleted then the error exception will contain metadata # to recreate using the last revision. raise err except nats.js.errors.KeyDeletedError as err: pa = await self.update( key, value, last=err.entry.revision, validate_keys=validate_keys, msg_ttl=msg_ttl ) return pa async def update( self, key: str, value: bytes, last: Optional[int] = None, validate_keys: bool = True, msg_ttl: Optional[float] = None, ) -> int: """ update will update the value if the latest revision matches. Note: TTL parameter is accepted for internal use by create(), but should not be used directly on update operations per NATS KV semantics. """ if validate_keys and not _is_key_valid(key): raise nats.js.errors.InvalidKeyError(key) hdrs = {} if not last: last = 0 hdrs[api.Header.EXPECTED_LAST_SUBJECT_SEQUENCE] = str(last) pa = None try: pa = await self._js.publish(f"{self._pre}{key}", value, headers=hdrs, msg_ttl=msg_ttl) except nats.js.errors.APIError as err: # Check for a BadRequest::KeyWrongLastSequenceError error code. if err.err_code == 10071: raise nats.js.errors.KeyWrongLastSequenceError(description=err.description) else: raise err return pa.seq async def delete( self, key: str, last: Optional[int] = None, validate_keys: bool = True, msg_ttl: Optional[float] = None ) -> bool: """ delete will place a delete marker and remove all previous revisions. :param key: The key to delete :param last: Expected last revision number (for optimistic concurrency) :param validate_keys: Whether to validate the key format :param msg_ttl: Optional TTL (time-to-live) in seconds for the delete marker """ if validate_keys and not _is_key_valid(key): raise nats.js.errors.InvalidKeyError(key) hdrs = {} hdrs[KV_OP] = KV_DEL if last and last > 0: hdrs[api.Header.EXPECTED_LAST_SUBJECT_SEQUENCE] = str(last) await self._js.publish(f"{self._pre}{key}", headers=hdrs, msg_ttl=msg_ttl) return True async def purge(self, key: str, msg_ttl: Optional[float] = None) -> bool: """ purge will remove the key and all revisions. :param key: The key to purge :param msg_ttl: Optional TTL (time-to-live) in seconds for the purge marker """ hdrs = {} hdrs[KV_OP] = KV_PURGE hdrs[api.Header.ROLLUP] = MSG_ROLLUP_SUBJECT await self._js.publish(f"{self._pre}{key}", headers=hdrs, msg_ttl=msg_ttl) return True async def purge_deletes(self, olderthan: int = 30 * 60) -> bool: """ purge will remove all current delete markers older. :param olderthan: time in seconds """ watcher = await self.watchall() delete_markers = [] async for update in watcher: if update is None: break if update.operation == KV_DEL or update.operation == KV_PURGE: delete_markers.append(update) for entry in delete_markers: keep = 0 subject = f"{self._pre}{entry.key}" duration = datetime.datetime.now(datetime.timezone.utc) - entry.created if olderthan > 0 and olderthan > duration.total_seconds(): keep = 1 await self._js.purge_stream(self._stream, subject=subject, keep=keep) return True async def status(self) -> BucketStatus: """ status retrieves the status and configuration of a bucket. """ info = await self._js.stream_info(self._stream) return KeyValue.BucketStatus(stream_info=info, bucket=self._name) class KeyWatcher: STOP_ITER = StopIterSentinel() def __init__(self, js): self._js = js self._updates: asyncio.Queue[KeyValue.Entry | None | StopIterSentinel] = asyncio.Queue(maxsize=256) self._sub = None self._pending: Optional[int] = None # init done means that the nil marker has been sent, # once this is sent it won't be sent anymore. self._init_done = False async def stop(self): """ stop will stop this watcher. """ await self._sub.unsubscribe() await self._updates.put(KeyValue.KeyWatcher.STOP_ITER) async def updates(self, timeout=5.0): """ updates fetches the next update from a watcher. """ try: return await asyncio.wait_for(self._updates.get(), timeout) except asyncio.TimeoutError: raise nats.errors.TimeoutError def __aiter__(self): return self async def __anext__(self): while True: entry = await self._updates.get() if isinstance(entry, StopIterSentinel): raise StopAsyncIteration return entry async def watchall(self, **kwargs) -> KeyWatcher: """ watchall returns a KeyValue watcher that matches all the keys. """ return await self.watch(">", **kwargs) async def keys(self, filters: List[str] = None, **kwargs) -> List[str]: """ Returns a list of the keys from a KeyValue store. Optionally filters the keys based on the provided filter list. """ watcher = await self.watchall( ignore_deletes=True, meta_only=True, ) keys = [] # Check consumer info and make sure filters are applied correctly try: consumer_info = await watcher._sub.consumer_info() if consumer_info and filters: # If NATS server < 2.10, filters might be ignored. if consumer_info.config.filter_subject != ">": logger.warning("Server may ignore filters if version is < 2.10.") except Exception as e: raise e async for key in watcher: # None entry is used to signal that there is no more info. if not key: break # Apply filters if any were provided if filters: if any(f in key.key for f in filters): keys.append(key.key) else: # No filters provided, append all keys keys.append(key.key) await watcher.stop() if not keys: raise nats.js.errors.NoKeysError return keys async def history(self, key: str) -> List[Entry]: """ history retrieves a list of the entries so far. """ watcher = await self.watch(key, include_history=True) entries = [] async for entry in watcher: # None entry is used to signal that there is no more info. if not entry: break entries.append(entry) await watcher.stop() if not entries: raise nats.js.errors.NoKeysError return entries async def watch( self, keys, headers_only=False, include_history=False, ignore_deletes=False, meta_only=False, inactive_threshold=None, ) -> KeyWatcher: """ watch will fire a callback when a key that matches the keys pattern is updated. The first update after starting the watch is None in case there are no pending updates. """ subject = f"{self._pre}{keys}" watcher = KeyValue.KeyWatcher(self) init_setup: asyncio.Future[bool] = asyncio.Future() async def watch_updates(msg): if not init_setup.done(): await asyncio.wait_for(init_setup, timeout=self._js._timeout) meta = msg.metadata op = None if msg.header and KV_OP in msg.header: op = msg.header.get(KV_OP) # keys() uses this if ignore_deletes: if op == KV_PURGE or op == KV_DEL: if meta.num_pending == 0 and not watcher._init_done: await watcher._updates.put(None) watcher._init_done = True return entry = KeyValue.Entry( bucket=self._name, key=msg.subject[len(self._pre) :], value=msg.data, revision=meta.sequence.stream, delta=meta.num_pending, created=meta.timestamp, operation=op, ) await watcher._updates.put(entry) # When there are no more updates send an empty marker # to signal that it is done, this will unblock iterators if meta.num_pending == 0 and (not watcher._init_done): await watcher._updates.put(None) watcher._init_done = True deliver_policy = None if not include_history: deliver_policy = api.DeliverPolicy.LAST_PER_SUBJECT # Cleanup watchers after 5 minutes of inactivity by default. if not inactive_threshold: inactive_threshold = 5 * 60 watcher._sub = await self._js.subscribe( subject, cb=watch_updates, ordered_consumer=True, deliver_policy=deliver_policy, headers_only=meta_only, inactive_threshold=inactive_threshold, ) await asyncio.sleep(0) # Check from consumer info what is the number of messages # awaiting to be consumed to send the initial signal marker. try: cinfo = await watcher._sub.consumer_info() watcher._pending = cinfo.num_pending # If no delivered and/or pending messages, then signal # that this is the start. # The consumer subscription will start receiving messages # so need to check those that have already made it. received = watcher._sub.delivered init_setup.set_result(True) if cinfo.num_pending == 0 and received == 0: await watcher._updates.put(None) watcher._init_done = True except Exception as err: init_setup.cancel() await watcher._sub.unsubscribe() raise err return watcher