feat(vscode): Add codelens provider w/ generate client command (#1710)

Co-authored-by: Joël Galeran <Jolg42@users.noreply.github.com>

packages/vscode/package.json

@@ -125,6 +125,26 @@
           ],
           "default": "off",
           "description": "Setting for logging between the VS Code extension and the language server."
+        },
+        "prisma.enableCodeLens": {
+          "type": "boolean",
+          "default": true
+        },
+        "prisma.scriptRunner": {
+          "type": "string",
+          "enum": [
+            "npx",
+            "pnpm",
+            "yarn"
+          ],
+          "default": "npx"
+        },
+        "prisma.schemaPath": {
+          "type": "string",
+          "examples": [
+            "/path/to/your/schema.prisma"
+          ],
+          "description": "If you have a Prisma schema file in a custom path, you will need to provide said path `/path/to/your/schema.prisma` to run generate"
         }
       }
     },
@@ -143,6 +163,21 @@
         "command": "prisma.filewatcherDisable",
         "title": "Disable the File Watcher functionality for Prisma Client.",
         "category": "Prisma"
+      },
+      {
+        "command": "prisma.enableCodeLens",
+        "title": "Enable CodeLens",
+        "category": "Prisma"
+      },
+      {
+        "command": "prisma.disableCodeLens",
+        "title": "Disable Codelens",
+        "category": "Prisma"
+      },
+      {
+        "command": "prisma.generate",
+        "title": "Generate",
+        "category": "Prisma"
       }
     ]
   },

packages/vscode/src/CodeLensProvider.ts

