delete S3 tables fixed

Author: vishnu

sqlmesh/core/engine_adapter/athena.py

@@ -138,6 +138,26 @@ def _get_data_objects(
             for row in df.itertuples()
         ]
 
+    def table_exists(self, table_name: TableName) -> bool:
+        from sqlmesh.core.engine_adapter.base import _get_data_object_cache_key
+        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:
+            # We don't use DESCRIBE because it fails with "Unsupported ddl with 2 catalogs"
+            # for cross-catalog queries in Athena.
+            # And since table_exists isn't run with the set_catalog decorator (which sets QueryExecutionContext),
+            # we must fallback to a query that works with fully qualified names or 
+            # uses the information_schema/limit 0. A limit 0 select works with fully qualified names in Athena.
+            self.execute(exp.select("1").from_(table).limit(0))
+            return True
+        except Exception:
+            return False
+
     def columns(
         self, table_name: TableName, include_pseudo_columns: bool = False
     ) -> t.Dict[str, exp.DataType]:
@@ -152,11 +172,81 @@ def columns(
             .where(exp.column("table_schema").eq(table.db), exp.column("table_name").eq(table.name))
             .order_by("ordinal_position")
         )
-        result = self.fetchdf(query, quote_identifiers=True)
-        return {
-            str(r.column_name): exp.DataType.build(str(r.data_type))
-            for r in result.itertuples(index=False)
-        }
+        
+        try:
+            result = self.fetchdf(query, quote_identifiers=True)
+            return {
+                str(r.column_name): exp.DataType.build(str(r.data_type))
+                for r in result.itertuples(index=False)
+            }
+        except Exception as e:
+            # If information_schema query fails, we fallback to DESCRIBE.
+            # But DESCRIBE with multiple catalogs fails in Athena, so we strip the catalog here
+            # and rely on the set_current_catalog mechanism (applied at the EngineAdapter method level)
+            # to set the catalog in the execution context.
+            describe_table = table.copy()
+            catalog = describe_table.catalog
+            current_catalog = self.get_current_catalog()
+            
+            if catalog and catalog != self._default_catalog:
+                describe_table.set("catalog", None)
+                if catalog != current_catalog:
+                    self.set_current_catalog(catalog)
+            
+            try:
+                self.execute(exp.Describe(this=describe_table, kind="TABLE"))
+                
+                from sqlmesh.core.engine_adapter.base import _decoded_str
+                import itertools
+                describe_output = self.cursor.fetchall()
+                return {
+                    # Note: MySQL  returns the column type as bytes.
+                    column_name: exp.DataType.build(_decoded_str(column_type), dialect=self.dialect)
+                    for column_name, column_type, *_ in itertools.takewhile(
+                        lambda t: not t[0].startswith("#"),
+                        describe_output,
+                    )
+                    if column_name and column_name.strip() and column_type and column_type.strip()
+                }
+            finally:
+                if catalog and catalog != self._default_catalog and current_catalog != catalog:
+                    if current_catalog is not None:
+                        self.set_current_catalog(current_catalog)
+
+    def _drop_object(
+        self,
+        name: TableName | SchemaName,
+        exists: bool = True,
+        kind: str = "TABLE",
+        cascade: bool = False,
+        **drop_args: t.Any,
+    ) -> None:
+        if cascade and kind.upper() in self.SUPPORTED_DROP_CASCADE_OBJECT_KINDS:
+            drop_args["cascade"] = cascade
+
+        target_table = exp.to_table(name).copy()
+        is_schema = kind.upper() == "SCHEMA"
+        catalog = target_table.db if is_schema else target_table.catalog
+        
+        if catalog and catalog != self._default_catalog:
+            if is_schema:
+                target_table.set("db", None)
+            else:
+                target_table.set("catalog", None)
+                
+            current_catalog = self.get_current_catalog()
+            if current_catalog != catalog:
+                self.set_current_catalog(catalog)
+            
+            try:
+                self.execute(exp.Drop(this=target_table, kind=kind, exists=exists, **drop_args))
+            finally:
+                if current_catalog is not None and current_catalog != catalog:
+                    self.set_current_catalog(current_catalog)
+        else:
+            self.execute(exp.Drop(this=target_table, kind=kind, exists=exists, **drop_args))
+            
+        self._clear_data_object_cache(name)
 
     def _create_schema(
         self,
@@ -167,18 +257,39 @@ def _create_schema(
         kind: str,
     ) -> None:
         schema = to_schema(schema_name)
-        if schema.catalog and schema.catalog != self._default_catalog:
-            logger.info(
-                "Skipping creation of schema '%s' because Athena does not support creating schemas in non-default catalogs.",
-                schema.sql(dialect=self.dialect),
-            )
-            return
 
         if location := self._table_location(table_properties=None, table=exp.to_table(schema_name)):
             # don't add extra LocationProperty's if one already exists
             if not any(p for p in properties if isinstance(p, exp.LocationProperty)):
                 properties.append(location)
 
+        if schema.catalog and schema.catalog != self._default_catalog:
+            target_schema = schema.copy()
+            catalog = target_schema.catalog
+            target_schema.set("catalog", None)
+            
+            current_catalog = self.get_current_catalog()
+            if current_catalog != catalog:
+                self.set_current_catalog(catalog)
+            
+            try:
+                self.execute(
+                    exp.Create(
+                        this=target_schema,
+                        kind=kind,
+                        exists=ignore_if_exists,
+                        properties=exp.Properties(expressions=properties),
+                    )
+                )
+            except Exception as e:
+                if not warn_on_error:
+                    raise
+                logger.warning("Failed to create %s '%s': %s", kind.lower(), schema_name, e)
+            finally:
+                if current_catalog is not None and current_catalog != catalog:
+                    self.set_current_catalog(current_catalog)
+            return
+
         return super()._create_schema(
             schema_name=schema_name,
             ignore_if_exists=ignore_if_exists,
@@ -187,6 +298,76 @@ def _create_schema(
             kind=kind,
         )
 
+    def _get_temp_table(
+        self, table: TableName, table_only: bool = False, quoted: bool = True
+    ) -> exp.Table:
+        """
+        Returns the name of the temp table that should be used for the given table name.
+        """
+        from sqlmesh.utils import random_id
+        
+        table = t.cast(exp.Table, exp.to_table(table).copy())
+        
+        # AWS S3 Tables (and Athena generally) prefer or require table names to start with a letter.
+        # S3 Tables specifically fail with: "The specified table name is not valid" if it starts with __temp_
+        table.set(
+            "this", exp.to_identifier(f"temp_{table.name}_{random_id(short=True)}", quoted=quoted)
+        )
+
+        if table_only:
+            table.set("db", None)
+            table.set("catalog", None)
+
+        return table
+
+    def _create_table(
+        self,
+        table_name_or_schema: t.Union[exp.Schema, TableName],
+        expression: t.Optional[exp.Expr],
+        exists: bool = True,
+        replace: bool = False,
+        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
+        table_description: t.Optional[str] = None,
+        column_descriptions: t.Optional[t.Dict[str, str]] = None,
+        table_kind: t.Optional[str] = None,
+        track_rows_processed: bool = True,
+        **kwargs: t.Any,
+    ) -> None:
+        table: exp.Table
+        if isinstance(table_name_or_schema, str):
+            table = exp.to_table(table_name_or_schema)
+        elif isinstance(table_name_or_schema, exp.Schema):
+            table = table_name_or_schema.this
+        else:
+            table = table_name_or_schema
+
+        catalog = table.catalog
+        current_catalog = self.get_current_catalog()
+        
+        # For non-CTAS CREATE TABLE in a non-default catalog, the catalog is stripped by _build_create_table_exp.
+        # We need to set the query execution context here.
+        if not expression and catalog and catalog != self._default_catalog:
+            if current_catalog != catalog:
+                self.set_current_catalog(catalog)
+        
+        try:
+            super()._create_table(
+                table_name_or_schema=table_name_or_schema,
+                expression=expression,
+                exists=exists,
+                replace=replace,
+                target_columns_to_types=target_columns_to_types,
+                table_description=table_description,
+                column_descriptions=column_descriptions,
+                table_kind=table_kind,
+                track_rows_processed=track_rows_processed,
+                **kwargs,
+            )
+        finally:
+            if not expression and catalog and catalog != self._default_catalog:
+                if current_catalog is not None and current_catalog != catalog:
+                    self.set_current_catalog(current_catalog)
+
     def _build_create_table_exp(
         self,
         table_name_or_schema: t.Union[exp.Schema, TableName],
@@ -240,8 +421,20 @@ def _build_create_table_exp(
             ]
             table_name_or_schema.args["expressions"] = filtered_expressions
 
+        create_table = table_name_or_schema.copy()
+        
+        # When creating a table without AS SELECT, Athena fails with "Unsupported ddl with 2 catalogs"
+        # if a custom catalog like s3tablescatalog/supply is provided in the CREATE TABLE statement.
+        # It requires the catalog to be provided via QueryExecutionContext instead.
+        # The set_catalog decorator (which calls set_current_catalog) passes it to the QueryExecutionContext.
+        # But we also need to strip it from the generated CREATE TABLE statement.
+        # Note: We must strip the catalog from the table in the schema if table_name_or_schema is a schema.
+        target_table = create_table.this if isinstance(create_table, exp.Schema) else create_table
+        if not expression and target_table.catalog and target_table.catalog != self._default_catalog:
+            target_table.set("catalog", None)
+
         return exp.Create(
-            this=table_name_or_schema,
+            this=create_table,
             kind=table_kind or "TABLE",
             replace=replace,
             exists=exists,
@@ -446,11 +639,29 @@ def _query_table_type_or_raise(self, table: exp.Table) -> TableType:
         """
         # Note: SHOW TBLPROPERTIES gets parsed by SQLGlot as an exp.Command anyway so we just use a string here
         # This also means we need to use dialect="hive" instead of dialect="athena" so that the identifiers get the correct quoting (backticks)
-        for row in self.fetchall(f"SHOW TBLPROPERTIES {table.sql(dialect='hive', identify=True)}"):
-            # This query returns a single column with values like 'EXTERNAL\tTRUE'
-            row_lower = row[0].lower()
-            if "external" in row_lower and "true" in row_lower:
-                return "hive"
+        target_table = table.copy()
+        if target_table.catalog and target_table.catalog != self._default_catalog:
+            catalog = target_table.catalog
+            target_table.set("catalog", None)
+            
+            current_catalog = self.get_current_catalog()
+            if current_catalog != catalog:
+                self.set_current_catalog(catalog)
+            
+            try:
+                for row in self.fetchall(f"SHOW TBLPROPERTIES {target_table.sql(dialect='hive', identify=True)}"):
+                    row_lower = row[0].lower()
+                    if "external" in row_lower and "true" in row_lower:
+                        return "hive"
+            finally:
+                if current_catalog is not None and current_catalog != catalog:
+                    self.set_current_catalog(current_catalog)
+        else:
+            for row in self.fetchall(f"SHOW TBLPROPERTIES {target_table.sql(dialect='hive', identify=True)}"):
+                # This query returns a single column with values like 'EXTERNAL\tTRUE'
+                row_lower = row[0].lower()
+                if "external" in row_lower and "true" in row_lower:
+                    return "hive"
         return "iceberg"
 
     def _is_hive_partitioned_table(self, table: exp.Table) -> bool:
@@ -700,5 +911,10 @@ def _boto3_client(self, name: str) -> t.Any:
             **conn._client_kwargs,
         )  # type: ignore
 
+    def set_current_catalog(self, catalog: str) -> None:
+        self.connection.catalog_name = catalog
+        if hasattr(self.cursor, "_catalog_name"):
+            self.cursor._catalog_name = catalog
+
     def get_current_catalog(self) -> t.Optional[str]:
         return self.connection.catalog_name