From b658dc75484a4b58f6109707171c4886be05e574 Mon Sep 17 00:00:00 2001
From: Myles Scolnick <myles@marimo.io>
Date: Mon, 27 Oct 2025 12:05:08 -0400
Subject: [PATCH] improvement: handle nested sql statements for sql smart cells

---
 .../src/__tests__/sql-parser.test.ts          | 61 +++++++++++++++++++
 .../smart-cells/src/parsers/sql-parser.ts     |  7 ++-
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/packages/smart-cells/src/__tests__/sql-parser.test.ts b/packages/smart-cells/src/__tests__/sql-parser.test.ts
index 72a51a0aee3..009d5aa8f94 100644
--- a/packages/smart-cells/src/__tests__/sql-parser.test.ts
+++ b/packages/smart-cells/src/__tests__/sql-parser.test.ts
@@ -76,6 +76,51 @@ ORDER BY price DESC`.trim(),
       expect(metadata.engine).toBe("sqlite");
     });
 
+    it("should handle f-strings with complex expressions containing quotes", () => {
+      const pythonCode = `_df = mo.sql(
+    f"""
+     SELECT
+        id AS idid,
+        value as valval
+    FROM
+        sample_data
+    WHERE
+        id IN ({",".join(df["id"][0:2].to_list())})
+    """
+)`;
+      const { code, metadata } = parser.transformIn(pythonCode);
+      expect(code).toBe(
+        `SELECT
+    id AS idid,
+    value as valval
+FROM
+    sample_data
+WHERE
+    id IN ({",".join(df["id"][0:2].to_list())})`.trim(),
+      );
+      expect(metadata.dataframeName).toBe("_df");
+    });
+
+    it("should handle f-strings with method calls containing strings", () => {
+      const pythonCode = `_df = mo.sql(f"""SELECT * FROM table WHERE col = {get_value("test")}""")`;
+      const { code, metadata } = parser.transformIn(pythonCode);
+      expect(code).toBe(`SELECT * FROM table WHERE col = {get_value("test")}`);
+      expect(metadata.dataframeName).toBe("_df");
+    });
+
+    it("should handle f-strings with nested brackets and quotes", () => {
+      const pythonCode = `result = mo.sql(f"""
+        SELECT * FROM users
+        WHERE id IN ({str(data["items"][0]["id"])})
+      """)`;
+      const { code, metadata } = parser.transformIn(pythonCode);
+      expect(code).toBe(
+        `SELECT * FROM users
+WHERE id IN ({str(data["items"][0]["id"])})`.trim(),
+      );
+      expect(metadata.dataframeName).toBe("result");
+    });
+
     it("should handle empty SQL string", () => {
       const pythonCode = 'next_df = mo.sql("")';
       const { code, offset } = parser.transformIn(pythonCode);
@@ -185,6 +230,22 @@ ORDER BY price DESC`.trim(),
         '# multiple\n# comments\ndf = mo.sql("")',
         'df = mo.sql("""SELECT 1""", output=True)',
         'df = mo.sql("""SELECT 1""", engine=postgres)',
+        // Complex f-string with quotes inside
+        `_df = mo.sql(
+    f"""
+     SELECT
+        id AS idid,
+        value as valval
+    FROM
+        sample_data
+    WHERE
+        id IN ({",".join(df["id"][0:2].to_list())})
+    """
+)`,
+        // F-string with method call containing string
+        `_df = mo.sql(f"""SELECT * FROM table WHERE col = {get_value("test")}""")`,
+        // F-string with nested brackets and quotes
+        `result = mo.sql(f"""SELECT * FROM users WHERE id IN ({str(data["items"][0]["id"])})""")`,
       ];
 
       for (const pythonCode of validCases) {
diff --git a/packages/smart-cells/src/parsers/sql-parser.ts b/packages/smart-cells/src/parsers/sql-parser.ts
index 9b39798cc18..d331eac9471 100644
--- a/packages/smart-cells/src/parsers/sql-parser.ts
+++ b/packages/smart-cells/src/parsers/sql-parser.ts
@@ -222,7 +222,8 @@ function parseSQLStatement(code: string): SQLParseInfo | null {
         }
 
         if (!isMoSql) {
-          return null;
+          // Skip non-mo.sql calls (e.g., embedded expressions in f-strings)
+          continue;
         }
 
         // Move to arguments
@@ -260,6 +261,10 @@ function parseSQLStatement(code: string): SQLParseInfo | null {
             break;
           }
         }
+
+        // Break after processing the mo.sql call to avoid processing
+        // embedded expressions in f-strings
+        break;
       }
     }