sqlmesh 0.217.1.dev1__py3-none-any.whl → 0.227.2.dev20__py3-none-any.whl

sqlmesh/core/engine_adapter/bigquery.py

@@ -8,9 +8,10 @@ from sqlglot import exp, parse_one
 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 (
-    InsertOverwriteWithMergeMixin,
     ClusteredByMixin,
+    GrantsFromInfoSchemaMixin,
     RowDiffMixin,
     TableAlterClusterByOperation,
 )
@@ -20,6 +21,7 @@ from sqlmesh.core.engine_adapter.shared import (
     DataObjectType,
     SourceQuery,
     set_catalog,
+    InsertOverwriteStrategy,
 )
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.schema_diff import TableAlterOperation, NestedSupport
@@ -39,7 +41,7 @@ if t.TYPE_CHECKING:
     from google.cloud.bigquery.table import Table as BigQueryTable
 
     from sqlmesh.core._typing import SchemaName, SessionProperties, TableName
-    from sqlmesh.core.engine_adapter._typing import BigframeSession, DF, Query
+    from sqlmesh.core.engine_adapter._typing import BigframeSession, DCL, DF, GrantsConfig, Query
     from sqlmesh.core.engine_adapter.base import QueryOrDF
 
 
@@ -54,7 +56,7 @@ NestedFieldsDict = t.Dict[str, t.List[NestedField]]
 
 
 @set_catalog()
-class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, RowDiffMixin):
+class BigQueryEngineAdapter(ClusteredByMixin, RowDiffMixin, GrantsFromInfoSchemaMixin):
     """
     BigQuery Engine Adapter using the `google-cloud-bigquery` library's DB API.
     """
@@ -64,10 +66,16 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
     SUPPORTS_TRANSACTIONS = False
     SUPPORTS_MATERIALIZED_VIEWS = True
     SUPPORTS_CLONING = True
+    SUPPORTS_GRANTS = True
+    CURRENT_USER_OR_ROLE_EXPRESSION: exp.Expression = exp.func("session_user")
+    SUPPORTS_MULTIPLE_GRANT_PRINCIPALS = True
+    USE_CATALOG_IN_GRANTS = True
+    GRANT_INFORMATION_SCHEMA_TABLE_NAME = "OBJECT_PRIVILEGES"
     MAX_TABLE_COMMENT_LENGTH = 1024
     MAX_COLUMN_COMMENT_LENGTH = 1024
     SUPPORTS_QUERY_EXECUTION_TRACKING = True
     SUPPORTED_DROP_CASCADE_OBJECT_KINDS = ["SCHEMA"]
