improvements

Author: idanasulin2706

superclient/__init__.py

@@ -1,2 +1,11 @@
+from importlib.metadata import PackageNotFoundError, version as _pkg_version
+
+try:
+    __version__ = _pkg_version("superclient")
+except PackageNotFoundError:
+    # Fallback for when the package isn't installed (e.g. running from source without editable install)
+    __version__ = "0.0.0"
+
+# Initialize the Superstream agent on import
 from .agent import initialize as _init
 _init() 

superclient/agent/__init__.py

@@ -5,7 +5,7 @@
 import sys
 from typing import Any, Dict, Optional, Tuple
 
-from ..logger import get_logger, set_debug_enabled
+from ..util.logger import get_logger, set_debug_enabled
 from ..util.config import get_env_vars, is_disabled
 from ..util.network import get_host_info
 from .interceptor import patch_kafka_python, patch_aiokafka, patch_confluent
@@ -24,7 +24,7 @@
 if is_disabled():
     logger.warn("Superstream functionality disabled via SUPERSTREAM_DISABLED")
 
-_VERSION = "1.0.0"
+# Preserve reference to built-in import function
 _original_import = builtins.__import__
 
 # ---------------------------------------------------------------------------

superclient/agent/clients.py

@@ -4,72 +4,85 @@
 import os
 from typing import Any, Dict
 
-from ..logger import get_logger
+from ..util.logger import get_logger
+from .. import __version__ as _VERSION
+from ..util.config import (
+    copy_client_configuration_properties,
+    convert_to_dot_syntax,
+    get_original_config,
+    mask_sensitive,
+)
 
 logger = get_logger("agent.clients")
 
-_VERSION = "0.1.0"
 _SUPERLIB_PREFIX = "superstreamlib-"
 
-def mask_sensitive(k: str, v: Any) -> Any:
-    """Mask sensitive configuration values."""
-    return "[MASKED]" if "password" in k.lower() or "sasl.jaas.config" in k.lower() else v
-
-def copy_security(src: Dict[str, Any], dst: Dict[str, Any]):
-    """Copy security-related configuration from source to destination."""
-    keys = [
-        "security.protocol",
-        "sasl.mechanism",
-        "sasl.jaas.config",
-        "ssl.keystore.password",
-        "ssl.truststore.password",
-        "ssl.key.password",
-        "client.dns.lookup",
-    ]
-    for k in keys:
-        if k in src and k not in dst:
-            dst[k] = src[k]
-
-def internal_send_clients(bootstrap: str, base_cfg: Dict[str, Any], payload: bytes) -> None:
-    """Send payload to superstream.clients using available Kafka library."""
-    # Attempt kafka-python first
-    try:
-        import kafka  # type: ignore
-
-        cfg = {
-            "bootstrap_servers": bootstrap,
-            "client_id": _SUPERLIB_PREFIX + "client-reporter",
-            "compression_type": "zstd",
-            "batch_size": 16_384,
-            "linger_ms": 1000,
-        }
-        copy_security(base_cfg, cfg)
-        prod = kafka.KafkaProducer(**{k.replace(".", "_"): v for k, v in cfg.items()})
-        prod.send("superstream.clients", payload)
-        prod.flush()
-        prod.close()
+def _create_producer_kafka_python(bootstrap: str, base_cfg: Dict[str, Any]):
+    import kafka  # type: ignore
+
+    cfg = {
+        "bootstrap.servers": bootstrap,
+        "client.id": _SUPERLIB_PREFIX + "client-reporter",
+        "compression.type": "zstd",
+        "batch.size": 16_384,
+        "linger.ms": 1000,
+    }
+    copy_client_configuration_properties(base_cfg, cfg)
+    kafka_cfg = {k.replace(".", "_"): v for k, v in cfg.items()}
+    return kafka.KafkaProducer(**kafka_cfg)
+
+
+def _create_producer_confluent(bootstrap: str, base_cfg: Dict[str, Any]):
+    from confluent_kafka import Producer as _CProducer  # type: ignore
+
+    cfg = {
+        "bootstrap.servers": bootstrap,
+        "client.id": _SUPERLIB_PREFIX + "client-reporter",
+        "compression.type": "zstd",
+        "batch.size": 16384,
+        "linger.ms": 1000,
+    }
+    copy_client_configuration_properties(base_cfg, cfg)
+    return _CProducer(cfg)
+
+
+_PRODUCER_BUILDERS = {
+    "kafka-python": _create_producer_kafka_python,
+    "confluent": _create_producer_confluent,
+    # aiokafka is async; for heartbeat reporting we use kafka-python builder as fallback
+    "aiokafka": _create_producer_kafka_python,
+}
+
+
+def internal_send_clients(
+    bootstrap: str,
+    base_cfg: Dict[str, Any],
+    payload: bytes,
+    lib_name: str,
+) -> None:
+    """Send payload to superstream.clients using the *same* Kafka library as the application.
+
+    In the unlikely event of failure we silently swallow exceptions – reporting
+    must never interrupt the user application – but we log for troubleshooting.
+    """
+
+    builder = _PRODUCER_BUILDERS.get(lib_name)
+    if builder is None:
+        logger.debug("Unknown Kafka library '{}', skipping client report", lib_name)
         return
