only cache data objects mantained by sqlmesh

Author: izeigerman

sqlmesh/core/engine_adapter/base.py

@@ -1046,7 +1046,8 @@ def create_table_like(
             target_table_name: The name of the table to create. Can be fully qualified or just table name.
             source_table_name: The name of the table to base the new table on.
         """
-        self.create_table(target_table_name, self.columns(source_table_name), exists=exists)
+        self._create_table_like(target_table_name, source_table_name, exists=exists, **kwargs)
+        self._clear_data_object_cache(target_table_name)
 
     def clone_table(
         self,
@@ -2271,24 +2272,34 @@ def rename_table(
                     "Tried to rename table across catalogs which is not supported"
                 )
         self._rename_table(old_table_name, new_table_name)
+        self._clear_data_object_cache(old_table_name)
+        self._clear_data_object_cache(new_table_name)
 
-    def get_data_object(self, target_name: TableName) -> t.Optional[DataObject]:
+    def get_data_object(
+        self, target_name: TableName, safe_to_cache: bool = False
+    ) -> t.Optional[DataObject]:
         target_table = exp.to_table(target_name)
         existing_data_objects = self.get_data_objects(
-            schema_(target_table.db, target_table.catalog), {target_table.name}
+            schema_(target_table.db, target_table.catalog),
+            {target_table.name},
+            safe_to_cache=safe_to_cache,
         )
         if existing_data_objects:
             return existing_data_objects[0]
         return None
 
     def get_data_objects(
-        self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
+        self,
+        schema_name: SchemaName,
+        object_names: t.Optional[t.Set[str]] = None,
+        safe_to_cache: bool = False,
     ) -> t.List[DataObject]:
         """Lists all data objects in the target schema.
 
         Args:
             schema_name: The name of the schema to list data objects from.
             object_names: If provided, only return data objects with these names.
+            safe_to_cache: Whether it is safe to cache the results of this call.
 
         Returns:
             A list of data objects in the target schema.
@@ -2323,31 +2334,36 @@ def get_data_objects(
                     object_names_list[i : i + self.DATA_OBJECT_FILTER_BATCH_SIZE]
                     for i in range(0, len(object_names_list), self.DATA_OBJECT_FILTER_BATCH_SIZE)
                 ]
-                fetched_objects = [
-                    obj
-                    for batch in batches
-                    for obj in self._get_data_objects(schema_name, set(batch))
-                ]
-
-                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
 
-                fetched_object_names = {obj.name for obj in fetched_objects}
-                for missing_name in missing_names - fetched_object_names:
-                    cache_key = _get_data_object_cache_key(
-                        target_schema.catalog, target_schema.db, missing_name
-                    )
-                    self._data_object_cache[cache_key] = None
+                fetched_objects = []
+                fetched_object_names = set()
+                for batch in batches:
+                    objects = self._get_data_objects(schema_name, set(batch))
+                    for obj in objects:
+                        if safe_to_cache:
+                            cache_key = _get_data_object_cache_key(
+                                obj.catalog, obj.schema_name, obj.name
+                            )
+                            self._data_object_cache[cache_key] = obj
+                        fetched_objects.append(obj)
+                        fetched_object_names.add(obj.name)
+
+                if safe_to_cache:
+                    for missing_name in missing_names - fetched_object_names:
+                        cache_key = _get_data_object_cache_key(
+                            target_schema.catalog, target_schema.db, missing_name
+                        )
+                        self._data_object_cache[cache_key] = None
 
                 return cached_objects + fetched_objects
 
             return cached_objects
 
         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
+        if safe_to_cache:
+            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(
@@ -2951,6 +2967,15 @@ def _create_column_comments(
                     exc_info=True,
                 )
 
+    def _create_table_like(
+        self,
+        target_table_name: TableName,
+        source_table_name: TableName,
+        exists: bool,
+        **kwargs: t.Any,
+    ) -> None:
+        self.create_table(target_table_name, self.columns(source_table_name), exists=exists)
+
     def _rename_table(
         self,
         old_table_name: TableName,
@@ -3017,5 +3042,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 = catalog or ""
-    return f"{catalog}.{schema_name}.{object_name}"
+    catalog = f"{catalog}." if catalog else ""
+    return f"{catalog}{schema_name}.{object_name}"

sqlmesh/core/engine_adapter/clickhouse.py

@@ -224,7 +224,7 @@ def _insert_overwrite_by_condition(
         target_columns_to_types = target_columns_to_types or self.columns(target_table)
 
         temp_table = self._get_temp_table(target_table)
-        self._create_table_like(temp_table, target_table)
+        self.create_table_like(temp_table, target_table)
 
         # REPLACE BY KEY: extract kwargs if present
         dynamic_key = kwargs.get("dynamic_key")
@@ -456,7 +456,11 @@ def insert_overwrite_by_partition(
         )
 
     def _create_table_like(
-        self, target_table_name: TableName, source_table_name: TableName
+        self,
+        target_table_name: TableName,
+        source_table_name: TableName,
+        exists: bool,
+        **kwargs: t.Any,
     ) -> None:
         """Create table with identical structure as source table"""
         self.execute(
@@ -632,16 +636,15 @@ def _drop_object(
             kind: What kind of object to drop. Defaults to TABLE
             **drop_args: Any extra arguments to set on the Drop expression
         """
-        self.execute(
-            exp.Drop(
-                this=exp.to_table(name),
-                kind=kind,
-                exists=exists,
-                cluster=exp.OnCluster(this=exp.to_identifier(self.cluster))
-                if self.engine_run_mode.is_cluster
-                else None,
-                **drop_args,
-            )
+        super()._drop_object(
+            name=name,
+            exists=exists,
+            kind=kind,
+            cascade=cascade,
+            cluster=exp.OnCluster(this=exp.to_identifier(self.cluster))
+            if self.engine_run_mode.is_cluster
+            else None,
+            **drop_args,
         )
 
     def _build_partitioned_by_exp(

sqlmesh/core/engine_adapter/mysql.py

@@ -164,11 +164,11 @@ def _create_column_comments(
                         exc_info=True,
                     )
 
-    def create_table_like(
+    def _create_table_like(
         self,
         target_table_name: TableName,
         source_table_name: TableName,
-        exists: bool = True,
+        exists: bool,
         **kwargs: t.Any,
     ) -> None:
         self.execute(

sqlmesh/core/engine_adapter/postgres.py

@@ -79,11 +79,11 @@ def _fetch_native_df(
             self._connection_pool.commit()
         return df
 
-    def create_table_like(
+    def _create_table_like(
         self,
         target_table_name: TableName,
         source_table_name: TableName,
-        exists: bool = True,
+        exists: bool,
         **kwargs: t.Any,
     ) -> None:
         self.execute(

sqlmesh/core/engine_adapter/spark.py

@@ -402,14 +402,16 @@ def get_current_database(self) -> str:
             return self.spark.catalog.currentDatabase()
         return self.fetchone(exp.select(exp.func("current_database")))[0]  # type: ignore
 
-    def get_data_object(self, target_name: TableName) -> t.Optional[DataObject]:
+    def get_data_object(
+        self, target_name: TableName, safe_to_cache: bool = False
+    ) -> t.Optional[DataObject]:
         target_table = exp.to_table(target_name)
         if isinstance(target_table.this, exp.Dot) and target_table.this.expression.name.startswith(
             f"{self.BRANCH_PREFIX}{self.WAP_PREFIX}"
         ):
             # Exclude the branch name
             target_table.set("this", target_table.this.this)
-        return super().get_data_object(target_table)
+        return super().get_data_object(target_table, safe_to_cache=safe_to_cache)
 
     def create_state_table(
         self,

sqlmesh/core/snapshot/evaluator.py

@@ -1564,7 +1564,9 @@ def _get_data_objects_in_schema(
             gateway: t.Optional[str] = None,
         ) -> t.List[DataObject]:
             logger.info("Listing data objects in schema %s", schema.sql())
-            return self.get_adapter(gateway).get_data_objects(schema, object_names)
+            return self.get_adapter(gateway).get_data_objects(
+                schema, object_names, safe_to_cache=True
+            )
 
         with self.concurrent_context():
             existing_objects: t.List[DataObject] = []