Skip to content

Commit ee80bd1

Browse files
authored
improvement: handle nested sql statements for sql smart cells (#6947)
Fixes #6882. This handles nested SQL in the SQL parsing logic for sql smart cells: e.g. ```sql SELECT id AS idid, value as valval FROM sample_data WHERE id IN ({",".join(df["id"][0:2].to_list())}) ```
1 parent 7c5309d commit ee80bd1

File tree

2 files changed

+67
-1
lines changed

2 files changed

+67
-1
lines changed

packages/smart-cells/src/__tests__/sql-parser.test.ts

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,51 @@ ORDER BY price DESC`.trim(),
7676
expect(metadata.engine).toBe("sqlite");
7777
});
7878

79+
it("should handle f-strings with complex expressions containing quotes", () => {
80+
const pythonCode = `_df = mo.sql(
81+
f"""
82+
SELECT
83+
id AS idid,
84+
value as valval
85+
FROM
86+
sample_data
87+
WHERE
88+
id IN ({",".join(df["id"][0:2].to_list())})
89+
"""
90+
)`;
91+
const { code, metadata } = parser.transformIn(pythonCode);
92+
expect(code).toBe(
93+
`SELECT
94+
id AS idid,
95+
value as valval
96+
FROM
97+
sample_data
98+
WHERE
99+
id IN ({",".join(df["id"][0:2].to_list())})`.trim(),
100+
);
101+
expect(metadata.dataframeName).toBe("_df");
102+
});
103+
104+
it("should handle f-strings with method calls containing strings", () => {
105+
const pythonCode = `_df = mo.sql(f"""SELECT * FROM table WHERE col = {get_value("test")}""")`;
106+
const { code, metadata } = parser.transformIn(pythonCode);
107+
expect(code).toBe(`SELECT * FROM table WHERE col = {get_value("test")}`);
108+
expect(metadata.dataframeName).toBe("_df");
109+
});
110+
111+
it("should handle f-strings with nested brackets and quotes", () => {
112+
const pythonCode = `result = mo.sql(f"""
113+
SELECT * FROM users
114+
WHERE id IN ({str(data["items"][0]["id"])})
115+
""")`;
116+
const { code, metadata } = parser.transformIn(pythonCode);
117+
expect(code).toBe(
118+
`SELECT * FROM users
119+
WHERE id IN ({str(data["items"][0]["id"])})`.trim(),
120+
);
121+
expect(metadata.dataframeName).toBe("result");
122+
});
123+
79124
it("should handle empty SQL string", () => {
80125
const pythonCode = 'next_df = mo.sql("")';
81126
const { code, offset } = parser.transformIn(pythonCode);
@@ -185,6 +230,22 @@ ORDER BY price DESC`.trim(),
185230
'# multiple\n# comments\ndf = mo.sql("")',
186231
'df = mo.sql("""SELECT 1""", output=True)',
187232
'df = mo.sql("""SELECT 1""", engine=postgres)',
233+
// Complex f-string with quotes inside
234+
`_df = mo.sql(
235+
f"""
236+
SELECT
237+
id AS idid,
238+
value as valval
239+
FROM
240+
sample_data
241+
WHERE
242+
id IN ({",".join(df["id"][0:2].to_list())})
243+
"""
244+
)`,
245+
// F-string with method call containing string
246+
`_df = mo.sql(f"""SELECT * FROM table WHERE col = {get_value("test")}""")`,
247+
// F-string with nested brackets and quotes
248+
`result = mo.sql(f"""SELECT * FROM users WHERE id IN ({str(data["items"][0]["id"])})""")`,
188249
];
189250

190251
for (const pythonCode of validCases) {

packages/smart-cells/src/parsers/sql-parser.ts

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,8 @@ function parseSQLStatement(code: string): SQLParseInfo | null {
222222
}
223223

224224
if (!isMoSql) {
225-
return null;
225+
// Skip non-mo.sql calls (e.g., embedded expressions in f-strings)
226+
continue;
226227
}
227228

228229
// Move to arguments
@@ -260,6 +261,10 @@ function parseSQLStatement(code: string): SQLParseInfo | null {
260261
break;
261262
}
262263
}
264+
265+
// Break after processing the mo.sql call to avoid processing
266+
// embedded expressions in f-strings
267+
break;
263268
}
264269
}
265270

0 commit comments

Comments
 (0)