Repository: briefercloud/briefer (language: python)

Issue:
run only selected parts of SQL when page setting is enabled



Discussion / hints:

Fix patch:
diff --git a/apps/api/src/v1/workspaces/workspace/documents/document/settings.ts b/apps/api/src/v1/workspaces/workspace/documents/document/settings.ts
index dc1efc9f..e206daaf 100644
--- a/apps/api/src/v1/workspaces/workspace/documents/document/settings.ts
+++ b/apps/api/src/v1/workspaces/workspace/documents/document/settings.ts
@@ -6,7 +6,8 @@ import prisma, { recoverFromNotFound } from '@briefer/database'
 import { broadcastDocument } from '../../../../../websocket/workspace/documents.js'
 
 const DocumentSettings = z.object({
-  runUnexecutedBlocks: z.boolean(),
+  runUnexecutedBlocks: z.boolean().optional(),
+  runSQLSelection: z.boolean().optional(),
 })
 
 export default function settingsRouter(socketServer: IOServer) {
@@ -21,13 +22,13 @@ export default function settingsRouter(socketServer: IOServer) {
       return
     }
 
-    const runUnexecutedBlocks = body.data.runUnexecutedBlocks
+    const { runUnexecutedBlocks, runSQLSelection } = body.data
 
     try {
       const doc = await recoverFromNotFound(
         prisma().document.update({
           where: { id: documentId },
-          data: { runUnexecutedBlocks },
+          data: { runUnexecutedBlocks, runSQLSelection },
         })
       )
       if (!doc) {
diff --git a/apps/api/src/yjs/v2/executors/blocks/sql.ts b/apps/api/src/yjs/v2/executors/blocks/sql.ts
index bbd912d0..d9922a4d 100644
--- a/apps/api/src/yjs/v2/executors/blocks/sql.ts
+++ b/apps/api/src/yjs/v2/executors/blocks/sql.ts
@@ -9,7 +9,7 @@ import {
 } from '@briefer/editor'
 import PQueue from 'p-queue'
 import * as Y from 'yjs'
-import {
+import prisma, {
   getCredentialsInfo,
   getDatabaseURL,
   listDataSources,
@@ -94,6 +94,7 @@ export type SQLEffects = {
   listDataSources: typeof listDataSources
   listDataFrames: typeof listDataFrames
   editWithAI: typeof editWithAI
+  documentHasRunSQLSelectionEnabled: (id: string) => Promise<boolean>
 }
 
 type RunninQuery = {
@@ -106,7 +107,8 @@ export interface ISQLExecutor {
   runQuery(
     block: Y.XmlElement<SQLBlock>,
     tr: Y.Transaction,
-    isSuggestion: boolean
+    isSuggestion: boolean,
+    isRunAll: boolean
   ): Promise<void>
   abortQuery(block: Y.XmlElement<SQLBlock>): Promise<void>
   renameDataFrame(block: Y.XmlElement<SQLBlock>): Promise<void>
@@ -155,7 +157,8 @@ export class SQLExecutor implements ISQLExecutor {
   public async runQuery(
     block: Y.XmlElement<SQLBlock>,
     tr: Y.Transaction,
-    isSuggestion: boolean
+    isSuggestion: boolean,
+    isRunAll: boolean
   ) {
     this.events.sqlRun(EventContext.fromYTransaction(tr))
     const abortController = new AbortController()
@@ -186,8 +189,16 @@ export class SQLExecutor implements ISQLExecutor {
             'executing query'
           )
 
-          const { dataSourceId, dataframeName, isFileDataSource } =
-            getSQLAttributes(block, this.blocks)
+          const {
+            id: blockId,
+            aiSuggestions,
+            source,
+            configuration,
+            dataSourceId,
+            dataframeName,
+            isFileDataSource,
+            selectedCode,
+          } = getSQLAttributes(block, this.blocks)
           if ((!dataSourceId && !isFileDataSource) || !dataframeName) {
             return
           }
@@ -212,15 +223,17 @@ export class SQLExecutor implements ISQLExecutor {
 
           block.removeAttribute('result')
 
-          const {
-            id: blockId,
-            aiSuggestions,
-            source,
-            configuration,
-          } = getSQLAttributes(block, this.blocks)
-
-          const actualSource =
-            (isSuggestion ? aiSuggestions : source)?.toJSON()?.trim() ?? ''
+          let actualSource =
+            (isSuggestion ? aiSuggestions : source)?.toJSON().trim() ?? ''
+          if (!isRunAll && selectedCode) {
+            const hasRunSQLSelection =
+              await this.effects.documentHasRunSQLSelectionEnabled(
+                this.documentId
+              )
+            if (hasRunSQLSelection) {
+              actualSource = selectedCode.trim()
+            }
+          }
 
           let resultType: RunQueryResult['type'] | 'empty-query' = 'empty-query'
           if (actualSource !== '') {
@@ -552,6 +565,13 @@ export class SQLExecutor implements ISQLExecutor {
         listDataSources,
         listDataFrames,
         editWithAI,
+        documentHasRunSQLSelectionEnabled: (id: string) =>
+          prisma()
+            .document.findFirst({
+              where: { id },
+              select: { runSQLSelection: true },
+            })
+            .then((doc) => doc?.runSQLSelection ?? false),
       },
       events
     )
diff --git a/apps/api/src/yjs/v2/executors/run-all.ts b/apps/api/src/yjs/v2/executors/run-all.ts
index 53b6a898..d66c6485 100644
--- a/apps/api/src/yjs/v2/executors/run-all.ts
+++ b/apps/api/src/yjs/v2/executors/run-all.ts
@@ -400,7 +400,7 @@ export class RunAllExecutor implements IRunAllExecutor {
     await switchBlockType(block, {
       onRichText: () => Promise.resolve(),
       onPython: (block) => this.executors.python.run(block, tr, false),
-      onSQL: (block) => this.executors.sql.runQuery(block, tr, false),
+      onSQL: (block) => this.executors.sql.runQuery(block, tr, false, true),
       onVisualization: (block) => this.executors.visualization.run(block, tr),
       onInput: (block) => this.executors.input.saveValue(block),
       onDropdownInput: (block) => this.executors.dropdownInput.saveValue(block),
diff --git a/apps/api/src/yjs/v2/observers/blocks/sql.ts b/apps/api/src/yjs/v2/observers/blocks/sql.ts
index a966ed21..50831ac0 100644
--- a/apps/api/src/yjs/v2/observers/blocks/sql.ts
+++ b/apps/api/src/yjs/v2/observers/blocks/sql.ts
@@ -147,11 +147,11 @@ export class SQLObserver implements ISQLObserver {
           block.setAttribute('status', 'idle')
           break
         case 'running':
-          await this.executor.runQuery(block, tr, false)
+          await this.executor.runQuery(block, tr, false, false)
           block.setAttribute('status', 'idle')
           break
         case 'running-suggestion':
-          await this.executor.runQuery(block, tr, true)
+          await this.executor.runQuery(block, tr, true, false)
           block.setAttribute('status', 'idle')
           break
         case 'idle':
diff --git a/apps/web/src/components/PageSettingsPanel.tsx b/apps/web/src/components/PageSettingsPanel.tsx
index 5178e990..13b7c85f 100644
--- a/apps/web/src/components/PageSettingsPanel.tsx
+++ b/apps/web/src/components/PageSettingsPanel.tsx
@@ -1,5 +1,4 @@
 import * as Y from 'yjs'
-import { useState } from 'react'
 import { Switch } from '@headlessui/react'
 import { ChevronDoubleRightIcon } from '@heroicons/react/24/outline'
 import { Transition } from '@headlessui/react'
@@ -56,6 +55,12 @@ export default function PageSettingsPanel(props: Props) {
               enabled={document?.runUnexecutedBlocks ?? false}
               onToggle={api.toggleRunUnexecutedBlocks}
             />
+            <PageSettingToggle
+              name="Run selected SQL only"
+              description="Whether Briefer should only run selected code when a SQL block is executed."
+              enabled={document?.runSQLSelection ?? false}
+              onToggle={api.toggleRunSQLSelection}
+            />
           </div>
         </div>
       </Transition>
diff --git a/apps/web/src/components/v2Editor/CodeEditor/index.tsx b/apps/web/src/components/v2Editor/CodeEditor/index.tsx
index 23f70d7b..15290184 100644
--- a/apps/web/src/components/v2Editor/CodeEditor/index.tsx
+++ b/apps/web/src/components/v2Editor/CodeEditor/index.tsx
@@ -17,6 +17,16 @@ import useEditorAwareness from '@/hooks/useEditorAwareness'
 import { useSQLExtension } from './sql'
 import { usePythonExtension } from './python'
 
+function getSelectedCodeFromState(e: EditorState): string | null {
+  const { from, to } = e.selection.main
+  const selection = e.sliceDoc(from, to)
+  if (selection === '') {
+    return null
+  }
+
+  return selection
+}
+
 function createTextSync(source: Y.Text) {
   const plugin = ViewPlugin.fromClass(
     class YTextSync {
@@ -172,6 +182,32 @@ const pythonCompartment = new Compartment()
 const readonlyCompartment = new Compartment()
 const themeCompartment = new Compartment()
 const keymapsCompartment = new Compartment()
+const watchSelectionCompartment = new Compartment()
+
+const createWatchSelectionPlugin = (
+  callback: (selectedCode: string | null) => void
+) =>
+  ViewPlugin.fromClass(
+    class WatchSelection {
+      private selectedCode: string | null = null
+      constructor(view: EditorView) {
+        this.selectedCode = getSelectedCodeFromState(view.state)
+      }
+
+      public update(update: ViewUpdate) {
+        if (!update.selectionSet) {
+          return
+        }
+
+        const selectedCode = getSelectedCodeFromState(update.state)
+        if (selectedCode === this.selectedCode) {
+          return
+        }
+        this.selectedCode = selectedCode
+        callback(selectedCode)
+      }
+    }
+  )
 
 interface Props {
   workspaceId: string
@@ -186,6 +222,7 @@ interface Props {
   diff?: Y.Text
   dataSourceId?: string | null
   disabled: boolean
+  onSelectionChanged?: (selectedCode: string | null) => void
 }
 export function CodeEditor(props: Props) {
   const [editorState, editorAPI] = useEditorAwareness()
@@ -211,7 +248,13 @@ export function CodeEditor(props: Props) {
       return
     }
 
-    function getExtensions(source: Y.Text, disabled: boolean) {
+    function getExtensions(
+      source: Y.Text,
+      {
+        disabled,
+        watchSelection,
+      }: { disabled: boolean; watchSelection: boolean }
+    ) {
       return [
         keymapsCompartment.of(
           brieferKeyMaps({
@@ -235,6 +278,15 @@ export function CodeEditor(props: Props) {
         ),
         createTextSync(source),
         themeCompartment.of(materialLight(disabled)),
+        ...(watchSelection
+          ? [
+              watchSelectionCompartment.of(
+                createWatchSelectionPlugin(
+                  props.onSelectionChanged ?? (() => {})
+                )
+              ),
+            ]
+          : []),
       ]
     }
 
@@ -267,7 +319,10 @@ export function CodeEditor(props: Props) {
     function initializeEditorView(parent: Element) {
       const selection = getSelection()
       const config = {
-        extensions: getExtensions(props.source, props.disabled),
+        extensions: getExtensions(props.source, {
+          disabled: props.disabled,
+          watchSelection: true,
+        }),
         doc: props.source.toString(),
         selection,
       }
@@ -315,12 +370,18 @@ export function CodeEditor(props: Props) {
       destroyCurrent()
 
       const a = {
-        extensions: getExtensions(props.source, props.disabled),
+        extensions: getExtensions(props.source, {
+          disabled: props.disabled,
+          watchSelection: false,
+        }),
         doc: props.source.toString(),
         selection,
       }
       const b = {
-        extensions: getExtensions(diff, props.disabled),
+        extensions: getExtensions(diff, {
+          disabled: props.disabled,
+          watchSelection: true,
+        }),
         doc: diff.toString(),
       }
 
@@ -451,6 +512,24 @@ export function CodeEditor(props: Props) {
     onRunInsertBlock,
   ])
 
+  useEffect(() => {
+    const effect = watchSelectionCompartment.reconfigure(
+      createWatchSelectionPlugin(props.onSelectionChanged ?? (() => {}))
+    )
+
+    if (viewRef.current) {
+      viewRef.current.dispatch({
+        effects: effect,
+      })
+    }
+
+    if (mergeRef.current) {
+      mergeRef.current.view.a.dispatch({
+        effects: effect,
+      })
+    }
+  }, [viewRef, mergeRef, props.onSelectionChanged])
+
   useEffect(() => {
     if (
       editorState.cursorBlockId === props.blockId &&
diff --git a/apps/web/src/components/v2Editor/customBlocks/sql/index.tsx b/apps/web/src/components/v2Editor/customBlocks/sql/index.tsx
index 5bf286fa..4072ac9c 100644
--- a/apps/web/src/components/v2Editor/customBlocks/sql/index.tsx
+++ b/apps/web/src/components/v2Editor/customBlocks/sql/index.tsx
@@ -9,7 +9,7 @@ import {
   BookOpenIcon,
 } from '@heroicons/react/20/solid'
 import { CopyToClipboard } from 'react-copy-to-clipboard'
-import { useCallback, useEffect, useMemo, useState } from 'react'
+import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
 import * as Y from 'yjs'
 import {
   type SQLBlock,
@@ -115,6 +115,13 @@ function SQLBlock(props: Props) {
     })
   }, [props.block])
 
+  const onSQLSelectionChanged = useCallback(
+    (selectedCode: string | null) => {
+      props.block.setAttribute('selectedCode', selectedCode)
+    },
+    [props.block]
+  )
+
   const onRun = useCallback(() => {
     props.onRun(props.block)
   }, [props.onRun, props.block])
@@ -531,6 +538,7 @@ function SQLBlock(props: Props) {
                 diff={aiSuggestions ?? undefined}
                 dataSourceId={dataSourceId}
                 disabled={statusIsDisabled}
+                onSelectionChanged={onSQLSelectionChanged}
               />
             </div>
           </div>
diff --git a/apps/web/src/hooks/useDocument.ts b/apps/web/src/hooks/useDocument.ts
index 5937cbf4..20c9075f 100644
--- a/apps/web/src/hooks/useDocument.ts
+++ b/apps/web/src/hooks/useDocument.ts
@@ -6,6 +6,7 @@ type API = {
   setIcon: (icon: string) => Promise<void>
   publish: () => Promise<void>
   toggleRunUnexecutedBlocks: () => Promise<void>
+  toggleRunSQLSelection: () => Promise<void>
 }
 
 type UseDocument = [
@@ -59,12 +60,34 @@ function useDocument(workspaceId: string, documentId: string): UseDocument {
     api.updateDocumentSettings,
   ])
 
+  const toggleRunSQLSelection = useCallback(async () => {
+    const newRunSQLSelection = !document?.runSQLSelection
+    try {
+      await api.updateDocumentSettings(documentId, {
+        runSQLSelection: newRunSQLSelection,
+      })
+    } catch (err) {
+      alert('Failed to update document settings')
+    }
+  }, [
+    workspaceId,
+    documentId,
+    document?.runSQLSelection,
+    api.updateDocumentSettings,
+  ])
+
   return useMemo(
     () => [
       { document, loading, publishing },
-      { setIcon, publish, toggleRunUnexecutedBlocks },
+      { setIcon, publish, toggleRunUnexecutedBlocks, toggleRunSQLSelection },
     ],
-    [loading, setIcon, publish, toggleRunUnexecutedBlocks]
+    [
+      loading,
+      setIcon,
+      publish,
+      toggleRunUnexecutedBlocks,
+      toggleRunSQLSelection,
+    ]
   )
 }
 
diff --git a/apps/web/src/hooks/useDocuments.tsx b/apps/web/src/hooks/useDocuments.tsx
index 0bd740c4..453555b1 100644
--- a/apps/web/src/hooks/useDocuments.tsx
+++ b/apps/web/src/hooks/useDocuments.tsx
@@ -73,7 +73,8 @@ function upsertDocumentInMemory(
       clock: 0,
       appClock: 0,
       userAppClock: {},
-      runUnexecutedBlocks: true,
+      runUnexecutedBlocks: false,
+      runSQLSelection: false,
       hasDashboard: false,
     })
   }
@@ -206,7 +207,7 @@ type API = {
   ) => Promise<void>
   updateDocumentSettings: (
     id: string,
-    settings: { runUnexecutedBlocks: boolean }
+    settings: { runUnexecutedBlocks?: boolean; runSQLSelection?: boolean }
   ) => Promise<void>
   publish: (id: string) => Promise<void>
 }
@@ -738,7 +739,10 @@ export function useDocuments(workspaceId: string): UseDocuments {
   )
 
   const updateDocumentSettings = useCallback(
-    async (id: string, settings: { runUnexecutedBlocks: boolean }) => {
+    async (
+      id: string,
+      settings: { runUnexecutedBlocks?: boolean; runSQLSelection?: boolean }
+    ) => {
       const document = documents.find((doc) => doc.id === id)
       if (!document) {
         return
diff --git a/packages/database/prisma/migrations/20241106182901_add_run_sql_selection_to_document_table/migration.sql b/packages/database/prisma/migrations/20241106182901_add_run_sql_selection_to_document_table/migration.sql
new file mode 100644
index 00000000..8eb76533
--- /dev/null
+++ b/packages/database/prisma/migrations/20241106182901_add_run_sql_selection_to_document_table/migration.sql
@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "Document" ADD COLUMN     "runSQLSelection" BOOLEAN NOT NULL DEFAULT false;
diff --git a/packages/database/prisma/schema.prisma b/packages/database/prisma/schema.prisma
index 434c5653..961ee19f 100644
--- a/packages/database/prisma/schema.prisma
+++ b/packages/database/prisma/schema.prisma
@@ -37,6 +37,7 @@ model Document {
   reusableComponents ReusableComponent[]
 
   runUnexecutedBlocks Boolean @default(false)
+  runSQLSelection     Boolean @default(false)
 }
 
 model YjsDocument {
diff --git a/packages/editor/src/blocks/sql.ts b/packages/editor/src/blocks/sql.ts
index 508755f5..5e3d2879 100644
--- a/packages/editor/src/blocks/sql.ts
+++ b/packages/editor/src/blocks/sql.ts
@@ -36,6 +36,7 @@ export type SQLBlock = BaseBlock<BlockType.SQL> & {
     | 'edit-with-ai-running'
     | 'fix-with-ai-requested'
     | 'fix-with-ai-running'
+  selectedCode: string | null
   dataframeName: DataframeName
   dataSourceId: string | null
   isFileDataSource: boolean
@@ -74,6 +75,7 @@ export const makeSQLBlock = (
     type: BlockType.SQL,
     title: '',
     source: new Y.Text(opts?.source ?? ''),
+    selectedCode: null,
     status: 'idle',
     dataframeName: getDataframeName(blocks),
     dataSourceId: opts?.dataSourceId ?? null,
@@ -107,6 +109,7 @@ export function getSQLAttributes(
     ...getBaseAttributes(block),
     source: getSQLSource(block),
     status: getAttributeOr(block, 'status', 'idle'),
+    selectedCode: getAttributeOr(block, 'selectedCode', null),
     dataframeName: getAttributeOr(
       block,
       'dataframeName',
@@ -148,6 +151,7 @@ export function duplicateSQLBlock(
     ...duplicateBaseAttributes(newId, prevAttributes),
     source: duplicateYText(prevAttributes.source),
     status: options?.noState ? 'idle' : prevAttributes.status,
+    selectedCode: prevAttributes.selectedCode,
     dataframeName: clone(prevAttributes.dataframeName),
     dataSourceId: prevAttributes.dataSourceId
       ? options?.datasourceMap?.get(prevAttributes.dataSourceId) ??
diff --git a/packages/editor/src/component.ts b/packages/editor/src/component.ts
index 09902370..f998f468 100644
--- a/packages/editor/src/component.ts
+++ b/packages/editor/src/component.ts
@@ -55,6 +55,7 @@ type SQLComponentAttrs = Omit<
   | 'index'
   | 'type'
   | 'status'
+  | 'selectedCode'
   | 'result'
   | 'lastQuery'
   | 'lastQueryTime'


Test patch:
diff --git a/apps/api/src/yjs/v2/executors/blocks/sql.test.ts b/apps/api/src/yjs/v2/executors/blocks/sql.test.ts
index 436f66d0..c082e7a3 100644
--- a/apps/api/src/yjs/v2/executors/blocks/sql.test.ts
+++ b/apps/api/src/yjs/v2/executors/blocks/sql.test.ts
@@ -42,6 +42,7 @@ describe('SQLExecutor', () => {
       listDataSources: jest.fn(),
       listDataFrames: jest.fn(),
       editWithAI: jest.fn(),
+      documentHasRunSQLSelectionEnabled: jest.fn(),
     }
     events = {
       aiUsage: jest.fn(),
@@ -131,6 +132,7 @@ describe('SQLExecutor', () => {
       await blocksExecutor.runQuery(
         block,
         new Y.Transaction(ydoc, {}, true),
+        false,
         false
       )
 
@@ -194,7 +196,8 @@ describe('SQLExecutor', () => {
       await blocksExecutor.runQuery(
         block,
         new Y.Transaction(ydoc, {}, true),
-        true
+        true,
+        false
       )
 
       expect(effects.makeSQLQuery).toHaveBeenCalledWith(
@@ -237,6 +240,7 @@ describe('SQLExecutor', () => {
       const runningPromise = blocksExecutor.runQuery(
         block,
         new Y.Transaction(ydoc, {}, true),
+        false,
         false
       )
       await blocksExecutor.abortQuery(block)
@@ -280,6 +284,7 @@ describe('SQLExecutor', () => {
         block,
 
         new Y.Transaction(ydoc, {}, true),
+        false,
         false
       )
 
diff --git a/apps/api/src/yjs/v2/observers/blocks/sql.test.ts b/apps/api/src/yjs/v2/observers/blocks/sql.test.ts
index a544af59..257477e2 100644
--- a/apps/api/src/yjs/v2/observers/blocks/sql.test.ts
+++ b/apps/api/src/yjs/v2/observers/blocks/sql.test.ts
@@ -210,7 +210,12 @@ describe.only('SQLObserver', () => {
             tr
           )
 
-          expect(executorMock.runQuery).toHaveBeenCalledWith(block, tr, false)
+          expect(executorMock.runQuery).toHaveBeenCalledWith(
+            block,
+            tr,
+            false,
+            false
+          )
           expect(block.getAttribute('status')).toEqual('idle')
         })
 
@@ -227,7 +232,12 @@ describe.only('SQLObserver', () => {
             tr
           )
 
-          expect(executorMock.runQuery).toHaveBeenCalledWith(block, tr, false)
+          expect(executorMock.runQuery).toHaveBeenCalledWith(
+            block,
+            tr,
+            false,
+            false
+          )
           expect(block.getAttribute('status')).toEqual('idle')
         })
       })
@@ -246,7 +256,12 @@ describe.only('SQLObserver', () => {
             tr
           )
 
-          expect(executorMock.runQuery).toHaveBeenCalledWith(block, tr, true)
+          expect(executorMock.runQuery).toHaveBeenCalledWith(
+            block,
+            tr,
+            true,
+            false
+          )
           expect(block.getAttribute('status')).toEqual('idle')
         })
 
@@ -263,7 +278,12 @@ describe.only('SQLObserver', () => {
             tr
           )
 
-          expect(executorMock.runQuery).toHaveBeenCalledWith(block, tr, true)
+          expect(executorMock.runQuery).toHaveBeenCalledWith(
+            block,
+            tr,
+            true,
+            false
+          )
           expect(block.getAttribute('status')).toEqual('idle')
         })
       })