primary key expression (#3031)

<!-- CURSOR_SUMMARY -->
> [!NOTE]
> Adds first-class primary key expression support across OLAP pipelines,
generating/reading PRIMARY KEY in DDL, normalizing for diffs, and
exposing config in TS/Python SDKs with tests and docs.
> 
> - **OLAP/ClickHouse core**:
> - Add `primary_key_expression` to `Table`/ClickHouse models and
protobuf; include in (de)serialization and introspection.
> - Generate PRIMARY KEY in `create_table_query` (uses expression or
column flags) and expose in list_tables via SQL parsing
(`extract_primary_key_from_create_table`).
> - Normalize PKs (strip spaces/backticks/outer parens) and compare in
diffs; PK changes trigger drop+create.
> - **Code generation**:
> - TS/Python generators output
`primaryKeyExpression`/`primary_key_expression` in table configs and
switch key-wrapping logic to use it.
> - **SDK/Internal config**:
> - TS/Python SDKs add `primaryKeyExpression`/`primary_key_expression`
to `OlapTable` config and internal representations.
> - **Tests**:
> - Extensive unit tests for PK parsing/normalization/diff behavior and
DDL generation; e2e template tests validating PRIMARY KEY and ORDER BY
in DDL.
> - **Docs**:
> - New/updated docs detailing primary key expressions, constraints, and
examples in TS/Python, including ORDER BY prefix requirement.
> 
> <sup>Written by [Cursor
Bugbot](https://cursor.com/dashboard?tab=bugbot) for commit
3bdd55e. This will update automatically
on new commits. Configure
[here](https://cursor.com/dashboard?tab=bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

Author: phiSgr

apps/framework-cli-e2e/test/templates.test.ts

@@ -292,6 +292,54 @@ const createTemplateTestSuite = (config: TemplateTestConfig) => {
       }
     });
 
+    it("should include PRIMARY KEY expression in DDL when configured", async function () {
+      if (config.isTestsVariant) {
+        // Test 1: Primary key with hash function
+        const ddl1 = await getTableDDL("PrimaryKeyExpressionTest", "local");
+        const primaryKeyPattern =
+          config.language === "typescript" ?
+            "PRIMARY KEY (userId, cityHash64(eventId))"
+          : "PRIMARY KEY (user_id, cityHash64(event_id))";
+        const orderByPattern =
+          config.language === "typescript" ?
+            "ORDER BY (userId, cityHash64(eventId), timestamp)"
+          : "ORDER BY (user_id, cityHash64(event_id), timestamp)";
+
+        if (!ddl1.includes(primaryKeyPattern)) {
+          throw new Error(
+            `PRIMARY KEY expression not found in PrimaryKeyExpressionTest DDL. Expected: ${primaryKeyPattern}. DDL: ${ddl1}`,
+          );
+        }
+        if (!ddl1.includes(orderByPattern)) {
+          throw new Error(
+            `ORDER BY expression not found in PrimaryKeyExpressionTest DDL. Expected: ${orderByPattern}. DDL: ${ddl1}`,
+          );
+        }
+
+        // Test 2: Primary key with different ordering
+        const ddl2 = await getTableDDL("PrimaryKeyOrderingTest", "local");
+        const primaryKeyPattern2 =
+          config.language === "typescript" ?
+            "PRIMARY KEY productId"
+          : "PRIMARY KEY product_id";
+        const orderByPattern2 =
+          config.language === "typescript" ?
+            "ORDER BY (productId, category, brand)"
+          : "ORDER BY (product_id, category, brand)";
+
+        if (!ddl2.includes(primaryKeyPattern2)) {
+          throw new Error(
+            `PRIMARY KEY expression not found in PrimaryKeyOrderingTest DDL. Expected: ${primaryKeyPattern2}. DDL: ${ddl2}`,
+          );
+        }
+        if (!ddl2.includes(orderByPattern2)) {
+          throw new Error(
+            `ORDER BY expression not found in PrimaryKeyOrderingTest DDL. Expected: ${orderByPattern2}. DDL: ${ddl2}`,
+          );
+        }
+      }
+    });
+
     it("should generate FixedString types in DDL including type aliases", async function () {
       if (config.isTestsVariant && config.language === "python") {
         const ddl = await getTableDDL("FixedStringTest", "local");

apps/framework-cli-e2e/test/utils/schema-definitions.ts

@@ -421,6 +421,25 @@ export const TYPESCRIPT_TEST_SCHEMAS: ExpectedTableSchema[] = [
       { name: "payloadBasic", type: "JSON(count Int64, name String)" },
     ],
   },
+  // Primary Key Expression Tests
+  {
+    tableName: "PrimaryKeyExpressionTest",
+    columns: [
+      { name: "userId", type: "String" },
+      { name: "eventId", type: "String" },
+      { name: "timestamp", type: /DateTime\('UTC'\)/ },
+      { name: "category", type: "String" },
+    ],
+  },
+  {
+    tableName: "PrimaryKeyOrderingTest",
+    columns: [
+      { name: "productId", type: "String" },
+      { name: "category", type: "String" },
+      { name: "brand", type: "String" },
+      { name: "timestamp", type: /DateTime\('UTC'\)/ },
+    ],
+  },
 ];
 
 // ============ PYTHON TEMPLATE SCHEMA DEFINITIONS ============
@@ -805,6 +824,25 @@ export const PYTHON_TEST_SCHEMAS: ExpectedTableSchema[] = [
       { name: "payload_basic", type: "JSON(count Int64, name String)" },
     ],
   },
+  // Primary Key Expression Tests
+  {
+    tableName: "PrimaryKeyExpressionTest",
+    columns: [
+      { name: "user_id", type: "String" },
+      { name: "event_id", type: "String" },
+      { name: "timestamp", type: /DateTime\('UTC'\)/ },
+      { name: "category", type: "String" },
+    ],
+  },
+  {
+    tableName: "PrimaryKeyOrderingTest",
+    columns: [
+      { name: "product_id", type: "String" },
+      { name: "category", type: "String" },
+      { name: "brand", type: "String" },
+      { name: "timestamp", type: /DateTime\('UTC'\)/ },
+    ],
+  },
 ];
 
 // ============ HELPER FUNCTIONS ============

