From 997ee9deb9d25f48ce2b20c78c13ea22798fe8c4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Malcolm=20Nihl=C3=A9n?= <malcolm.nihlen@gmail.com>
Date: Tue, 16 Aug 2022 17:51:40 +0200
Subject: [PATCH] fix: don't break on nested segments (#62)

---
 index.test.js | 25 +++++++++++++++++++++++++
 lib/index.js  |  2 +-
 test/index.js |  6 +++++-
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/index.test.js b/index.test.js
index caf7909..de18053 100644
--- 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
index f3f9bdf..b0ceadb 100644
--- 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
index 41c6088..c0f492f 100644
--- 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;'))