fix copy options serialization

Author: emrberk

src/parser/ast.ts

@@ -627,6 +627,8 @@ export interface CopyOption extends AstNode {
   type: "copyOption"
   key: string
   value?: string | number | boolean | string[]
+  /** When true, the string value originated from a string literal and should be quoted in toSql. */
+  quoted?: boolean
 }
 
 export interface CheckpointStatement extends AstNode {

src/parser/cst-types.d.ts

@@ -1452,6 +1452,7 @@ export type CopyOptionCstChildren = {
   DataPageSize?: IToken[];
   StatisticsEnabled?: IToken[];
   ParquetVersion?: IToken[];
+  NumberLiteral?: IToken[];
   RawArrayEncoding?: IToken[];
 };
 

src/parser/parser.ts

@@ -2746,7 +2746,10 @@ class QuestDBParser extends CstParser {
       {
         ALT: () => {
           this.CONSUME(ParquetVersion)
-          this.SUBRULE2(this.stringOrIdentifier)
+          this.OR5([
+            { ALT: () => this.CONSUME(NumberLiteral) },
+            { ALT: () => this.SUBRULE2(this.stringOrIdentifier) },
+          ])
         },
       },
       {

src/parser/toSql.ts

@@ -1268,7 +1268,7 @@ function copyOptionToSql(opt: AST.CopyOption): string {
   if (opt.value === true) return `${opt.key} TRUE`
   if (opt.value === false) return `${opt.key} FALSE`
   if (typeof opt.value === "string")
-    return `${opt.key} ${escapeString(opt.value)}`
+    return `${opt.key} ${opt.quoted ? escapeString(opt.value) : opt.value}`
   if (Array.isArray(opt.value)) return `${opt.key} ${opt.value.join(" ")}`
   return `${opt.key} ${opt.value}`
 }

src/parser/visitor.ts

@@ -2298,16 +2298,26 @@ class QuestDBVisitor extends BaseVisitor {
 
     if (ctx.booleanLiteral) {
       result.value = this.visit(ctx.booleanLiteral) as boolean
+    } else if (ctx.NumberLiteral) {
+      // PARQUET_VERSION with bare number literal (e.g., PARQUET_VERSION 2)
+      result.value = parseInt(ctx.NumberLiteral[0].image, 10)
     } else if (ctx.stringOrIdentifier) {
       result.value = this.extractMaybeString(ctx.stringOrIdentifier[0])
+      // Mark as quoted when the stringOrIdentifier resolved to a string literal
+      const soiChildren = (ctx.stringOrIdentifier[0] as CstNode).children
+      if (soiChildren.StringLiteral) {
+        result.quoted = true
+      }
     } else if (ctx.StringLiteral) {
       result.value = ctx.StringLiteral[0].image.slice(1, -1)
+      result.quoted = true
     } else if (ctx.expression) {
       const expr = this.visit(ctx.expression) as AST.Expression
       if (expr?.type === "literal" && expr.literalType === "number") {
         result.value = expr.value as number
       } else if (expr?.type === "literal" && expr.literalType === "string") {
         result.value = expr.value as string
+        result.quoted = true
       } else if (expr?.type === "literal") {
         result.value = expr.raw ?? String(expr.value ?? "")
       } else {

tests/parser.test.ts

@@ -3210,6 +3210,74 @@ orders PIVOT (sum(amount) FOR status IN ('open'))`
         expect(result2.ast[0].type).toBe(result1.ast[0].type)
       })
     }
+
+    it("COPY TO round-trip: PARTITION_BY value is NOT quoted", () => {
+      const sql = "COPY trades TO '/export/trades' WITH PARTITION_BY MONTH"
+      const result = parseToAst(sql)
+      expect(result.errors).toHaveLength(0)
+
+      const roundtrip = toSql(result.ast[0])
+      expect(roundtrip).toContain("PARTITION_BY MONTH")
+      expect(roundtrip).not.toContain("PARTITION_BY 'MONTH'")
+
+      const result2 = parseToAst(roundtrip)
+      expect(result2.errors).toHaveLength(0)
+    })
+
+    it("COPY TO round-trip: FORMAT, PARTITION_BY and COMPRESSION_CODEC are NOT quoted", () => {
+      const sql =
+        "COPY trades TO '/export/trades' WITH FORMAT PARQUET " +
+        "PARTITION_BY MONTH COMPRESSION_CODEC ZSTD"
+      const result = parseToAst(sql)
+      expect(result.errors).toHaveLength(0)
+
+      const roundtrip = toSql(result.ast[0])
+      expect(roundtrip).toContain("FORMAT PARQUET")
+      expect(roundtrip).toContain("PARTITION_BY MONTH")
+      expect(roundtrip).toContain("COMPRESSION_CODEC ZSTD")
+
+      const result2 = parseToAst(roundtrip)
+      expect(result2.errors).toHaveLength(0)
+    })
+
+    it("COPY TO round-trip: PARQUET_VERSION accepts bare number literal", () => {
+      const sql =
+        "COPY trades TO '/export/trades' WITH FORMAT PARQUET PARQUET_VERSION 2"
+      const result = parseToAst(sql)
+      expect(result.errors).toHaveLength(0)
+
+      const roundtrip = toSql(result.ast[0])
+      expect(roundtrip).toContain("PARQUET_VERSION 2")
+
+      const result2 = parseToAst(roundtrip)
+      expect(result2.errors).toHaveLength(0)
+    })
+
+    it("COPY FROM round-trip: string literal options are still quoted", () => {
+      const sql = "COPY trades FROM '/data/trades.csv' WITH DELIMITER ','"
+      const result = parseToAst(sql)
+      expect(result.errors).toHaveLength(0)
+
+      const roundtrip = toSql(result.ast[0])
+      expect(roundtrip).toContain("DELIMITER ','")
+
+      const result2 = parseToAst(roundtrip)
+      expect(result2.errors).toHaveLength(0)
+    })
+
+    it("COPY TO with subquery and all Parquet options", () => {
+      const sql =
+        "COPY (SELECT * FROM trades WHERE timestamp IN '2024') TO '/export/trades' " +
+        "WITH FORMAT PARQUET PARTITION_BY MONTH COMPRESSION_CODEC ZSTD COMPRESSION_LEVEL 3 " +
+        "ROW_GROUP_SIZE 100000 DATA_PAGE_SIZE 1048576 STATISTICS_ENABLED true PARQUET_VERSION 2 " +
+        "RAW_ARRAY_ENCODING true"
+      const result = parseToAst(sql)
+      expect(result.errors).toHaveLength(0)
+
+      const roundtrip = toSql(result.ast[0])
+      const result2 = parseToAst(roundtrip)
+      expect(result2.errors).toHaveLength(0)
+    })
   })
 
   // ===========================================================================