Pretty sql (#509)

Make it look nicer, `WITH` broke it.
QuesmaOrg · Jul 10, 2024 · 42ad685 · 42ad685
1 parent 4590a4b
commit 42ad685
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 5 deletions.
diff --git a/quesma/util/sql_pretty_fmt.go b/quesma/util/sql_pretty_fmt.go
@@ -34,11 +34,7 @@ func SqlPrettyPrint(sqlData []byte) string {
 
 		// Skip original whitespace
 		if token.Type == sqllexer.WS {
-			if strings.ContainsRune(token.Value, '\n') {
-				continue
-			} else {
-				token.Value = " "
-			}
+			token.Value = " "
 		}
 
 		// Add new line if needed
@@ -47,6 +43,14 @@ func SqlPrettyPrint(sqlData []byte) string {
 			lineLength = 0
 			isBreakIndent = false
 		}
+		if token.Value == "WITH" {
+			if len(stack) > 0 {
+				sb.WriteString("\n\n")
+				lineLength = 0
+			}
+			isBreakIndent = false
+			stack = []string{token.Value}
+		}
 		if token.Value == "SELECT" {
 			if len(stack) > 0 && stack[len(stack)-1] == "SELECT" {
 				stack = stack[:len(stack)-1]
@@ -55,6 +59,9 @@ func SqlPrettyPrint(sqlData []byte) string {
 					lineLength = 0
 				}
 			}
+			if len(stack) > 0 && stack[len(stack)-1] == "WITH" {
+				stack = stack[:len(stack)-1]
+			}
 			if len(stack) > 0 {
 				subQueryIndent += 1
 			}

diff --git a/quesma/util/sql_pretty_fmt_test.go b/quesma/util/sql_pretty_fmt_test.go
@@ -107,3 +107,29 @@ ORDER BY count() DESC`
 	sqlFormatted := SqlPrettyPrint([]byte(expect))
 	assert.Equal(t, expect, sqlFormatted)
 }
+
+func TestSqlWith(t *testing.T) {
+	sql := `SELECT count() FROM "kibana_sample_data_ecommerce"
+WITH subQuery_1 AS (SELECT "animalType" AS "subQuery_1_1", count() AS "subQuery_1_cnt" FROM "default"."animal_index" WHERE ("date">=parseDateTime64BestEffort('2024-04-17T08:53:18.456Z') AND "date"<=parseDateTime64BestEffort('2024-07-10T08:53:18.456Z')) GROUP BY "animalType" ORDER BY count() DESC, "animalType" LIMIT 5) SELECT "animalType", "zooName", count() FROM "default"."animal_index" INNER JOIN "subQuery_1" ON "animalType" = "subQuery_1_1" WHERE ("date">=parseDateTime64BestEffort('2024-04-17T08:53:18.456Z') AND "date"<=parseDateTime64BestEffort('2024-07-10T08:53:18.456Z')) GROUP BY "animalType", "zooName", subQuery_1_cnt ORDER BY subQuery_1_cnt DESC, "animalType", count() DESC, "zooName" LIMIT 6`
+	expect := `SELECT count()
+FROM "kibana_sample_data_ecommerce"
+
+WITH subQuery_1 AS (
+  SELECT "animalType" AS "subQuery_1_1", count() AS "subQuery_1_cnt"
+  FROM "default"."animal_index"
+  WHERE ("date">=parseDateTime64BestEffort('2024-04-17T08:53:18.456Z') AND
+    "date"<=parseDateTime64BestEffort('2024-07-10T08:53:18.456Z'))
+  GROUP BY "animalType"
+  ORDER BY count() DESC, "animalType"
+  LIMIT 5)
+SELECT "animalType", "zooName", count()
+FROM "default"."animal_index" INNER JOIN "subQuery_1" ON "animalType" =
+  "subQuery_1_1"
+WHERE ("date">=parseDateTime64BestEffort('2024-04-17T08:53:18.456Z') AND "date"
+  <=parseDateTime64BestEffort('2024-07-10T08:53:18.456Z'))
+GROUP BY "animalType", "zooName", subQuery_1_cnt
+ORDER BY subQuery_1_cnt DESC, "animalType", count() DESC, "zooName"
+LIMIT 6`
+	sqlFormatted := SqlPrettyPrint([]byte(sql))
+	assert.Equal(t, expect, sqlFormatted)
+}