Refactor Kafka consumer and producer setup

kafka/consumer.py

@@ -3,13 +3,24 @@
 import time
 import argparse
 from typing import Optional
-from confluent_kafka import Consumer, Producer
+
+try:
+    from confluent_kafka import Consumer, Producer
+    KAFKA_LIB = "confluent"
+except ImportError:
+    try:
+        from kafka import KafkaConsumer, KafkaProducer
+        KAFKA_LIB = "python"
+    except ImportError:
+        print("Error: Neither confluent-kafka nor kafka-python is installed.")
+        sys.exit(1)
 
 try:
     import httpx
 except ImportError:
     httpx = None
 
+
 def get_scorer(mode: str, api_url: str):
     if mode == "direct":
         from api.schemas import TransactionIn
@@ -36,6 +47,59 @@ def api_score(tx_dict: dict) -> tuple[Optional[dict], Optional[str]]:
                 return None, str(e)
         return api_score
 
+
+def get_kafka_consumer(servers: str, group_id: str):
+    """Create a Kafka consumer, preferring confluent-kafka."""
+    if KAFKA_LIB == "confluent":
+        return Consumer({
+            "bootstrap.servers": servers,
+            "group.id": group_id,
+            "auto.offset.reset": "earliest",
+            "enable.auto.commit": False
+        })
+    else:
+        print("[warning] Using kafka-python consumer fallback — some features may differ.")
+        return KafkaConsumer(
+            bootstrap_servers=servers.split(","),
+            group_id=group_id,
+            auto_offset_reset="earliest",
+            enable_auto_commit=False,
+            value_deserializer=lambda m: m  # raw bytes, same as confluent
+        )
+
+
+def get_kafka_producer(servers: str):
+    """Create a Kafka producer, preferring confluent-kafka."""
+    if KAFKA_LIB == "confluent":
+        return Producer({"bootstrap.servers": servers})
+    else:
+        # NOTE: No value_serializer here — we send raw bytes explicitly.
+        # DLQ receives raw bytes (from upstream), scored-transactions receives
+        # JSON-encoded bytes. A serializer would break the DLQ path.
+        return KafkaProducer(bootstrap_servers=servers.split(","))
+
+
+def produce_dlq(producer, raw_val: bytes):
+    """Send a malformed message to the DLQ. Raises on failure so caller does not commit."""
+    if KAFKA_LIB == "confluent":
+        producer.produce("fraud-dlq", value=raw_val)
+        producer.poll(0)
+    else:
+        producer.send("fraud-dlq", value=raw_val)
+
+
+def produce_scored(producer, res_dict: dict):
+    """Publish a scored transaction. Raises on failure so caller does not commit."""
+    key = str(res_dict.get("transaction_id", "")).encode("utf-8")
+    payload = json.dumps(res_dict).encode("utf-8")
+    if KAFKA_LIB == "confluent":
+        producer.produce("scored-transactions", key=key, value=payload)
+        producer.poll(0)
+    else:
+        # Send raw bytes — KafkaProducer has no serializer configured
+        producer.send("scored-transactions", key=key, value=payload)
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--mode", choices=["direct", "api"], default="direct")
@@ -44,92 +108,123 @@ def main():
     parser.add_argument("--api-url", type=str, default="http://localhost:8000")
     args = parser.parse_args()
 
-    consumer = Consumer({
-        "bootstrap.servers": args.servers,
-        "group.id": "fraud-radar-consumer-v1",
-        "auto.offset.reset": "earliest",
-        "enable.auto.commit": False
-    })
-
-    producer = Producer({
-        "bootstrap.servers": args.servers
-    })
-
-    consumer.subscribe(["transactions"])
+    consumer = get_kafka_consumer(args.servers, "fraud-radar-consumer-v1")
+    producer = get_kafka_producer(args.servers)
+
+    if KAFKA_LIB == "confluent":
+        consumer.subscribe(["transactions"])
+    else:
+        consumer.subscribe(["transactions"])
+
     score_fn = get_scorer(args.mode, args.api_url)
-    
-    print(f"[consumer] listening on {args.servers} <- transactions (mode={args.mode})")
-    
+
+    print(f"[consumer] listening on {args.servers} <- transactions (mode={args.mode}, lib={KAFKA_LIB})")
+
     processed = 0
     flagged = 0
     errors = 0
     start_time = time.time()
     latencies = []
-    
+
     try:
         while True:
