address comments

Author: izeigerman

sqlmesh/core/engine_adapter/base.py

@@ -1470,8 +1470,14 @@ def columns(
         }
 
     def table_exists(self, table_name: TableName) -> bool:
+        table = exp.to_table(table_name)
+        data_object_cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name)
+        if data_object_cache_key in self._data_object_cache:
+            logger.debug("Table existence cache hit: %s", data_object_cache_key)
+            return self._data_object_cache[data_object_cache_key] is not None
+
         try:
-            self.execute(exp.Describe(this=exp.to_table(table_name), kind="TABLE"))
+            self.execute(exp.Describe(this=table, kind="TABLE"))
             return True
         except Exception:
             return False
@@ -2301,11 +2307,13 @@ def get_data_objects(
                     target_schema.catalog, target_schema.db, name
                 )
                 if cache_key in self._data_object_cache:
+                    logger.debug("Data object cache hit: %s", cache_key)
                     data_object = self._data_object_cache[cache_key]
                     # If the object is none, then the table was previously looked for but not found
                     if data_object:
                         cached_objects.append(data_object)
                 else:
+                    logger.debug("Data object cache miss: %s", cache_key)
                     missing_names.add(name)
 
             # Fetch missing objects from database
@@ -2321,7 +2329,6 @@ def get_data_objects(
                     for obj in self._get_data_objects(schema_name, set(batch))
                 ]
 
-                # Cache the fetched objects
                 for obj in fetched_objects:
                     cache_key = _get_data_object_cache_key(obj.catalog, obj.schema_name, obj.name)
                     self._data_object_cache[cache_key] = obj
