feat: add experiment statement gutter highlight (#155)

* add experiment statement gutter highlight

* initial highlight code

* change how we select statement

package.json

@@ -94,7 +94,6 @@
     "react-resizable-panels": "^1.0.9",
     "sonner": "^1.4.41",
     "sql-formatter": "^15.3.2",
-    "sql-query-identifier": "^2.6.0",
     "tailwind-merge": "^2.2.2",
     "tailwindcss-animate": "^1.0.7",
     "zod": "^3.22.4"

src/components/gui/sql-editor/index.tsx

@@ -21,6 +21,7 @@ import { sqliteDialect } from "@/drivers/sqlite/sqlite-dialect";
 import { functionTooltip } from "./function-tooltips";
 import sqliteFunctionList from "@/drivers/sqlite/function-tooltip.json";
 import { toast } from "sonner";
+import SqlStatementHighlightPlugin from "./statement-highlight";
 import { SupportedDialect } from "@/drivers/base-driver";
 
 interface SqlEditorProps {
@@ -142,10 +143,17 @@ const SqlEditor = forwardRef<ReactCodeMirrorRef, SqlEditorProps>(
       }
 
       return [
+        EditorView.baseTheme({
+          "& .cm-line": {
+            borderLeft: "3px solid transparent",
+            paddingLeft: "10px",
+          },
+        }),
         keyExtensions,
         sqlDialect,
         tooltipExtension,
         tableNameHighlightPlugin,
+        SqlStatementHighlightPlugin,
         EditorView.updateListener.of((state) => {
           const pos = state.state.selection.main.head;
           const line = state.state.doc.lineAt(pos);

src/components/gui/sql-editor/statement-highlight.test.ts

@@ -0,0 +1,155 @@
+import { MySQL, SQLite } from "@codemirror/lang-sql";
+import { EditorState } from "@codemirror/state";
+import { splitSqlQuery } from "./statement-highlight";
+
+function sqlite(code: string) {
+  const state = EditorState.create({ doc: code, extensions: [SQLite] });
+  return splitSqlQuery(state).map((p) => p.text);
+}
+
+function mysql(code: string) {
+  const state = EditorState.create({ doc: code, extensions: [MySQL] });
+  return splitSqlQuery(state).map((p) => p.text);
+}
+
+describe("split sql statements", () => {
+  test("should parse a query with different statements in a single line", () => {
+    expect(
+      sqlite(
+        `INSERT INTO Persons (PersonID, Name) VALUES (1, 'Jack');SELECT * FROM Persons`
+      )
+    ).toEqual([
+      `INSERT INTO Persons (PersonID, Name) VALUES (1, 'Jack');`,
+      `SELECT * FROM Persons`,
+    ]);
+  });
+
+  test("should identify a query with different statements in multiple lines", () => {
+    expect(
+      sqlite(`
+        INSERT INTO Persons (PersonID, Name) VALUES (1, 'Jack');
+        SELECT * FROM Persons';
+      `)
+    ).toEqual([
+      `INSERT INTO Persons (PersonID, Name) VALUES (1, 'Jack');`,
+      `SELECT * FROM Persons';\n      `,
+    ]);
+  });
+
+  test("sholud be able to split statement with BEGIN and END", () => {
+    expect(
+      sqlite(`CREATE TABLE customer(
+  cust_id INTEGER PRIMARY KEY,
+  cust_name TEXT,
+  cust_addr TEXT
+);
+
+-- some comment here that should be ignore
+
+
+CREATE VIEW customer_address AS
+   SELECT cust_id, cust_addr FROM customer;
+CREATE TRIGGER cust_addr_chng
+INSTEAD OF UPDATE OF cust_addr ON customer_address
+BEGIN
+  UPDATE customer SET cust_addr=NEW.cust_addr
+   WHERE cust_id=NEW.cust_id;
+END ;`)
+    ).toEqual([
+      `CREATE TABLE customer(\n  cust_id INTEGER PRIMARY KEY,\n  cust_name TEXT,\n  cust_addr TEXT\n);`,
+      `CREATE VIEW customer_address AS\n   SELECT cust_id, cust_addr FROM customer;`,
+      `CREATE TRIGGER cust_addr_chng\nINSTEAD OF UPDATE OF cust_addr ON customer_address\nBEGIN\n  UPDATE customer SET cust_addr=NEW.cust_addr\n   WHERE cust_id=NEW.cust_id;\nEND ;`,
+    ]);
+  });
+
+  test("should be able to split statement with BEGIN and END and CONDITION inside", () => {
+    expect(
+      mysql(`CREATE TRIGGER upd_check BEFORE UPDATE ON account
+FOR EACH ROW
+BEGIN
+    IF NEW.amount < 0 THEN
+        SET NEW.amount = 0;
+    ELSEIF NEW.amount > 100 THEN
+        SET NEW.amount = 100;
+    END IF;
+END; SELECT * FROM hello`)
+    ).toEqual([
+      `CREATE TRIGGER upd_check BEFORE UPDATE ON account\nFOR EACH ROW\nBEGIN\n    IF NEW.amount < 0 THEN\n        SET NEW.amount = 0;\n    ELSEIF NEW.amount > 100 THEN\n        SET NEW.amount = 100;\n    END IF;\nEND;`,
+      "SELECT * FROM hello",
+    ]);
+  });
+
+  test("should be able to split statement with BEGIN with no end", () => {
+    expect(
+      mysql(`SELECT * FROM outerbase; CREATE TRIGGER upd_check BEFORE UPDATE ON account
+FOR EACH ROW
+BEGIN
+    IF NEW.amount < 0 THEN
+        SET NEW.amount = 0;
+    ELSEIF NEW.amount > 100 THEN
+        SET NEW.amount = 100;`)
+    ).toEqual([
+      "SELECT * FROM outerbase;",
+      `CREATE TRIGGER upd_check BEFORE UPDATE ON account
+FOR EACH ROW
+BEGIN
+    IF NEW.amount < 0 THEN
+        SET NEW.amount = 0;
+    ELSEIF NEW.amount > 100 THEN
+        SET NEW.amount = 100;`,
+    ]);
+  });
+
+  test("should be able to split TRIGGER without begin", () => {
+    expect(
+      mysql(`create trigger hire_log after insert on employees 
+for each row insert into hiring values (new.id, current_time());
+    
+insert into employees (first_name, last_name) values ("Tim", "Sehn");`)
+    ).toEqual([
+      `create trigger hire_log after insert on employees \nfor each row insert into hiring values (new.id, current_time());`,
+      `insert into employees (first_name, last_name) values ("Tim", "Sehn");`,
+    ]);
+  });
+
+  test("should be able to split nested BEGIN", () => {
+    expect(
+      mysql(
+        `CREATE PROCEDURE procCreateCarTable
+IS
+BEGIN
+  BEGIN
+    EXECUTE IMMEDIATE 'DROP TABLE CARS';
+    EXCEPTION WHEN OTHERS THEN NULL;
+    EXECUTE IMMEDIATE 'CREATE TABLE CARS (ID VARCHAR2(1), NAME VARCHAR2(10), TITLE    
+    VARCHAR2(10))';
+  END;
+  BEGIN
+    EXECUTE IMMEDIATE 'DROP TABLE TRUCKS';
+    EXCEPTION WHEN OTHERS THEN NULL;
+    EXECUTE IMMEDIATE 'CREATE TABLE TRUCKS (ID VARCHAR2(1), NAME VARCHAR2(10), TITLE 
+    VARCHAR2(10))';
+  END;
+END; SELECT * FROM outeerbase;`
+      )
+    ).toEqual([
+      `CREATE PROCEDURE procCreateCarTable
+IS
+BEGIN
+  BEGIN
+    EXECUTE IMMEDIATE 'DROP TABLE CARS';
+    EXCEPTION WHEN OTHERS THEN NULL;
+    EXECUTE IMMEDIATE 'CREATE TABLE CARS (ID VARCHAR2(1), NAME VARCHAR2(10), TITLE    
+    VARCHAR2(10))';
+  END;
+  BEGIN
+    EXECUTE IMMEDIATE 'DROP TABLE TRUCKS';
+    EXCEPTION WHEN OTHERS THEN NULL;
+    EXECUTE IMMEDIATE 'CREATE TABLE TRUCKS (ID VARCHAR2(1), NAME VARCHAR2(10), TITLE 
+    VARCHAR2(10))';
+  END;
+END;`,
+      "SELECT * FROM outeerbase;",
+    ]);
+  });
+});

src/components/gui/sql-editor/statement-highlight.ts

@@ -0,0 +1,172 @@
+import {
+  Decoration,
+  EditorState,
+  EditorView,
+  StateField,
+  Range,
+} from "@uiw/react-codemirror";
+import { syntaxTree } from "@codemirror/language";
+import { SyntaxNode } from "@lezer/common";
+
+const statementLineHighlight = Decoration.line({
+  class: "cm-highlight-statement",
+});
+
+export interface StatementSegment {
+  from: number;
+  to: number;
+  text: string;
+}
+
+function toNodeString(state: EditorState, node: SyntaxNode) {
+  return state.doc.sliceString(node.from, node.to);
+}
+
+function isRequireEndStatement(state: EditorState, node: SyntaxNode): number {
+  const ptr = node.firstChild;
+  if (!ptr) return 0;
+
+  // Majority of the query will fall in SELECT, INSERT, UPDATE, DELETE
+  const firstKeyword = toNodeString(state, ptr).toLowerCase();
+  if (firstKeyword === "select") return 0;
+  if (firstKeyword === "insert") return 0;
+  if (firstKeyword === "update") return 0;
+  if (firstKeyword === "delete") return 0;
+
+  const keywords = node.getChildren("Keyword");
+  if (keywords.length === 0) return 0;
+
+  return keywords.filter(
+    (k) => toNodeString(state, k).toLowerCase() === "begin"
+  ).length;
+}
+
+function isEndStatement(state: EditorState, node: SyntaxNode) {
+  let ptr = node.firstChild;
+  if (!ptr) return false;
+  if (toNodeString(state, ptr).toLowerCase() !== "end") return false;
+
+  ptr = ptr.nextSibling;
+  if (!ptr) return false;
+  if (toNodeString(state, ptr) !== ";") return false;
+
+  return true;
+}
+
+export function splitSqlQuery(
+  state: EditorState,
+  generateText: boolean = true
+): StatementSegment[] {
+  const topNode = syntaxTree(state).topNode;
+
+  // Get all the statements
+  let needEndStatementCounter = 0;
+  const statements = topNode.getChildren("Statement");
+
+  if (statements.length === 0) return [];
+
+  const statementGroups: SyntaxNode[][] = [];
+  let accumulateNodes: SyntaxNode[] = [];
+  let i = 0;
+
+  for (; i < statements.length; i++) {
+    const statement = statements[i];
+    needEndStatementCounter += isRequireEndStatement(state, statement);
+
+    if (needEndStatementCounter) {
+      accumulateNodes.push(statement);
+    } else {
+      statementGroups.push([statement]);
+    }
+
+    if (needEndStatementCounter && isEndStatement(state, statement)) {
+      needEndStatementCounter--;
+      if (needEndStatementCounter === 0) {
+        statementGroups.push(accumulateNodes);
+        accumulateNodes = [];
+      }
+    }
+  }
+
+  if (accumulateNodes.length > 0) {
+    statementGroups.push(accumulateNodes);
+  }
+
+  return statementGroups.map((r) => ({
+    from: r[0].from,
+    to: r[r.length - 1].to,
+    text: generateText
+      ? state.doc.sliceString(r[0].from, r[r.length - 1].to)
+      : "",
+  }));
+}
+
+export function resolveToNearestStatement(
+  state: EditorState
+): { from: number; to: number } | null {
+  // Breakdown and grouping the statement
+  const cursor = state.selection.main.from;
+  const statements = splitSqlQuery(state, false);
+
+  if (statements.length === 0) return null;
+
+  // Check if our current cursor is within any statement
+  let i = 0;
+  for (; i < statements.length; i++) {
+    const statement = statements[i];
+    if (cursor < statement.from) break;
+    if (cursor > statement.to) continue;
+    if (cursor >= statement.from && cursor <= statement.to) return statement;
+  }
+
+  if (i === 0) return statements[0];
+  if (i === statements.length) return statements[i - 1];
+
+  const cursorLine = state.doc.lineAt(cursor).number;
+  const topLine = state.doc.lineAt(statements[i - 1].to).number;
+  const bottomLine = state.doc.lineAt(statements[i].from).number;
+
+  if (cursorLine - topLine >= bottomLine - cursorLine) {
+    return statements[i];
+  } else {
+    return statements[i - 1];
+  }
+}
+function getDecorationFromState(state: EditorState) {
+  const statement = resolveToNearestStatement(state);
+
+  if (!statement) return Decoration.none;
+
+  // Get the line of the node
+  const fromLineNumber = state.doc.lineAt(statement.from).number;
+  const toLineNumber = state.doc.lineAt(statement.to).number;
+
+  const d: Range<Decoration>[] = [];
+  for (let i = fromLineNumber; i <= toLineNumber; i++) {
+    d.push(statementLineHighlight.range(state.doc.line(i).from));
+  }
+
+  return Decoration.set(d);
+}
+
+const SqlStatementStateField = StateField.define({
+  create(state) {
+    return getDecorationFromState(state);
+  },
+
+  update(_, tr) {
+    return getDecorationFromState(tr.state);
+  },
+
+  provide: (f) => EditorView.decorations.from(f),
+});
+
+const SqlStatementTheme = EditorView.baseTheme({
+  ".cm-highlight-statement": {
+    borderLeft: "3px solid #ff9ff3 !important",
+  },
+});
+
+const SqlStatementHighlightPlugin = [SqlStatementStateField, SqlStatementTheme];
+
+export default SqlStatementHighlightPlugin;