From b8703ea6a113bc3c4fd4fe7d309802ca8632673f Mon Sep 17 00:00:00 2001
From: "Pedro S. Lopez" <pedroslopez@me.com>
Date: Thu, 12 May 2022 22:38:30 -0400
Subject: [PATCH] update CDK airbyte protocol models to fix master build
 (#12829)

---
 .../airbyte_cdk/models/airbyte_protocol.py    | 31 ++++++++++++++++---
 1 file changed, 27 insertions(+), 4 deletions(-)

diff --git a/airbyte-cdk/python/airbyte_cdk/models/airbyte_protocol.py b/airbyte-cdk/python/airbyte_cdk/models/airbyte_protocol.py
index e309b7f00d8d..5e904cf70c4e 100644
--- a/airbyte-cdk/python/airbyte_cdk/models/airbyte_protocol.py
+++ b/airbyte-cdk/python/airbyte_cdk/models/airbyte_protocol.py
@@ -36,11 +36,16 @@ class Config:
     namespace: Optional[str] = Field(None, description="the namespace of this record's stream")
 
 
-class AirbyteStateMessage(BaseModel):
-    class Config:
-        extra = Extra.allow
+class AirbyteStateType(Enum):
+    GLOBAL = "GLOBAL"
+    PER_STREAM = "PER_STREAM"
+
 
-    data: Dict[str, Any] = Field(..., description="the state data")
+class AirbyteStateBlob(BaseModel):
+    pass
+
+    class Config:
+        extra = Extra.forbid
 
 
 class Level(Enum):
@@ -157,6 +162,14 @@ class OAuthConfigSpecification(BaseModel):
     )
 
 
+class AirbyteStreamState(BaseModel):
+    class Config:
+        extra = Extra.forbid
+
+    name: Optional[str] = Field(None, description="Stream name")
+    state: Optional[AirbyteStateBlob] = None
+
+
 class AirbyteTraceMessage(BaseModel):
     class Config:
         extra = Extra.allow
@@ -244,6 +257,16 @@ class Config:
     )
 
 
+class AirbyteStateMessage(BaseModel):
+    class Config:
+        extra = Extra.allow
+
+    state_type: Optional[AirbyteStateType] = None
+    data: Optional[Dict[str, Any]] = Field(None, description="(Deprecated) the state data")
+    global_: Optional[AirbyteStateBlob] = Field(None, alias="global")
+    streams: Optional[List[AirbyteStreamState]] = None
+
+
 class AirbyteCatalog(BaseModel):
     class Config:
         extra = Extra.allow