Edit

[EDIT] FILE: kv.py

# 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

CANCEL