feat: added enum completion (#96)

* feat: added enum completion

* chore: remove console.log

* chore: added guard clause on enumSchema.length

* chore: included package.json

* chore: moved replaceAll to top level

* chore: improved operator checking logic

* feat: handle enum inside the IN (...)

---------

Co-authored-by: Visal .In <[email protected]>

Author: rin-yato

package.json

@@ -68,6 +68,7 @@
     "electron-updater": "^5.3.0",
     "mysql2": "^3.5.0",
     "node-sql-parser": "^4.7.0",
+    "query-master-lang-sql": "^1.0.2",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
     "react-router-dom": "^6.8.1",

src/renderer/components/CodeEditor/SqlCodeEditor.tsx

@@ -6,20 +6,43 @@ import {
   acceptCompletion,
   completionStatus,
   startCompletion,
+  autocompletion,
+  CompletionContext,
+  CompletionResult,
 } from '@codemirror/autocomplete';
 import { defaultKeymap, insertTab } from '@codemirror/commands';
 import { keymap } from '@codemirror/view';
-import { SQLConfig, sql, MySQL } from '@codemirror/lang-sql';
-import { Ref, forwardRef } from 'react';
+import { Ref, forwardRef, useCallback } from 'react';
 import useCodeEditorTheme from './useCodeEditorTheme';
+import type { EnumSchema } from 'renderer/screens/DatabaseScreen/QueryWindow';
+import { SyntaxNode } from '@lezer/common';
+import {
+  SQLConfig,
+  sql,
+  MySQL,
+  genericCompletion,
+  keywordCompletionSource,
+  schemaCompletionSource,
+} from 'query-master-lang-sql';
+import handleCustomSqlAutoComplete from './handleCustomSqlAutoComplete';
 
 const SqlCodeEditor = forwardRef(function SqlCodeEditor(
-  props: ReactCodeMirrorProps & { schema: SQLConfig['schema'] },
+  props: ReactCodeMirrorProps & {
+    schema: SQLConfig['schema'];
+    enumSchema: EnumSchema;
+  },
   ref: Ref<ReactCodeMirrorRef>
 ) {
-  const { schema, ...codeMirrorProps } = props;
+  const { schema, enumSchema, ...codeMirrorProps } = props;
   const theme = useCodeEditorTheme();
 
+  const enumCompletion = useCallback(
+    (context: CompletionContext, tree: SyntaxNode): CompletionResult | null => {
+      return handleCustomSqlAutoComplete(context, tree, enumSchema);
+    },
+    [enumSchema]
+  );
+
   return (
     <CodeMirror
       ref={ref}
@@ -50,7 +73,13 @@ const SqlCodeEditor = forwardRef(function SqlCodeEditor(
         ]),
         sql({
           dialect: MySQL,
-          schema,
+        }),
+        autocompletion({
+          override: [
+            keywordCompletionSource(MySQL),
+            schemaCompletionSource({ schema }),
+            genericCompletion(enumCompletion),
+          ],
         }),
       ]}
       {...codeMirrorProps}

src/renderer/components/CodeEditor/handleCustomSqlAutoComplete.ts

@@ -0,0 +1,99 @@
+import { SyntaxNode } from '@lezer/common';
+import { CompletionContext, CompletionResult } from '@codemirror/autocomplete';
+import { EnumSchema } from 'renderer/screens/DatabaseScreen/QueryWindow';
+
+function getNodeString(context: CompletionContext, node: SyntaxNode) {
+  return context.state.doc.sliceString(node.from, node.to);
+}
+
+function allowNodeWhenSearchForIdentify(
+  context: CompletionContext,
+  node: SyntaxNode
+) {
+  if (node.type.name === 'Operator') return true;
+  if (node.type.name === 'Keyword') {
+    return ['IN'].includes(getNodeString(context, node).toUpperCase());
+  }
+  return false;
+}
+
+function searchForIdentifier(
+  context: CompletionContext,
+  node: SyntaxNode
+): string | null {
+  let currentNode = node.prevSibling;
+  while (currentNode) {
+    if (['CompositeIdentifier', 'Identifier'].includes(currentNode.type.name)) {
+      return getNodeString(context, currentNode);
+    } else if (!allowNodeWhenSearchForIdentify(context, node)) {
+      return null;
+    }
+
+    currentNode = node.prevSibling;
+  }
+
+  return null;
+}
+
+function handleEnumAutoComplete(
+  context: CompletionContext,
+  node: SyntaxNode,
+  enumSchema: EnumSchema
+): CompletionResult | null {
+  let currentNode = node;
+
+  if (currentNode.type.name !== 'String') {
+    return null;
+  }
+
+  // This will handle
+  // SELECT * FROM tblA WHERE tblA.colA IN (....)
+  if (currentNode?.parent?.type?.name === 'Parens') {
+    currentNode = currentNode.parent;
+  }
+
+  if (!currentNode.prevSibling) return null;
+
+  // Let search for identifer
+  const identifier = searchForIdentifier(context, currentNode.prevSibling);
+  if (!identifier) return null;
+
+  const [table, column] = identifier.replaceAll('`', '').split('.');
+  if (!table) return null;
+
+  const enumValues = enumSchema.find((tempEnum) => {
+    if (column) {
+      // normally column will be enum
+      return tempEnum.column === column;
+    } else {
+      // when column is a keyword, the node will be counted as 2
+      // so there will be no column, the enum will be in table variable instead
+      return tempEnum.column === table;
+    }
+  })?.values;
+
+  if (!enumValues) return null;
+
+  const options: CompletionResult['options'] = enumValues.map((value) => ({
+    label: value,
+    displayLabel: value,
+    type: 'keyword',
+  }));
+
+  return {
+    from: node.from + 1,
+    to: node.to - 1,
+    options,
+  };
+}
+
+export default function handleCustomSqlAutoComplete(
+  context: CompletionContext,
+  tree: SyntaxNode,
+  enumSchema: EnumSchema
+): CompletionResult | null {
+  // dont run if there is no enumSchema
+  if (enumSchema.length === 0) return null;
+
+  return handleEnumAutoComplete(context, tree, enumSchema);
+}

src/renderer/screens/DatabaseScreen/QueryWindow.tsx

@@ -21,6 +21,12 @@ import { transformResultHeaderUseSchema } from 'libs/TransformResult';
 import { SqlStatementResult } from 'libs/SqlRunnerManager';
 import { EditorState } from '@codemirror/state';
 
+export type EnumSchema = Array<{
+  table: string;
+  column: string;
+  values: string[];
+}>;
+
 interface QueryWindowProps {
   initialSql?: string;
   initialRun?: boolean;
@@ -34,13 +40,16 @@ export default function QueryWindow({
   tabKey,
 }: QueryWindowProps) {
   const editorRef = useRef<ReactCodeMirrorRef>(null);
-  const { selectedTab, setTabData, saveWindowTabHistory } = useWindowTab();
+
+  const [loading, setLoading] = useState(false);
+  const [queryKeyCounter, setQueryKeyCounter] = useState(0);
+  const [result, setResult] = useState<SqlStatementResult[]>([]);
+
   const { runner } = useSqlExecute();
   const { showErrorDialog } = useDialog();
-  const [result, setResult] = useState<SqlStatementResult[]>([]);
-  const [queryKeyCounter, setQueryKeyCounter] = useState(0);
-  const [loading, setLoading] = useState(false);
   const { schema, currentDatabase } = useSchema();
+  const { selectedTab, setTabData, saveWindowTabHistory } = useWindowTab();
+
   const codeMirrorSchema = useMemo(() => {
     return currentDatabase && schema
       ? Object.values(schema[currentDatabase].tables).reduce(
@@ -52,6 +61,27 @@ export default function QueryWindow({
         )
       : {};
   }, [schema, currentDatabase]);
+
+  const enumSchema = useMemo(() => {
+    if (!schema || !currentDatabase) return [];
+
+    const results: EnumSchema = [];
+
+    for (const table of Object.values(schema[currentDatabase].tables)) {
+      for (const column of Object.values(table.columns)) {
+        if (column.dataType === 'enum') {
+          results.push({
+            table: table.name,
+            column: column.name,
+            values: column.enumValues || [],
+          });
+        }
+      }
+    }
+
+    return results;
+  }, [schema, currentDatabase]);
+
   const [code, setCode] = useState(initialSql || '');
 
   const { handleContextMenu } = useContextMenu(() => {
@@ -237,6 +267,7 @@ export default function QueryWindow({
             }}
             height="100%"
             schema={codeMirrorSchema}
+            enumSchema={enumSchema}
           />
         </div>
       </div>