fix(sqllab): clean comments within quotes

apache · May 3, 2023 · 167d6a2 · 167d6a2
1 parent 594d3e0
commit 167d6a2
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 3 deletions.
diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.js b/superset-frontend/src/SqlLab/actions/sqlLab.js
@@ -357,12 +357,59 @@ export function fetchQueryResults(query, displayLimit) {
   };
 }
 
+const quotes = '\'"`'.split('');
+const quotedBlockHash = shortid.generate();
+const quotedBlockMatch = new RegExp(`${quotedBlockHash}:\\d+:`, 'g');
+
+function splitByQuotedBlock(str) {
+  const chunks = [];
+  let currentQuote = '';
+  let chunkStart = 0;
+
+  str.split('').forEach((currentChar, i) => {
+    if (
+      (currentQuote && currentChar === currentQuote) ||
+      (!currentQuote && quotes.includes(currentChar))
+    ) {
+      if (currentQuote) {
+        chunks.push(str.substring(chunkStart, i + 1));
+        chunkStart = i + 1;
+        currentQuote = '';
+      } else {
+        chunks.push(str.substring(chunkStart, i));
+        chunkStart = i;
+        currentQuote = currentChar;
+      }
+    }
+  });
+
+  if (chunkStart < str.length) {
+    const lastChunk = str.substring(chunkStart);
+    if (lastChunk) {
+      chunks.push(lastChunk);
+    }
+  }
+
+  return chunks;
+}
+
 export function cleanSqlComments(sql) {
   if (!sql) return '';
   // it sanitizes the following comment block groups
   // group 1 -> /* */
   // group 2 -> --
-  return sql.replace(/(--.*?$|\/\*[\s\S]*?\*\/)\n?/gm, '\n').trim();
+  const chunks = splitByQuotedBlock(sql);
+  return chunks
+    .map((chunk, index) =>
+      quotes.includes(chunk[0]) ? `${quotedBlockHash}:${index}:` : chunk,
+    )
+    .join('')
+    .replace(/(--.*?$|\/\*[\s\S]*?\*\/)\n?/gm, '\n')
+    .replace(
+      quotedBlockMatch,
+      quotedBlock => chunks[quotedBlock.match(/:\d+/)[0].substring(1)],
+    )
+    .trim();
 }
 
 export function runQuery(query) {

diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.test.js b/superset-frontend/src/SqlLab/actions/sqlLab.test.js
@@ -312,7 +312,16 @@ describe('async actions', () => {
     const makeRequest = () => {
       const request = actions.runQuery({
         ...query,
-        sql: '/*\nSELECT * FROM\n */\nSELECT 213--, {{ds}}\n/*\n{{new_param1}}\n{{new_param2}}*/\n\nFROM table',
+        sql: `/*
+  SELECT * FROM
+*/
+SELECT 213--, {{ds}} "quote out"
+/*
+{{new_param1}}
+{{new_param2}}*/
+
+FROM table
+WHERE value = '--"NULL"--' --{{test_param}}`,
       });
       return request(dispatch, () => initialState);
     };
@@ -326,7 +335,7 @@ describe('async actions', () => {
         expect(fetchMock.calls(runQueryEndpoint)).toHaveLength(1);
         expect(
           JSON.parse(fetchMock.calls(runQueryEndpoint)[0][1].body).sql,
-        ).toEqual('SELECT 213\n\n\nFROM table');
+        ).toEqual(`SELECT 213\n\n\nFROM table\nWHERE value = '--"NULL"--'`);
       });
     });
   });