ensure that table aliases are not in reserved keywords

defog-ai · Jun 4, 2024 · 0c5c4e6 · 0c5c4e6
1 parent 5de773d
commit 0c5c4e6
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/utils/gen_prompt.py b/utils/gen_prompt.py
@@ -60,16 +60,17 @@ def generate_aliases(table_names: list) -> str:
     Generate aliases for table names
     """
     aliases = {}
+    reserved_keywords = ['all', 'and', 'any', 'as', 'asc', 'do', 'end', 'for', 'in', 'is', 'not', 'to']
     for table_name in table_names:
         alias = table_name[0]
-        if alias in aliases.values() and "_" in table_name:
+        if (alias in aliases.values() and "_" in table_name) or alias.lower() in reserved_keywords:
             alias = table_name.split("_")[0] + table_name.split("_")[1]
-        if alias in aliases.values():
+        if alias in aliases.values() or alias.lower() in reserved_keywords:
             alias = table_name[:2]
-        if alias in aliases.values():
+        if alias in aliases.values() or alias.lower() in reserved_keywords:
             alias = table_name[:3]
         num = 2
-        while alias in aliases.values():
+        while (alias in aliases.values() or alias.lower() in reserved_keywords):
             alias = table_name[0] + str(num)
             num += 1