+    INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.MERGE
 
     SCHEMA_DIFFER_KWARGS = {
         "compatible_types": {
@@ -168,17 +176,18 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
         )
 
         def query_factory() -> Query:
-            if bigframes_pd and isinstance(df, bigframes_pd.DataFrame):
-                df.to_gbq(
+            ordered_df = df[list(source_columns_to_types)]
+            if bigframes_pd and isinstance(ordered_df, bigframes_pd.DataFrame):
+                ordered_df.to_gbq(
                     f"{temp_bq_table.project}.{temp_bq_table.dataset_id}.{temp_bq_table.table_id}",
                     if_exists="replace",
                 )
             elif not self.table_exists(temp_table):
                 # Make mypy happy
-                assert isinstance(df, pd.DataFrame)
+                assert isinstance(ordered_df, pd.DataFrame)
                 self._db_call(self.client.create_table, table=temp_bq_table, exists_ok=False)
                 result = self.__load_pandas_to_table(
-                    temp_bq_table, df, source_columns_to_types, replace=False
+                    temp_bq_table, ordered_df, source_columns_to_types, replace=False
                 )
                 if result.errors:
                     raise SQLMeshError(result.errors)
@@ -742,6 +751,12 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
             )
 
     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:
@@ -753,6 +768,28 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
         except NotFound:
             return False
 
+    def get_table_last_modified_ts(self, table_names: t.List[TableName]) -> t.List[int]:
+        from sqlmesh.utils.date import to_timestamp
+
+        datasets_to_tables: t.DefaultDict[str, t.List[str]] = defaultdict(list)
+        for table_name in table_names:
+            table = exp.to_table(table_name)
+            datasets_to_tables[table.db].append(table.name)
+
+        results = []
+
+        for dataset, tables in datasets_to_tables.items():
+            query = (
+                f"SELECT TIMESTAMP_MILLIS(last_modified_time) FROM `{dataset}.__TABLES__` WHERE "
+            )
+            for i, table_name in enumerate(tables):
+                query += f"TABLE_ID = '{table_name}'"
+                if i < len(tables) - 1:
+                    query += " OR "
+            results.extend(self.fetchall(query))
+
+        return [to_timestamp(row[0]) for row in results]
+
     def _get_table(self, table_name: TableName) -> BigQueryTable:
         """
         Returns a BigQueryTable object for the given table name.
@@ -1295,6 +1332,108 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
     def _session_id(self, value: t.Any) -> None:
         self._connection_pool.set_attribute("session_id", value)
 
+    def _get_current_schema(self) -> str:
+        raise NotImplementedError("BigQuery does not support current schema")
+
+    def _get_bq_dataset_location(self, project: str, dataset: str) -> str:
+        return self._db_call(self.client.get_dataset, dataset_ref=f"{project}.{dataset}").location
+
+    def _get_grant_expression(self, table: exp.Table) -> exp.Expression:
+        if not table.db:
+            raise ValueError(
+                f"Table {table.sql(dialect=self.dialect)} does not have a schema (dataset)"
+            )
+        project = table.catalog or self.get_current_catalog()
+        if not project:
+            raise ValueError(
+                f"Table {table.sql(dialect=self.dialect)} does not have a catalog (project)"
+            )
+
+        dataset = table.db
+        table_name = table.name
+        location = self._get_bq_dataset_location(project, dataset)
+
+        # https://cloud.google.com/bigquery/docs/information-schema-object-privileges
+        # OBJECT_PRIVILEGES is a project-level INFORMATION_SCHEMA view with regional qualifier
+        object_privileges_table = exp.to_table(
+            f"`{project}`.`region-{location}`.INFORMATION_SCHEMA.{self.GRANT_INFORMATION_SCHEMA_TABLE_NAME}",
+            dialect=self.dialect,
+        )
+        return (
+            exp.select("privilege_type", "grantee")
+            .from_(object_privileges_table)
+            .where(
+                exp.and_(
+                    exp.column("object_schema").eq(exp.Literal.string(dataset)),
+                    exp.column("object_name").eq(exp.Literal.string(table_name)),
+                    # Filter out current_user
+                    # BigQuery grantees format: "user:email" or "group:name"
+                    exp.func("split", exp.column("grantee"), exp.Literal.string(":"))[
+                        exp.func("OFFSET", exp.Literal.number("1"))
+                    ].neq(self.CURRENT_USER_OR_ROLE_EXPRESSION),
+                )
+            )
+        )
+
+    @staticmethod
+    def _grant_object_kind(table_type: DataObjectType) -> str:
+        if table_type == DataObjectType.VIEW:
+            return "VIEW"
+        if table_type == DataObjectType.MATERIALIZED_VIEW:
+            # We actually need to use "MATERIALIZED VIEW" here even though it's not listed
+            # as a supported resource_type in the BigQuery DCL doc:
+            # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-control-language
+            return "MATERIALIZED VIEW"
+        return "TABLE"
+
+    def _dcl_grants_config_expr(
+        self,
+        dcl_cmd: t.Type[DCL],
+        table: exp.Table,
+        grants_config: GrantsConfig,
+        table_type: DataObjectType = DataObjectType.TABLE,
+    ) -> t.List[exp.Expression]:
+        expressions: t.List[exp.Expression] = []
+        if not grants_config:
+            return expressions
+
+        # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-control-language
+
+        def normalize_principal(p: str) -> str:
+            if ":" not in p:
+                raise ValueError(f"Principal '{p}' missing a prefix label")
+
+            # allUsers and allAuthenticatedUsers special groups that are cas-sensitive and must start with "specialGroup:"
+            if p.endswith("allUsers") or p.endswith("allAuthenticatedUsers"):
+                if not p.startswith("specialGroup:"):
+                    raise ValueError(
+                        f"Special group principal '{p}' must start with 'specialGroup:' prefix label"
+                    )
+                return p
+
+            label, principal = p.split(":", 1)
+            # always lowercase principals
+            return f"{label}:{principal.lower()}"
+
+        object_kind = self._grant_object_kind(table_type)
+        for privilege, principals in grants_config.items():
+            if not principals:
+                continue
+
+            noramlized_principals = [exp.Literal.string(normalize_principal(p)) for p in principals]
+            args: t.Dict[str, t.Any] = {
+                "privileges": [exp.GrantPrivilege(this=exp.to_identifier(privilege, quoted=True))],
+                "securable": table.copy(),
+                "principals": noramlized_principals,
+            }
+
+            if object_kind:
+                args["kind"] = exp.Var(this=object_kind)
+
+            expressions.append(dcl_cmd(**args))  # type: ignore[arg-type]
+
+        return expressions
+
 
 class _ErrorCounter:
     """

sqlmesh/core/engine_adapter/clickhouse.py

@@ -112,8 +112,9 @@ class ClickhouseEngineAdapter(EngineAdapterWithIndexSupport, LogicalMergeMixin):
                     storage_format=exp.var("MergeTree"),
                     **kwargs,
                 )
