kafka_cdc decorator

Author: omgbebebe

src/mindwm/knfunc/decorators.py

@@ -12,8 +12,9 @@
 from fastapi import FastAPI, Request, Response, status
 from fastapi.responses import JSONResponse
 from mindwm import logging
-from mindwm.model.events import (CloudEvent, IoDocumentEvent, LLMAnswerEvent,
-                                 TouchEvent)
+from mindwm.model.events import (CloudEvent, IoDocumentEvent, KafkaCdcEvent,
+                                 LLMAnswerEvent, TouchEvent)
+from mindwm.model.graph import KafkaCdc
 from neontology import auto_constrain, init_neontology
 from opentelemetry import trace
 from opentelemetry._logs import set_logger_provider
@@ -250,3 +251,83 @@ async def inner(**kwargs):
         return res
 
     return wrapper
+
+
+def kafka_cdc(func):
+
+    @app.post("/")
+    async def wrapper(r: Request, response: Response):
+        func_sig = inspect.signature(func)
+        xx = [p.annotation for p in func_sig.parameters.values()]
+        kwargs = dict(func_sig.parameters)
+        b = await r.body()
+        logger.debug(f"headers: {r.headers}")
+        logger.debug(f"body: {b}")
+        cdc_obj = KafkaCdc.model_validate_json(b)
+        logger.info(f"cdc_obj: {cdc_obj}")
+        if 'traceparent' in r.headers.keys():
+            carrier = r.headers.get('traceparent')
+
+        if 'obj' in kwargs:
+            kwargs['obj'] = cdc_obj
+
+        match cdc_obj.payload.type:
+            case 'relationship':
+                carrier = cdc_obj.payload.traceparent
+            case 'node':
+                if cdc_obj.meta.operation != 'deleted':
+                    carrier = cdc_obj.payload.after.properties.traceparent
+                else:
+                    carrier = cdc_obj.payload.before.properties.traceparent
+
+        @with_trace(carrier=r.headers)
+        @wraps(func)
+        async def inner(**kwargs):
+
+            value = await func(**kwargs)
+
+            logger.debug(f"return value: {value}")
+
+            headers = {}
+            if carrier:
+                headers = {"traceparent": carrier}
+
+            if not value:
+                return Response(status_code=status.HTTP_200_OK,
+                                headers=headers)
+            else:
+                context_name = os.environ.get('CONTEXT_NAME', 'NO_CONTEXT')
+
+                obj_ev = CloudEvent.make_obj_event(value)
+                match cdc_obj.payload.type:
+                    case "node":
+                        if cdc_obj.meta.operation != 'deleted':
+                            event_type = cdc_obj.payload.after.labels[0].lower(
+                            )
+                        else:
+                            event_type = cdc_obj.payload.before.labels[
+                                0].lower()
+                    case "relationship":
+                        event_type = cdc_obj.payload.label.lower()
+                    case _:
+                        event_type = "UNKNOWN"
+
+                #logger.info(f"ctx: {span.context}")
+                ev = CloudEvent(
+                    id=uuid4().hex,
+                    source=
+                    f"mindwm.{context_name}.graph.{ cdc_obj.payload.type.lower() }",
+                    subject=f"{ cdc_obj.meta.operation }",
+                    type=event_type,
+                    data=obj_ev.model_dump(),
+                    traceparent=carrier)
+
+                headers['content-type'] = 'application/cloudevents+json'
+                logger.debug(f"response headers: {headers}")
+                logger.debug(f"response event: {ev}")
+                return Response(content=ev.model_dump_json(), headers=headers)
+
+        res = await inner(**kwargs)
+        return res
+
+    return wrapper

src/mindwm/model/events.py

@@ -2,6 +2,7 @@
 
 from pydantic import BaseModel, Field
 
+from .graph import KafkaCdc
 from .objects import IoDocument, LLMAnswer, ShowMessage, Touch, TypeText
 
 
@@ -45,18 +46,24 @@ class TypeTextEvent(BaseEvent):
     type: Literal['typetextevent'] = 'typetextevent'
 
 
-Obj = TypeVar("Obj", IoDocument, Touch, LLMAnswer, ShowMessage, TypeText)
+class KafkaCdcEvent(BaseEvent):
+    data: KafkaCdc
+    type: Literal['kafkacdcevent'] = 'kafkacdcevent'
+
+
+Obj = TypeVar("Obj", IoDocument, Touch, LLMAnswer, ShowMessage, TypeText,
+              KafkaCdc)
 
 ObjEvent = TypeVar("ObjEvent", IoDocumentEvent, TouchEvent, LLMAnswerEvent,
-                   ShowMessageEvent, TypeTextEvent)
+                   ShowMessageEvent, TypeTextEvent, KafkaCdcEvent)
 
 
 class CloudEvent(BaseEvent):
     id: str
     source: str
     specversion: str = "1.0"
     data: Annotated[Union[IoDocumentEvent, TouchEvent, LLMAnswerEvent,
-                          ShowMessageEvent, TypeTextEvent],
+                          ShowMessageEvent, TypeTextEvent, KafkaCdcEvent],
                     Field(discriminator='type')]
     type: Optional[str] = None
     datacontenttype: Optional[str] = None
@@ -76,6 +83,10 @@ class CloudEvent(BaseEvent):
         Field(min_length=1,
               description="a comma-delimited list of key-value pairs")]] = None
     knativearrivaltime: Optional[str] = None
+    key: Optional[str] = None
+    knativekafkaoffset: Optional[int] = None
+    knativekafkapartition: Optional[int] = None
+    partitionkey: Optional[str] = None
 
     @classmethod
     def make_obj_event(cls, obj: Type[Obj]) -> Type[ObjEvent]:
