Fix autocomplete alias resolution for qualified tables

Author: Maxteabag

sqlit/domains/query/completion/core.py

@@ -426,6 +426,23 @@ def build_alias_map(refs: list[TableRef], known_tables: list[str]) -> dict[str,
     Only includes aliases for tables that exist in known_tables.
     """
     known_lower = {t.lower() for t in known_tables}
+
+    def strip_identifier(value: str) -> str:
+        value = value.strip()
+        if len(value) >= 2 and value[0] == value[-1] and value[0] in {"`", '"'}:
+            return value[1:-1]
+        if value.startswith("[") and value.endswith("]") and len(value) >= 2:
+            return value[1:-1]
+        return value
+
+    # Add unqualified table names for qualified entries (db.schema.table, schema.table, etc.).
+    for table in known_tables:
+        parts = [strip_identifier(part) for part in table.split(".")]
+        # Take the last non-empty component as the unqualified table name.
+        for part in reversed(parts):
+            if part:
+                known_lower.add(part.lower())
+                break
     alias_map: dict[str, str] = {}
 
     for ref in refs:

sqlit/domains/query/ui/mixins/autocomplete_suggestions.py

@@ -37,10 +37,27 @@ def _build_alias_map(self: AutocompleteMixinHost, text: str) -> dict[str, str]:
         table_refs = extract_table_refs(text)
         known_tables = set(t.lower() for t in self._schema_cache.get("tables", []))
         known_tables.update(t.lower() for t in self._schema_cache.get("views", []))
+        table_metadata = getattr(self, "_table_metadata", {}) or {}
+        known_metadata = {key.lower() for key in table_metadata.keys()}
 
         alias_map: dict[str, str] = {}
+
+        def is_known(name: str) -> bool:
+            lowered = name.lower()
+            return lowered in known_tables or lowered in known_metadata
+
         for ref in table_refs:
-            if ref.alias and ref.name.lower() in known_tables:
+            if not ref.alias:
+                continue
+
+            # Prefer schema-qualified name when present (db.table or schema.table).
+            if ref.schema:
+                qualified = f"{ref.schema}.{ref.name}"
+                if is_known(qualified):
+                    alias_map[ref.alias.lower()] = qualified
+                    continue
+
+            if is_known(ref.name):
                 alias_map[ref.alias.lower()] = ref.name
         return alias_map
 

tests/integration/test_autocomplete_mysql_columns.py

@@ -0,0 +1,155 @@
+"""Integration tests for MySQL autocomplete column suggestions."""
+
+from __future__ import annotations
+
+import os
+import tempfile
+
+import pytest
+
+from sqlit.domains.shell.app.main import SSMSTUI
+from tests.fixtures.mysql import MYSQL_HOST, MYSQL_PASSWORD, MYSQL_PORT, MYSQL_USER
+from tests.helpers import ConnectionConfig
+from tests.integration.browsing_base import wait_for_condition
+
+
+EXTRA_DB_NAME = os.environ.get("MYSQL_AUTOCOMPLETE_EXTRA_DB", "test_sqlit_autocomplete_other")
+
+
+@pytest.fixture
+def temp_config_dir():
+    """Create a temporary config directory for tests."""
+    with tempfile.TemporaryDirectory(prefix="sqlit-test-") as tmpdir:
+        original = os.environ.get("SQLIT_CONFIG_DIR")
+        os.environ["SQLIT_CONFIG_DIR"] = tmpdir
+        yield tmpdir
+        if original:
+            os.environ["SQLIT_CONFIG_DIR"] = original
+        else:
+            os.environ.pop("SQLIT_CONFIG_DIR", None)
+
+
+@pytest.fixture
+def mysql_extra_db(mysql_server_ready: bool):
+    """Create a second user database to force multi-db autocomplete mode."""
+    if not mysql_server_ready:
+        pytest.skip("MySQL is not available")
+
+    try:
+        import pymysql
+    except ImportError:
+        pytest.skip("PyMySQL is not installed")
+
+    conn = pymysql.connect(
+        host=MYSQL_HOST,
+        port=MYSQL_PORT,
+        user=MYSQL_USER,
+        password=MYSQL_PASSWORD,
+        database="mysql",
+        connect_timeout=10,
+    )
+    cursor = conn.cursor()
+    cursor.execute(f"CREATE DATABASE IF NOT EXISTS `{EXTRA_DB_NAME}`")
+    conn.commit()
+    conn.close()
+
+    yield EXTRA_DB_NAME
+
+    conn = pymysql.connect(
+        host=MYSQL_HOST,
+        port=MYSQL_PORT,
+        user=MYSQL_USER,
+        password=MYSQL_PASSWORD,
+        database="mysql",
+        connect_timeout=10,
+    )
+    cursor = conn.cursor()
+    cursor.execute(f"DROP DATABASE IF EXISTS `{EXTRA_DB_NAME}`")
+    conn.commit()
+    conn.close()
+
+
+def _set_cursor_to_end(app: SSMSTUI) -> None:
+    text = app.query_input.text
+    lines = text.split("\n")
+    app.query_input.cursor_location = (len(lines) - 1, len(lines[-1]) if lines else 0)
+
+
+@pytest.mark.asyncio
+async def test_mysql_alias_autocomplete_shows_columns_multi_db(
+    mysql_server_ready: bool,
+    mysql_db: str,
+    mysql_extra_db: str,  # noqa: ARG001 - fixture used for setup/teardown
+    temp_config_dir: str,  # noqa: ARG001 - fixture used for isolation
+) -> None:
+    """Autocomplete should resolve alias columns even with multi-db schema cache."""
+    if not mysql_server_ready:
+        pytest.skip("MySQL is not available")
+
+    config = ConnectionConfig(
+        name="test-mysql-autocomplete",
+        db_type="mysql",
+        server=MYSQL_HOST,
+        port=str(MYSQL_PORT),
+        database="",
+        username=MYSQL_USER,
+        password=MYSQL_PASSWORD,
+    )
+
+    app = SSMSTUI()
+    async with app.run_test(size=(120, 40)) as pilot:
+        await pilot.pause(0.1)
+
+        app.connections = [config]
+        app.refresh_tree()
+        await pilot.pause(0.1)
+
+        await wait_for_condition(
+            pilot,
+            lambda: len(app.object_tree.root.children) > 0,
+            timeout_seconds=5.0,
+            description="tree to be populated with connections",
+        )
+
+        app.connect_to_server(config)
+        await wait_for_condition(
+            pilot,
+            lambda: app.current_connection is not None,
+            timeout_seconds=15.0,
+            description="connection to be established",
+        )
+
+        # Ensure schema cache is available for autocomplete in headless tests.
+        load_schema_async = getattr(app, "_load_schema_cache_async", None)
+        if callable(load_schema_async):
+            await load_schema_async()
+
+        await wait_for_condition(
+            pilot,
+            lambda: (
+                "test_users" in getattr(app, "_table_metadata", {})
+                or f"{mysql_db}.test_users" in getattr(app, "_table_metadata", {})
+            ),
+            timeout_seconds=20.0,
+            description="table metadata to be loaded",
+        )
+
+        app.query_input.text = f"SELECT * FROM {mysql_db}.test_users u WHERE u."
+        _set_cursor_to_end(app)
+        app._trigger_autocomplete(app.query_input)
+
+        expected = {"id", "name", "email"}
+
+        await wait_for_condition(
+            pilot,
+            lambda: bool(app._schema_cache.get("columns", {}).get("test_users")),
+            timeout_seconds=10.0,
+            description="autocomplete columns to load",
+        )
+
+        text = app.query_input.text
+        cursor_loc = app.query_input.cursor_location
+        cursor_pos = app._location_to_offset(text, cursor_loc)
+        suggestions = app._get_autocomplete_suggestions(text, cursor_pos)
+        assert expected.issubset({item.lower() for item in suggestions})
+        assert "Loading..." not in suggestions

tests/unit/test_autocomplete_alias_multi_db.py

@@ -0,0 +1,34 @@
+"""Tests for autocomplete alias handling in multi-database schemas."""
+
+from __future__ import annotations
+
+from sqlit.domains.query.ui.mixins.autocomplete_suggestions import AutocompleteSuggestionsMixin
+
+
+class _DummyAutocompleteHost(AutocompleteSuggestionsMixin):
+    def __init__(self, schema_cache: dict) -> None:
+        self._schema_cache = schema_cache
+        self._columns_loading = set()
+        self.loaded: list[str] = []
+
+    def _load_columns_for_table(self, table_name: str, *, allow_retry: bool = True) -> None:  # noqa: ARG002
+        self.loaded.append(table_name)
+
+
+def test_alias_column_loads_table_when_tables_are_qualified() -> None:
+    """Alias-based column lookup should resolve to the table name, not the alias."""
+    host = _DummyAutocompleteHost(
+        {
+            "tables": ["db.test_users"],
+            "views": [],
+            "columns": {},
+            "procedures": [],
+        }
+    )
+    host._table_metadata = {"test_users": ("", "test_users", "db")}
+
+    sql = "SELECT * FROM test_users u WHERE u."
+    suggestions = host._get_autocomplete_suggestions(sql, len(sql))
+
+    assert suggestions == ["Loading..."]
+    assert host.loaded == ["test_users"]