+                ordered_df = df[list(source_columns_to_types)]
 
-                self.cursor.client.insert_df(temp_table.sql(dialect=self.dialect), df=df)
+                self.cursor.client.insert_df(temp_table.sql(dialect=self.dialect), df=ordered_df)
 
             return exp.select(*self._casted_columns(target_columns_to_types, source_columns)).from_(
                 temp_table
@@ -223,7 +224,7 @@ class ClickhouseEngineAdapter(EngineAdapterWithIndexSupport, LogicalMergeMixin):
         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")
@@ -455,7 +456,11 @@ class ClickhouseEngineAdapter(EngineAdapterWithIndexSupport, LogicalMergeMixin):
         )
 
     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(
@@ -631,16 +636,15 @@ class ClickhouseEngineAdapter(EngineAdapterWithIndexSupport, LogicalMergeMixin):
             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/databricks.py

@@ -5,7 +5,9 @@ import typing as t
 from functools import partial
 
 from sqlglot import exp
+
 from sqlmesh.core.dialect import to_schema
+from sqlmesh.core.engine_adapter.mixins import GrantsFromInfoSchemaMixin
 from sqlmesh.core.engine_adapter.shared import (
     CatalogSupport,
     DataObject,
@@ -28,12 +30,16 @@ if t.TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class DatabricksEngineAdapter(SparkEngineAdapter):
+class DatabricksEngineAdapter(SparkEngineAdapter, GrantsFromInfoSchemaMixin):
     DIALECT = "databricks"
     INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.REPLACE_WHERE
     SUPPORTS_CLONING = True
     SUPPORTS_MATERIALIZED_VIEWS = True
     SUPPORTS_MATERIALIZED_VIEW_SCHEMA = True
+    SUPPORTS_GRANTS = True
+    USE_CATALOG_IN_GRANTS = True
+    # Spark has this set to false for compatibility when mixing with Trino but that isn't a concern with Databricks
+    QUOTE_IDENTIFIERS_IN_VIEWS = True
     SCHEMA_DIFFER_KWARGS = {
         "support_positional_add": True,
         "nested_support": NestedSupport.ALL,
@@ -149,6 +155,28 @@ class DatabricksEngineAdapter(SparkEngineAdapter):
     def catalog_support(self) -> CatalogSupport:
         return CatalogSupport.FULL_SUPPORT
 
+    @staticmethod
+    def _grant_object_kind(table_type: DataObjectType) -> str:
+        if table_type == DataObjectType.VIEW:
+            return "VIEW"
+        if table_type == DataObjectType.MATERIALIZED_VIEW:
+            return "MATERIALIZED VIEW"
+        return "TABLE"
+
+    def _get_grant_expression(self, table: exp.Table) -> exp.Expression:
+        # We only care about explicitly granted privileges and not inherited ones
+        # if this is removed you would see grants inherited from the catalog get returned
+        expression = super()._get_grant_expression(table)
+        expression.args["where"].set(
+            "this",
+            exp.and_(
+                expression.args["where"].this,
+                exp.column("inherited_from").eq(exp.Literal.string("NONE")),
+                wrap=False,
+            ),
+        )
+        return expression
+
     def _begin_session(self, properties: SessionProperties) -> t.Any:
         """Begin a new session."""
         # Align the different possible connectors to a single catalog
@@ -266,7 +294,9 @@ class DatabricksEngineAdapter(SparkEngineAdapter):
                 exp.column("table_catalog").as_("catalog"),
                 exp.case(exp.column("table_type"))
                 .when(exp.Literal.string("VIEW"), exp.Literal.string("view"))
-                .when(exp.Literal.string("MATERIALIZED_VIEW"), exp.Literal.string("view"))
+                .when(
+                    exp.Literal.string("MATERIALIZED_VIEW"), exp.Literal.string("materialized_view")
+                )
                 .else_(exp.Literal.string("table"))
                 .as_("type"),
             )
@@ -297,6 +327,7 @@ class DatabricksEngineAdapter(SparkEngineAdapter):
         target_table_name: TableName,
         source_table_name: TableName,
         replace: bool = False,
+        exists: bool = True,
         clone_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
         **kwargs: t.Any,
     ) -> None:
@@ -363,3 +394,20 @@ class DatabricksEngineAdapter(SparkEngineAdapter):
             expressions.append(clustered_by_exp)
             properties = exp.Properties(expressions=expressions)
         return properties
+
+    def _build_column_defs(
+        self,
+        target_columns_to_types: t.Dict[str, exp.DataType],
+        column_descriptions: t.Optional[t.Dict[str, str]] = None,
+        is_view: bool = False,
+        materialized: bool = False,
+    ) -> t.List[exp.ColumnDef]:
+        # Databricks requires column types to be specified when adding column comments
+        # in CREATE MATERIALIZED VIEW statements. Override is_view to False to force
+        # column types to be included when comments are present.
+        if is_view and materialized and column_descriptions:
+            is_view = False
+
+        return super()._build_column_defs(
+            target_columns_to_types, column_descriptions, is_view, materialized
+        )

sqlmesh/core/engine_adapter/fabric.py

@@ -10,23 +10,17 @@ from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_resul
 from sqlmesh.core.engine_adapter.mssql import MSSQLEngineAdapter
 from sqlmesh.core.engine_adapter.shared import (
     InsertOverwriteStrategy,
-    SourceQuery,
 )
-from sqlmesh.core.engine_adapter.base import EngineAdapter
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.connection_pool import ConnectionPool
+from sqlmesh.core.schema_diff import TableAlterOperation
+from sqlmesh.utils import random_id
 
 
-if t.TYPE_CHECKING:
-    from sqlmesh.core._typing import TableName
-
-
-from sqlmesh.core.engine_adapter.mixins import LogicalMergeMixin
-
 logger = logging.getLogger(__name__)
 
 
-class FabricEngineAdapter(LogicalMergeMixin, MSSQLEngineAdapter):
+class FabricEngineAdapter(MSSQLEngineAdapter):
     """
     Adapter for Microsoft Fabric.
     """
@@ -58,26 +52,6 @@ class FabricEngineAdapter(LogicalMergeMixin, MSSQLEngineAdapter):
     def _target_catalog(self, value: t.Optional[str]) -> None:
         self._connection_pool.set_attribute("target_catalog", value)
 
-    def _insert_overwrite_by_condition(
-        self,
-        table_name: TableName,
-        source_queries: t.List[SourceQuery],
-        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
-        where: t.Optional[exp.Condition] = None,
-        insert_overwrite_strategy_override: t.Optional[InsertOverwriteStrategy] = None,
-        **kwargs: t.Any,
-    ) -> None:
-        # Override to avoid MERGE statement which isn't fully supported in Fabric
-        return EngineAdapter._insert_overwrite_by_condition(
-            self,
-            table_name=table_name,
-            source_queries=source_queries,
-            target_columns_to_types=target_columns_to_types,
-            where=where,
-            insert_overwrite_strategy_override=InsertOverwriteStrategy.DELETE_INSERT,
-            **kwargs,
-        )
-
     @property
     def api_client(self) -> FabricHttpClient:
         # the requests Session is not guaranteed to be threadsafe
@@ -181,6 +155,113 @@ class FabricEngineAdapter(LogicalMergeMixin, MSSQLEngineAdapter):
                 f"Unable to switch catalog to {catalog_name}, catalog ended up as {catalog_after_switch}"
             )
 