apps/framework-cli/src/cli/local_webserver.rs

@@ -3564,6 +3564,7 @@ mod tests {
             database: None,
             table_ttl_setting: None,
             cluster_name: None,
+            primary_key_expression: None,
         }
     }
 

apps/framework-cli/src/cli/routines/migrate.rs

@@ -781,6 +781,7 @@ mod tests {
             table_settings: None,
             table_ttl_setting: None,
             cluster_name: None,
+            primary_key_expression: None,
         }
     }
 

apps/framework-cli/src/cli/routines/seed_data.rs

@@ -614,6 +614,7 @@ mod tests {
             table_settings: None,
             table_ttl_setting: None,
             cluster_name: None,
+            primary_key_expression: None,
         }
     }
 

apps/framework-cli/src/framework/core/infra_reality_checker.rs

@@ -533,6 +533,7 @@ mod tests {
             database: None,
             table_ttl_setting: None,
             cluster_name: None,
+            primary_key_expression: None,
         }
     }
 

apps/framework-cli/src/framework/core/infrastructure/table.rs

@@ -302,6 +302,10 @@ pub struct Table {
     /// Optional cluster name for ON CLUSTER support in ClickHouse
     #[serde(skip_serializing_if = "Option::is_none", default)]
     pub cluster_name: Option<String>,
+    /// Optional PRIMARY KEY expression (overrides column-level primary_key flags when specified)
+    /// Allows for complex primary keys using functions or different column ordering
+    #[serde(skip_serializing_if = "Option::is_none", default)]
+    pub primary_key_expression: Option<String>,
 }
 
 impl Table {
@@ -402,6 +406,75 @@ impl Table {
             .collect()
     }
 
+    /// Returns a normalized representation of the primary key for comparison purposes.
+    ///
+    /// This handles both:
+    /// - `primary_key_expression`: Uses the expression directly
+    /// - Column-level `primary_key` flags: Builds an expression from column names
+    ///
+    /// The result is normalized (trimmed, spaces removed, backticks removed, and outer
+    /// parentheses stripped for single-element tuples) to enable semantic comparison.
+    /// For example:
+    /// - `primary_key_expression: Some("(foo, bar)")` returns "(foo,bar)"
+    /// - Columns foo, bar with `primary_key: true` returns "(foo,bar)"
+    /// - `primary_key_expression: Some("foo")` returns "foo"
+    /// - `primary_key_expression: Some("(foo)")` returns "foo" (outer parens stripped)
+    /// - Single column foo with `primary_key: true` returns "foo"
+    pub fn normalized_primary_key_expr(&self) -> String {
+        let expr = if let Some(ref pk_expr) = self.primary_key_expression {
+            // Use the explicit primary_key_expression
+            pk_expr.clone()
+        } else {
+            // Build from column-level primary_key flags
+            let pk_cols = self.primary_key_columns();
+            if pk_cols.is_empty() {
+                String::new()
+            } else if pk_cols.len() == 1 {
+                pk_cols[0].to_string()
+            } else {
+                format!("({})", pk_cols.join(", "))
+            }
+        };
+
+        // Normalize: trim, remove backticks, remove spaces
+        let mut normalized = expr
+            .trim()
+            .trim_matches('`')
+            .replace('`', "")
+            .replace(" ", "");
+
+        // Strip outer parentheses if this is a single-element tuple
+        // E.g., "(col)" -> "col", "(cityHash64(col))" -> "cityHash64(col)"
+        // But keep "(col1,col2)" as-is
+        if normalized.starts_with('(') && normalized.ends_with(')') {
+            // Check if there are any top-level commas (not inside nested parentheses)
+            let inner = &normalized[1..normalized.len() - 1];
+            let has_top_level_comma = {
+                let mut depth = 0;
+                let mut found_comma = false;
+                for ch in inner.chars() {
+                    match ch {
+                        '(' => depth += 1,
+                        ')' => depth -= 1,
+                        ',' if depth == 0 => {
+                            found_comma = true;
+                            break;
+                        }
+                        _ => {}
+                    }
+                }
+                found_comma
+            };
+
+            // If no top-level comma, it's a single-element tuple - strip outer parens
+            if !has_top_level_comma {
+                normalized = inner.to_string();
+            }
+        }
+
+        normalized
+    }
+
     pub fn order_by_with_fallback(&self) -> OrderBy {
         // table (in infra map created by older version of moose) may leave order_by unspecified,
         // but the implicit order_by from primary keys can be the same
@@ -473,6 +546,7 @@ impl Table {
             table_settings: self.table_settings.clone().unwrap_or_default(),
             table_ttl_setting: self.table_ttl_setting.clone(),
             cluster_name: self.cluster_name.clone(),
+            primary_key_expression: self.primary_key_expression.clone(),
             metadata: MessageField::from_option(self.metadata.as_ref().map(|m| {
                 infrastructure_map::Metadata {
                     description: m.description.clone().unwrap_or_default(),
@@ -581,6 +655,7 @@ impl Table {
             database: proto.database,
             table_ttl_setting: proto.table_ttl_setting,
             cluster_name: proto.cluster_name,
+            primary_key_expression: proto.primary_key_expression,
         }
     }
 }
@@ -1647,6 +1722,7 @@ mod tests {
             database: None,
             table_ttl_setting: None,
             cluster_name: None,
+            primary_key_expression: None,
         };
         assert_eq!(table1.id(DEFAULT_DATABASE_NAME), "local_users");
 
@@ -1776,6 +1852,7 @@ mod tests {
             database: None,
             table_ttl_setting: None,
             cluster_name: None,
+            primary_key_expression: None,
         };
 
         // Target table from code: explicit order_by that matches primary key
@@ -1799,6 +1876,7 @@ mod tests {
             database: None,
             table_ttl_setting: None,
             cluster_name: None,
+            primary_key_expression: None,
         };
 
         // These should be equal because:

apps/framework-cli/src/framework/core/infrastructure_map.rs

@@ -3065,6 +3065,7 @@ mod tests {
             database: None,
             table_ttl_setting: None,
             cluster_name: None,
+            primary_key_expression: None,
         };
 
         let after = Table {
@@ -3121,6 +3122,7 @@ mod tests {
             database: None,
             table_ttl_setting: None,
             cluster_name: None,
+            primary_key_expression: None,
         };
 
         let diff = compute_table_columns_diff(&before, &after);
@@ -3297,6 +3299,7 @@ mod diff_tests {
             database: None,
             table_ttl_setting: None,
             cluster_name: None,
+            primary_key_expression: None,
         }
     }
 

apps/framework-cli/src/framework/core/partial_infrastructure_map.rs

@@ -273,6 +273,9 @@ struct PartialTable {
     /// Optional cluster name for ON CLUSTER support
     #[serde(default)]
     pub cluster: Option<String>,
+    /// Optional PRIMARY KEY expression (overrides column-level primary_key flags when specified)
+    #[serde(default, alias = "primary_key_expression")]
+    pub primary_key_expression: Option<String>,
 }
 
 /// Represents a topic definition from user code before it's converted into a complete [`Topic`].
@@ -743,6 +746,7 @@ impl PartialInfrastructureMap {
                     table_ttl_setting,
                     database: partial_table.database.clone(),
                     cluster_name: partial_table.cluster.clone(),
+                    primary_key_expression: partial_table.primary_key_expression.clone(),
                 };
                 Ok((table.id(default_database), table))
             })

apps/framework-cli/src/framework/core/plan.rs

@@ -527,6 +527,7 @@ mod tests {
             database: None,
             table_ttl_setting: None,
             cluster_name: None,
+            primary_key_expression: None,
         }
     }