fixes

Author: idanasulin2706

superclient/agent/interceptor.py

@@ -5,7 +5,7 @@
 
 from ..util.logger import get_logger
 from ..util.config import get_topics_list, is_disabled
-from .metadata import fetch_metadata, optimal_cfg, _DEFAULTS
+from .metadata import fetch_metadata_sync, optimal_cfg, _DEFAULTS
 from ..core.reporter import send_clients_msg
 from ..core.manager import normalize_bootstrap
 from .tracker import ProducerTracker, Heartbeat
@@ -57,7 +57,7 @@ def init_patch(self, *args, **kwargs):
         try:
             # Get topics and metadata for optimization
             topics_env = get_topics_list()
-            metadata = fetch_metadata(bootstrap, orig_cfg, "kafka-python")
+            metadata = fetch_metadata_sync(bootstrap, orig_cfg, "kafka-python")
 
             # Check if Superstream is active for this cluster
             error_msg = ""
@@ -185,7 +185,7 @@ def init_patch(self, *args, **kwargs):
 
         try:
             topics_env = get_topics_list()
-            metadata = fetch_metadata(bootstrap, orig_cfg, "aiokafka")
+            metadata = fetch_metadata_sync(bootstrap, orig_cfg, "aiokafka")
             error_msg = ""
             if metadata is None:
                 error_msg = "[ERR-304] Failed to fetch metadata for producer with client id {}: Unable to connect to Superstream service".format(client_id)
@@ -294,7 +294,7 @@ def init_patch(self, conf: Dict[str, Any], *args, **kwargs):
 
         try:
             topics_env = get_topics_list()
-            metadata = fetch_metadata(bootstrap, conf, "confluent")
+            metadata = fetch_metadata_sync(bootstrap, conf, "confluent")
             error_msg = ""
             if metadata is None:
                 error_msg = "[ERR-304] Failed to fetch metadata for producer with client id {}: Unable to connect to Superstream service".format(client_id)

superclient/agent/metadata.py

@@ -57,7 +57,7 @@ def _create_consumer_confluent(bootstrap: str, base_cfg: Dict[str, Any]):
     "confluent": _create_consumer_confluent,
 }
 
-def fetch_metadata(
+async def fetch_metadata(
     bootstrap: str,
     cfg: Dict[str, Any],
     lib_name: Literal["kafka-python", "aiokafka", "confluent"],
@@ -103,7 +103,7 @@ def fetch_metadata(
             if msg and msg.value():
                 return json.loads(msg.value().decode())
 
-        else:  # kafka-python / aiokafka share similar API
+        elif lib_name == "kafka-python":
             import kafka as _kafka  # type: ignore
 
             tp = _kafka.TopicPartition(topic, 0)
@@ -126,10 +126,46 @@ def fetch_metadata(
             for batch in recs.values():
                 for rec in batch:
                     return json.loads(rec.value.decode())
+
+        elif lib_name == "aiokafka":
+            # aiokafka uses its own TopicPartition and async API
+            from aiokafka import TopicPartition  # type: ignore
+
+            tp = TopicPartition(topic, 0)
+            consumer.assign([tp])
+            
+            # Get the end offset safely using aiokafka's API
+            end_offsets = await consumer.end_offsets([tp])
+            end = end_offsets.get(tp, 0)
+            
+            if end == 0:
+                logger.error(
+                    "[ERR-202] Unable to retrieve optimizations data from Superstream – topic empty."
+                )
+                consumer.close()
+                return None
+                
+            consumer.seek(tp, end - 1)
+            recs = await consumer.getmany(timeout_ms=5000)
+            consumer.close()
+            for batch in recs.values():
+                for rec in batch:
+                    return json.loads(rec.value.decode())
     except Exception as exc:
         logger.error("[ERR-203] Failed to fetch metadata: {}", exc)
     return None
 
+def fetch_metadata_sync(
+    bootstrap: str,
+    cfg: Dict[str, Any],
+    lib_name: Literal["kafka-python", "aiokafka", "confluent"],
+) -> Optional[Dict[str, Any]]:
+    """Synchronous wrapper for fetch_metadata."""
+    import asyncio
+    
+    # Run the async function synchronously for all libraries
+    return asyncio.run(fetch_metadata(bootstrap, cfg, lib_name))
+
 def optimal_cfg(metadata: Optional[Dict[str, Any]], topics: list[str], orig: Dict[str, Any], lib_name: str) -> tuple[Dict[str, Any], str]:
     """Compute optimal configuration based on metadata and topics.
     

superclient/util/config.py

@@ -31,9 +31,31 @@ def get_topics_list() -> List[str]:
 def mask_sensitive(k: str, v: Any) -> Any:
     """Mask sensitive configuration values."""
     sensitive_patterns = [
+        # Existing patterns
         "password", "sasl.jaas.config", "basic.auth.user.info",
         "ssl.key", "ssl.cert", "ssl.truststore", "ssl.keystore",
-        "sasl.kerberos.keytab", "sasl.kerberos.principal"
+        "sasl.kerberos.keytab", "sasl.kerberos.principal",
+        
+        # Additional SSL properties
+        "ssl.cafile", "ssl.certfile", "ssl.keyfile",
+        "ssl.certificate.location", "ssl.certificate.pem",
+        "ssl.ca.location", "ssl.ca.pem",
+        "ssl.ca.certificate.stores", "ssl.crl.location",
+        "ssl.providers", "ssl.context",
+        
+        # Additional SASL properties
+        "sasl.username", "sasl.password",
+        "sasl.plain.username", "sasl.plain.password",
+        "sasl.oauthbearer.config", "sasl.oauthbearer.client.secret",
+        "sasl.oauthbearer.extensions",
+        
+        # OAuth callback configurations
+        "sasl.oauth.token.provider", "oauth_cb", "sasl_oauth_token_provider",
+        
+        # Library-specific variations
+        "sasl_plain_username", "sasl_plain_password",
+        "ssl_cafile", "ssl_certfile", "ssl_keyfile",
+        "sasl_oauth_token_provider"
     ]
     return "[MASKED]" if any(pattern in k.lower() for pattern in sensitive_patterns) else v
 
@@ -129,6 +151,13 @@ def copy_client_configuration_properties(src: Dict[str, Any], dst: Dict[str, Any
             lib_key = _JAVA_TO_LIB_MAPPING[lib_name].get(java_key, java_key)
             if lib_key in src and lib_key not in dst:
                 dst[lib_key] = src[lib_key]
+    
+    # Debug log to show config comparison
+    # Mask sensitive data before logging
+    src_masked = {k: mask_sensitive(k, v) for k, v in src.items()}
+    dst_masked = {k: mask_sensitive(k, v) for k, v in dst.items()}
+    
+    logger.debug("copy_client_configuration_properties - Source config: {}, Destination config: {}", src_masked, dst_masked)
 
 # ---------------------------------------------------------------------------
 # Field name mapping between Java-style and library-specific representations