From ef2dbb0b6ade7987cdb9def693bde182cacb7360 Mon Sep 17 00:00:00 2001
From: Younghoon Kim <yh.kim@databricks.com>
Date: Fri, 22 Nov 2024 14:50:32 -0800
Subject: [PATCH] apply single quote escape

---
 src/compile/data/timeunit.ts       | 13 +++++--------
 src/util.ts                        |  4 ++++
 test/compile/data/timeunit.test.ts | 30 ++++++++++++++++++++++++++++++
 3 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/src/compile/data/timeunit.ts b/src/compile/data/timeunit.ts
index 7e4bda8b49..54bc9e33a8 100644
--- a/src/compile/data/timeunit.ts
+++ b/src/compile/data/timeunit.ts
@@ -10,7 +10,7 @@ import {
   normalizeTimeUnit
 } from '../../timeunit';
 import {TimeUnitTransform} from '../../transform';
-import {Dict, duplicate, entries, hash, isEmpty, replacePathInField, vals} from '../../util';
+import {Dict, duplicate, entries, escapeSingleQuotes, hash, isEmpty, replacePathInField, vals} from '../../util';
 import {ModelWithField, isUnitModel} from '../model';
 import {DataFlowNode} from './dataflow';
 import {isRectBasedMark} from '../../mark';
@@ -38,10 +38,7 @@ export class TimeUnitNode extends DataFlowNode {
     return new TimeUnitNode(null, duplicate(this.timeUnits));
   }
 
-  constructor(
-    parent: DataFlowNode,
-    private timeUnits: Dict<TimeUnitComponent>
-  ) {
+  constructor(parent: DataFlowNode, private timeUnits: Dict<TimeUnitComponent>) {
     super(parent);
   }
 
@@ -218,7 +215,7 @@ function offsetExpr({timeUnit, field, reverse}: {timeUnit: TimeUnitParams; field
   const smallestUnit = getSmallestTimeUnitPart(unit);
   const {part, step} = getDateTimePartAndStep(smallestUnit, timeUnit.step);
   const offsetFn = utc ? 'utcOffset' : 'timeOffset';
-  const expr = `${offsetFn}('${part}', datum['${field}'], ${reverse ? -step : step})`;
+  const expr = `${offsetFn}('${part}', datum['${escapeSingleQuotes(field)}'], ${reverse ? -step : step})`;
   return expr;
 }
 
@@ -228,8 +225,8 @@ function offsetedRectFormulas(
   timeUnit: TimeUnitParams
 ): VgFormulaTransform[] {
   if (rectBandPosition !== undefined && rectBandPosition !== 0.5) {
-    const startExpr = `datum['${startField}']`;
-    const endExpr = `datum['${endField}']`;
+    const startExpr = `datum['${escapeSingleQuotes(startField)}']`;
+    const endExpr = `datum['${escapeSingleQuotes(endField)}']`;
     return [
       {
         type: 'formula',
diff --git a/src/util.ts b/src/util.ts
index 5159e1fed0..ebfbe29755 100644
--- a/src/util.ts
+++ b/src/util.ts
@@ -297,6 +297,10 @@ function escapePathAccess(string: string) {
   return string.replace(/(\[|\]|\.|'|")/g, '\\$1');
 }
 
+export function escapeSingleQuotes(value: string) {
+  return value.replaceAll("'", "\\'");
+}
+
 /**
  * Replaces path accesses with access to non-nested field.
  * For example, `foo["bar"].baz` becomes `foo\\.bar\\.baz`.
diff --git a/test/compile/data/timeunit.test.ts b/test/compile/data/timeunit.test.ts
index 85cce49580..9a9e029755 100644
--- a/test/compile/data/timeunit.test.ts
+++ b/test/compile/data/timeunit.test.ts
@@ -82,6 +82,36 @@ describe('compile/data/timeunit', () => {
       ]);
     });
 
+    it('should return a unit transform for bar with bandPosition=0 and escaped field name', () => {
+      const model = parseUnitModel({
+        data: {values: []},
+        mark: 'bar',
+        encoding: {
+          x: {field: "\\'a\\'", type: 'temporal', timeUnit: 'month', bandPosition: 0}
+        }
+      });
+
+      expect(assembleFromEncoding(model)).toEqual([
+        {
+          type: 'timeunit',
+          field: "\\'a\\'",
+          as: ["month_'a'", "month_'a'_end"],
+          units: ['month']
+        },
+        {
+          type: 'formula',
+          expr: "0.5 * timeOffset('month', datum['month_\\'a\\''], -1) + 0.5 * datum['month_\\'a\\'']",
+          as: `month_'a'_${OFFSETTED_RECT_START_SUFFIX}`
+        },
+
+        {
+          type: 'formula',
+          expr: "0.5 * datum['month_\\'a\\''] + 0.5 * datum['month_\\'a\\'_end']",
+          as: `month_'a'_${OFFSETTED_RECT_END_SUFFIX}`
+        }
+      ]);
+    });
+
     it('should return a timeunit transform with step for bar with bandPosition=0', () => {
       const model = parseUnitModel({
         data: {values: []},