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