Skip to content

Commit 1ccbe61

Browse files
ilyabocursoragent
ilyabo
and
cursoragent
authored
feat: Snowflake connector integration (#401)
Signed-off-by: Ilya Boyandin <[email protected]> Co-authored-by: Cursor Agent <[email protected]>
1 parent b44ac17 commit 1ccbe61

25 files changed

Lines changed: 2432 additions & 373 deletions

‎apps/sqlrooms-cli-ui/src/components/AssistantDrawer.tsx‎

Lines changed: 99 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,4 @@
1-
import {
2-
AiSettingsPanel,
3-
AnalysisResultsContainer,
4-
ModelSelector,
5-
PromptSuggestions,
6-
QueryControls,
7-
SessionControls,
8-
} from '@sqlrooms/ai';
1+
import {AiSettingsPanel, Chat} from '@sqlrooms/ai';
92
import {
103
Button,
114
Dialog,
@@ -37,6 +30,7 @@ export const AssistantDrawer: React.FC<{
3730
(s) => s.ai.config.currentSessionId || null,
3831
);
3932
const isDataAvailable = useRoomStore((state) => state.room.initialized);
33+
const updateProvider = useRoomStore((s) => s.aiSettings.updateProvider);
4034
const settingsPanelOpen = useDisclosure();
4135
const isAssistantOpen = useRoomStore((state) => state.isAssistantOpen);
4236
const setAssistantOpen = useRoomStore((state) => state.setAssistantOpen);
@@ -66,97 +60,106 @@ export const AssistantDrawer: React.FC<{
6660
</Button>
6761
</DrawerClose>
6862
</DrawerHeader>
69-
<div className="flex min-h-0 flex-1 flex-col gap-0 overflow-hidden p-4">
70-
<div className="mb-4 flex items-center justify-between gap-2">
71-
<SessionControls className="w-full" />
72-
{currentSessionId && (
73-
<Dialog
74-
open={settingsPanelOpen.isOpen}
75-
onOpenChange={(open) => {
76-
if (open) {
77-
settingsPanelOpen.onOpen();
78-
} else {
79-
settingsPanelOpen.onClose();
80-
}
81-
}}
82-
>
83-
<DialogTrigger asChild>
84-
<Button
85-
variant="outline"
86-
className="hover:bg-accent flex items-center justify-center transition-colors"
87-
title="Configuration"
88-
size="sm"
89-
>
90-
<Settings className="h-4 w-4" />
91-
</Button>
92-
</DialogTrigger>
93-
<DialogContent className="flex h-[80vh] w-[90vw] max-w-3xl flex-col overflow-hidden">
94-
<DialogHeader>
95-
<DialogTitle>AI Assistant Settings</DialogTitle>
96-
</DialogHeader>
97-
<Tabs
98-
defaultValue="providers"
99-
className="flex min-h-0 flex-1 flex-col"
100-
>
101-
<TabsList className="grid w-full shrink-0 grid-cols-3">
102-
<TabsTrigger value="providers">Providers</TabsTrigger>
103-
<TabsTrigger value="models">Models</TabsTrigger>
104-
<TabsTrigger value="parameters">Parameters</TabsTrigger>
105-
</TabsList>
106-
<TabsContent
107-
value="providers"
108-
className="flex-1 overflow-y-auto"
109-
>
110-
<AiSettingsPanel.ProvidersSettings />
111-
</TabsContent>
112-
<TabsContent
113-
value="models"
114-
className="flex-1 overflow-y-auto"
63+
<Chat.Root>
64+
<div className="flex min-h-0 flex-1 flex-col gap-0 overflow-hidden p-4">
65+
<div className="mb-4 flex items-center justify-between gap-2">
66+
<Chat.Sessions className="w-full" />
67+
{currentSessionId && (
68+
<Dialog
69+
open={settingsPanelOpen.isOpen}
70+
onOpenChange={(open) => {
71+
if (open) {
72+
settingsPanelOpen.onOpen();
73+
} else {
74+
settingsPanelOpen.onClose();
75+
}
76+
}}
77+
>
78+
<DialogTrigger asChild>
79+
<Button
80+
variant="outline"
81+
className="hover:bg-accent flex items-center justify-center transition-colors"
82+
title="Configuration"
83+
size="sm"
11584
>
116-
<AiSettingsPanel.ModelsSettings />
117-
</TabsContent>
118-
<TabsContent
119-
value="parameters"
120-
className="flex-1 overflow-y-auto"
85+
<Settings className="h-4 w-4" />
86+
</Button>
87+
</DialogTrigger>
88+
<DialogContent className="flex h-[80vh] w-[90vw] max-w-3xl flex-col overflow-hidden">
89+
<DialogHeader>
90+
<DialogTitle>AI Assistant Settings</DialogTitle>
91+
</DialogHeader>
92+
<Tabs
93+
defaultValue="providers"
94+
className="flex min-h-0 flex-1 flex-col"
12195
>
122-
<AiSettingsPanel.ModelParametersSettings />
123-
</TabsContent>
124-
</Tabs>
125-
</DialogContent>
126-
</Dialog>
127-
)}
128-
</div>
129-
<div className="print-container grow overflow-auto">
130-
{!currentSessionId ? (
131-
<div className="flex h-full w-full flex-col items-center justify-center">
132-
<p className="text-muted-foreground mt-4">
133-
No session selected
134-
</p>
135-
</div>
136-
) : isDataAvailable ? (
137-
<AnalysisResultsContainer key={currentSessionId} />
138-
) : (
139-
<div className="flex h-full w-full flex-col items-center justify-center">
140-
<SkeletonPane className="p-4" />
141-
<p className="text-muted-foreground mt-4">
142-
Loading database...
143-
</p>
144-
</div>
145-
)}
146-
</div>{' '}
147-
<PromptSuggestions.Container>
148-
<PromptSuggestions.Item text="What questions can I ask to get insights from my data?" />
149-
<PromptSuggestions.Item text="Show me a summary of the data" />
150-
<PromptSuggestions.Item text="What are the key trends?" />
151-
<PromptSuggestions.Item text="Help me understand the data structure" />
152-
</PromptSuggestions.Container>
153-
<QueryControls placeholder="What would you like to learn about the data?">
154-
<div className="flex items-center justify-end gap-2">
155-
<PromptSuggestions.VisibilityToggle />
156-
<ModelSelector />
96+
<TabsList className="grid w-full shrink-0 grid-cols-3">
97+
<TabsTrigger value="providers">Providers</TabsTrigger>
98+
<TabsTrigger value="models">Models</TabsTrigger>
99+
<TabsTrigger value="parameters">
100+
Parameters
101+
</TabsTrigger>
102+
</TabsList>
103+
<TabsContent
104+
value="providers"
105+
className="flex-1 overflow-y-auto"
106+
>
107+
<AiSettingsPanel.ProvidersSettings />
108+
</TabsContent>
109+
<TabsContent
110+
value="models"
111+
className="flex-1 overflow-y-auto"
112+
>
113+
<AiSettingsPanel.ModelsSettings />
114+
</TabsContent>
115+
<TabsContent
116+
value="parameters"
117+
className="flex-1 overflow-y-auto"
118+
>
119+
<AiSettingsPanel.ModelParametersSettings />
120+
</TabsContent>
121+
</Tabs>
122+
</DialogContent>
123+
</Dialog>
124+
)}
157125
</div>
158-
</QueryControls>
159-
</div>
126+
<div className="print-container grow overflow-auto">
127+
{!currentSessionId ? (
128+
<div className="flex h-full w-full flex-col items-center justify-center">
129+
<p className="text-muted-foreground mt-4">
130+
No session selected
131+
</p>
132+
</div>
133+
) : isDataAvailable ? (
134+
<Chat.Messages key={currentSessionId} />
135+
) : (
136+
<div className="flex h-full w-full flex-col items-center justify-center">
137+
<SkeletonPane className="p-4" />
138+
<p className="text-muted-foreground mt-4">
139+
Loading database...
140+
</p>
141+
</div>
142+
)}
143+
</div>{' '}
144+
<Chat.PromptSuggestions>
145+
<Chat.PromptSuggestions.Item text="What questions can I ask to get insights from my data?" />
146+
<Chat.PromptSuggestions.Item text="Show me a summary of the data" />
147+
<Chat.PromptSuggestions.Item text="What are the key trends?" />
148+
<Chat.PromptSuggestions.Item text="Help me understand the data structure" />
149+
</Chat.PromptSuggestions>
150+
<Chat.Composer placeholder="What would you like to learn about the data?">
151+
<Chat.InlineApiKeyInput
152+
onSaveApiKey={(provider, apiKey) => {
153+
updateProvider(provider, {apiKey});
154+
}}
155+
/>
156+
<div className="flex items-center justify-end gap-2">
157+
<Chat.PromptSuggestions.VisibilityToggle />
158+
<Chat.ModelSelector />
159+
</div>
160+
</Chat.Composer>
161+
</div>
162+
</Chat.Root>
160163
</div>
161164
</DrawerContent>
162165
</Drawer>

‎apps/sqlrooms-cli-ui/src/runtimeConfig.ts‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,28 @@ export type RuntimeConfig = {
44
llmProvider?: string;
55
llmModel?: string;
66
apiKey?: string;
7+
aiProviders?: Record<
8+
string,
9+
{
10+
baseUrl: string;
11+
apiKey: string;
12+
models: Array<{modelName: string}>;
13+
}
14+
>;
715
dbPath?: string;
816
metaNamespace?: string;
9-
postgresBridgeEnabled?: boolean;
17+
dbBridge?: {
18+
id: string;
19+
connections: Array<{
20+
id: string;
21+
engineId: string;
22+
title: string;
23+
runtimeSupport?: 'browser' | 'server' | 'both';
24+
requiresBridge?: boolean;
25+
bridgeId?: string;
26+
isCore?: boolean;
27+
}>;
28+
};
1029
};
1130

1231
export async function fetchRuntimeConfig(): Promise<RuntimeConfig> {

‎apps/sqlrooms-cli-ui/src/store.ts‎

Lines changed: 54 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ import {
1919
createCellsSlice,
2020
createDefaultCellRegistry,
2121
} from '@sqlrooms/cells';
22-
import {createHttpDbBridge} from '@sqlrooms/db';
2322
import {createWebSocketDuckDbConnector} from '@sqlrooms/duckdb';
2423
import {
2524
createNotebookSlice,
@@ -48,6 +47,7 @@ import {
4847
import {produce} from 'immer';
4948
import {z} from 'zod';
5049

50+
import {createHttpDbBridge, DbConnection} from '@sqlrooms/db';
5151
import {getDefaultScaffoldTree} from './helpers';
5252
import {LAYOUT} from './layout';
5353
import {fetchRuntimeConfig} from './runtimeConfig';
@@ -98,6 +98,14 @@ export type RoomState = RoomShellSliceState &
9898
};
9999

100100
export const runtimeConfig = await fetchRuntimeConfig();
101+
const runtimeAiProviders =
102+
(runtimeConfig.aiProviders as AiSettingsSliceConfig['providers']) || {};
103+
const defaultProviderFromConfig =
104+
runtimeConfig.llmProvider || Object.keys(runtimeAiProviders)[0] || 'openai';
105+
const defaultModelFromProvider =
106+
runtimeAiProviders[defaultProviderFromConfig]?.models?.[0]?.modelName;
107+
const defaultModelFromConfig =
108+
runtimeConfig.llmModel || defaultModelFromProvider || 'gpt-4o-mini';
101109

102110
const connector = createWebSocketDuckDbConnector({
103111
wsUrl: runtimeConfig.wsUrl || 'ws://localhost:4000',
@@ -113,6 +121,26 @@ connector.loadFile = async (file, desiredTableName, options) => {
113121
return baseLoadFile(file, desiredTableName, options);
114122
};
115123

124+
function getRuntimeBridgeConfig():
125+
| {
126+
id: string;
127+
connections: Array<{
128+
id: string;
129+
engineId: string;
130+
title: string;
131+
runtimeSupport?: 'browser' | 'server' | 'both';
132+
requiresBridge?: boolean;
133+
bridgeId?: string;
134+
isCore?: boolean;
135+
}>;
136+
}
137+
| undefined {
138+
if (runtimeConfig.dbBridge?.connections?.length) {
139+
return runtimeConfig.dbBridge;
140+
}
141+
return undefined;
142+
}
143+
116144
export const {roomStore, useRoomStore} = createRoomStore<RoomState>(
117145
persistSliceConfigs<RoomState>(
118146
{
@@ -121,7 +149,7 @@ export const {roomStore, useRoomStore} = createRoomStore<RoomState>(
121149
room: BaseRoomConfig,
122150
layout: LayoutConfig,
123151
ai: AiSliceConfig,
124-
aiSettings: AiSettingsSliceConfig,
152+
// aiSettings: AiSettingsSliceConfig,
125153
sqlEditor: SqlEditorSliceConfig,
126154
cells: CellsSliceConfig,
127155
notebook: NotebookSliceConfig,
@@ -203,14 +231,15 @@ export const {roomStore, useRoomStore} = createRoomStore<RoomState>(
203231
})(set, get, store),
204232

205233
...createAiSettingsSlice({
206-
config: {providers: {} as AiSettingsSliceConfig['providers']},
234+
config: {providers: runtimeAiProviders},
207235
})(set, get, store),
208236

209237
...createAiSlice({
210238
config: AiSliceConfig.parse({sessions: []}),
211-
defaultProvider: (runtimeConfig.llmProvider as any) || 'openai',
212-
defaultModel: runtimeConfig.llmModel || 'gpt-4o-mini',
213-
getApiKey: () => runtimeConfig.apiKey || '',
239+
defaultProvider: defaultProviderFromConfig as any,
240+
defaultModel: defaultModelFromConfig,
241+
getApiKey: (provider) =>
242+
runtimeAiProviders[provider]?.apiKey || runtimeConfig.apiKey || '',
214243
getBaseUrl: () => runtimeConfig.apiBaseUrl || '',
215244
getInstructions: () => createDefaultAiInstructions(store),
216245
tools: {
@@ -222,22 +251,24 @@ export const {roomStore, useRoomStore} = createRoomStore<RoomState>(
222251
),
223252
);
224253

225-
if (runtimeConfig.postgresBridgeEnabled) {
226-
const postgresBridgeId = 'postgres-http-bridge';
227-
roomStore.getState().db.connectors.registerBridge(
228-
createHttpDbBridge({
229-
id: postgresBridgeId,
230-
baseUrl: runtimeConfig.apiBaseUrl || '',
231-
runtimeSupport: 'server',
232-
}),
233-
);
234-
roomStore.getState().db.connectors.registerConnection({
235-
id: 'postgres',
236-
engineId: 'postgres',
237-
title: 'Postgres',
238-
runtimeSupport: 'server',
239-
requiresBridge: true,
240-
bridgeId: postgresBridgeId,
241-
isCore: false,
254+
const bridgeConfig = getRuntimeBridgeConfig();
255+
if (bridgeConfig?.connections.length) {
256+
const bridge = createHttpDbBridge({
257+
id: bridgeConfig.id,
258+
baseUrl: runtimeConfig.apiBaseUrl || '',
242259
});
260+
const state = roomStore.getState();
261+
state.db.connectors.registerBridge(bridge);
262+
for (const connection of bridgeConfig.connections) {
263+
const normalizedConnection: DbConnection = {
264+
id: connection.id,
265+
engineId: connection.engineId,
266+
title: connection.title || connection.id,
267+
runtimeSupport: connection.runtimeSupport || 'server',
268+
requiresBridge: connection.requiresBridge ?? true,
269+
bridgeId: connection.bridgeId || bridgeConfig.id,
270+
isCore: connection.isCore ?? false,
271+
};
272+
state.db.connectors.registerConnection(normalizedConnection);
273+
}
243274
}

‎packages/canvas/src/CanvasSlice.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -159,8 +159,8 @@ export function createCanvasSlice(
159159
get().cells.config.currentSheetId ||
160160
get().cells.config.sheetOrder[0];
161161
if (!sheetId) return;
162-
await get().cells.runAllCellsCascade(sheetId);
163-
await get().db.refreshTableSchemas();
162+
// don't await this - it will block the UI
163+
get().cells.runAllCellsCascade(sheetId);
164164
},
165165

166166
addNode: async ({

‎packages/cells/src/defaultCellRegistry.tsx‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,13 @@ export function createDefaultCellRegistry(): CellRegistry {
5050
});
5151
},
5252
runCell: async ({id, opts, get, set}) => {
53+
const ownerSheetId = findSheetIdForCell(get(), id);
54+
if (ownerSheetId) {
55+
// Recompute this SQL cell's dependencies on explicit run so graph state
56+
// stays current without paying the cost on every keystroke.
57+
await get().cells.updateEdgesFromSql(ownerSheetId, id);
58+
}
59+
5360
const controller = new AbortController();
5461
set((s) =>
5562
produce(s, (draft) => {

0 commit comments

Comments
 (0)