@@ -90,6 +101,8 @@ def make_obj_event(cls, obj: Type[Obj]) -> Type[ObjEvent]:
                 return ShowMessageEvent(data=obj)
             case TypeText():
                 return TypeTextEvent(data=obj)
+            case KafkaCdc():
+                return KafkaCdcEvent(data=obj)
             case _:
                 msg = f"unknown object type: {obj}"
                 raise TypeError(msg)

src/mindwm/model/graph.py

@@ -1,13 +1,18 @@
-from typing import ClassVar, Dict, List, TypeVar, Union
+from datetime import datetime
+from typing import (Annotated, Any, ClassVar, Dict, List, Literal, Optional,
+                    TypeVar, Union)
 
 from neontology import BaseNode, BaseRelationship
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from .objects import IoDocument
 
 
 class MindwmNode(BaseNode):
-    atime: int = 0
+    atime: Optional[int] = 0
+    created: Optional[datetime] = None
+    merged: Optional[datetime] = None
+    traceparent: Optional[str] = None
     #    atime: datetime = Field(
     #    default_factory=datetime.now
     #)
@@ -66,79 +71,85 @@ class Parameter(MindwmNode):
 
 
 # Relations
-class UserHasHost(BaseRelationship):
+class MindwmRelationship(BaseRelationship):
+    traceparent: Optional[str] = None
+    created: Optional[datetime] = None
+    merged: Optional[datetime] = None
+
+
+class UserHasHost(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_HOST"
     source: User
     target: Host
 
 
-class HostHasTmux(BaseRelationship):
+class HostHasTmux(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_TMUX"
     source: Host
     target: Tmux
 
 
-class TmuxHasTmuxSession(BaseRelationship):
+class TmuxHasTmuxSession(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_TMUX_SESSION"
     source: Tmux
     target: TmuxSession
 
 
-class TmuxSessionHasTmuxPane(BaseRelationship):
+class TmuxSessionHasTmuxPane(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_TMUX_PANE"
     source: TmuxSession
     target: TmuxPane
 
 
-class TmuxPaneHasIoDocument(BaseRelationship):
+class TmuxPaneHasIoDocument(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_IO_DOCUMENT"
     source: TmuxPane
     target: IoDocument
 
 
-class HostHasClipboard(BaseRelationship):
+class HostHasClipboard(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_CLIPBOARD"
     source: Host
     target: Clipboard
 
 
-class UserHasTmux(BaseRelationship):
+class UserHasTmux(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_TMUX"
     source: User
     target: Tmux
 
 
-class IoDocumentHasUser(BaseRelationship):
+class IoDocumentHasUser(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_USER"
     source: IoDocument
     target: User
 
 
-class UserHasParameter(BaseRelationship):
+class UserHasParameter(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_PARAMETER"
     source: User
     target: Parameter
 
 
-class HostHasParameter(BaseRelationship):
+class HostHasParameter(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_PARAMETER"
     source: Host
     target: Parameter
 
 
-class TmuxHasParameter(BaseRelationship):
+class TmuxHasParameter(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_PARAMETER"
     source: Tmux
     target: Parameter
 
 
-class TmuxSessionHasParameter(BaseRelationship):
+class TmuxSessionHasParameter(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_PARAMETER"
     source: TmuxSession
     target: Parameter
 
 
-class TmuxPaneHasParameter(BaseRelationship):
+class TmuxPaneHasParameter(MindwmRelationship):
     __relationshiptype__: ClassVar[str] = "HAS_PARAMETER"
     source: TmuxPane
     target: Parameter
@@ -155,19 +166,42 @@ class KafkaCdcMeta(BaseModel):
     source: Dict[str, str]
 
 
-Prop = TypeVar("Prop", User, Host)
+Prop = TypeVar("Prop", User, Host, Tmux, TmuxSession, TmuxPane, IoDocument,
+               Clipboard, Parameter)
+
+
+class KafkaCdcRelNode(BaseModel):
+    id: str
+    labels: List[str]
+    ids: Dict[str, str]
+
+
+class KafkaCdcRelProp(BaseModel):
+    # FIX: more pricise parametrization if needed
+    properties: Any
+
+
+class KafkaCdcRelation(BaseModel):
+    id: int
+    start: KafkaCdcRelNode
+    end: KafkaCdcRelNode
+    before: Optional[KafkaCdcRelProp] = None
+    after: Optional[KafkaCdcRelProp] = None
+    label: str
+    type: Literal['relationship'] = 'relationship'
+    traceparent: Optional[str] = None
 
 
-class KafkaCdcPayloadData(BaseModel):
+class KafkaCdcNodeData(BaseModel):
     properties: Prop
     labels: List[str]
 
 
-class KafkaCdcPayload(BaseModel):
+class KafkaCdcNode(BaseModel):
     id: int
-    type: str
-    before: KafkaCdcPayloadData
-    after: KafkaCdcPayloadData
+    type: Literal['node'] = 'node'
+    before: Optional[KafkaCdcNodeData] = None
+    after: Optional[KafkaCdcNodeData] = None
 
 
 class KafkaCdcSchema(BaseModel):
@@ -177,7 +211,8 @@ class KafkaCdcSchema(BaseModel):
 
 class KafkaCdc(BaseModel):
     meta: KafkaCdcMeta
-    payload: KafkaCdcPayload
+    payload: Annotated[Union[KafkaCdcNode, KafkaCdcRelation],
+                       Field(discriminator='type')]
     schema: KafkaCdcSchema
 
     def get_object_before(self):
@@ -187,4 +222,7 @@ def get_object_after(self):
         return self.payload.after.properties
 
     def get_object(self):
-        return self.get_object_after()
+        if self.payload.meta.operation != 'deleted':
+            return self.get_object_after()
+        else:
+            return self.get_object_before()