@@ -2336,7 +2343,12 @@ def get_data_objects(
                 return cached_objects + fetched_objects
 
             return cached_objects
-        return self._get_data_objects(schema_name)
+
+        fetched_objects = self._get_data_objects(schema_name)
+        for obj in fetched_objects:
+            cache_key = _get_data_object_cache_key(obj.catalog, obj.schema_name, obj.name)
+            self._data_object_cache[cache_key] = obj
+        return fetched_objects
 
     def fetchone(
         self,
@@ -2746,10 +2758,12 @@ def _to_sql(self, expression: exp.Expression, quote: bool = True, **kwargs: t.An
     def _clear_data_object_cache(self, table_name: t.Optional[TableName] = None) -> None:
         """Clears the cache entry for the given table name, or clears the entire cache if table_name is None."""
         if table_name is None:
+            logger.debug("Clearing entire data object cache")
             self._data_object_cache.clear()
         else:
             table = exp.to_table(table_name)
             cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name)
+            logger.debug("Clearing data object cache key: %s", cache_key)
             self._data_object_cache.pop(cache_key, None)
 
     def _get_data_objects(
@@ -3003,7 +3017,5 @@ def _decoded_str(value: t.Union[str, bytes]) -> str:
 
 def _get_data_object_cache_key(catalog: t.Optional[str], schema_name: str, object_name: str) -> str:
     """Returns a cache key for a data object based on its fully qualified name."""
-    catalog_part = catalog.lower() if catalog else ""
-    schema_part = schema_name.lower()
-    object_part = object_name.lower()
-    return f"{catalog_part}.{schema_part}.{object_part}"
+    catalog = catalog or ""
+    return f"{catalog}.{schema_name}.{object_name}"

sqlmesh/core/engine_adapter/base_postgres.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
 import typing as t
+import logging
 
 from sqlglot import exp
 
 from sqlmesh.core.dialect import to_schema
-from sqlmesh.core.engine_adapter import EngineAdapter
+from sqlmesh.core.engine_adapter.base import EngineAdapter, _get_data_object_cache_key
 from sqlmesh.core.engine_adapter.shared import (
     CatalogSupport,
     CommentCreationTable,
@@ -20,6 +21,9 @@
     from sqlmesh.core.engine_adapter._typing import QueryOrDF
 
 
+logger = logging.getLogger(__name__)
+
+
 class BasePostgresEngineAdapter(EngineAdapter):
     DEFAULT_BATCH_SIZE = 400
     COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
@@ -75,6 +79,10 @@ def table_exists(self, table_name: TableName) -> bool:
         Reference: https://github.com/aws/amazon-redshift-python-driver/blob/master/redshift_connector/cursor.py#L528-L553
         """
         table = exp.to_table(table_name)
+        data_object_cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name)
+        if data_object_cache_key in self._data_object_cache:
+            logger.debug("Table existence cache hit: %s", data_object_cache_key)
+            return self._data_object_cache[data_object_cache_key] is not None
 
         sql = (
             exp.select("1")

sqlmesh/core/engine_adapter/bigquery.py

@@ -8,6 +8,7 @@
 from sqlglot.transforms import remove_precision_parameterized_types
 
 from sqlmesh.core.dialect import to_schema
+from sqlmesh.core.engine_adapter.base import _get_data_object_cache_key
 from sqlmesh.core.engine_adapter.mixins import (
     ClusteredByMixin,
     RowDiffMixin,
@@ -744,6 +745,12 @@ def insert_overwrite_by_partition(
             )
 
     def table_exists(self, table_name: TableName) -> bool:
+        table = exp.to_table(table_name)
+        data_object_cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name)
+        if data_object_cache_key in self._data_object_cache:
+            logger.debug("Table existence cache hit: %s", data_object_cache_key)
+            return self._data_object_cache[data_object_cache_key] is not None
+
         try:
             from google.cloud.exceptions import NotFound
         except ModuleNotFoundError:

sqlmesh/core/engine_adapter/mssql.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import typing as t
+import logging
 
 from sqlglot import exp
 
@@ -13,6 +14,7 @@
     InsertOverwriteStrategy,
     MERGE_SOURCE_ALIAS,
     MERGE_TARGET_ALIAS,
+    _get_data_object_cache_key,
 )
 from sqlmesh.core.engine_adapter.mixins import (
     GetCurrentCatalogFromFunctionMixin,
@@ -36,6 +38,9 @@
     from sqlmesh.core.engine_adapter._typing import DF, Query, QueryOrDF
 
 
+logger = logging.getLogger(__name__)
+
+
 @set_catalog()
 class MSSQLEngineAdapter(
     EngineAdapterWithIndexSupport,
@@ -144,6 +149,10 @@ def build_var_length_col(
     def table_exists(self, table_name: TableName) -> bool:
         """MsSql doesn't support describe so we query information_schema."""
         table = exp.to_table(table_name)
+        data_object_cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name)
+        if data_object_cache_key in self._data_object_cache:
+            logger.debug("Table existence cache hit: %s", data_object_cache_key)
+            return self._data_object_cache[data_object_cache_key] is not None
 
         sql = (
             exp.select("1")

sqlmesh/core/engine_adapter/postgres.py

@@ -34,7 +34,7 @@ class PostgresEngineAdapter(
     HAS_VIEW_BINDING = True
     CURRENT_CATALOG_EXPRESSION = exp.column("current_catalog")
     SUPPORTS_REPLACE_TABLE = False
-    MAX_IDENTIFIER_LENGTH = 63
+    MAX_IDENTIFIER_LENGTH: t.Optional[int] = 63
     SUPPORTS_QUERY_EXECUTION_TRACKING = True
     SCHEMA_DIFFER_KWARGS = {
         "parameterized_type_defaults": {

tests/core/engine_adapter/test_base.py

@@ -3723,6 +3723,27 @@ def test_data_object_cache_get_data_objects(
     assert mock_get_data_objects.call_count == 1  # Should not increase
 
 
+def test_data_object_cache_get_data_objects_no_object_names(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+):
+    adapter = make_mocked_engine_adapter(EngineAdapter, patch_get_data_objects=False)
+
+    table1 = DataObject(catalog=None, schema="test_schema", name="table1", type="table")
+    table2 = DataObject(catalog=None, schema="test_schema", name="table2", type="table")
+
+    mock_get_data_objects = mocker.patch.object(
+        adapter, "_get_data_objects", return_value=[table1, table2]
+    )
+
+    result1 = adapter.get_data_objects("test_schema")
+    assert len(result1) == 2
+    assert mock_get_data_objects.call_count == 1
+
+    result2 = adapter.get_data_objects("test_schema", {"table1", "table2"})
+    assert len(result2) == 2
+    assert mock_get_data_objects.call_count == 1  # Should not increase
+
+
 def test_data_object_cache_get_data_object(
     make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
 ):

tests/dbt/test_transformation.py

@@ -122,10 +122,10 @@ def test_dbt_custom_materialization():
     selected_model = list(plan.selected_models)[0]
     assert selected_model == "model.sushi.custom_incremental_model"
 
-    qoery = "SELECT * FROM sushi.custom_incremental_model ORDER BY created_at"
+    query = "SELECT * FROM sushi.custom_incremental_model ORDER BY created_at"
     hook_table = "SELECT * FROM hook_table ORDER BY id"
     sushi_context.apply(plan)
-    result = sushi_context.engine_adapter.fetchdf(qoery)
+    result = sushi_context.engine_adapter.fetchdf(query)
     assert len(result) == 1
     assert {"created_at", "id"}.issubset(result.columns)
 
@@ -140,7 +140,7 @@ def test_dbt_custom_materialization():
     tomorrow = datetime.now() + timedelta(days=1)
     sushi_context.run(select_models=["sushi.custom_incremental_model"], execution_time=tomorrow)
 
-    result_after_run = sushi_context.engine_adapter.fetchdf(qoery)
+    result_after_run = sushi_context.engine_adapter.fetchdf(query)
     assert {"created_at", "id"}.issubset(result_after_run.columns)
 
     # this should have added new unique values for the new row