update schema after every query run

Author: m-abdelwahab

src/components/editor.tsx

@@ -10,6 +10,7 @@ import {
  editorSchemaAtom,
  queryAtom,
  queryResultAtom,
+ schemaAtom,
  showCommandPaletteAtom,
 } from '@/utils/atoms';
 import { useAtom } from 'jotai';
@@ -29,6 +30,7 @@ export const Editor = () => {
  );
  const [hasConfiguredDatabase] = useAtom(connectionStringAtom);
  const [editorSchema] = useAtom(editorSchemaAtom);
+ const [schema, setSchema] = useAtom(schemaAtom);
  const queryClient = useQueryClient();
 
  const { mutate, isLoading } = useMutation(
@@ -38,6 +40,7 @@ export const Editor = () => {
  onSuccess: (data) => {
  // @ts-ignore
  setQueryResult(data);
+ setSchema(data.databaseSchema);
  queryClient.invalidateQueries({ queryKey: ['schema'] });
  },
  onError: (error) => {

src/utils/query.ts

@@ -15,12 +15,32 @@ export const runQuery = async ({ connectionString, query }: queryOptions) => {
 
  const { rows, rowCount, fields } = await client.query(query);
 
+ const { rows: schemaRows } = await client.query(
+ ` SELECT table_name, array_agg(column_name || ' ' || data_type) AS columns_and_types FROM information_schema.columns WHERE table_schema = 'public' GROUP BY table_name ORDER BY table_name;`
+ );
+
+ // transform into CSV
+ const databaseSchema = schemaRows
+ .map((table) => {
+ const columns = table.columns_and_types
+ // @ts-ignore
+ .map((column) => {
+ const [name, type] = column.split(' ');
+ return `${name} ${type}`;
+ })
+ .join(', ');
+
+ return `${table.table_name} (${columns});`;
+ })
+ .join('\n');
+
  client.end();
 
  // get query execution time
  startTime = Date.now();
 
  return {
+ databaseSchema,
  rows: rows ?? [],
  rowCount: rowCount ?? 0,
  columns: fields?.map((field) => field.name) ?? [],