Merge pull request #47 from bruinenxyz/hungle/bru-1053-add-autocomple…

…te-to-sql-query-editor

BRU-1053 Add autocomplete to SQL query editor

backend/src/databases/databases.controller.ts

@@ -14,7 +14,11 @@ import { OrgGuard } from "@/auth/organizations.guard";
 import { TrackingInterceptor } from "@/interceptors/tracking/tracking.interceptor";
 import { ZodSerializerDto } from "nestjs-zod";
 import { ApiOkResponse } from "@nestjs/swagger";
-import { CleanDatabase, CleanDatabaseSchema } from "@/definitions";
+import {
+  CleanDatabase,
+  CleanDatabaseSchema,
+  ExternalColumn,
+} from "@/definitions";
 import {
   GetDatabaseResponseDto,
   GetAllDatabasesResponseDto,
@@ -74,7 +78,9 @@ export class DatabasesController {
 
   @Get("/:id/schemas")
   @UseGuards(OrgGuard("database"))
-  findAllSchemas(@Param("id") id: string): Promise<string[]> {
+  findAllSchemas(
+    @Param("id") id: string,
+  ): Promise<Record<string, Record<string, Record<string, ExternalColumn>>>> {
     return this.databasesService.findAllSchemas(id);
   }
 }

backend/src/databases/databases.service.ts

@@ -7,7 +7,12 @@ import {
 } from "@nestjs/common";
 import { ConfigService } from "@nestjs/config";
 import { HttpRequestContextService } from "@/shared/http-request-context/http-request-context.service";
-import { CreateDatabase, Database, UpdateDatabase } from "@/definitions";
+import {
+  CreateDatabase,
+  Database,
+  ExternalColumn,
+  UpdateDatabase,
+} from "@/definitions";
 import { AES } from "crypto-js";
 import * as CryptoJS from "crypto-js";
 import * as assert from "assert";
@@ -123,19 +128,19 @@ export class DatabasesService {
     return database;
   }
 
-  async findAllSchemas(databaseId: string): Promise<string[]> {
+  async findAllSchemas(
+    databaseId: string,
+  ): Promise<Record<string, Record<string, Record<string, ExternalColumn>>>> {
     const database = await this.findOne(databaseId);
 
     if (!database) {
       throw new NotFoundException("Database not found");
     }
 
     const schemas =
-      await this.postgresAdapterService.getAllDatabaseSchema(databaseId);
-
-    const systemSchemas = ["information_schema", "pg_catalog", "pg_toast"];
+      await this.postgresAdapterService.getAllTableSchema(databaseId);
 
-    return schemas.filter((schema) => !systemSchemas.includes(schema));
+    return schemas;
   }
 
   async updateDatabase(

backend/src/postgres-adapter/postgres-adapter.service.ts

@@ -3,6 +3,7 @@ import { Inject, Injectable, forwardRef } from "@nestjs/common";
 import { Pool } from "pg";
 import { ExternalColumn } from "@/definitions/postgres-adapter";
 import * as assert from "assert";
+import { SYSTEM_SCHEMAS } from "@/shared/constant/database";
 
 interface CreateDatabaseQueryDto {
   databaseId: string;
@@ -105,6 +106,8 @@ export class PostgresAdapterService {
       const columns = result.rows;
       const columnMap = {};
       for (const column of columns) {
+        if (SYSTEM_SCHEMAS.includes(column.table_schema)) continue;
+
         if (!columnMap[column.table_schema]) {
           columnMap[column.table_schema] = {};
         }
@@ -136,7 +139,9 @@ export class PostgresAdapterService {
         `SELECT schema_name FROM information_schema.schemata`,
       );
       const columns = result.rows;
-      const schemata = columns.map((column) => column.schema_name);
+      const schemata = columns
+        .map((column) => column.schema_name)
+        .filter((schema) => !SYSTEM_SCHEMAS.includes(schema));
 
       return schemata;
     } catch (e) {

backend/src/shared/constant/database.ts

@@ -0,0 +1 @@
+export const SYSTEM_SCHEMAS = ["information_schema", "pg_catalog", "pg_toast"];

frontend/src/app/databases/[databaseId]/database-details.tsx

@@ -111,7 +111,9 @@ const DatabaseDetails = ({ database }: { database: CleanDatabase }) => {
           <Divider />
           <div className="flex flex-row justify-between">
             <Text>Schema count</Text>
-            <Text className="ml-2 bp5-text-muted">{schemas!.length}</Text>
+            <Text className="ml-2 bp5-text-muted">
+              {_.keys(schemas).length}
+            </Text>
           </div>
           <div className="flex flex-row justify-between">
             <Text>Table count</Text>

frontend/src/app/databases/database-selector.tsx

@@ -7,6 +7,7 @@ import {
 import AddDatabase from "@/app/databases/add-database";
 import React, { useEffect, useState } from "react";
 import { useRouter } from "next/navigation";
+import _ from "lodash";
 
 const DatabaseSelector = ({
   selectedDatabase,
@@ -88,7 +89,7 @@ const DatabaseSelector = ({
           }
           popoverProps={{ usePortal: true }}
         >
-          {databaseSchemas?.map((schema) => (
+          {_.keys(databaseSchemas).map((schema) => (
             <MenuItem
               key={schema}
               text={schema}

frontend/src/components/query/query-editor.tsx

@@ -7,22 +7,72 @@ import {
   MonacoThemeLight,
 } from "@blueprintjs/monaco-editor-theme";
 import { useDarkModeContext } from "../context/dark-mode-context";
+import { LanguageIdEnum, setupLanguageFeatures } from "monaco-sql-languages";
+import { useSelectedDatabase } from "@/stores";
+import { useDatabaseSchemas } from "@/data/use-database";
+import _ from "lodash";
+import { ExternalColumn } from "@/definitions";
+
+interface SuggestionProps {
+  label: string;
+  kind: monaco.languages.CompletionItemKind;
+  documentation?: string;
+  insertText: string;
+  range: monaco.IRange;
+}
 
 function createDependencyProposals(
   range: monaco.IRange,
   monacoInstance: typeof monaco,
+  schemas?: Record<string, Record<string, Record<string, ExternalColumn>>>,
 ) {
   // returning a static list of proposals, not even looking at the prefix (filtering is done by the Monaco editor),
   // here you could do a server side lookup
-  return [
-    {
-      label: "table_name",
+  const suggestions: SuggestionProps[] = [];
+  const tables = new Set<string>();
+  const columns = new Set<string>();
+
+  _.keys(schemas).forEach((schema: string) => {
+    if (!schemas) return;
+
+    _.keys(schemas[schema]).forEach((table) => {
+      _.keys(schemas[schema][table]).forEach((column) => {
+        columns.add(column);
+      });
+
+      tables.add(table);
+    });
+
+    suggestions.push({
+      label: schema,
+      kind: monacoInstance.languages.CompletionItemKind.Module,
+      documentation: "Schema name",
+      insertText: schema,
+      range: range,
+    });
+  });
+
+  tables.forEach((table) => {
+    suggestions.push({
+      label: table,
       kind: monacoInstance.languages.CompletionItemKind.Field,
-      documentation: "The Lodash library exported as Node.js modules.",
-      insertText: "table_name",
+      documentation: "Table name",
+      insertText: table,
       range: range,
-    },
-  ];
+    });
+  });
+
+  columns.forEach((column) => {
+    suggestions.push({
+      label: column,
+      kind: monacoInstance.languages.CompletionItemKind.Variable,
+      documentation: "Column name",
+      insertText: column,
+      range: range,
+    });
+  });
+
+  return suggestions;
 }
 
 interface QueryEditorProps {
@@ -33,14 +83,16 @@ interface QueryEditorProps {
 function QueryEditor({ value, onChange }: QueryEditorProps) {
   const monacoInstance = useMonaco();
   const { darkMode, setDarkMode } = useDarkModeContext();
+  const [selectedDatabase, setSelectedDatabase] = useSelectedDatabase();
+  const {
+    data: schemas,
+    isLoading: isLoadingSchemas,
+    error: schemasError,
+  } = useDatabaseSchemas(selectedDatabase?.id);
 
   useEffect(() => {
     if (!monacoInstance || !window) return;
     // const monaco = require("monaco-editor");
-    const {
-      setupLanguageFeatures,
-      LanguageIdEnum,
-    } = require("monaco-sql-languages");
     // const require("monaco-sql-languages/out/esm/pgsql/pgsql.ts");
     // console.log(typeof setupLanguageFeatures);
     // monacoInstance.languages.setLanguageConfiguration("pgsql", conf);
@@ -59,7 +111,11 @@ function QueryEditor({ value, onChange }: QueryEditorProps) {
           endColumn: word.endColumn,
         };
         return {
-          suggestions: createDependencyProposals(range, monacoInstance),
+          suggestions: createDependencyProposals(
+            range,
+            monacoInstance,
+            schemas,
+          ),
         };
       },
     });

frontend/src/data/use-database.ts

@@ -1,5 +1,10 @@
 import useSWR, { useSWRConfig } from "swr";
-import { CleanDatabase, CreateDatabase, UpdateDatabase } from "@/definitions";
+import {
+  CleanDatabase,
+  CreateDatabase,
+  ExternalColumn,
+  UpdateDatabase,
+} from "@/definitions";
 import {
   backendCreate,
   backendGet,
@@ -36,10 +41,9 @@ export const useCreateDatabase = () => {
 };
 
 export const useDatabaseSchemas = (id?: string) => {
-  const { data, error, isLoading, isValidating, mutate } = useSWR<string[]>(
-    id ? `/databases/${id}/schemas` : null,
-    backendGet,
-  );
+  const { data, error, isLoading, isValidating, mutate } = useSWR<
+    Record<string, Record<string, Record<string, ExternalColumn>>>
+  >(id ? `/databases/${id}/schemas` : null, backendGet);
   return { data, error, isLoading, isValidating, mutate };
 };