-            msg = consumer.poll(1.0)
-            if msg is None:
-                continue
-            if msg.error():
-                print(f"Consumer error: {msg.error()}")
-                continue
-
+            if KAFKA_LIB == "confluent":
+                msg = consumer.poll(1.0)
+                if msg is None:
+                    continue
+                if msg.error():
+                    print(f"Consumer error: {msg.error()}")
+                    continue
+                raw_val = msg.value()
+            else:
+                # kafka-python: iterator-based with 1-second timeout via poll
+                msg_pack = consumer.poll(timeout_ms=1000)
+                if not msg_pack:
+                    continue
+                # poll returns a dict of TopicPartition -> list of ConsumerRecord
+                got_msg = False
+                for tp, records in msg_pack.items():
+                    for record in records:
+                        raw_val = record.value
+                        got_msg = True
+                        break
+                    if got_msg:
+                        break
+                if not got_msg:
+                    continue
+
             t0 = time.perf_counter()
-            raw_val = msg.value()
-
+
             try:
                 tx_dict = json.loads(raw_val.decode("utf-8"))
             except Exception:
-                # DLQ
-                producer.produce("fraud-dlq", value=raw_val)
-                producer.poll(0)
-                consumer.commit(msg)
+                # DLQ: do NOT commit if DLQ publish fails — message will be reprocessed
+                try:
+                    produce_dlq(producer, raw_val)
+                    if KAFKA_LIB == "confluent":
+                        consumer.commit(msg)
+                    else:
+                        consumer.commit_async()
+                except Exception as e:
+                    print(f"[warning] DLQ publish failed, skipping commit: {e}")
                 errors += 1
                 continue
-                
+
             # Score
             res_dict, err = score_fn(tx_dict)
             if err or not res_dict:
-                producer.produce("fraud-dlq", value=raw_val)
-                producer.poll(0)
-                consumer.commit(msg)
+                try:
+                    produce_dlq(producer, raw_val)
+                    if KAFKA_LIB == "confluent":
+                        consumer.commit(msg)
+                    else:
+                        consumer.commit_async()
+                except Exception as e:
+                    print(f"[warning] DLQ publish failed, skipping commit: {e}")
                 errors += 1
                 continue
-                
+
             # Publish to scored-transactions
             try:
-                key = str(res_dict.get("transaction_id", "")).encode("utf-8")
-                producer.produce("scored-transactions", key=key, value=json.dumps(res_dict).encode("utf-8"))
-                producer.poll(0)
-            except Exception:
-                pass
-
-            consumer.commit(msg)
-
+                produce_scored(producer, res_dict)
+            except Exception as e:
+                print(f"[warning] scored-transactions publish failed, skipping commit: {e}")
+                continue
+
+            # Only commit after successful downstream publish
+            if KAFKA_LIB == "confluent":
+                consumer.commit(msg)
+            else:
+                consumer.commit_async()
+
             latency = (time.perf_counter() - t0) * 1000.0
             latencies.append(latency)
             if len(latencies) > 100:
                 latencies.pop(0)
-                
+
             processed += 1
             if res_dict.get("is_flagged"):
                 flagged += 1
-                if args.print_flagged:
-                    print(f"🚨 FLAGGED  {res_dict['transaction_id']}  risk={res_dict['risk_score']:.3f}  level={res_dict['risk_level']}  ${res_dict['amount']:.2f}")
+            if args.print_flagged:
+                print(f"🚨 FLAGGED {res_dict['transaction_id']} risk={res_dict['risk_score']:.3f} level={res_dict['risk_level']} ${res_dict['amount']:.2f}")
 
             if processed % 100 == 0:
                 elapsed = time.time() - start_time
                 rate = processed / elapsed if elapsed > 0 else 0
                 avg_lat = sum(latencies) / len(latencies) if latencies else 0
                 pct = (flagged / processed) * 100
-                print(f"[metrics] processed={processed}  flagged={flagged} ({pct:.1f}%)  errors={errors}  rate={rate:.1f} tx/s  avg_latency={avg_lat:.1f}ms")
+                print(f"[metrics] processed={processed} flagged={flagged} ({pct:.1f}%) errors={errors} rate={rate:.1f} tx/s avg_latency={avg_lat:.1f}ms")
 
     except KeyboardInterrupt:
         pass
     finally:
-        consumer.close()
-        producer.flush()
+        if KAFKA_LIB == "confluent":
+            consumer.close()
+            producer.flush()
+        else:
+            consumer.close()
+            producer.flush()
+
 
 if __name__ == "__main__":
-    main()
+    main()