From 01d00f14c1146416c8747f6ee20c764134d18aba Mon Sep 17 00:00:00 2001
From: Brian Holmes <briangregoryholmes@gmail.com>
Date: Thu, 14 Dec 2023 15:20:39 -0500
Subject: [PATCH] adds autocomplate for models in the code editor

---
 web-common/src/features/models/selectors.ts   | 27 +++++++++++++++++++
 .../features/models/workspace/Editor.svelte   | 10 +++++++
 web-common/src/features/sources/selectors.ts  |  4 +--
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/web-common/src/features/models/selectors.ts b/web-common/src/features/models/selectors.ts
index 358083ae9e6..bb6049be320 100644
--- a/web-common/src/features/models/selectors.ts
+++ b/web-common/src/features/models/selectors.ts
@@ -11,9 +11,16 @@ import {
   getRuntimeServiceListFilesQueryKey,
   runtimeServiceListFiles,
   V1ListFilesResponse,
+  V1ProfileColumn,
 } 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 +38,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)
+      ),
+      (sourceColumnResponses) =>
+        sourceColumnResponses.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..75c61d356bf 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;
@@ -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