[SchemaRegistry] add async version of avro serializer (#21026)

fixes: #20709

sdk/schemaregistry/azure-schemaregistry-avroserializer/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Async version of `AvroSerializer` has been added under `azure.schemaregistry.serializer.avroserializer.aio`.
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/schemaregistry/azure-schemaregistry-avroserializer/azure/schemaregistry/serializer/avroserializer/_schema_registry_avro_serializer.py

@@ -29,10 +29,10 @@
     from backports.functools_lru_cache import lru_cache
 from io import BytesIO
 from typing import Any, Dict, Mapping
-import avro
 
 from ._constants import SCHEMA_ID_START_INDEX, SCHEMA_ID_LENGTH, DATA_START_INDEX
 from ._avro_serializer import AvroObjectSerializer
+from ._utils import parse_schema
 
 
 class AvroSerializer(object):
@@ -45,7 +45,7 @@ class AvroSerializer(object):
     :paramtype client: ~azure.schemaregistry.SchemaRegistryClient
     :keyword str group_name: Required. Schema group under which schema should be registered.
     :keyword bool auto_register_schemas: When true, register new schemas passed to serialize.
-     Otherwise, and by default, fail if it has not been pre-registered in the registry.
+     Otherwise, and by default, serialization will fail if the schema has not been pre-registered in the registry.
 
     """
 
@@ -89,8 +89,7 @@ def _get_schema_id(self, schema_name, schema_str, **kwargs):
 
         :param schema_name: Name of the schema
         :type schema_name: str
-        :param schema: Schema object
-        :type schema: avro.schema.Schema
+        :param str schema_str: Schema string
         :return: Schema Id
         :rtype: str
         """
@@ -114,11 +113,6 @@ def _get_schema(self, schema_id, **kwargs):
         ).schema_definition
         return schema_str
 