-    except Exception:
-        pass  # fallthrough
 
-    # Fallback to confluent-kafka if available
     try:
-        from confluent_kafka import Producer as _CProducer  # type: ignore
-
-        cfg = {
-            "bootstrap.servers": bootstrap,
-            "client.id": _SUPERLIB_PREFIX + "client-reporter",
-            "compression.type": "zstd",
-            "batch.size": 16384,
-            "linger.ms": 1000,
-        }
-        copy_security(base_cfg, cfg)
-        prod = _CProducer(cfg)
-        prod.produce("superstream.clients", payload)
-        prod.flush()
+        prod = builder(bootstrap, base_cfg)
+        # confluent-kafka Producer uses .produce, kafka-python uses .send
+        if hasattr(prod, "produce"):
+            prod.produce("superstream.clients", payload)
+            prod.flush()
+        else:
+            prod.send("superstream.clients", payload)
+            prod.flush()
+            prod.close()
     except Exception:
-        # As a last resort just log and drop – should never interrupt app
-        logger.debug("Failed to send clients message via all libraries")
+        logger.debug("Failed to send clients message via {}", lib_name)
 
 def get_host_info() -> tuple[str, str]:
     """Get hostname and IP address."""
@@ -84,22 +97,29 @@ def get_host_info() -> tuple[str, str]:
 def send_clients_msg(tracker: Any, error: str = "") -> None:
     """Send a message to the clients topic."""
     hostname, ip = get_host_info()
+    # Prepare configuration dictionaries in dot syntax with defaults merged
+    orig_cfg_dot = get_original_config(tracker.orig_cfg, tracker.library)
+    orig_cfg_masked = {k: mask_sensitive(k, v) for k, v in orig_cfg_dot.items()}
+
+    opt_cfg_dot = convert_to_dot_syntax(tracker.opt_cfg, tracker.library)
+    opt_cfg_masked = {k: mask_sensitive(k, v) for k, v in opt_cfg_dot.items()}
+
     msg_dict = {
         "client_id": tracker.client_id,
         "ip_address": ip,
         "type": "producer",
-        "message_type": "client_stats" if not error else "client_info",
+        "message_type": "client_stats",
         "version": _VERSION,
         "topics": sorted(tracker.topics),
-        "original_configuration": {k: mask_sensitive(k, v) for k, v in tracker.orig_cfg.items()},
-        "optimized_configuration": {k: mask_sensitive(k, v) for k, v in tracker.opt_cfg.items()},
+        "original_configuration": orig_cfg_masked,
+        "optimized_configuration": opt_cfg_masked,
         "environment_variables": {k: v for k, v in os.environ.items() if k.startswith("SUPERSTREAM_")},
         "hostname": hostname,
         "superstream_client_uid": tracker.uuid,
         "most_impactful_topic": tracker.determine_topic(),
-        "language": "Python",
+        "language": f"Python ({tracker.library})",
         "error": error,
     }
     payload = json.dumps(msg_dict).encode()
-    internal_send_clients(tracker.bootstrap, tracker.orig_cfg, payload)
+    internal_send_clients(tracker.bootstrap, tracker.orig_cfg, payload, tracker.library)
     logger.debug("Sent clients message for {}", tracker.client_id) 

superclient/agent/interceptor.py

@@ -3,9 +3,9 @@
 import os
 from typing import Any, Dict
 
-from ..logger import get_logger
+from ..util.logger import get_logger
 from ..util.config import get_topics_list, is_disabled
-from ..core.manager import fetch_metadata, optimal_cfg
+from .metadata import fetch_metadata, optimal_cfg
 from ..core.reporter import send_clients_msg
 from .tracker import ProducerTracker, Heartbeat
 
@@ -30,12 +30,17 @@ def init_patch(self, *args, **kwargs):
         bootstrap = orig_cfg.get("bootstrap_servers") or (args[0] if args else None)
         if not bootstrap:
             return orig_init(self, *args, **kwargs)
+        bootstrap = normalize_bootstrap(bootstrap)
         client_id = orig_cfg.get("client_id", "")
         if client_id.startswith(_SUPERLIB_PREFIX):
             return orig_init(self, *args, **kwargs)
         topics_env = get_topics_list()
-        metadata = fetch_metadata(bootstrap, orig_cfg)
-        opt_cfg = optimal_cfg(metadata, topics_env, orig_cfg)
+        metadata = fetch_metadata(bootstrap, orig_cfg, "kafka-python")
+        error_msg = ""
+        if metadata and not metadata.active:
+            error_msg = "[ERR-301] Superstream optimization is not active for this kafka cluster, please head to the Superstream console and activate it."
+            logger.error(error_msg)
+        opt_cfg = optimal_cfg(metadata, topics_env, orig_cfg) if metadata and metadata.active else {}
         for k, v in opt_cfg.items():
             snake = k.replace(".", "_")
             if kwargs.get(snake) != v:
@@ -64,13 +69,15 @@ def send_patch(inner, topic, *a, **kw):
             orig_close = self.close
 
             def close_patch(inner, *a, **kw):
-                tr.close()
-                Heartbeat.unregister_tracker(tr.uuid)
+                if not hasattr(self, "_superstream_closed"):
+                    self._superstream_closed = True
+                    tr.close()
+                    Heartbeat.unregister_tracker(tr.uuid)
                 return orig_close(*a, **kw)
 
             self.close = close_patch
             self._superstream_patch = True
-        send_clients_msg(tr)
+        send_clients_msg(tr, error_msg)
         logger.info("Successfully optimized producer configuration for {}", client_id)
 
     Producer.__init__ = init_patch
@@ -92,12 +99,17 @@ def init_patch(self, *args, **kwargs):
             bootstrap = args[0]
         if not bootstrap:
             return orig_init(self, *args, **kwargs)
+        bootstrap = normalize_bootstrap(bootstrap)
         client_id = orig_cfg.get("client_id", "")
         if client_id.startswith(_SUPERLIB_PREFIX):
             return orig_init(self, *args, **kwargs)
         topics_env = get_topics_list()
-        metadata = fetch_metadata(bootstrap, orig_cfg)
-        opt_cfg = optimal_cfg(metadata, topics_env, orig_cfg)
+        metadata = fetch_metadata(bootstrap, orig_cfg, "aiokafka")
+        error_msg = ""
+        if metadata and not metadata.active:
+            error_msg = "[ERR-301] Superstream optimization is not active for this kafka cluster, please head to the Superstream console and activate it."
+            logger.error(error_msg)
+        opt_cfg = optimal_cfg(metadata, topics_env, orig_cfg) if metadata and metadata.active else {}
         for k, v in opt_cfg.items():
             if kwargs.get(k) != v:
                 logger.debug("Overriding configuration: {} -> {}", k, v)
@@ -125,13 +137,15 @@ async def send_patch(inner, topic, *a, **kw):
             original_stop = self.stop
 
             async def stop_patch(inner, *a, **kw):
+                if not hasattr(self, "_superstream_closed"):
+                    self._superstream_closed = True
+                    tr.close()
+                    Heartbeat.unregister_tracker(tr.uuid)
                 await original_stop(*a, **kw)
-                tr.close()
-                Heartbeat.unregister_tracker(tr.uuid)
 
             self.stop = stop_patch
             self._superstream_patch = True
-        send_clients_msg(tr)
+        send_clients_msg(tr, error_msg)
         logger.info("Successfully optimized producer configuration for {}", client_id)
 
     Producer.__init__ = init_patch
@@ -149,12 +163,19 @@ def init_patch(self, conf: Dict[str, Any], *args, **kwargs):
             return orig_init(self, conf, *args, **kwargs)
         conf = dict(conf)
         bootstrap = conf.get("bootstrap.servers")
+        if not bootstrap:
+            return orig_init(self, conf, *args, **kwargs)
+        bootstrap = normalize_bootstrap(bootstrap)
         client_id = conf.get("client.id", "")
         if client_id.startswith(_SUPERLIB_PREFIX):
             return orig_init(self, conf, *args, **kwargs)
         topics_env = get_topics_list()
-        metadata = fetch_metadata(bootstrap, conf)
-        opt_cfg = optimal_cfg(metadata, topics_env, conf)
+        metadata = fetch_metadata(bootstrap, conf, "confluent")
+        error_msg = ""
+        if metadata and not metadata.active:
+            error_msg = "[ERR-301] Superstream optimization is not active for this kafka cluster, please head to the Superstream console and activate it."
+            logger.error(error_msg)
+        opt_cfg = optimal_cfg(metadata, topics_env, conf) if metadata and metadata.active else {}
         for k, v in opt_cfg.items():
             if conf.get(k) != v:
                 logger.debug("Overriding configuration: {} -> {}", k, v)
@@ -179,25 +200,18 @@ def produce_patch(inner, topic, *a, **kw):
                 return original_produce(topic, *a, **kw)
 
             self.produce = produce_patch
-            orig_flush = self.flush
-
-            def flush_patch(inner, *a, **kw):
-                tr.close()
-                Heartbeat.unregister_tracker(tr.uuid)
-                return orig_flush(*a, **kw)
-
-            self.flush = flush_patch
-            if hasattr(self, "close"):
-                orig_close = self.close
+            orig_close = self.close
 
-                def close_patch(inner, *a, **kw):
+            def close_patch(inner, *a, **kw):
+                if not hasattr(self, "_superstream_closed"):
+                    self._superstream_closed = True
                     tr.close()
                     Heartbeat.unregister_tracker(tr.uuid)
-                    return orig_close(*a, **kw)
+                return orig_close(*a, **kw)
 
-                self.close = close_patch
+            self.close = close_patch
             self._superstream_patch = True
-        send_clients_msg(tr)
+        send_clients_msg(tr, error_msg)
         logger.info("Successfully optimized producer configuration for {}", client_id)
 
     Producer.__init__ = init_patch