Merge remote-tracking branch 'origin/add-async-compat-schema-registry' into add-async-compat-schema-registry

Author: rohitsanj

src/confluent_kafka/schema_registry/_sync/avro.py

@@ -30,7 +30,7 @@
                SchemaRegistryClient)
 from confluent_kafka.serialization import (SerializationError,
                                            SerializationContext)
-from confluent_kafka.schema_registry.common import _ContextStringIO
+from confluent_kafka.schema_registry.common import _ContextStringIO, asyncinit
 from confluent_kafka.schema_registry.rule_registry import RuleRegistry
 from confluent_kafka.schema_registry.serde import BaseSerializer, BaseDeserializer, ParsedSchemaCache
 

src/confluent_kafka/schema_registry/avro.py

@@ -16,4 +16,5 @@
 # limitations under the License.
 
 from .common.avro import *
+from ._async.avro import *
 from ._sync.avro import *

src/confluent_kafka/schema_registry/common/__init__.py

@@ -89,3 +89,22 @@ def __enter__(self):
     def __exit__(self, *args):
         self.close()
         return False
+
+
+def asyncinit(cls):
+    """
+    Decorator to make a class async-initializable.
+    """
+    __new__ = cls.__new__
+
+    async def init(obj, *arg, **kwarg):
+        await obj.__init__(*arg, **kwarg)
+        return obj
+
+    def new(klass, *arg, **kwarg):
+        obj = __new__(klass)
+        coro = init(obj, *arg, **kwarg)
+        return coro
+
+    cls.__new__ = new
+    return cls

src/confluent_kafka/schema_registry/json_schema.py

@@ -16,4 +16,5 @@
 # limitations under the License.
 
 from .common.json_schema import *
+from ._async.json_schema import *
 from ._sync.json_schema import *

src/confluent_kafka/schema_registry/protobuf.py

@@ -16,4 +16,5 @@
 # limitations under the License.
 
 from .common.protobuf import *
+from ._async.protobuf import *
 from ._sync.protobuf import *

src/confluent_kafka/schema_registry/schema_registry_client.py

@@ -16,6 +16,7 @@
 # limitations under the License.
 
 from .common.schema_registry_client import *
+from ._async.schema_registry_client import *
 from ._sync.schema_registry_client import *
 
 from .error import SchemaRegistryError

src/confluent_kafka/schema_registry/serde.py

@@ -17,4 +17,5 @@
 #
 
 from .common.serde import *
+from ._async.serde import *
 from ._sync.serde import *

tests/common/_async/consumer.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2025 Confluent Inc.
+#
+# 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.
+#
+
+import asyncio
+
+from confluent_kafka.cimpl import Consumer
+from confluent_kafka.error import ConsumeError, KeyDeserializationError, ValueDeserializationError
+from confluent_kafka.serialization import MessageField, SerializationContext
+
+ASYNC_CONSUMER_POLL_INTERVAL_SECONDS: int = 0.2
+ASYNC_CONSUMER_POLL_INFINITE_TIMEOUT_SECONDS: int = -1
+
+class AsyncConsumer(Consumer):
+    def __init__(
+            self,
+            conf: dict,
+            loop: asyncio.AbstractEventLoop = None,
+            poll_interval_seconds: int = ASYNC_CONSUMER_POLL_INTERVAL_SECONDS
+        ):
+        super().__init__(conf)
+
+        self._loop = loop or asyncio.get_event_loop()
+        self._poll_interval = poll_interval_seconds
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        return await self.poll(None)
+
+    async def poll(self, timeout: int = -1):
+        timeout = None if timeout == -1 else timeout
+        async with asyncio.timeout(timeout):
+            while True:
+                # Zero timeout here is what makes it non-blocking
+                msg = super().poll(0)
+                if msg is not None:
+                    return msg
+                else:
+                    await asyncio.sleep(self._poll_interval)
+
+
+class TestAsyncDeserializingConsumer(AsyncConsumer):
+    def __init__(self, conf):
+        conf_copy = conf.copy()
+        self._key_deserializer = conf_copy.pop('key.deserializer', None)
+        self._value_deserializer = conf_copy.pop('value.deserializer', None)
+        super().__init__(conf_copy)
+
+    async def poll(self, timeout=-1):
+        msg = await super().poll(timeout)
+
+        if msg is None:
+            return None
+
+        if msg.error() is not None:
+            raise ConsumeError(msg.error(), kafka_message=msg)
+
+        ctx = SerializationContext(msg.topic(), MessageField.VALUE, msg.headers())
+        value = msg.value()
+        if self._value_deserializer is not None:
+            try:
+                value = await self._value_deserializer(value, ctx)
+            except Exception as se:
+                raise ValueDeserializationError(exception=se, kafka_message=msg)
+
+        key = msg.key()
+        ctx.field = MessageField.KEY
+        if self._key_deserializer is not None:
+            try:
+                key = await self._key_deserializer(key, ctx)
+            except Exception as se:
+                raise KeyDeserializationError(exception=se, kafka_message=msg)
+
+        msg.set_key(key)
+        msg.set_value(value)
+        return msg
+
+    def consume(self, num_messages=1, timeout=-1):
+        """
+        :py:func:`Consumer.consume` not implemented, use
+        :py:func:`DeserializingConsumer.poll` instead
+        """
+
+        raise NotImplementedError