-    @classmethod
-    @lru_cache(maxsize=128)
-    def _parse_schema(cls, schema):
-        return avro.schema.parse(schema)
-
     def serialize(self, value, **kwargs):
         # type: (Mapping[str, Any], Any) -> bytes
         """
@@ -137,7 +131,7 @@ def serialize(self, value, **kwargs):
         except KeyError as e:
             raise TypeError("'{}' is a required keyword.".format(e.args[0]))
 
-        cached_schema = AvroSerializer._parse_schema(raw_input_schema)
+        cached_schema = parse_schema(raw_input_schema)
         record_format_identifier = b"\0\0\0\0"
         schema_id = self._get_schema_id(cached_schema.fullname, str(cached_schema), **kwargs)
         data_bytes = self._avro_serializer.serialize(value, cached_schema)

sdk/schemaregistry/azure-schemaregistry-avroserializer/azure/schemaregistry/serializer/avroserializer/_utils.py

@@ -0,0 +1,13 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+try:
+    from functools import lru_cache
+except ImportError:
+    from backports.functools_lru_cache import lru_cache
+import avro
+
+@lru_cache(maxsize=128)
+def parse_schema(schema):
+    return avro.schema.parse(schema)

sdk/schemaregistry/azure-schemaregistry-avroserializer/azure/schemaregistry/serializer/avroserializer/aio/__init__.py

@@ -0,0 +1,30 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+from ._schema_registry_avro_serializer_async import AvroSerializer
+
+__all__ = [
+    "AvroSerializer"
+]

sdk/schemaregistry/azure-schemaregistry-avroserializer/azure/schemaregistry/serializer/avroserializer/aio/_async_lru.py

@@ -0,0 +1,230 @@
+# --------------------------------------------------------------------------
+# The MIT License
+#
+# Copyright (c) 2018 aio-libs team https://github.com/aio-libs/
+# Copyright (c) 2017 Ocean S. A. https://ocean.io/
+# Copyright (c) 2016-2017 WikiBusiness Corporation http://wikibusiness.org/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+# --------------------------------------------------------------------------
+# Copying over `async_lru.py`[https://github.com/aio-libs/async-lru/blob/master/async_lru.py]
+# from `aio-libs`[https://github.com/aio-libs/async-lru] for the following reasons:
+# 1. There has not been an official release of `async_lru` in 2 years.
+# 2. The last update to the library was a year ago, so it seems the library is
+#    not being actively maintained.
+
+import asyncio
+from collections import OrderedDict
+from functools import _CacheInfo, _make_key, partial, wraps
+
+
+__version__ = "1.0.2"
+
+__all__ = ("alru_cache",)
+
+
+def unpartial(fn):
+    while hasattr(fn, "func"):
+        fn = fn.func
+
+    return fn
+
+
+def _done_callback(fut, task):
+    if task.cancelled():
+        fut.cancel()
+        return
+
+    exc = task.exception()
+    if exc is not None:
+        fut.set_exception(exc)
+        return
+
+    fut.set_result(task.result())
+
+
+def _cache_invalidate(wrapped, typed, *args, **kwargs):
+    # pylint: disable=protected-access
+    key = _make_key(args, kwargs, typed)
+
+    exists = key in wrapped._cache
+
+    if exists:
+        wrapped._cache.pop(key)
+
+    return exists
+
+
+def _cache_clear(wrapped):
+    # pylint: disable=protected-access
+    wrapped.hits = wrapped.misses = 0
+    wrapped._cache = OrderedDict()
+    wrapped.tasks = set()
+
+
+def _open(wrapped):
+    if not wrapped.closed:
+        raise RuntimeError("alru_cache is not closed")
+
+    # pylint: disable=protected-access
+    was_closed = (
+        wrapped.hits == wrapped.misses == len(wrapped.tasks) == len(wrapped._cache) == 0
+    )
+
+    if not was_closed:
+        raise RuntimeError("alru_cache was not closed correctly")
+
+    wrapped.closed = False
+
+
+def _close(wrapped, *, cancel=False, return_exceptions=True):
+    if wrapped.closed:
+        raise RuntimeError("alru_cache is closed")
+
+    wrapped.closed = True
+
+    if cancel:
+        for task in wrapped.tasks:
+            if not task.done():  # not sure is it possible
+                task.cancel()
+
+    return _wait_closed(wrapped, return_exceptions=return_exceptions)
+
+
+async def _wait_closed(wrapped, *, return_exceptions):
+    wait_closed = asyncio.gather(*wrapped.tasks, return_exceptions=return_exceptions)
+
+    wait_closed.add_done_callback(partial(_close_waited, wrapped))
+
+    ret = await wait_closed
+
+    # hack to get _close_waited callback to be executed
+    await asyncio.sleep(0)
+
+    return ret
+
+
+def _close_waited(wrapped, _):
+    wrapped.cache_clear()
+
+
+def _cache_info(wrapped, maxsize):
+    # pylint: disable=protected-access
+    return _CacheInfo(
+        wrapped.hits,
+        wrapped.misses,
+        maxsize,
+        len(wrapped._cache),
+    )
+
+
+def __cache_touch(wrapped, key):
+    # pylint: disable=protected-access
+    try:
+        wrapped._cache.move_to_end(key)
+    except KeyError:  # not sure is it possible
+        pass
+
+
+def _cache_hit(wrapped, key):
+    wrapped.hits += 1
+    __cache_touch(wrapped, key)
+
+
+def _cache_miss(wrapped, key):
+    wrapped.misses += 1
+    __cache_touch(wrapped, key)
+
+
+def alru_cache(
+    fn=None,
+    maxsize=128,
+    typed=False,
+    *,
+    cache_exceptions=True,
+):
+    def wrapper(fn):
+        # pylint: disable=protected-access
+        _origin = unpartial(fn)
+
+        if not asyncio.iscoroutinefunction(_origin):
+            raise RuntimeError("Coroutine function is required, got {}".format(fn))
+
+        # functools.partialmethod support
+        if hasattr(fn, "_make_unbound_method"):
+            fn = fn._make_unbound_method()
+
+        @wraps(fn)
+        async def wrapped(*fn_args, **fn_kwargs):
+            if wrapped.closed:
+                raise RuntimeError("alru_cache is closed for {}".format(wrapped))
+
+            loop = asyncio.get_event_loop()
+
+            key = _make_key(fn_args, fn_kwargs, typed)
+
+            fut = wrapped._cache.get(key)
+
+            if fut is not None:
+                if not fut.done():
+                    _cache_hit(wrapped, key)
+                    return await asyncio.shield(fut)
+
+                exc = fut._exception
+
+                if exc is None or cache_exceptions:
+                    _cache_hit(wrapped, key)
+                    return fut.result()
+
+                # exception here and cache_exceptions == False
+                wrapped._cache.pop(key)
+
+            fut = loop.create_future()
+            task = loop.create_task(fn(*fn_args, **fn_kwargs))
+            task.add_done_callback(partial(_done_callback, fut))
+
+            wrapped.tasks.add(task)
+            task.add_done_callback(wrapped.tasks.remove)
+
+            wrapped._cache[key] = fut
+
+            if maxsize is not None and len(wrapped._cache) > maxsize:
+                wrapped._cache.popitem(last=False)
+
+            _cache_miss(wrapped, key)
+            return await asyncio.shield(fut)
+
+        _cache_clear(wrapped)
+        wrapped._origin = _origin
+        wrapped.closed = False
+        wrapped.cache_info = partial(_cache_info, wrapped, maxsize)
+        wrapped.cache_clear = partial(_cache_clear, wrapped)
+        wrapped.invalidate = partial(_cache_invalidate, wrapped, typed)
+        wrapped.close = partial(_close, wrapped)
+        wrapped.open = partial(_open, wrapped)
+
+        return wrapped
+
+    if fn is None:
+        return wrapper
+
+    if callable(fn) or hasattr(fn, "_make_unbound_method"):
+        return wrapper(fn)
+
+    raise NotImplementedError("{} decorating is not supported".format(fn))