fix aiokafka admin class

backend/deepchecks_monitoring/bgtasks/model_version_topic_delete.py

@@ -2,7 +2,6 @@
 from typing import Optional
 
 import pendulum as pdl
-from aiokafka.admin import AIOKafkaAdminClient
 from kafka.errors import KafkaError, UnknownTopicOrPartitionError
 from sqlalchemy import delete, select
 from sqlalchemy.dialects.postgresql import insert
@@ -15,6 +14,7 @@
 from deepchecks_monitoring.resources import ResourcesProvider
 from deepchecks_monitoring.schema_models import Model, ModelVersion
 from deepchecks_monitoring.utils import database
+from deepchecks_monitoring.utils.other import ExtendedAIOKafkaAdminClient
 
 __all__ = ['ModelVersionTopicDeletionWorker', 'insert_model_version_topic_delete_task']
 
@@ -36,7 +36,7 @@ class ModelVersionTopicDeletionWorker(BackgroundWorker):
 
     def __init__(self):
         self.lock = threading.Lock()
-        self.kafka_admin: Optional[AIOKafkaAdminClient] = None
+        self.kafka_admin: Optional[ExtendedAIOKafkaAdminClient] = None
 
     def queue_name(self) -> str:
         return QUEUE_NAME
@@ -48,7 +48,7 @@ async def run(self, task: 'Task', session: AsyncSession, resources_provider: Res
         if self.kafka_admin is None:
             with self.lock:
                 if self.kafka_admin is None:
-                    self.kafka_admin = AIOKafkaAdminClient(**resources_provider.kafka_settings.kafka_admin_params)
+                    self.kafka_admin = ExtendedAIOKafkaAdminClient(**resources_provider.kafka_settings.kafka_params)
                     await self.kafka_admin.start()
 
         # Backward compatibility, remove in next release and replace with:

backend/deepchecks_monitoring/config.py

@@ -79,16 +79,6 @@ def kafka_params(self):
             'metadata_max_age_ms': self.kafka_max_metadata_age
         }
 
-    @property
-    def kafka_admin_params(self):
-        """Get connection parameters for kafka admin."""
-        return {
-            'bootstrap_servers': self.kafka_host,
-            'security_protocol': self.kafka_security_protocol,
-            'ssl_context': create_ssl_context(),
-            'metadata_max_age_ms': self.kafka_max_metadata_age
-        }
-
 
 class DatabaseSettings(BaseDeepchecksSettings):
     """Database settings."""

backend/deepchecks_monitoring/utils/other.py

@@ -1,7 +1,11 @@
 """Represent global utility functions."""
 import typing as t
+from ssl import SSLContext
 
 import pendulum as pdl
+from aiokafka.admin import AIOKafkaAdminClient
+from aiokafka.admin import __version__ as aiokafka_version
+from aiokafka.client import AIOKafkaClient
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
 
 from deepchecks_monitoring.public_models import Organization, User, UserOAuthDTO
@@ -10,6 +14,39 @@
 __all__ = ['generate_random_user', 'generate_test_user', 'datetime_sample_formatter']
 
 
+class ExtendedAIOKafkaAdminClient(AIOKafkaAdminClient):  # pylint: disable=missing-class-docstring
+    # pylint: disable=super-init-not-called
+    def __init__(self, *, loop=None,
+                 bootstrap_servers: str = 'localhost',
+                 client_id: str = 'aiokafka-' + aiokafka_version,
+                 request_timeout_ms: int = 40000,
+                 connections_max_idle_ms: int = 540000,
+                 retry_backoff_ms: int = 100,
+                 metadata_max_age_ms: int = 300000,
+                 security_protocol: str = 'PLAINTEXT',
+                 ssl_context: t.Optional[SSLContext] = None,
+                 api_version: str = 'auto',
+                 sasl_mechanism='PLAIN',
+                 sasl_plain_username=None,
+                 sasl_plain_password=None):
+        self._closed = False
+        self._started = False
+        self._version_info = {}
+        self._request_timeout_ms = request_timeout_ms
+        self._client = AIOKafkaClient(
+            loop=loop, bootstrap_servers=bootstrap_servers,
+            client_id=client_id, metadata_max_age_ms=metadata_max_age_ms,
+            request_timeout_ms=request_timeout_ms,
+            retry_backoff_ms=retry_backoff_ms,
+            api_version=api_version,
+            ssl_context=ssl_context,
+            security_protocol=security_protocol,
+            connections_max_idle_ms=connections_max_idle_ms,
+            sasl_mechanism=sasl_mechanism,
+            sasl_plain_username=sasl_plain_username,
+            sasl_plain_password=sasl_plain_password)
+
+
 async def generate_random_user(session: AsyncSession, auth_jwt_secret: str, with_org: bool = True):
     """Generate a random user."""
     try: