Merge pull request #20 from jwills/jwills_like_escape

Handle the LIKE <pattern> ESCAPE <token> pattern in Presto queries from DBeaver

buenavista/bv_dialects.py

@@ -45,31 +45,51 @@ def _duckdb_command_handler(self, expression):
         entity = tokens[0].upper()
         if entity == "CATALOGS":
             q = "SELECT DISTINCT catalog_name as Catalog FROM information_schema.schemata"
-            if len(tokens) == 3:
-                q += " WHERE catalog_name LIKE " + tokens[2]
+            if len(tokens) >= 3:
+                if len(tokens) == 5:
+                    like = tokens[2].replace(tokens[4], "")
+                else:
+                    like = tokens[2]
+                q += " WHERE catalog_name LIKE " + like
             else:
                 return q
         elif entity == "SCHEMAS":
             q = "SELECT DISTINCT schema_name as Schema FROM information_schema.schemata"
             if len(tokens) > 1 and tokens[1].upper() == "FROM":
                 q += f" WHERE catalog_name = '{tokens[2]}'"
-                if len(tokens) == 5:
-                    q += " AND schema_name LIKE " + tokens[4]
+                if len(tokens) >= 5:
+                    if len(tokens) == 7:
+                        like = tokens[4].replace(tokens[6], "")
+                    else:
+                        like = tokens[4]
+                    q += " AND schema_name LIKE " + like
             else:
                 q += " WHERE catalog_name IN (SELECT current_database())"
-                if len(tokens) == 3:
-                    q += " AND schema_name LIKE " + tokens[2]
+                if len(tokens) >= 3:
+                    if len(tokens) == 5:
+                        like = tokens[2].replace(tokens[4], "")
+                    else:
+                        like = tokens[2]
+                    q += " AND schema_name LIKE " + like
             return q
         elif entity == "TABLES":
             q = "SELECT DISTINCT table_name as Table from information_schema.tables"
             if len(tokens) > 1 and tokens[1].upper() == "FROM":
                 q += f" WHERE schema_name = '{tokens[2]}'"
-                if len(tokens) == 5:
-                    q += " AND table_name LIKE " + tokens[4]
+                if len(tokens) >= 5:
+                    if len(tokens) == 7:
+                        like = tokens[4].replace(tokens[6], "")
+                    else:
+                        like = tokens[4]
+                    q += " AND table_name LIKE " + like
             else:
                 q += " WHERE catalog_name IN (SELECT current_schema())"
-                if len(tokens) == 3:
-                    q += " AND table_name LIKE " + tokens[2]
+                if len(tokens) >= 3:
+                    if len(tokens) == 5:
+                        like = tokens[2].replace(tokens[4], "")
+                    else:
+                        like = tokens[2]
+                    q += " AND table_name LIKE " + like
             return q
         elif entity == "COLUMNS" and tokens[1].upper() == "FROM":
             return f"DESCRIBE {tokens[2]}"

buenavista/examples/duckdb_http.py

@@ -1,4 +1,5 @@
 import os
+import re
 
 import duckdb
 from fastapi import FastAPI
@@ -8,7 +9,25 @@
 from ..http import main
 
 #### Rewriter setup/config
-rewriter = rewrite.Rewriter(bv_dialects.BVTrino(), bv_dialects.BVDuckDB())
+
+ESCAPE_PATTERN = re.compile(r"LIKE (\S+) ESCAPE (\S+)")
+
+
+def _escape_replace(match):
+    token = match.group(1)
+    escape = match.group(2).strip("'") or "\\"
+    return "LIKE " + token.replace(escape, "")
+
+
+class DuckDBHTTPRewriter(rewrite.Rewriter):
+    def rewrite(self, sql: str) -> str:
+        sql = super().rewrite(sql)
+        sql = ESCAPE_PATTERN.sub(_escape_replace, sql)
+        print("Rewritten to %s" % sql)
+        return sql
+
+
+rewriter = DuckDBHTTPRewriter(bv_dialects.BVTrino(), bv_dialects.BVDuckDB())
 
 
 @rewriter.relation("system.jdbc.tables")