@@ -0,0 +1,110 @@
+import * as vscode from 'vscode'
+import * as cp from 'child_process'
+
+/**
+ * CodelensProvider
+ */
+export class CodelensProvider implements vscode.CodeLensProvider {
+  private enabled: boolean
+
+  private generatorRegex: RegExp
+  private _onDidChangeCodeLenses: vscode.EventEmitter<void> = new vscode.EventEmitter<void>()
+  public readonly onDidChangeCodeLenses: vscode.Event<void> = this._onDidChangeCodeLenses.event
+
+  constructor() {
+    this.generatorRegex = /(generator +[a-zA-Z0-9]+ +{)/g
+    this.enabled = vscode.workspace.getConfiguration('prisma').get('enableCodeLens', true)
+
+    vscode.workspace.onDidChangeConfiguration((_) => {
+      this._onDidChangeCodeLenses.fire()
+    })
+  }
+
+  public provideCodeLenses(document: vscode.TextDocument, token: vscode.CancellationToken): vscode.CodeLens[] {
+    if (!this.enabled) {
+      return []
+    }
+
+    const codelenses = this.getCodeLensGenerateSchema(document, token)
+    return ([] as vscode.CodeLens[]).concat(...codelenses)
+  }
+
+  private getCodeLensGenerateSchema(document: vscode.TextDocument, token: vscode.CancellationToken): vscode.CodeLens[] {
+    const generatorRanges = this.getGeneratorRange(document, token)
+
+    const lenses = generatorRanges.map(
+      (range) =>
+        new vscode.CodeLens(range, {
+          title: 'Generate',
+          command: 'prisma.generate',
+          tooltip: `Run "prisma generate"`,
+          // ? (@druue) The arguments property does not seem to actually
+          // ? return an array of arguments. It would consistently
+          // ? return one singular string element, even when defined as:
+          // ? [this.scriptRunner, this.schemaPath]
+          // ?
+          // ? I've tried to understand why as there are usages in other
+          // ? codebases that do pass in multiple args so I have to imagine
+          // ? that it can work, but unsure.
+          // ? Reference: https://github.com/microsoft/vscode-extension-samples/blob/main/codelens-sample/
+          // ? arguments: [this.scriptRunner]
+        }),
+    )
+
+    return lenses
+  }
+
+  // ? (@druue) I really don't like finding it like this :|
+  private getGeneratorRange(document: vscode.TextDocument, _token: vscode.CancellationToken): vscode.Range[] {
+    const regex = new RegExp(this.generatorRegex)
+    const text = document.getText()
+
+    const ranges = []
+    let matches
+    while ((matches = regex.exec(text)) !== null) {
+      const line = document.lineAt(document.positionAt(matches.index).line)
+      const indexOf = line.text.indexOf(matches[0])
+      const position = new vscode.Position(line.lineNumber, indexOf)
+      const range = document.getWordRangeAtPosition(position, new RegExp(this.generatorRegex))
+      if (range) {
+        ranges.push(range)
+      }
+    }
+
+    return ranges
+  }
+}
+
+export function generateClient(_args: string) {
+  const prismaGenerateOutputChannel = vscode.window.createOutputChannel('Prisma Generate')
+  const rootPath = vscode.workspace.workspaceFolders?.[0].uri.path
+
+  const scriptRunner = vscode.workspace.getConfiguration('prisma').get('scriptRunner', 'npx')
+  const schemaPath: string | undefined = vscode.workspace.getConfiguration('prisma').get('schemaPath')
+
+  const pathArgFlag = ` --schema=${schemaPath}`
+  const cmd = `(cd ${rootPath} && ${scriptRunner} prisma generate${schemaPath ? pathArgFlag : ''})`
+
+  prismaGenerateOutputChannel.clear()
+  prismaGenerateOutputChannel.show(true)
+  prismaGenerateOutputChannel.appendLine(['Running prisma generate:', rootPath, cmd].join('\n- '))
+
+  const handleExec = (err: cp.ExecException | null, stdout: string, stderr: string) => {
+    try {
+      if (err) {
+        prismaGenerateOutputChannel.appendLine(err.message)
+        return
+      }
+      if (stdout) {
+        prismaGenerateOutputChannel.append(stdout)
+      }
+      if (stderr) {
+        prismaGenerateOutputChannel.append(stderr)
+      }
+    } catch (e) {
+      prismaGenerateOutputChannel.append(e as string)
+    }
+  }
+
+  cp.exec(cmd, (err, stdout, stderr) => handleExec(err, stdout, stderr))
+}

packages/vscode/src/plugins/prisma-language-server/index.ts

@@ -13,6 +13,7 @@ import {
   TextDocument,
   window,
   workspace,
+  languages,
 } from 'vscode'
 import {
   CodeAction as lsCodeAction,
@@ -35,6 +36,7 @@ import {
 import { PrismaVSCodePlugin } from '../types'
 import paths from 'env-paths'
 import FileWatcher from 'watcher'
+import { CodelensProvider, generateClient } from '../../CodeLensProvider'
 
 const packageJson = require('../../../../package.json') // eslint-disable-line
 
@@ -116,6 +118,9 @@ const plugin: PrismaVSCodePlugin = {
   enabled: () => true,
   activate: async (context) => {
     const isDebugOrTest = isDebugOrTestSession()
+    const codelensProvider = new CodelensProvider()
+
+    languages.registerCodeLensProvider('*', codelensProvider)
 
     setGenerateWatcher(!!workspace.getConfiguration('prisma').get('fileWatcher'))
 
@@ -213,11 +218,21 @@ const plugin: PrismaVSCodePlugin = {
         }
       }),
 
+      commands.registerCommand('prisma.generate', (args: string) => generateClient(args)),
+
       commands.registerCommand('prisma.restartLanguageServer', async () => {
         client = await restartClient(context, client, serverOptions, clientOptions)
         window.showInformationMessage('Prisma language server restarted.') // eslint-disable-line @typescript-eslint/no-floating-promises
       }),
 
+      commands.registerCommand('prisma.enableCodeLens', async () => {
+        await workspace.getConfiguration('prisma').update('enableCodeLens', true, true)
+      }),
+
+      commands.registerCommand('prisma.disableCodeLens', async () => {
+        await workspace.getConfiguration('prisma').update('enableCodeLens', false, true)
+      }),
+
       /* This command is part of the workaround for https://github.com/prisma/language-tools/issues/311 */
       commands.registerCommand('prisma.applySnippetWorkspaceEdit', applySnippetWorkspaceEdit()),