add python tests

callicles · callicles · commit d05b7e48d731 · 2025-12-04T21:25:13.000-05:00
diff --git a/apps/framework-cli-e2e/test/templates.test.ts b/apps/framework-cli-e2e/test/templates.test.ts
@@ -1770,6 +1770,145 @@ const createTemplateTestSuite = (config: TemplateTestConfig) => {
           expect(apiData).to.be.an("array");
         });
 
+        // Extra fields test for Python (ENG-1617)
+        // Tests that IngestApi accepts payloads with extra fields when the model has extra='allow'.
+        // Extra fields are passed through to streaming functions and stored in a JSON column.
+        //
+        // KEY CONCEPTS:
+        // - IngestApi/Stream with extra='allow': CAN accept variable fields
+        // - OlapTable: Requires fixed schema (ClickHouse needs to know columns)
+        // - Transform: Receives ALL fields via model_extra, outputs to fixed schema with JSON column
+        it("should pass extra fields to streaming function via Pydantic extra='allow' (PY)", async function () {
+          this.timeout(TIMEOUTS.TEST_SETUP_MS);
+
+          const userId = randomUUID();
+          const timestamp = new Date().toISOString();
+
+          // Send data with known fields plus arbitrary extra fields
+          await withRetries(
+            async () => {
+              const response = await fetch(
+                `${SERVER_CONFIG.url}/ingest/userEventIngestApi`,
+                {
+                  method: "POST",
+                  headers: { "Content-Type": "application/json" },
+                  body: JSON.stringify({
+                    // Known fields defined in the model (snake_case for Python)
+                    timestamp: timestamp,
+                    event_name: "page_view",
+                    user_id: userId,
+                    org_id: "org-123",
+                    // Extra fields - allowed by extra='allow', passed to streaming function
+                    customProperty: "custom-value",
+                    pageUrl: "/dashboard",
+                    sessionDuration: 120,
+                    nested: {
+                      level1: "value1",
+                      level2: { deep: "nested" },
+                    },
+                  }),
+                },
+              );
+              if (!response.ok) {
+                const text = await response.text();
+                throw new Error(`${response.status}: ${text}`);
+              }
+            },
+            { attempts: 5, delayMs: 500 },
+          );
+
+          // Wait for the transform to process and write to output table
+          await waitForDBWrite(
+            devProcess!,
+            "UserEventOutput",
+            1,
+            60_000,
+            "local",
+            `user_id = '${userId}'`,
+          );
+
+          // Verify the data was written correctly
+          const client = createClient(CLICKHOUSE_CONFIG);
+          const result = await client.query({
+            query: `
+              SELECT 
+                user_id,
+                event_name,
+                org_id,
+                properties
+              FROM local.UserEventOutput 
+              WHERE user_id = '${userId}'
+            `,
+            format: "JSONEachRow",
+          });
+
+          const rows: any[] = await result.json();
+          await client.close();
+
+          if (rows.length === 0) {
+            throw new Error(
+              `No data found in UserEventOutput for user_id ${userId}`,
+            );
+          }
+
+          const row = rows[0];
+
+          // Verify known fields are correctly passed through (snake_case for Python)
+          if (row.event_name !== "page_view") {
+            throw new Error(
+              `Expected event_name to be 'page_view', got '${row.event_name}'`,
+            );
+          }
+          if (row.org_id !== "org-123") {
+            throw new Error(
+              `Expected org_id to be 'org-123', got '${row.org_id}'`,
+            );
+          }
+
+          // Verify extra fields are stored in the properties JSON column
+          if (row.properties === undefined) {
+            throw new Error("Expected properties JSON column to exist");
+          }
+
+          // Parse properties if it's a string (ClickHouse may return JSON as string)
+          const properties =
+            typeof row.properties === "string" ?
+              JSON.parse(row.properties)
+            : row.properties;
+
+          // Verify extra fields were received by streaming function via model_extra
+          if (properties.customProperty !== "custom-value") {
+            throw new Error(
+              `Expected properties.customProperty to be 'custom-value', got '${properties.customProperty}'. Properties: ${JSON.stringify(properties)}`,
+            );
+          }
+          if (properties.pageUrl !== "/dashboard") {
+            throw new Error(
+              `Expected properties.pageUrl to be '/dashboard', got '${properties.pageUrl}'`,
+            );
+          }
+          // Note: ClickHouse JSON may return numbers as strings
+          if (Number(properties.sessionDuration) !== 120) {
+            throw new Error(
+              `Expected properties.sessionDuration to be 120, got '${properties.sessionDuration}'`,
+            );
+          }
+          if (
+            !properties.nested ||
+            properties.nested.level1 !== "value1" ||
+            !properties.nested.level2 ||
+            properties.nested.level2.deep !== "nested"
+          ) {
+            throw new Error(
+              `Expected nested object to be preserved, got '${JSON.stringify(properties.nested)}'`,
+            );
+          }
+
+          console.log(
+            "✅ Extra fields test passed (Python) - extra fields received by streaming function via model_extra and stored in properties column",
+          );
+        });
+
         // DateTime precision test for Python
         it("should preserve microsecond precision with clickhouse_datetime64 annotations via streaming transform (PY)", async function () {
           this.timeout(TIMEOUTS.TEST_SETUP_MS);
diff --git a/apps/framework-cli-e2e/test/utils/schema-definitions.ts b/apps/framework-cli-e2e/test/utils/schema-definitions.ts
@@ -936,6 +936,19 @@ export const PYTHON_TEST_SCHEMAS: ExpectedTableSchema[] = [
       },
     ],
   },
