From 50a047e7e24a19fcf7d78401f67c4d344151313d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gunnar=20Sv=20Sigurbj=C3=B6rnsson?=
 <gunnar.sigurbjornsson@gmail.com>
Date: Fri, 27 Sep 2024 12:24:08 +0000
Subject: [PATCH] feat: Support raw string s in Rust

---
 src/extension.ts           | 25 +++++++++++++++++++++----
 test/testdata/multiline.rs |  1 +
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/extension.ts b/src/extension.ts
index 69db3e2..39343d3 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -3,7 +3,8 @@ import sqlLint from 'sql-lint';
 import * as configuration from './configuration';
 
 export const PHP_SQL = '<<<SQL';
-export const SQL_START_REGEX = /(?<token>"""|"|'''|'|`)--\s*sql/;
+export const SQL_START_REGEX =
+    /(?<token>r(?<hashes>#*)?"|"""|"|'''|'|`)--\s*sql/;
 
 async function checkRange(
     log: vscode.OutputChannel,
@@ -60,6 +61,7 @@ export async function refreshDiagnostics(
     let startRangePosition = -1;
     let sqlStringBound = '';
     let sqlStartLineIndex = -1;
+    let hashes = ''; // For Rust raw strings
 
     if (
         configuration.get<boolean>('lintSQLFiles') &&
@@ -88,15 +90,29 @@ export async function refreshDiagnostics(
         if (sqlStartLineIndex === -1) {
             if ((match = SQL_START_REGEX.exec(lineOfText)) !== null) {
                 startRangePosition = match.index + match.groups!.token.length;
-                sqlStringBound = match.groups!.token;
                 sqlStartLineIndex = lineIndex;
+                if (match.groups!.hashes !== undefined) {
+                    hashes = match.groups!.hashes;
+                    sqlStringBound = `"${hashes}`;
+                } else {
+                    sqlStringBound = match.groups!.token;
+                }
             } else if ((phpPatternStart = lineOfText.indexOf(PHP_SQL)) !== -1) {
                 startRangePosition = phpPatternStart + PHP_SQL.length;
                 sqlStringBound = 'SQL;';
                 sqlStartLineIndex = lineIndex;
             }
-        } else if (sqlStringBound !== '') {
-            let endSqlIndex = lineOfText.indexOf(sqlStringBound);
+        }
+        if (sqlStringBound !== '') {
+            let endSqlIndex = -1;
+            if (lineIndex === sqlStartLineIndex) {
+                endSqlIndex = lineOfText.indexOf(
+                    sqlStringBound,
+                    startRangePosition,
+                );
+            } else {
+                endSqlIndex = lineOfText.indexOf(sqlStringBound);
+            }
             if (endSqlIndex !== -1) {
                 sqlStringCnt += 1;
                 const range = new vscode.Range(
@@ -109,6 +125,7 @@ export async function refreshDiagnostics(
                 diagnostics.push(...subDiagnostics);
                 sqlStartLineIndex = -1;
                 sqlStringBound = '';
+                hashes = '';
             }
         }
     }
diff --git a/test/testdata/multiline.rs b/test/testdata/multiline.rs
index aa63332..db04070 100644
--- a/test/testdata/multiline.rs
+++ b/test/testdata/multiline.rs
@@ -4,4 +4,5 @@ select * from book;
 "#;
 
     let _another = r#"--sql; select * from book;"#;
+    let _another = r##"--sql; select * from "book";"##;
 }