+    def alter_table(
+        self, alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]]
+    ) -> None:
+        """
+        Applies alter expressions to a table. Fabric has limited support for ALTER TABLE,
+        so this method implements a workaround for column type changes.
+        This method is self-contained and sets its own catalog context.
+        """
+        if not alter_expressions:
+            return
+
+        # Get the target table from the first expression to determine the correct catalog.
+        first_op = alter_expressions[0]
+        expression = first_op.expression if isinstance(first_op, TableAlterOperation) else first_op
+        if not isinstance(expression, exp.Alter) or not expression.this.catalog:
+            # Fallback for unexpected scenarios
+            logger.warning(
+                "Could not determine catalog from alter expression, executing with current context."
+            )
+            super().alter_table(alter_expressions)
+            return
+
+        target_catalog = expression.this.catalog
+        self.set_current_catalog(target_catalog)
+
+        with self.transaction():
+            for op in alter_expressions:
+                expression = op.expression if isinstance(op, TableAlterOperation) else op
+
+                if not isinstance(expression, exp.Alter):
+                    self.execute(expression)
+                    continue
+
+                for action in expression.actions:
+                    table_name = expression.this
+
+                    table_name_without_catalog = table_name.copy()
+                    table_name_without_catalog.set("catalog", None)
+
+                    is_type_change = isinstance(action, exp.AlterColumn) and action.args.get(
+                        "dtype"
+                    )
+
+                    if is_type_change:
+                        column_to_alter = action.this
+                        new_type = action.args["dtype"]
+                        temp_column_name_str = f"{column_to_alter.name}__{random_id(short=True)}"
+                        temp_column_name = exp.to_identifier(temp_column_name_str)
+
+                        logger.info(
+                            "Applying workaround for column '%s' on table '%s' to change type to '%s'.",
+                            column_to_alter.sql(),
+                            table_name.sql(),
+                            new_type.sql(),
+                        )
+
+                        # Step 1: Add a temporary column.
+                        add_column_expr = exp.Alter(
+                            this=table_name_without_catalog.copy(),
+                            kind="TABLE",
+                            actions=[
+                                exp.ColumnDef(this=temp_column_name.copy(), kind=new_type.copy())
+                            ],
+                        )
+                        add_sql = self._to_sql(add_column_expr)
+                        self.execute(add_sql)
+
+                        # Step 2: Copy and cast data.
+                        update_sql = self._to_sql(
+                            exp.Update(
+                                this=table_name_without_catalog.copy(),
+                                expressions=[
+                                    exp.EQ(
+                                        this=temp_column_name.copy(),
+                                        expression=exp.Cast(
+                                            this=column_to_alter.copy(), to=new_type.copy()
+                                        ),
+                                    )
+                                ],
+                            )
+                        )
+                        self.execute(update_sql)
+
+                        # Step 3: Drop the original column.
+                        drop_sql = self._to_sql(
+                            exp.Alter(
+                                this=table_name_without_catalog.copy(),
+                                kind="TABLE",
+                                actions=[exp.Drop(this=column_to_alter.copy(), kind="COLUMN")],
+                            )
+                        )
+                        self.execute(drop_sql)
+
+                        # Step 4: Rename the temporary column.
+                        old_name_qualified = f"{table_name_without_catalog.sql(dialect=self.dialect)}.{temp_column_name.sql(dialect=self.dialect)}"
+                        new_name_unquoted = column_to_alter.sql(
+                            dialect=self.dialect, identify=False
+                        )
+                        rename_sql = f"EXEC sp_rename '{old_name_qualified}', '{new_name_unquoted}', 'COLUMN'"
+                        self.execute(rename_sql)
+                    else:
+                        # For other alterations, execute directly.
+                        direct_alter_expr = exp.Alter(
+                            this=table_name_without_catalog.copy(), kind="TABLE", actions=[action]
+                        )
+                        self.execute(direct_alter_expr)
+
 
 class FabricHttpClient:
     def __init__(self, tenant_id: str, workspace_id: str, client_id: str, client_secret: str):