From 3576169a8b48610803cae2f015a22863a0411c1e Mon Sep 17 00:00:00 2001 From: Peter Adams <18162810+Maxteabag@users.noreply.github.com> Date: Fri, 10 Jul 2026 16:32:43 +0200 Subject: [PATCH] Make Postgres autocomplete suggestions executable for mixed-case names Postgres folds unquoted identifiers to lowercase, so accepting autocomplete for PascalCase tables or camelCase columns produced SQL that failed against real schemas. The fix keeps raw schema metadata for lookup and lazy loading, but formats dropdown insertion text through a dialect hook. The default hook preserves existing behavior; Postgres quotes only identifiers that are not safe lowercase unquoted names. Constraint: PostgreSQL quoted identifiers preserve case while unquoted identifiers fold to lowercase Rejected: Quote every autocomplete identifier | too noisy and would degrade normal lowercase workflows Rejected: Rewrite schema cache entries as quoted names | risks breaking raw metadata lookup and column loading Confidence: high Scope-risk: moderate Directive: Keep autocomplete display formatting separate from raw schema metadata unless lookup tests are updated alongside it Tested: uv run ruff check sqlit/domains/connections/providers/adapters/base.py sqlit/domains/connections/providers/postgresql/base.py sqlit/domains/query/ui/mixins/autocomplete_suggestions.py tests/integration/test_autocomplete_postgresql_schema.py Tested: uv run mypy sqlit/domains/connections/providers/adapters/base.py sqlit/domains/connections/providers/postgresql/base.py sqlit/domains/query/ui/mixins/autocomplete_suggestions.py Tested: uv run --extra postgres pytest tests/integration/test_autocomplete_postgresql_schema.py tests/unit/test_autocomplete_multidb.py tests/unit/test_autocomplete_alias_multi_db.py tests/unit/sql_completion tests/test_autocomplete_database_modes.py tests/ui/test_autocomplete_dropdown.py tests/integration/test_autocomplete_mysql_columns.py -q Not-tested: Full provider integration matrix beyond autocomplete-focused tests --- .../connections/providers/adapters/base.py | 9 ++ .../connections/providers/postgresql/base.py | 9 ++ .../ui/mixins/autocomplete_suggestions.py | 54 ++++++- .../test_autocomplete_postgresql_schema.py | 150 ++++++++++++++++++ 4 files changed, 218 insertions(+), 4 deletions(-) diff --git a/sqlit/domains/connections/providers/adapters/base.py b/sqlit/domains/connections/providers/adapters/base.py index 87770b1f..c906ff09 100644 --- a/sqlit/domains/connections/providers/adapters/base.py +++ b/sqlit/domains/connections/providers/adapters/base.py @@ -417,6 +417,15 @@ def quote_identifier(self, name: str) -> str: """Quote an identifier (table name, column name, etc.).""" pass + def format_autocomplete_identifier(self, name: str) -> str: + """Format an identifier for insertion from autocomplete. + + Most dialects keep the raw name for backwards-compatible suggestions. + Dialects with case-sensitive quoted identifiers can override this to + return executable SQL for names that need quoting. + """ + return name + def qualified_name(self, database: str | None, schema: str | None, name: str) -> str: """Build a quoted qualified identifier, skipping empty segments. diff --git a/sqlit/domains/connections/providers/postgresql/base.py b/sqlit/domains/connections/providers/postgresql/base.py index c5140eca..349b2287 100644 --- a/sqlit/domains/connections/providers/postgresql/base.py +++ b/sqlit/domains/connections/providers/postgresql/base.py @@ -2,6 +2,7 @@ from __future__ import annotations +import re from dataclasses import replace from typing import TYPE_CHECKING, Any @@ -17,6 +18,8 @@ if TYPE_CHECKING: from sqlit.domains.connections.domain.config import ConnectionConfig +_POSTGRES_SAFE_UNQUOTED_IDENTIFIER = re.compile(r"^[a-z_][a-z0-9_]*$") + class PostgresBaseAdapter(CursorBasedAdapter): """Base class for PostgreSQL-compatible databases (PostgreSQL, CockroachDB). @@ -119,6 +122,12 @@ def quote_identifier(self, name: str) -> str: escaped = name.replace('"', '""') return f'"{escaped}"' + def format_autocomplete_identifier(self, name: str) -> str: + """Quote autocomplete identifiers that PostgreSQL would otherwise fold.""" + if _POSTGRES_SAFE_UNQUOTED_IDENTIFIER.fullmatch(name): + return name + return self.quote_identifier(name) + def build_select_query(self, table: str, limit: int, database: str | None = None, schema: str | None = None) -> str: """Build SELECT LIMIT query for PostgreSQL.""" schema = schema or "public" diff --git a/sqlit/domains/query/ui/mixins/autocomplete_suggestions.py b/sqlit/domains/query/ui/mixins/autocomplete_suggestions.py index 94ef4e7e..126c14f6 100644 --- a/sqlit/domains/query/ui/mixins/autocomplete_suggestions.py +++ b/sqlit/domains/query/ui/mixins/autocomplete_suggestions.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +from typing import cast from sqlit.domains.query.completion import ( SuggestionType, @@ -16,6 +17,46 @@ class AutocompleteSuggestionsMixin: """Mixin providing SQL autocomplete suggestion logic.""" + def _format_autocomplete_identifier(self: AutocompleteMixinHost, name: str) -> str: + provider = getattr(self, "current_provider", None) + formatter = getattr(getattr(provider, "dialect", None), "format_autocomplete_identifier", None) + if callable(formatter): + return cast(str, formatter(name)) + return name + + def _format_autocomplete_name(self: AutocompleteMixinHost, name: str) -> str: + if any(marker in name for marker in ('"', "`", "[", "]")): + return name + table_metadata = getattr(self, "_table_metadata", {}) or {} + metadata = table_metadata.get(name.lower()) + if metadata: + schema_name, object_name, _database = metadata + provider = getattr(self, "current_provider", None) + default_schema = getattr(getattr(provider, "capabilities", None), "default_schema", None) + formatted_object = AutocompleteSuggestionsMixin._format_autocomplete_identifier(self, object_name) + if schema_name and schema_name != default_schema and name.lower().startswith(f"{schema_name.lower()}."): + formatted_schema = AutocompleteSuggestionsMixin._format_autocomplete_identifier(self, schema_name) + return f"{formatted_schema}.{formatted_object}" + return formatted_object + if "." not in name: + return AutocompleteSuggestionsMixin._format_autocomplete_identifier(self, name) + return ".".join( + AutocompleteSuggestionsMixin._format_autocomplete_identifier(self, part) + for part in name.split(".") + ) + + def _format_autocomplete_columns( + self: AutocompleteMixinHost, + columns: dict[str, list[str]], + ) -> dict[str, list[str]]: + return { + table_key: [ + AutocompleteSuggestionsMixin._format_autocomplete_identifier(self, column) + for column in column_names + ] + for table_key, column_names in columns.items() + } + def _get_current_word(self: AutocompleteMixinHost, text: str, cursor_pos: int) -> str: """Get the word currently being typed at cursor position.""" before_cursor = text[:cursor_pos] @@ -64,8 +105,13 @@ def is_known(name: str) -> bool: def _get_autocomplete_suggestions(self: AutocompleteMixinHost, text: str, cursor_pos: int) -> list[str]: """Get autocomplete suggestions using the SQL completion engine.""" # Build schema data for get_completions - tables = self._schema_cache.get("tables", []) + self._schema_cache.get("views", []) - columns = self._schema_cache.get("columns", {}) + raw_tables = self._schema_cache.get("tables", []) + self._schema_cache.get("views", []) + raw_columns = self._schema_cache.get("columns", {}) + tables = [ + AutocompleteSuggestionsMixin._format_autocomplete_name(self, table) + for table in raw_tables + ] + columns = AutocompleteSuggestionsMixin._format_autocomplete_columns(self, raw_columns) procedures = self._schema_cache.get("procedures", []) # First check if we need to lazy-load columns before calling get_completions @@ -81,7 +127,7 @@ def needs_column_load(key: str) -> bool: been loaded yet. Returning Loading... for an unknown key would wedge forever because the loader skips unknown tables silently. """ - if key in columns: + if key in raw_columns: return False if key not in table_metadata: return False @@ -102,7 +148,7 @@ def needs_column_load(key: str) -> bool: scope = suggestion.table_scope if scope: scope_lower = scope.lower() - table_key = alias_map.get(scope_lower, scope_lower) + table_key = alias_map.get(scope_lower, scope_lower).lower() if needs_column_load(table_key) and table_key not in loading: self._load_columns_for_table(table_key) diff --git a/tests/integration/test_autocomplete_postgresql_schema.py b/tests/integration/test_autocomplete_postgresql_schema.py index fcc481ca..5782b016 100644 --- a/tests/integration/test_autocomplete_postgresql_schema.py +++ b/tests/integration/test_autocomplete_postgresql_schema.py @@ -79,10 +79,84 @@ def postgres_schema_table(postgres_server_ready: bool, postgres_db: str): conn.close() +@pytest.fixture +def postgres_mixed_case_table(postgres_server_ready: bool, postgres_db: str): + """Create mixed-case PostgreSQL identifiers that require double quotes.""" + if not postgres_server_ready: + pytest.skip("PostgreSQL is not available") + + try: + import psycopg2 + except ImportError: + pytest.skip("psycopg2 is not installed") + + conn = psycopg2.connect( + host=POSTGRES_HOST, + port=POSTGRES_PORT, + database=postgres_db, + user=POSTGRES_USER, + password=POSTGRES_PASSWORD, + connect_timeout=10, + ) + conn.autocommit = True + cursor = conn.cursor() + cursor.execute('DROP TABLE IF EXISTS "UserProfile" CASCADE') + cursor.execute(""" + CREATE TABLE "UserProfile" ( + id INTEGER PRIMARY KEY, + "createdAt" TIMESTAMP NOT NULL, + "displayName" TEXT NOT NULL + ) + """) + cursor.execute(""" + INSERT INTO "UserProfile" (id, "createdAt", "displayName") + VALUES (1, NOW(), 'Ada') + """) + conn.close() + + yield + + conn = psycopg2.connect( + host=POSTGRES_HOST, + port=POSTGRES_PORT, + database=postgres_db, + user=POSTGRES_USER, + password=POSTGRES_PASSWORD, + connect_timeout=10, + ) + conn.autocommit = True + cursor = conn.cursor() + cursor.execute('DROP TABLE IF EXISTS "UserProfile" CASCADE') + conn.close() + + def _suggestions(app: SSMSTUI, sql: str) -> list[str]: return app._get_autocomplete_suggestions(sql, len(sql)) +def _assert_postgres_requires_quotes_for_mixed_case(postgres_db: str) -> None: + """Sanity-check the real PostgreSQL behavior behind the autocomplete regression.""" + psycopg2 = pytest.importorskip("psycopg2") + + conn = psycopg2.connect( + host=POSTGRES_HOST, + port=POSTGRES_PORT, + database=postgres_db, + user=POSTGRES_USER, + password=POSTGRES_PASSWORD, + connect_timeout=10, + ) + conn.autocommit = True + cursor = conn.cursor() + cursor.execute('SELECT "createdAt", "displayName" FROM "UserProfile"') + assert cursor.fetchone()[1] == "Ada" + + with pytest.raises(psycopg2.Error): + cursor.execute("SELECT createdAt FROM UserProfile") + + conn.close() + + @pytest.mark.asyncio async def test_postgresql_schema_table_autocomplete( postgres_server_ready: bool, @@ -156,3 +230,79 @@ async def test_postgresql_schema_table_autocomplete( assert {"id", "greeting"}.issubset({item.lower() for item in alias_columns}) assert "Loading..." not in alias_columns + + +@pytest.mark.asyncio +async def test_postgresql_mixed_case_identifier_autocomplete_quotes_suggestions( + postgres_server_ready: bool, + postgres_db: str, + postgres_mixed_case_table, + temp_config_dir: str, +) -> None: + """Autocomplete should emit executable quoted identifiers for mixed-case PostgreSQL names.""" + if not postgres_server_ready: + pytest.skip("PostgreSQL is not available") + + _assert_postgres_requires_quotes_for_mixed_case(postgres_db) + + config = ConnectionConfig( + name="test-postgres-quoted-autocomplete", + db_type="postgresql", + server=POSTGRES_HOST, + port=str(POSTGRES_PORT), + database=postgres_db, + username=POSTGRES_USER, + password=POSTGRES_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 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", + ) + + 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: "userprofile" in getattr(app, "_table_metadata", {}), + timeout_seconds=20.0, + description="mixed-case table metadata to be loaded", + ) + + table_suggestions = _suggestions(app, "SELECT * FROM User") + assert '"UserProfile"' in table_suggestions + assert "UserProfile" not in table_suggestions + + alias_columns = _suggestions(app, 'SELECT * FROM "UserProfile" u WHERE u.') + if alias_columns == ["Loading..."]: + await wait_for_condition( + pilot, + lambda: bool(app._schema_cache.get("columns", {}).get("userprofile")), + timeout_seconds=10.0, + description="mixed-case table columns to load", + ) + alias_columns = _suggestions(app, 'SELECT * FROM "UserProfile" u WHERE u.') + + assert '"createdAt"' in alias_columns + assert '"displayName"' in alias_columns + assert "createdAt" not in alias_columns + assert "displayName" not in alias_columns + assert "Loading..." not in alias_columns