fix: wh query types (#3951)

* fix: update query type * fix: select query type
rudderlabs · Oct 9, 2023 · 1e415e0 · 1e415e0
1 parent c94c953
commit 1e415e0
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 19 deletions.
diff --git a/warehouse/utils/querytype.go b/warehouse/utils/querytype.go
@@ -12,27 +12,27 @@ var (
 
 func init() {
  regexes := []string{
- "(?P<SELECT>SELECT)",
- "(?P<UPDATE>UPDATE\\s+.*SET)",
- "(?P<DELETE_FROM>DELETE\\s+FROM)",
- "(?P<INSERT_INTO>INSERT\\s+INTO)",
+ "(WITH.*\\(.*)*(?P<SELECT>SELECT)",
+ "(?P<UPDATE>UPDATE.*SET)",
+ "(?P<DELETE_FROM>DELETE.*FROM)",
+ "(?P<INSERT_INTO>INSERT.*INTO)",
  "(?P<COPY>COPY)",
- "(?P<MERGE_INTO>MERGE\\s+INTO)",
- "(?P<CREATE_TEMP_TABLE>CREATE\\s+TEMP(?:ORARY)*\\s+TABLE)",
- "(?P<CREATE_DATABASE>CREATE\\s+DATABASE)",
- "(?P<CREATE_SCHEMA>CREATE\\s+SCHEMA)",
- "(?P<CREATE_TABLE>(?:IF\\s+NOT\\s+EXISTS\\s+.*)*CREATE\\s+(?:OR\\s+REPLACE\\s+)*TABLE)",
- "(?P<CREATE_INDEX>CREATE\\s+INDEX)",
- "(?P<ALTER_TABLE>ALTER\\s+TABLE)",
- "(?P<ALTER_SESSION>ALTER\\s+SESSION)",
- "(?P<DROP_TABLE>(?:IF\\s+.*)*DROP\\s+TABLE)",
- "(?P<SHOW_TABLES>SHOW\\s+TABLES)",
- "(?P<SHOW_PARTITIONS>SHOW\\s+PARTITIONS)",
- "(?P<DESCRIBE_TABLE>DESCRIBE\\s+(?:QUERY\\s+)*TABLE)",
- "(?P<SET_TO>SET\\s+.*\\s+TO)",
+ "(?P<MERGE_INTO>MERGE.*INTO)",
+ "(?P<CREATE_TEMP_TABLE>CREATE.*TEMP(?:ORARY)*.*TABLE)",
+ "(?P<CREATE_DATABASE>CREATE.*DATABASE)",
+ "(?P<CREATE_SCHEMA>CREATE.*SCHEMA)",
+ "(?P<CREATE_TABLE>(?:IF.*NOT.*EXISTS.*)*CREATE.*(?:OR.*REPLACE.*)*TABLE)",
+ "(?P<CREATE_INDEX>CREATE.*INDEX)",
+ "(?P<ALTER_TABLE>ALTER.*TABLE)",
+ "(?P<ALTER_SESSION>ALTER.*SESSION)",
+ "(?P<DROP_TABLE>(?:IF.*)*DROP.*TABLE)",
+ "(?P<SHOW_TABLES>SHOW.*TABLES)",
+ "(?P<SHOW_PARTITIONS>SHOW.*PARTITIONS)",
+ "(?P<DESCRIBE_TABLE>DESCRIBE.*(?:QUERY.*)*TABLE)",
+ "(?P<SET_TO>SET.*TO)",
  }
 
- queryTypeRegex = regexp.MustCompile(`^(?i)\s*(` + strings.Join(regexes, "|") + `)\s+`)
+ queryTypeRegex = regexp.MustCompile(`^(?ism)\s*(` + strings.Join(regexes, "|") + `)\s+`)
 
  queryTypeIndexes = make(map[int]string)
  for i, n := range queryTypeRegex.SubexpNames() {

diff --git a/warehouse/utils/querytype_test.go b/warehouse/utils/querytype_test.go
@@ -15,7 +15,11 @@ func TestGetQueryType(t *testing.T) {
  }{
  {"select 1", "Select * from table", "SELECT", true},
  {"select 2", "\t\n\n \t\n\n seLeCt * from table", "SELECT", true},
- {"update", "\t\n\n \t\n\n UpDaTe something SET some_column = 'x'", "UPDATE", true},
+ {"select 3", "\n\t\tWITH load_files as (\n\t\t SELECT\n\t\t\tlocation,\n\t\t\tmetadata,\n\t\t\t " +
+ "FROM\n\t\t\tt1\n\t\t WHERE\n\t\t\ta IN (1,2,3) \n\t\t)\n\t\tSELECT\n\t\t location,\n\t\t " +
+ "metadata\n\t\tFROM\n\t\t t2\n\t\tWHERE\n\t\t x = 1\n\t\t;\n", "SELECT", true},
+ {"update 1", "\t\n\n \t\n\n UpDaTe something SET some_column = 'x'", "UPDATE", true},
+ {"update 2", "\n\t\tUPDATE\n\t\t t1\n\t\tSET\n\t\t a=$2,b=$3\n\t\tWHERE\n\t\t id = $1;", "UPDATE", true},
  {"delete", "\t\n\n \t\n\n DeLeTe FROm something", "DELETE_FROM", true},
  {"insert", "\t\n\n \t\n\n InSerT INTO something", "INSERT_INTO", true},
  {"copy", "\t\n\n \t\n\n cOpY t1 from t2", "COPY", true},