From 59a557ce86014ecba777ae78b5090675a1804c09 Mon Sep 17 00:00:00 2001
From: Kenneth Shaw <ken@usql.app>
Date: Sat, 26 Oct 2024 12:47:26 +0700
Subject: [PATCH] Fixing issues with escaping single characters

---
 stmt/parse.go      |  9 ++++-----
 stmt/parse_test.go |  2 +-
 stmt/stmt.go       | 10 ++++++----
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/stmt/parse.go b/stmt/parse.go
index e1971516ac..f5f42c664a 100644
--- a/stmt/parse.go
+++ b/stmt/parse.go
@@ -177,11 +177,11 @@ func readVarName(r []rune, i, end int, q rune) (int, rune) {
 // A command is defined as the first non-blank text after \, followed by
 // parameters up to either the next \ or a control character (for example, \n):
 func readCommand(r []rune, i, end int) (int, int) {
-command:
 	// find end of command
+	c, next := rune(0), rune(0)
+command:
 	for ; i < end; i++ {
-		next := grab(r, i+1, end)
-		switch {
+		switch next = grab(r, i+1, end); {
 		case next == 0:
 			return end, end
 		case next == '\\' || unicode.IsControl(next):
@@ -196,8 +196,7 @@ command:
 params:
 	// find end of params
 	for ; i < end; i++ {
-		c, next := r[i], grab(r, i+1, end)
-		switch {
+		switch c, next = r[i], grab(r, i+1, end); {
 		case next == 0:
 			return cmd, end
 		case quote == 0 && (c == '\'' || c == '"' || c == '`'):
diff --git a/stmt/parse_test.go b/stmt/parse_test.go
index 53780bcfeb..047a6c7a02 100644
--- a/stmt/parse_test.go
+++ b/stmt/parse_test.go
@@ -537,7 +537,7 @@ func TestSubstituteVar(t *testing.T) {
 	for i, test := range tests {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
 			z := []rune(test.s)
-			y, l := substituteVar(z, test.v, test.sub)
+			y, l := test.v.Substitute(z, test.sub, true)
 			if sl := len([]rune(test.sub)); test.v.Len != sl {
 				t.Errorf("expected v.Len to be %d, got: %d", sl, test.v.Len)
 			}
diff --git a/stmt/stmt.go b/stmt/stmt.go
index e731287b46..ff6608f7ab 100644
--- a/stmt/stmt.go
+++ b/stmt/stmt.go
@@ -71,7 +71,7 @@ func (b *Stmt) RawString() string {
 	i, s, z := 0, string(b.Buf), new(bytes.Buffer)
 	// deinterpolate vars
 	for _, v := range b.Vars {
-		if !v.Defined {
+		if !v.Defined && v.Quote != '\\' {
 			continue
 		}
 		if len(s) > i {
@@ -222,8 +222,6 @@ parse:
 		case b.quote != 0 || b.multilineComment || b.balanceCount != 0:
 		// skip escaped backslash, semicolon, colon
 		case c == '\\' && (next == '\\' || next == ';' || next == ':'):
-			// FIXME: the below works, but it may not make sense to keep this enabled.
-			// FIXME: also, the behavior is slightly different than psql
 			v := &Var{
 				I:     i,
 				End:   i + 2,
@@ -234,6 +232,7 @@ parse:
 			if b.r, b.rlen = v.Substitute(b.r, string(next), false); b.Len != 0 {
 				v.I += b.Len + 1
 			}
+			i++
 		// start of command
 		case c == '\\':
 			// parse command and params end positions
@@ -373,7 +372,10 @@ func (v *Var) String() string {
 
 // Substitute substitutes part of r, with s.
 func (v *Var) Substitute(r []rune, s string, ok bool) ([]rune, int) {
-	if v.Quote == '?' {
+	switch v.Quote {
+	case '\\':
+		s = "\\" + s
+	case '?':
 		s = trueFalse(ok)
 	}
 	sr, rcap := []rune(s), cap(r)