+  // Extra fields test table (ENG-1617)
+  // Extra fields from Pydantic's extra='allow' are stored in properties JSON column
+  {
+    tableName: "UserEventOutput",
+    columns: [
+      { name: "timestamp", type: /DateTime\('UTC'\)/ },
+      { name: "event_name", type: "String" },
+      { name: "user_id", type: "String" },
+      { name: "org_id", type: "Nullable(String)" },
+      { name: "project_id", type: "Nullable(String)" },
+      { name: "properties", type: "JSON" },
+    ],
+  },
 ];
 
 // ============ HELPER FUNCTIONS ============
diff --git a/packages/ts-moose-lib/src/dmv2/sdk/ingestApi.ts b/packages/ts-moose-lib/src/dmv2/sdk/ingestApi.ts
@@ -39,7 +39,12 @@ export class IngestApi<T> extends TypedBase<T, IngestConfig<T>> {
    */
   constructor(name: string, config?: IngestConfig<T>);
 
-  /** @internal **/
+  /**
+   * @internal
+   * Note: `validators` parameter is a positional placeholder (always undefined for IngestApi).
+   * It exists because TypedBase has validators as the 5th param, and we need to pass
+   * allowExtraFields as the 6th param. IngestApi doesn't use validators.
+   */
   constructor(
     name: string,
     config: IngestConfig<T>,
diff --git a/packages/ts-moose-lib/src/dmv2/sdk/stream.ts b/packages/ts-moose-lib/src/dmv2/sdk/stream.ts
@@ -228,7 +228,12 @@ export class Stream<T> extends TypedBase<T, StreamConfig<T>> {
    */
   constructor(name: string, config?: StreamConfig<T>);
 
-  /** @internal **/
+  /**
+   * @internal
+   * Note: `validators` parameter is a positional placeholder (always undefined for Stream).
+   * It exists because TypedBase has validators as the 5th param, and we need to pass
+   * allowExtraFields as the 6th param. Stream doesn't use validators.
+   */
   constructor(
     name: string,
     config: StreamConfig<T>,
diff --git a/templates/python-tests/src/ingest/models.py b/templates/python-tests/src/ingest/models.py
@@ -3,7 +3,8 @@
 from moose_lib import Point, Ring, LineString, MultiLineString, Polygon, MultiPolygon
 from moose_lib import Key, IngestPipeline, IngestPipelineConfig, StringToEnumMixin, clickhouse_default, ClickHouseCodec, OlapTable, \
     OlapConfig, MergeTreeEngine, ReplacingMergeTreeEngine, AggregatingMergeTreeEngine, simple_aggregated, \
-    ClickhouseSize, UInt8, UInt16, UInt32, UInt64, Int8, Int16, Int32, Int64, Float32, Float64, ClickhousePrecision
+    ClickhouseSize, UInt8, UInt16, UInt32, UInt64, Int8, Int16, Int32, Int64, Float32, Float64, ClickhousePrecision, \
+    IngestApi, IngestConfigWithDestination
 from datetime import datetime, date
 from typing import Optional, Annotated, Any
 from pydantic import BaseModel, BeforeValidator, ConfigDict
@@ -778,3 +779,78 @@ class MaterializedTest(BaseModel):
         dead_letter_queue=True
     )
 )
+
+
+# =======Extra Fields Test (ENG-1617)=========
+# Tests the ability to accept arbitrary payload fields using Pydantic's extra='allow'
+# This is the Python equivalent of TypeScript's index signatures
+#
+# KEY CONCEPTS:
+# - IngestApi and Stream: CAN allow extra fields (accept variable payloads)
+# - OlapTable: Requires fixed schema (ClickHouse needs to know the columns)
+# - Transform functions: Receive ALL fields (including extras)
+#   and must output to a fixed schema for OlapTable storage
+#
+# DATA FLOW:
+#   IngestApi (variable) → Stream (variable) → Transform → Stream (fixed) → OlapTable (fixed)
+
+class UserEventInput(BaseModel):
+    """Input type that allows arbitrary additional fields.
+    
+    Known fields (timestamp, event_name, user_id) are validated.
+    Unknown fields are passed through to streaming functions via extra='allow'.
+    """
+    model_config = ConfigDict(extra='allow')  # Allow arbitrary extra fields
+    
+    timestamp: datetime
+    event_name: str
+    user_id: Key[str]
+    # Optional known fields
+    org_id: Optional[str] = None
+    project_id: Optional[str] = None
+
+
+class UserEventOutput(BaseModel):
+    """Output type with a FIXED schema for OlapTable storage.
+    
+    Extra fields from the input are stored in a JSON column.
+    OlapTable requires fixed columns - ClickHouse needs to know the schema.
+    """
+    timestamp: datetime
+    event_name: str
+    user_id: Key[str]
+    org_id: Optional[str] = None
+    project_id: Optional[str] = None
+    # JSON column for extra properties
+    # This is how you persist variable fields to ClickHouse
+    properties: dict[str, Any]
+
+
+# Input stream for raw events (accepts variable fields via extra='allow')
+user_event_input_stream = Stream[UserEventInput](
+    "UserEventInput",
+    StreamConfig(destination=None)
+)
+
+# IngestApi accepting arbitrary payload fields
+user_event_ingest_api = IngestApi[UserEventInput](
+    "userEventIngestApi",
+    IngestConfigWithDestination(
+        destination=user_event_input_stream,
+        version="0.1"
+    )
+)
+
+# Output table with fixed schema (JSON column stores variable data)
+user_event_output_table = OlapTable[UserEventOutput](
+    "UserEventOutput",
+    OlapConfig(
+        order_by_fields=["user_id", "timestamp"]
+    )
+)
+
+# Stream for processed events (fixed schema)
+user_event_output_stream = Stream[UserEventOutput](
+    "UserEventOutput",
+    StreamConfig(destination=user_event_output_table)
+)
diff --git a/templates/python-tests/src/ingest/transforms.py b/templates/python-tests/src/ingest/transforms.py
@@ -146,3 +146,61 @@ def datetime_precision_transform(input_data: DateTimePrecisionTestData) -> DateT
     destination=datetime_precision_output_stream,
     transformation=datetime_precision_transform,
 )
+
+
+# ============================================================================
+# Extra Fields Transform (ENG-1617)
+# Demonstrates how streaming functions receive extra fields from Pydantic's extra='allow'
+# ============================================================================
+
+from src.ingest.models import (
+    user_event_input_stream,
+    user_event_output_stream,
+    UserEventInput,
+    UserEventOutput,
+)
+
+
+def user_event_transform(input_data: UserEventInput) -> UserEventOutput:
+    """Transform for extra fields types.
+    
+    KEY POINT: The streaming function receives ALL fields from the input,
+    including extra fields allowed by Pydantic's `model_config = ConfigDict(extra='allow')`.
+    
+    Since OlapTable requires a fixed schema, we extract known fields and store
+    extra fields in a JSON column (`properties`).
+    """
+    # Get known fields
+    known_fields = {
+        "timestamp": input_data.timestamp,
+        "event_name": input_data.event_name,
+        "user_id": input_data.user_id,
+        "org_id": input_data.org_id,
+        "project_id": input_data.project_id,
+    }
+    
+    # Get extra fields using model_extra (Pydantic v2)
+    # This contains all fields that were not defined in the model schema
+    extra_fields = input_data.model_extra or {}
+    
+    # Log to demonstrate that extra fields ARE received by the streaming function
+    print("Extra fields transform received:", {
+        "known_fields": known_fields,
+        "extra_fields": extra_fields,  # These came through extra='allow'!
+    })
+    
+    return UserEventOutput(
+        timestamp=input_data.timestamp,
+        event_name=input_data.event_name,
+        user_id=input_data.user_id,
+        org_id=input_data.org_id,
+        project_id=input_data.project_id,
+        # Store extra fields in JSON column for persistence to ClickHouse
+        properties=extra_fields,
+    )
+
+
+user_event_input_stream.add_transform(
+    destination=user_event_output_stream,
+    transformation=user_event_transform,
+)
