<<< add column names from all tables present after the from keyword

Author: gdementen

larray_editor/editor.py

@@ -445,6 +445,7 @@ class SQLWidget(QTextEdit):
         "LIMIT", "OFFSET", "UNION", "ALL",
         "VALUES", "SET", "CASE", "WHEN", "THEN", "ELSE", "END"
     ]
+    SQL_KEYWORDS_SET = set(SQL_KEYWORDS)
 
     def __init__(self, editor_window):
         import polars as pl
@@ -477,15 +478,42 @@ def __init__(self, editor_window):
         self.sql_context = pl.SQLContext(eager=False)
         self.update_completer_options({})
 
-    def update_completer_options(self, data, selected=None):
-        data = self._filter_data_for_sql(data)
-        if selected is not None and self._handled_by_polars_sql(selected):
-            data['self'] = selected
-            col_names = selected.collect_schema().names()
+    def update_completer_options(self, data=None, selected=None):
+        if data is not None:
+            data = self._filter_data_for_sql(data)
+            if selected is not None and self._handled_by_polars_sql(selected):
+                data['self'] = selected
+            self.data = data
+            self.sql_context.register_many(data)
         else:
-            col_names = []
-        self.sql_context.register_many(data)
+            data = self.data
+
+        if 'self' in data:
+            table_names_to_fetch_columns = ['self']
+        else:
+            table_names_to_fetch_columns = []
+
         table_names = [k for k, v in data.items()]
+
+        # extract table names from the current FROM clause
+        query_text = self.toPlainText()
+        m = re.search(r'\s+FROM\s+(\S+)', query_text, re.IGNORECASE)
+        if m:
+            after_from = m.group(1)
+            # try any identifier found in the query after the FROM keyword
+            # there will probably be false positives if a column has the same
+            # name as another table but that should be rare
+            from_tables = [word for word in after_from.split()
+                           if word not in self.SQL_KEYWORDS_SET and word in data]
+            if from_tables:
+                table_names_to_fetch_columns = from_tables
+
+        # add column names from all the used tables or self, if present
+        col_names_set = set()
+        for table_name in table_names_to_fetch_columns:
+            col_names_set.update(set(data[table_name].collect_schema().names()))
+        col_names = sorted(col_names_set)
+
         logger.debug(f"available tables for SQL queries: {table_names}")
         logger.debug(f"available columns for SQL queries: {col_names}")
         completions = table_names + col_names + self.SQL_KEYWORDS
@@ -520,6 +548,7 @@ def insert_completion(self, completion):
         at_end = cursor.position() == len(self.toPlainText())
         cursor.insertText(completion + (' ' if at_end else ''))
         self.setTextCursor(cursor)
+        self.update_completer_options()
 
     def keyPressEvent(self, event):
         completer_popup = self.completer.popup()
@@ -574,14 +603,17 @@ def keyPressEvent(self, event):
                     self.clear()
                 return
         super().keyPressEvent(event)
+        self.update_completer_options()
         # we need to compute the prefix *after* the keypress event has been
-        # handled so that the prefix contains the last key stroke
+        # handled so that the prefix contains the last keystroke
         prefix = self.get_word_prefix()
         if prefix:
             self.completer.setCompletionPrefix(prefix)
             num_completion = self.completer.completionCount()
-            if (0 < num_completion <= self.completer.maxVisibleItems() and
-                    not completer_popup.isVisible()):
+            # we must show the popup even if it is already visible, because
+            # the number of completions might have changed and the popup size
+            # needs to be updated
+            if 0 < num_completion <= self.completer.maxVisibleItems():
                 self.show_autocomplete_popup()
             elif num_completion == 0 and completer_popup.isVisible():
                 completer_popup.hide()
@@ -625,6 +657,7 @@ def search_and_recall_history(self, direction: int):
             if self.history[index].startswith(prefix):
                 self.history_index = index
                 self.setPlainText(self.history[self.history_index])
+                self.update_completer_options()
                 cursor.setPosition(cursor_pos)
                 self.setTextCursor(cursor)
                 return True