fix: don't break on nested segments (#62)

scriptcoded · Aug 16, 2022 · 997ee9d · 997ee9d
1 parent 97d911e
commit 997ee9d
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 2 deletions.
diff --git a/index.test.js b/index.test.js
@@ -99,6 +99,21 @@ describe('unicode', () => {
     expect(hlUni("SELECT COUNT(id), `id`, `username` FROM `users` WHERE `email` = 'test@example.com' AND `foo` = 'BAR' OR 1=1"))
       .toBe("[keyword]SELECT[clear] [function]COUNT[clear][bracket]([clear]id[bracket])[clear][special],[clear] [string]`id`[clear][special],[clear] [string]`username`[clear] [keyword]FROM[clear] [string]`users`[clear] [keyword]WHERE[clear] [string]`email`[clear] [special]=[clear] [string]'test@example.com'[clear] [keyword]AND[clear] [string]`foo`[clear] [special]=[clear] [string]'BAR'[clear] [keyword]OR[clear] [number]1[clear][special]=[clear][number]1[clear]")
   })
+
+  it('query with identifiers without apostrophes', () => {
+    expect(hlUni('SELECT id FROM users'))
+      .toBe('[keyword]SELECT[clear] id [keyword]FROM[clear] users')
+  })
+
+  it('query with nested segments (minus in string)', () => {
+    expect(hlUni('DROP PROCEDURE IF EXISTS `some-database`.`some-table`;'))
+      .toBe('[keyword]DROP[clear] [keyword]PROCEDURE[clear] [keyword]IF[clear] [keyword]EXISTS[clear] [string]`some-database`[clear].[string]`some-table`[clear][special];[clear]')
+  })
+
+  it('multiple queries', () => {
+    expect(hlUni('SELECT * FROM a;SELECT * FROM b;'))
+      .toBe('[keyword]SELECT[clear] [special]*[clear] [keyword]FROM[clear] a[special];[clear][keyword]SELECT[clear] [special]*[clear] [keyword]FROM[clear] b[special];[clear]')
+  })
 })
 
 describe('html', () => {
@@ -181,6 +196,16 @@ describe('html', () => {
     expect(hlHtml('SELECT id FROM users'))
       .toBe('<span class="sql-hl-keyword">SELECT</span> id <span class="sql-hl-keyword">FROM</span> users')
   })
+
+  it('query with nested segments (minus in string)', () => {
+    expect(hlHtml('DROP PROCEDURE IF EXISTS `some-database`.`some-table`;'))
+      .toBe('<span class="sql-hl-keyword">DROP</span> <span class="sql-hl-keyword">PROCEDURE</span> <span class="sql-hl-keyword">IF</span> <span class="sql-hl-keyword">EXISTS</span> <span class="sql-hl-string">`some-database`</span>.<span class="sql-hl-string">`some-table`</span><span class="sql-hl-special">;</span>')
+  })
+
+  it('multiple queries', () => {
+    expect(hlHtml('SELECT * FROM a;SELECT * FROM b;'))
+      .toBe('<span class="sql-hl-keyword">SELECT</span> <span class="sql-hl-special">*</span> <span class="sql-hl-keyword">FROM</span> a<span class="sql-hl-special">;</span><span class="sql-hl-keyword">SELECT</span> <span class="sql-hl-special">*</span> <span class="sql-hl-keyword">FROM</span> b<span class="sql-hl-special">;</span>')
+  })
 })
 
 describe('getSegments', () => {

diff --git a/lib/index.js b/lib/index.js
@@ -87,7 +87,7 @@ function getSegments (sqlString) {
   const segments = []
   let upperBound = 0
   for (let i = 0; i < sortedMatches.length; i++) {
-    if (sortedMatches[i].start < upperBound) { break }
+    if (sortedMatches[i].start < upperBound) { continue }
 
     // If no match, add a default segment
     if (sortedMatches[i].start > upperBound) {

diff --git a/test/index.js b/test/index.js
@@ -3,7 +3,7 @@
  * that. It's just to see that it's working in the console.
  */
 
-const { highlight } = require('../index')
+const { highlight } = require('../lib')
 
 console.log(highlight("SELECT COUNT(id), COUNT(id), `id`, `username` FROM `users` WHERE `email` = 'test@example.com' AND `something` = 'oke' AND 1=1"))
 console.log(highlight('SELECT "users".* FROM "users"'))
@@ -21,3 +21,7 @@ console.log(highlight('SELECT id FROM listings WHERE status = "not available"'))
 
 console.log(highlight('SELECT \'{"json_index":"json_value"}\' AS test;'))
 console.log(highlight('SELECT "This is a \\"text\\" test" AS text;'))
+
+console.log(highlight('DROP PROCEDURE IF EXISTS `some-database`.`some-table`;'))
+
+console.log(highlight('SELECT * FROM a;SELECT * FROM b;'))