From 46e8a61027b633194d44ce8c110b1843da80cecf Mon Sep 17 00:00:00 2001
From: Han Wang <hwang@lyft.com>
Date: Sun, 7 Aug 2022 06:18:34 +0000
Subject: [PATCH] Add more fsql highlight scenarios

---
 package.json |  2 +-
 src/index.ts |  2 ++
 src/utils.ts | 24 +++++++++++++++++++++++-
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/package.json b/package.json
index 433bf96..8be9817 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "fugue-jupyter",
-  "version": "0.1.5",
+  "version": "0.1.6",
   "description": "Jupyterlab Extension for Fugue",
   "keywords": [
     "jupyter",
diff --git a/src/index.ts b/src/index.ts
index 66a17e1..c08ecf5 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -10,6 +10,7 @@ import { IEditorTracker } from '@jupyterlab/fileeditor';
 import {
   cellMagicExtractor,
   markerExtractor,
+  fsqlBlockExtractor,
   sqlCodeMirrorModesFor,
   registerCodeMirrorFor
 } from './utils';
@@ -126,6 +127,7 @@ const plugin: JupyterFrontEndPlugin<void> = {
     // and into a virtual document which is then passed to the sql-language-server
     // for code completion evaluation
     lspExtractorsMgr.register(markerExtractor('fsql'), 'python');
+    lspExtractorsMgr.register(fsqlBlockExtractor('fsql'), 'python');
     lspExtractorsMgr.register(cellMagicExtractor('fsql'), 'python');
     console.log('fugue-jupyter LSP extractors registered');
   }
diff --git a/src/utils.ts b/src/utils.ts
index 27d5b31..809567a 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -3,7 +3,7 @@ import { Mode } from 'codemirror';
 import { ICodeMirror } from '@jupyterlab/codemirror';
 
 function cell_magic(language: string) {
-  return `%%${language}`;
+  return `%%${language}.*`;
 }
 function start(language: string) {
   return `--start-${language}`;
@@ -11,6 +11,12 @@ function start(language: string) {
 function end(language: string) {
   return `--end-${language}`;
 }
+function fsql_start() {
+  return `fsql[\\s\\S]*\\([\\s\\S]*\\"\\"\\"`;
+}
+function fsql_end() {
+  return `\\"\\"\\"`;
+}
 
 const BEGIN = '(?:^|\n)';
 
@@ -28,6 +34,12 @@ export function sqlCodeMirrorModesFor(
       parseDelimiters: true,
       mode: sqlMode
     },
+    {
+      open: (RegExp(`${fsql_start()}`) as unknown) as string,
+      close: (RegExp(`${fsql_end()}`) as unknown) as string,
+      parseDelimiters: false,
+      mode: sqlMode
+    },
     {
       open: (RegExp(`${cell_magic(language)}`) as unknown) as string,
       close: '__A MARKER THAT WILL NEVER BE MATCHED__', // Cell magic: capture chars till the end of the cell
@@ -59,6 +71,16 @@ export function markerExtractor(language: string): RegExpForeignCodeExtractor {
   });
 }
 
+export function fsqlBlockExtractor(language: string): RegExpForeignCodeExtractor {
+  return new RegExpForeignCodeExtractor({
+    language: language,
+    pattern: `${fsql_start()}.*?\n([^]*?)${fsql_end()}`,
+    foreign_capture_groups: [1],
+    is_standalone: true,
+    file_extension: language
+  });
+}
+
 function set(str: string) {
   const obj: any = {},
     words = str.split(' ');