fix: add proper catalog support to engine adapter (#1559)

Author: eakmanrq

sqlmesh/core/_typing.py

@@ -6,3 +6,4 @@
 
 if t.TYPE_CHECKING:
     TableName = t.Union[str, exp.Table]
+    SchemaName = t.Union[str, exp.Table]

sqlmesh/core/dialect.py

@@ -839,3 +839,35 @@ def transform_values(
             yield exp.func("PARSE_JSON", f"'{value}'")
         else:
             yield value
+
+
+def to_schema(sql_path: str | exp.Table) -> exp.Table:
+    if isinstance(sql_path, exp.Table) and sql_path.this is None:
+        return sql_path
+    table = exp.to_table(sql_path.copy() if isinstance(sql_path, exp.Table) else sql_path)
+    table.set("catalog", table.args.get("db"))
+    table.set("db", table.args.get("this"))
+    table.set("this", None)
+    return table
+
+
+def schema_(
+    db: exp.Identifier | str,
+    catalog: t.Optional[exp.Identifier | str] = None,
+    quoted: t.Optional[bool] = None,
+) -> exp.Table:
+    """Build a Schema.
+
+    Args:
+        db: Database name.
+        catalog: Catalog name.
+        quoted: Whether to force quotes on the schema's identifiers.
+
+    Returns:
+        The new Schema instance.
+    """
+    return exp.Table(
+        this=None,
+        db=exp.to_identifier(db, quoted=quoted) if db else None,
+        catalog=exp.to_identifier(catalog, quoted=quoted) if catalog else None,
+    )

sqlmesh/core/engine_adapter/base.py

@@ -25,18 +25,23 @@
 from sqlglot.helper import ensure_list
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 
-from sqlmesh.core.dialect import add_table, select_from_values_for_batch_range
-from sqlmesh.core.engine_adapter.shared import DataObject
+from sqlmesh.core.dialect import (
+    add_table,
+    schema_,
+    select_from_values_for_batch_range,
+    to_schema,
+)
+from sqlmesh.core.engine_adapter.shared import DataObject, set_catalog
 from sqlmesh.core.model.kind import TimeColumn
 from sqlmesh.core.schema_diff import SchemaDiffer
 from sqlmesh.utils import double_escape
 from sqlmesh.utils.connection_pool import create_connection_pool
 from sqlmesh.utils.date import TimeLike, make_inclusive, to_ts
-from sqlmesh.utils.errors import SQLMeshError
+from sqlmesh.utils.errors import SQLMeshError, UnsupportedCatalogOperationError
 from sqlmesh.utils.pandas import columns_to_types_from_df
 
 if t.TYPE_CHECKING:
-    from sqlmesh.core._typing import TableName
+    from sqlmesh.core._typing import SchemaName, TableName
     from sqlmesh.core.engine_adapter._typing import (
         DF,
         PySparkDataFrame,
@@ -76,6 +81,28 @@ def requires_condition(self) -> bool:
         return self.is_replace_where or self.is_delete_insert
 
 
+class CatalogSupport(Enum):
+    UNSUPPORTED = 1
+    REQUIRES_SET_CATALOG = 2
+    FULL_SUPPORT = 3
+
+    @property
+    def is_unsupported(self) -> bool:
+        return self == CatalogSupport.UNSUPPORTED
+
+    @property
+    def is_requires_set_catalog(self) -> bool:
+        return self == CatalogSupport.REQUIRES_SET_CATALOG
+
+    @property
+    def is_full_support(self) -> bool:
+        return self == CatalogSupport.FULL_SUPPORT
+
+    @property
+    def is_supported(self) -> bool:
+        return self.is_requires_set_catalog or self.is_full_support
+
+
 class SourceQuery:
     def __init__(
         self,
@@ -115,7 +142,6 @@ class EngineAdapter:
 
     DIALECT = ""
     DEFAULT_BATCH_SIZE = 10000
-    DEFAULT_SQL_GEN_KWARGS: t.Dict[str, str | bool | int] = {}
     ESCAPE_JSON = False
     SUPPORTS_TRANSACTIONS = True
     SUPPORTS_INDEXES = False
@@ -125,6 +151,7 @@ class EngineAdapter:
     SUPPORTS_CLONING = False
     SCHEMA_DIFFER = SchemaDiffer()
     SUPPORTS_TUPLE_IN = True
+    CATALOG_SUPPORT = CatalogSupport.UNSUPPORTED
 
     def __init__(
         self,
@@ -267,6 +294,14 @@ def close(self) -> t.Any:
         """Closes all open connections and releases all allocated resources."""
         self._connection_pool.close_all()
 
+    def get_current_catalog(self) -> t.Optional[str]:
+        """Returns the catalog name of the current connection."""
+        raise NotImplementedError()
+
+    def set_current_catalog(self, catalog: str) -> None:
+        """Sets the catalog name of the current connection."""
+        raise NotImplementedError()
+
     def replace_query(
         self,
         table_name: TableName,
@@ -516,7 +551,10 @@ def clone_table(
                 this=exp.to_table(target_table_name),
                 kind="TABLE",
                 replace=replace,
-                clone=exp.Clone(this=exp.to_table(source_table_name), **(clone_kwargs or {})),
+                clone=exp.Clone(
+                    this=exp.to_table(source_table_name),
+                    **(clone_kwargs or {}),
+                ),
                 **kwargs,
             )
         )
@@ -625,18 +663,18 @@ def create_view(
                 )
             )
 
+    @set_catalog()
     def create_schema(
         self,
-        schema_name: str,
-        catalog_name: t.Optional[str] = None,
+        schema_name: SchemaName,
         ignore_if_exists: bool = True,
         warn_on_error: bool = True,
     ) -> None:
         """Create a schema from a name or qualified table name."""
         try:
             self.execute(
                 exp.Create(
-                    this=exp.table_(schema_name, catalog_name),
+                    this=to_schema(schema_name),
                     kind="SCHEMA",
                     exists=ignore_if_exists,
                 )
@@ -646,13 +684,16 @@ def create_schema(
                 raise
             logger.warning("Failed to create schema '%s': %s", schema_name, e)
 
+    @set_catalog()
     def drop_schema(
-        self, schema_name: str, ignore_if_not_exists: bool = True, cascade: bool = False
+        self,
+        schema_name: SchemaName,
+        ignore_if_not_exists: bool = True,
+        cascade: bool = False,
     ) -> None:
-        """Drop a schema from a name or qualified table name."""
         self.execute(
             exp.Drop(
-                this=exp.table_(schema_name.split(".")[0]),
+                this=to_schema(schema_name),
                 kind="SCHEMA",
                 exists=ignore_if_not_exists,
                 cascade=cascade,
@@ -672,6 +713,7 @@ def drop_view(
             )
         )
 
+    @set_catalog()
     def columns(
         self, table_name: TableName, include_pseudo_columns: bool = False
     ) -> t.Dict[str, exp.DataType]:
@@ -687,6 +729,7 @@ def columns(
             )
         }
 
+    @set_catalog()
     def table_exists(self, table_name: TableName) -> bool:
         try:
             self.execute(exp.Describe(this=exp.to_table(table_name), kind="TABLE"))
@@ -926,6 +969,8 @@ def scd_type_2(
         columns_to_types = columns_to_types or self.columns(target_table)
         if valid_from_name not in columns_to_types or valid_to_name not in columns_to_types:
             columns_to_types = self.columns(target_table)
+        if not columns_to_types:
+            raise SQLMeshError(f"Could not get columns_to_types. Does {target_table} exist?")
         if updated_at_name not in columns_to_types:
             raise SQLMeshError(
                 f"Column {updated_at_name} not found in {target_table}. Table must contain an `updated_at` timestamp for SCD Type 2"
@@ -1154,11 +1199,20 @@ def merge(
                     match_expressions=[when_matched, when_not_matched],
                 )
 
+    @set_catalog()
     def rename_table(
         self,
         old_table_name: TableName,
         new_table_name: TableName,
     ) -> None:
+        new_table = exp.to_table(new_table_name)
+        if new_table.catalog:
+            old_table = exp.to_table(old_table_name)
+            catalog = old_table.catalog or self.get_current_catalog()
+            if catalog != new_table.catalog:
+                raise UnsupportedCatalogOperationError(
+                    "Tried to rename table across catalogs which is not supported"
+                )
         self.execute(exp.rename_table(old_table_name, new_table_name))
 
     def fetchone(
@@ -1307,7 +1361,7 @@ def temp_table(
         with self.transaction():
             table = self._get_temp_table(name)
             if table.db:
-                self.create_schema(table.db)
+                self.create_schema(schema_(table.args["db"], table.args.get("catalog")))
             self._create_table_from_source_queries(
                 table, source_queries, columns_to_types, exists=True, **kwargs
             )
@@ -1346,7 +1400,6 @@ def _to_sql(self, expression: exp.Expression, quote: bool = True, **kwargs: t.An
             "dialect": self.dialect,
             "pretty": False,
             "comments": False,
-            **self.DEFAULT_SQL_GEN_KWARGS,
             **self.sql_gen_kwargs,
             **kwargs,
         }
@@ -1356,13 +1409,10 @@ def _to_sql(self, expression: exp.Expression, quote: bool = True, **kwargs: t.An
 
         return expression.sql(**sql_gen_kwargs)  # type: ignore
 
-    def _get_data_objects(
-        self, schema_name: str, catalog_name: t.Optional[str] = None
-    ) -> t.List[DataObject]:
+    def _get_data_objects(self, schema_name: SchemaName) -> t.List[DataObject]:
         """
         Returns all the data objects that exist in the given schema and optionally catalog.
         """
-
         raise NotImplementedError()
 
     def _get_temp_table(

sqlmesh/core/engine_adapter/base_postgres.py

@@ -4,12 +4,12 @@
 
 from sqlglot import exp
 
-from sqlmesh.core.engine_adapter.mixins import EngineAdapter
+from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
 from sqlmesh.utils.errors import SQLMeshError
 
 if t.TYPE_CHECKING:
-    from sqlmesh.core._typing import TableName
+    from sqlmesh.core._typing import SchemaName, TableName
     from sqlmesh.core.engine_adapter.base import QueryOrDF
 
 
@@ -92,13 +92,11 @@ def create_view(
                 **create_kwargs,
             )
 
-    def _get_data_objects(
-        self, schema_name: str, catalog_name: t.Optional[str] = None
-    ) -> t.List[DataObject]:
+    def _get_data_objects(self, schema_name: SchemaName) -> t.List[DataObject]:
         """
         Returns all the data objects that exist in the given schema and optionally catalog.
         """
-        catalog_name = f"'{catalog_name}'" if catalog_name else "NULL"
+        catalog_name = f"'{self.get_current_catalog()}'"
         query = f"""
             SELECT
                 {catalog_name} AS catalog_name,

sqlmesh/core/engine_adapter/bigquery.py

@@ -9,7 +9,8 @@
 from sqlglot.helper import ensure_list
 from sqlglot.transforms import remove_precision_parameterized_types
 
-from sqlmesh.core.engine_adapter.base import SourceQuery
+from sqlmesh.core.dialect import to_schema
+from sqlmesh.core.engine_adapter.base import CatalogSupport, SourceQuery
 from sqlmesh.core.engine_adapter.mixins import InsertOverwriteWithMergeMixin
 from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
 from sqlmesh.core.node import IntervalUnit
@@ -26,7 +27,7 @@
     from google.cloud.bigquery.job.base import _AsyncJob as BigQueryQueryResult
     from google.cloud.bigquery.table import Table as BigQueryTable
 
-    from sqlmesh.core._typing import TableName
+    from sqlmesh.core._typing import SchemaName, TableName
     from sqlmesh.core.engine_adapter._typing import DF, Query
     from sqlmesh.core.engine_adapter.base import QueryOrDF
 
@@ -44,6 +45,7 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin):
     SUPPORTS_TRANSACTIONS = False
     SUPPORTS_MATERIALIZED_VIEWS = True
     SUPPORTS_CLONING = True
+    CATALOG_SUPPORT = CatalogSupport.FULL_SUPPORT
 
     # SQL is not supported for adding columns to structs: https://cloud.google.com/bigquery/docs/managing-table-schemas#api_1
     # Can explore doing this with the API in the future
@@ -131,10 +133,17 @@ def _end_session(self) -> None:
     def _is_session_active(self) -> bool:
         return self._session_id is not None
 
+    def get_current_catalog(self) -> t.Optional[str]:
+        """Returns the catalog name of the current connection."""
+        return self.client.project
+
+    def set_current_catalog(self, catalog: str) -> None:
+        """Sets the catalog name of the current connection."""
+        self.client.project = catalog
+
     def create_schema(
         self,
-        schema_name: str,
-        catalog_name: t.Optional[str] = None,
+        schema_name: SchemaName,
         ignore_if_exists: bool = True,
         warn_on_error: bool = True,
     ) -> None:
@@ -144,7 +153,6 @@ def create_schema(
         try:
             super().create_schema(
                 schema_name,
-                catalog_name=catalog_name,
                 ignore_if_exists=ignore_if_exists,
                 warn_on_error=False,
             )
@@ -535,18 +543,16 @@ def execute(
             self.cursor._set_rowcount(query_results)
             self.cursor._set_description(query_results.schema)
 
-    def _get_data_objects(
-        self, schema_name: str, catalog_name: t.Optional[str] = None
-    ) -> t.List[DataObject]:
+    def _get_data_objects(self, schema_name: SchemaName) -> t.List[DataObject]:
         """
         Returns all the data objects that exist in the given schema and optionally catalog.
         """
         from google.api_core.exceptions import NotFound
         from google.cloud.bigquery import DatasetReference
 
-        dataset_ref = DatasetReference(
-            project=catalog_name or self.client.project, dataset_id=schema_name
-        )
+        schema = to_schema(schema_name)
+        catalog_name = schema.catalog or self.get_current_catalog()
+        dataset_ref = DatasetReference(project=catalog_name, dataset_id=schema.db)
         try:
             return [
                 DataObject(