From 2c0360b153685903d5a3bf5a47d02e53acb1fcbb Mon Sep 17 00:00:00 2001 From: TangoEnSkai <21152231+TangoEnSkai@users.noreply.github.com> Date: Sun, 23 Aug 2026 12:57:15 +0900 Subject: [PATCH 1/2] fix: cache get_foreign_keys() to stop redundant DESCRIBE TABLE EXTENDED MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit get_foreign_keys() was the only _describe_table_extended()-backed reflection method without @reflection.cache. SQLAlchemy threads an info_cache dict through every Inspector call so that repeated reflection of the same table costs one round-trip rather than one per call, but a dialect method only participates when it is decorated. Undecorated, every get_foreign_keys() call issued a fresh DESCRIBE TABLE EXTENDED `catalog`.`schema`.`table` against the warehouse, however many times the same table was reflected through the same Inspector. For an application that reflects on each request, that is a warehouse query per table per request. get_pk_constraint() sits directly above it, uses the same helper, and is decorated; get_table_names(), get_view_names(), has_table(), get_schema_names() and get_table_comment() are decorated too. This was an omission rather than a deliberate exclusion. Note the cache key includes fn.__name__, so this does not make get_pk_constraint() and get_foreign_keys() share one DESCRIBE within a single pass — that duplication is separate and would need the helper itself to be cached. Add unit tests covering the cache hit, that the cache is keyed per table, that a direct call passing no info_cache is unaffected, and a regression guard on the already-correct get_pk_constraint(). Resolves #72 Signed-off-by: TangoEnSkai <21152231+TangoEnSkai@users.noreply.github.com> --- CHANGELOG.md | 4 ++ src/databricks/sqlalchemy/base.py | 1 + tests/test_local/test_reflection_cache.py | 76 +++++++++++++++++++++++ 3 files changed, 81 insertions(+) create mode 100644 tests/test_local/test_reflection_cache.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e744830..fc20143 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Release History +# Unreleased + +- Fix: Cache `get_foreign_keys()` per table per reflection pass, so reflecting a schema no longer issues a redundant `DESCRIBE TABLE EXTENDED` per table on top of the one `get_pk_constraint()` already made (fixes #72) + # 2.0.10 (2026-06-18) - Fix: Quote bind parameter names containing non-identifier characters (e.g. hyphens, backticks) so columns and parameters with special characters bind correctly (databricks/databricks-sqlalchemy#60 by @msrathore-db) diff --git a/src/databricks/sqlalchemy/base.py b/src/databricks/sqlalchemy/base.py index bcdd6a8..1526b4e 100644 --- a/src/databricks/sqlalchemy/base.py +++ b/src/databricks/sqlalchemy/base.py @@ -253,6 +253,7 @@ def get_pk_constraint( # TODO: figure out how to return sqlalchemy.interfaces in a way that mypy respects return build_pk_dict(pk_name, pk_constraint_string) # type: ignore + @reflection.cache def get_foreign_keys( self, connection, table_name, schema=None, **kw ) -> List[ReflectedForeignKeyConstraint]: diff --git a/tests/test_local/test_reflection_cache.py b/tests/test_local/test_reflection_cache.py new file mode 100644 index 0000000..868ae2c --- /dev/null +++ b/tests/test_local/test_reflection_cache.py @@ -0,0 +1,76 @@ +from unittest.mock import patch + +from databricks.sqlalchemy import DatabricksDialect + + +class TestReflectionCache: + """Reflection methods backed by DESCRIBE TABLE EXTENDED must honour info_cache. + + SQLAlchemy passes an ``info_cache`` dict through every ``Inspector`` call so + that one reflection pass issues one round-trip per table. A dialect method + only participates when it is decorated with ``@reflection.cache``. + """ + + def _dialect_and_cache(self): + return DatabricksDialect(), {} + + def test_get_foreign_keys_is_cached(self): + dialect, info_cache = self._dialect_and_cache() + + with patch.object( + dialect, "_describe_table_extended", return_value=[] + ) as mock_dte: + first = dialect.get_foreign_keys( + None, "some_table", schema="some_schema", info_cache=info_cache + ) + second = dialect.get_foreign_keys( + None, "some_table", schema="some_schema", info_cache=info_cache + ) + + assert mock_dte.call_count == 1 + assert first == second + + def test_get_foreign_keys_cache_is_per_table(self): + dialect, info_cache = self._dialect_and_cache() + + with patch.object( + dialect, "_describe_table_extended", return_value=[] + ) as mock_dte: + dialect.get_foreign_keys( + None, "table_one", schema="some_schema", info_cache=info_cache + ) + dialect.get_foreign_keys( + None, "table_two", schema="some_schema", info_cache=info_cache + ) + + assert mock_dte.call_count == 2 + + def test_get_foreign_keys_without_info_cache_is_not_cached(self): + """Direct calls that pass no info_cache keep their existing behaviour.""" + + dialect, _ = self._dialect_and_cache() + + with patch.object( + dialect, "_describe_table_extended", return_value=[] + ) as mock_dte: + dialect.get_foreign_keys(None, "some_table", schema="some_schema") + dialect.get_foreign_keys(None, "some_table", schema="some_schema") + + assert mock_dte.call_count == 2 + + def test_get_pk_constraint_is_cached(self): + """Guards the already-correct sibling against regression.""" + + dialect, info_cache = self._dialect_and_cache() + + with patch.object( + dialect, "_describe_table_extended", return_value=[] + ) as mock_dte: + dialect.get_pk_constraint( + None, "some_table", schema="some_schema", info_cache=info_cache + ) + dialect.get_pk_constraint( + None, "some_table", schema="some_schema", info_cache=info_cache + ) + + assert mock_dte.call_count == 1 From c6c60f2fcd8cfee19a62ec01e1340c40b4702c46 Mon Sep 17 00:00:00 2001 From: TangoEnSkai <21152231+TangoEnSkai@users.noreply.github.com> Date: Sun, 23 Aug 2026 14:22:18 +0900 Subject: [PATCH 2/2] fix: cache get_columns() to stop a round-trip on every reflection call get_columns() was the second dialect method missing @reflection.cache, so it never participated in the cache SQLAlchemy threads through Inspector.get_columns(): col_defs = self.dialect.get_columns( conn, table_name, schema, info_cache=self.info_cache, **kw ) The info_cache kwarg arrived, landed in **kwargs, and was discarded. Measured against main, three calls sharing one info_cache produced three GetColumns round-trips and left info_cache empty, where the decorated get_pk_constraint() produced one. The cost is not always a single statement: when cur.columns() returns an empty list, get_columns() follows up with DESCRIBE TABLE EXTENDED to tell a column-less table from a missing one, so an uncached call on such a table is two round-trips, repeated every time. SQLAlchemy's own SQLite, PostgreSQL and MySQL dialects all decorate get_columns(), so this matches upstream practice rather than inventing one. Caching is safe with respect to Inspector._instantiate_types(), which mutates the returned column dicts in place but is guarded by 'if not isinstance(coltype, TypeEngine)' and is therefore a no-op on an already-instantiated cached list. get_indexes() remains undecorated deliberately: it returns the EMPTY_INDEX constant without touching the server. Resolves #75 Signed-off-by: TangoEnSkai <21152231+TangoEnSkai@users.noreply.github.com> --- CHANGELOG.md | 3 +- src/databricks/sqlalchemy/base.py | 1 + tests/test_local/test_reflection_cache.py | 69 +++++++++++++++++++++-- 3 files changed, 68 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fc20143..f55b7b6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,8 @@ # Unreleased -- Fix: Cache `get_foreign_keys()` per table per reflection pass, so reflecting a schema no longer issues a redundant `DESCRIBE TABLE EXTENDED` per table on top of the one `get_pk_constraint()` already made (fixes #72) +- Fix: Cache `get_foreign_keys()` so repeated reflection of a table through the same `Inspector` issues one `DESCRIBE TABLE EXTENDED` instead of one per call (fixes #72) +- Fix: Cache `get_columns()` so repeated reflection of a table through the same `Inspector` issues one `GetColumns` round-trip instead of one per call (fixes #75) # 2.0.10 (2026-06-18) diff --git a/src/databricks/sqlalchemy/base.py b/src/databricks/sqlalchemy/base.py index 1526b4e..8c2f36d 100644 --- a/src/databricks/sqlalchemy/base.py +++ b/src/databricks/sqlalchemy/base.py @@ -136,6 +136,7 @@ def create_connect_args(self, url): return [], kwargs + @reflection.cache def get_columns( self, connection, table_name, schema=None, **kwargs ) -> List[ReflectedColumn]: diff --git a/tests/test_local/test_reflection_cache.py b/tests/test_local/test_reflection_cache.py index 868ae2c..13f4429 100644 --- a/tests/test_local/test_reflection_cache.py +++ b/tests/test_local/test_reflection_cache.py @@ -1,14 +1,32 @@ -from unittest.mock import patch +from contextlib import contextmanager +from unittest.mock import Mock, patch from databricks.sqlalchemy import DatabricksDialect +@contextmanager +def _counting_cursor(self, connection): + """Stand in for get_connection_cursor, recording each columns() call.""" + + cursor = Mock() + + def columns(**kwargs): + _counting_cursor.calls.append(kwargs["table_name"]) + result = Mock() + result.fetchall.return_value = [] + return result + + cursor.columns.side_effect = columns + yield cursor + + class TestReflectionCache: - """Reflection methods backed by DESCRIBE TABLE EXTENDED must honour info_cache. + """Reflection methods must honour the info_cache SQLAlchemy threads through. SQLAlchemy passes an ``info_cache`` dict through every ``Inspector`` call so - that one reflection pass issues one round-trip per table. A dialect method - only participates when it is decorated with ``@reflection.cache``. + that reflecting the same table repeatedly costs one round-trip rather than + one per call. A dialect method only participates when it is decorated with + ``@reflection.cache``. """ def _dialect_and_cache(self): @@ -58,6 +76,49 @@ def test_get_foreign_keys_without_info_cache_is_not_cached(self): assert mock_dte.call_count == 2 + def _get_columns_calls(self, dialect, info_cache, times, table_name="some_table"): + """Call get_columns `times` times, returning the server round-trip counts.""" + + _counting_cursor.calls = [] + + with patch.object( + DatabricksDialect, "get_connection_cursor", _counting_cursor + ), patch.object( + dialect, "_describe_table_extended", return_value=[] + ) as mock_dte: + for _ in range(times): + kwargs = {} if info_cache is None else {"info_cache": info_cache} + dialect.get_columns(None, table_name, None, **kwargs) + + return len(_counting_cursor.calls), mock_dte.call_count + + def test_get_columns_is_cached(self): + dialect, info_cache = self._dialect_and_cache() + dialect.catalog = "some_catalog" + dialect.schema = "some_schema" + + column_calls, describe_calls = self._get_columns_calls( + dialect, info_cache, times=3 + ) + + assert column_calls == 1 + # An empty columns() result makes get_columns fall back to + # DESCRIBE TABLE EXTENDED to tell a column-less table from a missing + # one, so an uncached call costs two round-trips, not one. + assert describe_calls == 1 + + def test_get_columns_without_info_cache_is_not_cached(self): + """Direct calls that pass no info_cache keep their existing behaviour.""" + + dialect, _ = self._dialect_and_cache() + dialect.catalog = "some_catalog" + dialect.schema = "some_schema" + + column_calls, describe_calls = self._get_columns_calls(dialect, None, times=3) + + assert column_calls == 3 + assert describe_calls == 3 + def test_get_pk_constraint_is_cached(self): """Guards the already-correct sibling against regression."""