tests/common/_async/producer.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2025 Confluent Inc.
+#
+# 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 confluent_kafka.cimpl import Producer
+import inspect
+import asyncio
+
+from confluent_kafka.error import KeySerializationError, ValueSerializationError
+from confluent_kafka.serialization import MessageField, SerializationContext
+
+ASYNC_PRODUCER_POLL_INTERVAL: int = 0.2
+
+class AsyncProducer(Producer):
+    def __init__(
+            self,
+            conf: dict,
+            loop: asyncio.AbstractEventLoop = None,
+            poll_interval: int = ASYNC_PRODUCER_POLL_INTERVAL
+        ):
+        super().__init__(conf)
+
+        self._loop = loop or asyncio.get_event_loop()
+        self._poll_interval = poll_interval
+
+        self._poll_task = None
+        self._waiters: int = 0
+
+    async def produce(
+            self, topic, value=None, key=None, partition=-1,
+            on_delivery=None, timestamp=0, headers=None
+    ):
+        fut = self._loop.create_future()
+        self._waiters += 1
+        try:
+            if self._poll_task is None or self._poll_task.done():
+                self._poll_task = asyncio.create_task(self._poll_dr(self._poll_interval))
+
+            def wrapped_on_delivery(err, msg):
+                if on_delivery is not None:
+                    if inspect.iscoroutinefunction(on_delivery):
+                        asyncio.run_coroutine_threadsafe(
+                            on_delivery(err, msg),
+                            self._loop
+                        )
+                    else:
+                        self._loop.call_soon_threadsafe(on_delivery, err, msg)
+
+                if err:
+                    self._loop.call_soon_threadsafe(fut.set_exception, err)
+                else:
+                    self._loop.call_soon_threadsafe(fut.set_result, msg)
+
+            super().produce(
+                topic, 
+                value, 
+                key, 
+                headers=headers, 
+                partition=partition, 
+                timestamp=timestamp, 
+                on_delivery=wrapped_on_delivery
+            )
+            return await fut
+        finally:
+            self._waiters -= 1
+
+    async def _poll_dr(self, interval: int):
+        """Poll delivery reports at interval seconds"""
+        while self._waiters:
+            super().poll(0)
+            await asyncio.sleep(interval)
+
+
+class TestAsyncSerializingProducer(AsyncProducer):
+    def __init__(self, conf):
+        conf_copy = conf.copy()
+
+        self._key_serializer = conf_copy.pop('key.serializer', None)
+        self._value_serializer = conf_copy.pop('value.serializer', None)
+
+        super(TestAsyncSerializingProducer, self).__init__(conf_copy)
+
+    async def produce(self, topic, key=None, value=None, partition=-1,
+                on_delivery=None, timestamp=0, headers=None):
+        ctx = SerializationContext(topic, MessageField.KEY, headers)
+        if self._key_serializer is not None:
+            try:
+                key = await self._key_serializer(key, ctx)
+            except Exception as se:
+                raise KeySerializationError(se)
+        ctx.field = MessageField.VALUE
+        if self._value_serializer is not None:
+            try:
+                value = await self._value_serializer(value, ctx)
+            except Exception as se:
+                raise ValueSerializationError(se)
+
+        return await super().produce(topic, value, key,
+                                                 headers=headers,
+                                                 partition=partition,
+                                                 timestamp=timestamp,
+                                                 on_delivery=on_delivery)