From 567993e2769f6ff6414ae788589eb6afc1fbefd3 Mon Sep 17 00:00:00 2001
From: Brian Holmes <120223836+briangregoryholmes@users.noreply.github.com>
Date: Thu, 14 Dec 2023 18:59:10 -0500
Subject: [PATCH] suggest models in the code editor (#3705)

* adds autocomplate for models in the code editor

* remove unused type import

* change schema initialization, rename variable in useAllModelColumns
---
 web-common/src/features/models/selectors.ts   | 26 +++++++++++++++++++
 .../features/models/workspace/Editor.svelte   | 14 ++++++++--
 web-common/src/features/sources/selectors.ts  |  4 +--
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/web-common/src/features/models/selectors.ts b/web-common/src/features/models/selectors.ts
index 358083ae9e6..a8c27a8f7f5 100644
--- a/web-common/src/features/models/selectors.ts
+++ b/web-common/src/features/models/selectors.ts
@@ -14,6 +14,12 @@ import {
 } from "@rilldata/web-common/runtime-client";
 import type { QueryClient } from "@tanstack/query-core";
 import { TIMESTAMPS } from "../../lib/duckdb-data-types";
+import type { Readable } from "svelte/motion";
+import { derived } from "svelte/store";
+import {
+  createTableColumnsWithName,
+  type TableColumnsWithName,
+} from "../sources/selectors";
 
 export function useModels(instanceId: string) {
   return useFilteredResources(instanceId, ResourceKind.Model);
@@ -31,6 +37,26 @@ export function useModel(instanceId: string, name: string) {
   return useResource(instanceId, name, ResourceKind.Model);
 }
 
+export function useAllModelColumns(
+  queryClient: QueryClient,
+  instanceId: string
+): Readable<Array<TableColumnsWithName>> {
+  return derived([useModels(instanceId)], ([allModels], set) => {
+    if (!allModels.data?.length) {
+      set([]);
+      return;
+    }
+
+    derived(
+      allModels.data.map((r) =>
+        createTableColumnsWithName(queryClient, instanceId, r.meta.name.name)
+      ),
+      (modelColumnResponses) =>
+        modelColumnResponses.filter((res) => !!res.data).map((res) => res.data)
+    ).subscribe(set);
+  });
+}
+
 export async function getModelNames(
   queryClient: QueryClient,
   instanceId: string
diff --git a/web-common/src/features/models/workspace/Editor.svelte b/web-common/src/features/models/workspace/Editor.svelte
index 08383d006fe..2b9affc21f5 100644
--- a/web-common/src/features/models/workspace/Editor.svelte
+++ b/web-common/src/features/models/workspace/Editor.svelte
@@ -55,6 +55,7 @@
   import { editorTheme } from "../../../components/editor/theme";
   import { runtime } from "../../../runtime-client/runtime-store";
   import { useAllSourceColumns } from "../../sources/selectors";
+  import { useAllModelColumns } from "../selectors";
 
   export let content: string;
   export let editorHeight = 0;
@@ -89,11 +90,11 @@
       "select from where group by all having order limit sample unnest with window qualify values filter exclude replace like ilike glob as case when then else end in cast left join on not desc asc sum union",
   });
 
+  const schema: { [table: string]: string[] } = {};
+
   // Autocomplete: source tables
-  let schema: { [table: string]: string[] };
   $: allSourceColumns = useAllSourceColumns(queryClient, $runtime?.instanceId);
   $: if ($allSourceColumns?.length) {
-    schema = {};
     for (const sourceTable of $allSourceColumns) {
       const sourceIdentifier = sourceTable?.tableName;
       schema[sourceIdentifier] =
@@ -101,6 +102,15 @@
     }
   }
 
+  //Auto complete: model tables
+  $: allModelColumns = useAllModelColumns(queryClient, $runtime?.instanceId);
+  $: if ($allModelColumns?.length) {
+    for (const modelTable of $allModelColumns) {
+      const modelIdentifier = modelTable?.tableName;
+      schema[modelIdentifier] = modelTable.profileColumns?.map((c) => c.name);
+    }
+  }
+
   function getTableNameFromFromClause(
     sql: string,
     schema: { [table: string]: string[] }
diff --git a/web-common/src/features/sources/selectors.ts b/web-common/src/features/sources/selectors.ts
index d2188bc628a..57029976a6f 100644
--- a/web-common/src/features/sources/selectors.ts
+++ b/web-common/src/features/sources/selectors.ts
@@ -88,7 +88,7 @@ export function useIsLocalFileConnector(
   );
 }
 
-type TableColumnsWithName = {
+export type TableColumnsWithName = {
   tableName: string;
   profileColumns: Array<V1ProfileColumn>;
 };
@@ -116,7 +116,7 @@ export function useAllSourceColumns(
 /**
  * Fetches columns and adds the table name. By using the selector the results will be cached.
  */
-function createTableColumnsWithName(
+export function createTableColumnsWithName(
   queryClient: QueryClient,
   instanceId: string,
   tableName: string