refactor doris

Author: xinge-ji

sqlmesh/core/engine_adapter/base.py

@@ -2555,22 +2555,23 @@ def _get_data_objects(
         """
         raise NotImplementedError()
 
+    def _get_temp_table_name(self, table: TableName) -> str:
+        table_obj = exp.to_table(table)
+        return f"__temp_{table_obj.name}_{random_id(short=True)}"
+
     def _get_temp_table(
         self,
         table: TableName,
         table_only: bool = False,
         quoted: bool = True,
-        start_with: str = "__",
     ) -> exp.Table:
         """
         Returns the name of the temp table that should be used for the given table name.
         """
         table = t.cast(exp.Table, exp.to_table(table).copy())
         table.set(
             "this",
-            exp.to_identifier(
-                f"{start_with}temp_{table.name}_{random_id(short=True)}", quoted=quoted
-            ),
+            exp.to_identifier(self._get_temp_table_name(table), quoted=quoted),
         )
 
         if table_only:

sqlmesh/core/engine_adapter/doris.py

@@ -21,6 +21,7 @@
     set_catalog,
 )
 from sqlmesh.core.schema_diff import SchemaDiffer
+from sqlmesh.utils import random_id
 from sqlmesh.utils.errors import (
     SQLMeshError,
 )
@@ -34,7 +35,9 @@
 
 
 @set_catalog()
-class DorisEngineAdapter(LogicalMergeMixin, PandasNativeFetchDFSupportMixin, NonTransactionalTruncateMixin):
+class DorisEngineAdapter(
+    LogicalMergeMixin, PandasNativeFetchDFSupportMixin, NonTransactionalTruncateMixin
+):
     DIALECT = "doris"
     DEFAULT_BATCH_SIZE = 200
     SUPPORTS_TRANSACTIONS = False  # Doris doesn't support transactions
@@ -224,7 +227,9 @@ def _create_materialized_view(
         import pandas as pd
 
         if isinstance(query_or_df, pd.DataFrame):
-            values: t.List[t.Tuple[t.Any, ...]] = list(query_or_df.itertuples(index=False, name=None))
+            values: t.List[t.Tuple[t.Any, ...]] = list(
+                query_or_df.itertuples(index=False, name=None)
+            )
             target_columns_to_types, source_columns = self._columns_to_types(
                 query_or_df, target_columns_to_types, source_columns
             )
@@ -304,12 +309,16 @@ def _create_materialized_view(
                     doris_inline_clauses.append(f"BUILD {build_value}")
                 refresh = view_properties.get("refresh")
                 if refresh is not None:
-                    refresh_value = refresh.this if isinstance(refresh, exp.Literal) else str(refresh)
+                    refresh_value = (
+                        refresh.this if isinstance(refresh, exp.Literal) else str(refresh)
+                    )
                     doris_inline_clauses.append(f"REFRESH {refresh_value}")
                 refresh_trigger = view_properties.get("refresh_trigger")
                 if refresh_trigger is not None:
                     refresh_trigger_value = (
-                        refresh_trigger.this if isinstance(refresh_trigger, exp.Literal) else str(refresh_trigger)
+                        refresh_trigger.this
+                        if isinstance(refresh_trigger, exp.Literal)
+                        else str(refresh_trigger)
                     )
                     doris_inline_clauses.append(str(refresh_trigger_value))
                 # KEY / DUPLICATE KEY clause
@@ -339,7 +348,9 @@ def _create_materialized_view(
                     doris_inline_clauses.append(f"DUPLICATE KEY ({', '.join(cols)})")
                 # COMMENT clause
                 if table_description:
-                    doris_inline_clauses.append(f"COMMENT '{self._truncate_table_comment(table_description)}'")
+                    doris_inline_clauses.append(
+                        f"COMMENT '{self._truncate_table_comment(table_description)}'"
+                    )
             # PARTITION BY (inline for Doris MV)
             if partitioned_by:
                 part_cols = ", ".join(
@@ -349,7 +360,9 @@ def _create_materialized_view(
                             if isinstance(col, exp.Expression)
                             else exp.to_column(col).sql(dialect=self.dialect, identify=True)
                         )
-                        for col in (partitioned_by if isinstance(partitioned_by, list) else [partitioned_by])
+                        for col in (
+                            partitioned_by if isinstance(partitioned_by, list) else [partitioned_by]
+                        )
                     ]
                 )
                 doris_inline_clauses.append(f"PARTITION BY ({part_cols})")
@@ -398,7 +411,9 @@ def drop_view(
         # Remove cascade from kwargs as Doris doesn't support it
         if materialized and kwargs.get("view_properties"):
             view_properties = kwargs.pop("view_properties")
-            if view_properties.get("materialized_type") == "SYNC" and view_properties.get("source_table"):
+            if view_properties.get("materialized_type") == "SYNC" and view_properties.get(
+                "source_table"
+            ):
                 # Format the source table name properly for Doris
                 source_table = view_properties.get("source_table")
                 if isinstance(source_table, exp.Table):
@@ -432,18 +447,24 @@ def create_table_like(
             )
         )
 
-    def _create_table_comment(self, table_name: TableName, table_comment: str, table_kind: str = "TABLE") -> None:
+    def _create_table_comment(
+        self, table_name: TableName, table_comment: str, table_kind: str = "TABLE"
+    ) -> None:
         table_sql = exp.to_table(table_name).sql(dialect=self.dialect, identify=True)
 
-        self.execute(f'ALTER TABLE {table_sql} MODIFY COMMENT "{self._truncate_table_comment(table_comment)}"')
+        self.execute(
+            f'ALTER TABLE {table_sql} MODIFY COMMENT "{self._truncate_table_comment(table_comment)}"'
+        )
 
     def _build_create_comment_column_exp(
         self, table: exp.Table, column_name: str, column_comment: str, table_kind: str = "TABLE"
     ) -> exp.Comment | str:
         table_sql = table.sql(dialect=self.dialect, identify=True)
         return f'ALTER TABLE {table_sql} MODIFY COLUMN {column_name} COMMENT "{self._truncate_column_comment(column_comment)}"'
 
-    def delete_from(self, table_name: TableName, where: t.Optional[t.Union[str, exp.Expression]] = None) -> None:
+    def delete_from(
+        self, table_name: TableName, where: t.Optional[t.Union[str, exp.Expression]] = None
+    ) -> None:
         """
         Delete from a table.
 
@@ -491,7 +512,11 @@ def _find_subquery_in_condition(
 
     def _is_subquery_expression(self, expr: exp.Expression) -> bool:
         """Check if expression contains a subquery."""
-        return "query" in expr.args and expr.args["query"] and isinstance(expr.args["query"], exp.Subquery)
+        return (
+            "query" in expr.args
+            and expr.args["query"]
+            and isinstance(expr.args["query"], exp.Subquery)
+        )
 
     def _execute_delete_with_subquery(
         self, table_name: TableName, subquery_info: t.Tuple[exp.Expression, exp.Expression, bool]
@@ -561,7 +586,9 @@ def _create_table_from_columns(
         # Convert primary_key to unique_key for Doris (Doris doesn't support primary keys)
         if primary_key and "unique_key" not in table_properties:
             # Represent as a Tuple of columns to match downstream handling
-            table_properties["unique_key"] = exp.Tuple(expressions=[exp.to_column(col) for col in primary_key])
+            table_properties["unique_key"] = exp.Tuple(
+                expressions=[exp.to_column(col) for col in primary_key]
+            )
 
         # Update kwargs with the modified table_properties
         kwargs["table_properties"] = table_properties
@@ -606,7 +633,9 @@ def to_raw_sql(expr: t.Union[exp.Literal, exp.Var, str, t.Any]) -> exp.Var:
         if partitions:
             if isinstance(partitions, exp.Tuple):
                 create_expressions = [
-                    exp.Var(this=e.this, quoted=False) if isinstance(e, exp.Literal) else to_raw_sql(e)
+                    exp.Var(this=e.this, quoted=False)
+                    if isinstance(e, exp.Literal)
+                    else to_raw_sql(e)
                     for e in partitions.expressions
                 ]
             elif isinstance(partitions, exp.Literal):
@@ -645,13 +674,19 @@ def _build_table_properties_exp(
                 # Extract column names from Tuple expressions
                 column_names = []
                 for expr in unique_key.expressions:
-                    if isinstance(expr, exp.Column) and hasattr(expr, "this") and hasattr(expr.this, "this"):
+                    if (
+                        isinstance(expr, exp.Column)
+                        and hasattr(expr, "this")
+                        and hasattr(expr.this, "this")
+                    ):
                         column_names.append(str(expr.this.this))
                     elif hasattr(expr, "this"):
                         column_names.append(str(expr.this))
                     else:
                         column_names.append(str(expr))
-                properties.append(exp.UniqueKeyProperty(expressions=[exp.to_column(k) for k in column_names]))
+                properties.append(
+                    exp.UniqueKeyProperty(expressions=[exp.to_column(k) for k in column_names])
+                )
             elif isinstance(unique_key, exp.Column):
                 # Handle as single column
                 if hasattr(unique_key, "this") and hasattr(unique_key.this, "this"):
@@ -669,26 +704,38 @@ def _build_table_properties_exp(
                 # Extract column names from Tuple expressions
                 column_names = []
                 for expr in duplicate_key.expressions:
-                    if isinstance(expr, exp.Column) and hasattr(expr, "this") and hasattr(expr.this, "this"):
+                    if (
+                        isinstance(expr, exp.Column)
+                        and hasattr(expr, "this")
+                        and hasattr(expr.this, "this")
+                    ):
                         column_names.append(str(expr.this.this))
                     elif hasattr(expr, "this"):
                         column_names.append(str(expr.this))
                     else:
                         column_names.append(str(expr))
-                properties.append(exp.DuplicateKeyProperty(expressions=[exp.to_column(k) for k in column_names]))
+                properties.append(
+                    exp.DuplicateKeyProperty(expressions=[exp.to_column(k) for k in column_names])
+                )
             elif isinstance(duplicate_key, exp.Column):
                 # Handle as single column
                 if hasattr(duplicate_key, "this") and hasattr(duplicate_key.this, "this"):
                     column_name = str(duplicate_key.this.this)
                 else:
                     column_name = str(duplicate_key.this)
-                properties.append(exp.DuplicateKeyProperty(expressions=[exp.to_column(column_name)]))
+                properties.append(
+                    exp.DuplicateKeyProperty(expressions=[exp.to_column(column_name)])
+                )
             elif isinstance(duplicate_key, str):
-                properties.append(exp.DuplicateKeyProperty(expressions=[exp.to_column(duplicate_key)]))
+                properties.append(
+                    exp.DuplicateKeyProperty(expressions=[exp.to_column(duplicate_key)])
+                )
 
         if table_description:
             properties.append(
-                exp.SchemaCommentProperty(this=exp.Literal.string(self._truncate_table_comment(table_description)))
+                exp.SchemaCommentProperty(
+                    this=exp.Literal.string(self._truncate_table_comment(table_description))
+                )
             )
 
         # Handle partitioning
@@ -700,10 +747,14 @@ def _build_table_properties_exp(
                     # Handle literal strings like "RANGE(col)" or "LIST(col)"
                     if isinstance(expr, exp.Literal) and getattr(expr, "is_string", False):
                         text = str(expr.this)
-                        match = re.match(r"^\s*(RANGE|LIST)\s*\((.*?)\)\s*$", text, flags=re.IGNORECASE)
+                        match = re.match(
+                            r"^\s*(RANGE|LIST)\s*\((.*?)\)\s*$", text, flags=re.IGNORECASE
+                        )
                         if match:
                             inner = match.group(2)
-                            inner_cols = [c.strip().strip("`") for c in inner.split(",") if c.strip()]
+                            inner_cols = [
+                                c.strip().strip("`") for c in inner.split(",") if c.strip()
+                            ]
                             for col in inner_cols:
                                 normalized_partitioned_by.append(exp.to_column(col))
                             continue
@@ -720,7 +771,11 @@ def _build_table_properties_exp(
                 key_cols_set = set()
                 if isinstance(unique_key, exp.Tuple):
                     for expr in unique_key.expressions:
-                        if isinstance(expr, exp.Column) and hasattr(expr, "this") and hasattr(expr.this, "this"):
+                        if (
+                            isinstance(expr, exp.Column)
+                            and hasattr(expr, "this")
+                            and hasattr(expr.this, "this")
+                        ):
                             key_cols_set.add(str(expr.this.this))
                         elif hasattr(expr, "this"):
                             key_cols_set.add(str(expr.this))
@@ -801,12 +856,16 @@ def _build_table_properties_exp(
                         elif isinstance(expr.expression, exp.Array):
                             # Handle expressions array
                             distributed_info[key] = [
-                                str(e.this) for e in expr.expression.expressions if hasattr(e, "this")
+                                str(e.this)
+                                for e in expr.expression.expressions
+                                if hasattr(e, "this")
                             ]
                         elif isinstance(expr.expression, exp.Tuple):
                             # Handle expressions tuple (array of strings)
                             distributed_info[key] = [
-                                str(e.this) for e in expr.expression.expressions if hasattr(e, "this")
+                                str(e.this)
+                                for e in expr.expression.expressions
+                                if hasattr(e, "this")
                             ]
                         else:
                             distributed_info[key] = str(expr.expression)
@@ -859,13 +918,19 @@ def _build_table_properties_exp(
                     )
                     properties.append(prop)
         else:
-            unique_key_property = next((prop for prop in properties if isinstance(prop, exp.UniqueKeyProperty)), None)
+            unique_key_property = next(
+                (prop for prop in properties if isinstance(prop, exp.UniqueKeyProperty)), None
+            )
             if unique_key_property:
                 # Use the first column from unique_key as the distribution key
                 if unique_key_property.expressions:
                     first_col = unique_key_property.expressions[0]
-                    column_name = str(first_col.this) if hasattr(first_col, "this") else str(first_col)
-                    logger.info(f"[Doris] Adding default distributed_by using unique_key column: {column_name}")
+                    column_name = (
+                        str(first_col.this) if hasattr(first_col, "this") else str(first_col)
+                    )
+                    logger.info(
+                        f"[Doris] Adding default distributed_by using unique_key column: {column_name}"
+                    )
                     properties.append(
                         exp.DistributedByProperty(
                             expressions=[exp.to_column(column_name)],
@@ -882,13 +947,9 @@ def _build_table_properties_exp(
             return exp.Properties(expressions=properties)
         return None
 
-    def _get_temp_table(
-        self, table: TableName, table_only: bool = False, quoted: bool = True, start_with: str = "__"
-    ) -> exp.Table:
-        """
-        Returns the name of the temp table that should be used for the given table name.
-        """
-        return super()._get_temp_table(table, table_only, quoted, start_with="")
+    def _get_temp_table_name(self, table: TableName) -> str:
+        table_obj = exp.to_table(table)
+        return f"temp_{table_obj.name}_{random_id(short=True)}"
 
     def _properties_to_expressions(self, properties: t.Dict[str, t.Any]) -> t.List[exp.Expression]:
         """Convert a dictionary of properties to a list of exp.Property expressions."""