Repository: TobikoData/sqlmesh
Pull request #3201 (2024-09-30 03:35:12)

Feat: Support INCREMENTAL_BY_TIME_RANGE models on Athena/Hive
ref: issue #1315 


Patch:
diff --git a/sqlmesh/core/config/common.py b/sqlmesh/core/config/common.py
index ab07133d..819740c4 100644
--- a/sqlmesh/core/config/common.py
+++ b/sqlmesh/core/config/common.py
@@ -4,7 +4,7 @@ import typing as t
 from enum import Enum
 
 from sqlmesh.utils import classproperty
-from sqlmesh.utils.errors import ConfigError, SQLMeshError
+from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.pydantic import field_validator
 
 
@@ -86,17 +86,3 @@ variables_validator = field_validator(
     mode="before",
     check_fields=False,
 )(_variables_validator)
-
-
-def validate_s3_location(value: str, error_type: t.Type[Exception] = SQLMeshError) -> str:
-    if not value.startswith("s3://"):
-        raise error_type(f"Location '{value}' must be a s3:// URI")
-
-    if not value.endswith("/"):
-        value += "/"
-
-    # To avoid HIVE_METASTORE_ERROR: S3 resource path length must be less than or equal to 700.
-    if len(value) > 700:
-        raise error_type(f"Location '{value}' cannot be more than 700 characters")
-
-    return value
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index ed033b35..50e9cf17 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -18,7 +18,6 @@ from sqlmesh.core.config.base import BaseConfig
 from sqlmesh.core.config.common import (
     concurrent_tasks_validator,
     http_headers_validator,
-    validate_s3_location,
 )
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.utils.errors import ConfigError
@@ -29,6 +28,7 @@ from sqlmesh.utils.pydantic import (
     model_validator_v1_args,
     field_validator_v1_args,
 )
+from sqlmesh.utils.aws import validate_s3_uri
 
 if sys.version_info >= (3, 9):
     from typing import Literal
@@ -1555,11 +1555,13 @@ class AthenaConnectionConfig(ConnectionConfig):
             raise ConfigError("At least one of work_group or s3_staging_dir must be set")
 
         if s3_staging_dir:
-            values["s3_staging_dir"] = validate_s3_location(s3_staging_dir, error_type=ConfigError)
+            values["s3_staging_dir"] = validate_s3_uri(
+                s3_staging_dir, base=True, error_type=ConfigError
+            )
 
         if s3_warehouse_location:
-            values["s3_warehouse_location"] = validate_s3_location(
-                s3_warehouse_location, error_type=ConfigError
+            values["s3_warehouse_location"] = validate_s3_uri(
+                s3_warehouse_location, base=True, error_type=ConfigError
             )
 
         return values
diff --git a/sqlmesh/core/engine_adapter/athena.py b/sqlmesh/core/engine_adapter/athena.py
index 4058ca5e..24ceda4b 100644
--- a/sqlmesh/core/engine_adapter/athena.py
+++ b/sqlmesh/core/engine_adapter/athena.py
@@ -4,7 +4,7 @@ import typing as t
 import logging
 from sqlglot import exp
 from sqlmesh.core.dialect import to_schema
-from sqlmesh.core.config.common import validate_s3_location
+from sqlmesh.utils.aws import validate_s3_uri, parse_s3_uri
 from sqlmesh.core.engine_adapter.mixins import PandasNativeFetchDFSupportMixin
 from sqlmesh.core.engine_adapter.trino import TrinoEngineAdapter
 from sqlmesh.core.node import IntervalUnit
@@ -16,6 +16,8 @@ from sqlmesh.core.engine_adapter.shared import (
     DataObjectType,
     CommentCreationTable,
     CommentCreationView,
+    SourceQuery,
+    InsertOverwriteStrategy,
 )
 
 if t.TYPE_CHECKING:
@@ -53,9 +55,17 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin):
     @s3_warehouse_location.setter
     def s3_warehouse_location(self, value: t.Optional[str]) -> None:
         if value:
-            value = validate_s3_location(value)
+            value = validate_s3_uri(value, base=True)
         self._s3_warehouse_location = value
 
+    @property
+    def s3_warehouse_location_or_raise(self) -> str:
+        # this makes tests easier to write without extra null checks to keep mypy happy
+        if location := self.s3_warehouse_location:
+            return location
+
+        raise SQLMeshError("s3_warehouse_location was expected to be populated; it isnt")
+
     def create_state_table(
         self,
         table_name: str,
@@ -296,14 +306,27 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin):
         return None
 
     def _truncate_table(self, table_name: TableName) -> None:
-        table = exp.to_table(table_name)
-        # Athena doesnt support TRUNCATE TABLE. The closest thing is DELETE FROM <table> but it only works on Iceberg
-        self.execute(f"DELETE FROM {table.sql(dialect=self.dialect, identify=True)}")
+        if isinstance(table_name, str):
+            table_name = exp.to_table(table_name)
+
+        # Truncating an Iceberg table is just DELETE FROM <table>
+        if self._query_table_type(table_name) == "iceberg":
+            return self.delete_from(table_name, exp.true())
+
+        # Truncating a partitioned Hive table is dropping all partitions and deleting the data from S3
+        if self._is_hive_partitioned_table(table_name):
+            self._clear_partition_data(table_name, exp.true())
+        elif s3_location := self._query_table_s3_location(table_name):
+            # Truncating a non-partitioned Hive table is clearing out all data in its Location
+            self._clear_s3_location(s3_location)
 
     def _table_type(
-        self, table_type: t.Optional[str] = None
+        self, table_format: t.Optional[str] = None
     ) -> t.Union[t.Literal["hive"], t.Literal["iceberg"]]:
-        if table_type and table_type.lower() == "iceberg":
+        """
+        Interpret the "table_format" property to check if this is a Hive or an Iceberg table
+        """
+        if table_format and table_format.lower() == "iceberg":
             return "iceberg"
 
         # if we cant detect any indication of Iceberg, this is a Hive table
@@ -311,21 +334,29 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin):
 
     @lru_cache()
     def _query_table_type(
-        self, table_name: TableName
+        self, table: exp.Table
     ) -> t.Union[t.Literal["hive"], t.Literal["iceberg"]]:
         """
         Hit the DB to check if this is a Hive or an Iceberg table
         """
-        table_name = exp.to_table(table_name)
         # 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_name.sql(dialect='hive')}"):
+        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"
         return "iceberg"
 
+    def _is_hive_partitioned_table(self, table: exp.Table) -> bool:
+        try:
+            self._list_partitions(table=table, where=None, limit=1)
+            return True
+        except Exception as e:
+            if "TABLE_NOT_FOUND" in str(e):
+                return False
+            raise e
+
     def _table_location_or_raise(
         self, table_properties: t.Optional[t.Dict[str, exp.Expression]], table: exp.Table
     ) -> exp.LocationProperty:
@@ -359,7 +390,7 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin):
         else:
             return None
 
-        full_uri = validate_s3_location(os.path.join(base_uri, table.text("this") or ""))
+        full_uri = validate_s3_uri(os.path.join(base_uri, table.text("this") or ""), base=True)
         return exp.LocationProperty(this=exp.Literal.string(full_uri))
 
     def _find_matching_columns(
@@ -373,7 +404,105 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin):
                 matches.append((key, match_dtype))
         return matches
 
+    def _insert_overwrite_by_time_partition(
+        self,
+        table_name: TableName,
+        source_queries: t.List[SourceQuery],
+        columns_to_types: t.Dict[str, exp.DataType],
+        where: exp.Condition,
+        **kwargs: t.Any,
+    ) -> None:
+        if isinstance(table_name, str):
+            table_name = exp.to_table(table_name)
+
+        table_type = self._query_table_type(table_name)
+
+        if table_type == "iceberg":
+            # Iceberg tables work as expected, we can use the default behaviour
+            return super()._insert_overwrite_by_time_partition(
+                table_name, source_queries, columns_to_types, where, **kwargs
+            )
+
+        # For Hive tables, we need to drop all the partitions covered by the query and delete the data from S3
+        self._clear_partition_data(table_name, where)
+
+        # Now the data is physically gone, we can continue with inserting a new partition
+        return super()._insert_overwrite_by_time_partition(
+            table_name,
+            source_queries,
+            columns_to_types,
+            where,
+            insert_overwrite_strategy_override=InsertOverwriteStrategy.INTO_IS_OVERWRITE,  # since we already cleared the data
+            **kwargs,
+        )
+
+    def _clear_partition_data(self, table: exp.Table, where: t.Optional[exp.Condition]) -> None:
+        if partitions_to_drop := self._list_partitions(table, where):
+            for _, s3_location in partitions_to_drop:
+                logger.debug(
+                    f"Clearing S3 location for '{table.sql(dialect=self.dialect)}': {s3_location}"
+                )
+                self._clear_s3_location(s3_location)
+
+            partition_values = [k for k, _ in partitions_to_drop]
+            logger.debug(
+                f"Dropping partitions for '{table.sql(dialect=self.dialect)}' from metastore: {partition_values}"
+            )
+            self._drop_partitions_from_metastore(table, partition_values)
+
+    def _list_partitions(
+        self,
+        table: exp.Table,
+        where: t.Optional[exp.Condition] = None,
+        limit: t.Optional[int] = None,
+    ) -> t.List[t.Tuple[t.List[str], str]]:
+        # Use Athena's magic "$partitions" metadata table to identify the partitions to drop
+        # Doing it this way allows us to use SQL to filter the partition list
+        partition_table_name = table.copy()
+        partition_table_name.this.replace(
+            exp.to_identifier(f"{table.name}$partitions", quoted=True)
+        )
+
+        query = exp.select("*").from_(partition_table_name).where(where)
+        if limit:
+            query = query.limit(limit)
+
+        partition_values = [list(r) for r in self.fetchall(query, quote_identifiers=True)]
+
+        if partition_values:
+            response = self._glue_client.batch_get_partition(
+                DatabaseName=table.db,
+                TableName=table.name,
+                PartitionsToGet=[{"Values": v} for v in partition_values],
+            )
+            return sorted(
+                [(p["Values"], p["StorageDescriptor"]["Location"]) for p in response["Partitions"]]
+            )
+
+        return []
+
+    def _query_table_s3_location(self, table: exp.Table) -> str:
+        response = self._glue_client.get_table(DatabaseName=table.db, Name=table.name)
+
+        # Athena wont let you create a table without a location, so *theoretically* this should never be empty
+        if location := response.get("Table", {}).get("StorageDescriptor", {}).get("Location", None):
+            return location
+
+        raise SQLMeshError(f"Table {table} has no location set in the metastore!")
+
+    def _drop_partitions_from_metastore(
+        self, table: exp.Table, partition_values: t.List[t.List[t.Any]]
+    ) -> None:
+        self._glue_client.batch_delete_partition(
+            DatabaseName=table.db,
+            TableName=table.name,
+            PartitionsToDelete=[{"Values": v} for v in partition_values],
+        )
+
     def delete_from(self, table_name: TableName, where: t.Union[str, exp.Expression]) -> None:
+        if isinstance(table_name, str):
+            table_name = exp.to_table(table_name)
+
         table_type = self._query_table_type(table_name)
 
         # If Iceberg, DELETE operations work as expected
@@ -382,14 +511,50 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin):
 
         # If Hive, DELETE is an error
         if table_type == "hive":
-            # However, if the table is empty, we can make DELETE a no-op
+            # However, if there are no actual records to delete, we can make DELETE a no-op
             # This simplifies a bunch of calling code that just assumes DELETE works (which to be fair is a reasonable assumption since it does for every other engine)
             empty_check = (
-                exp.select("*").from_(table_name).limit(1)
+                exp.select("*").from_(table_name).where(where).limit(1)
             )  # deliberately not count(*) because we want the engine to stop as soon as it finds a record
             if len(self.fetchall(empty_check)) > 0:
-                # TODO: in future, if SQLMesh adds support for explicit partition management, we may
-                # be able to covert the DELETE query into an ALTER TABLE DROP PARTITION assuming the WHERE clause fully covers the partition bounds
-                raise SQLMeshError("Cannot delete from non-empty Hive table")
+                raise SQLMeshError("Cannot delete individual records from a Hive table")
 
         return None
+
+    def _clear_s3_location(self, s3_uri: str) -> None:
+        s3 = self._s3_client
+
+        bucket, key = parse_s3_uri(s3_uri)
+        if not key.endswith("/"):
+            key = f"{key}/"
+
+        keys_to_delete = []
+
+        # note: uses Delimiter=/ to prevent stepping into folders
+        # the assumption is that all the files in a partition live directly at the partition `Location`
+        for page in s3.get_paginator("list_objects_v2").paginate(
+            Bucket=bucket, Prefix=key, Delimiter="/"
+        ):
+            # list_objects_v2() returns 1000 keys per page so that lines up nicely with delete_objects() being able to delete 1000 keys at a time
+            keys_to_delete.append([item["Key"] for item in page.get("Contents", [])])
+
+        for chunk in keys_to_delete:
+            s3.delete_objects(Bucket=bucket, Delete={"Objects": [{"Key": k} for k in chunk]})
+
+    @property
+    def _glue_client(self) -> t.Any:
+        return self._boto3_client("glue")
+
+    @property
+    def _s3_client(self) -> t.Any:
+        return self._boto3_client("s3")
+
+    def _boto3_client(self, name: str) -> t.Any:
+        # use the client factory from PyAthena which is already configured with the correct AWS details
+        conn = self.connection
+        return conn.session.client(
+            name,
+            region_name=conn.region_name,
+            config=conn.config,
+            **conn._client_kwargs,
+        )  # type: ignore
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 251f501c..b6f76806 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -1207,7 +1207,21 @@ class EngineAdapter:
             low=low,
             high=high,
         )
-        self._insert_overwrite_by_condition(table_name, source_queries, columns_to_types, where)
+        return self._insert_overwrite_by_time_partition(
+            table_name, source_queries, columns_to_types, where, **kwargs
+        )
+
+    def _insert_overwrite_by_time_partition(
+        self,
+        table_name: TableName,
+        source_queries: t.List[SourceQuery],
+        columns_to_types: t.Dict[str, exp.DataType],
+        where: exp.Condition,
+        **kwargs: t.Any,
+    ) -> None:
+        return self._insert_overwrite_by_condition(
+            table_name, source_queries, columns_to_types, where
+        )
 
     def _values_to_sql(
         self,
diff --git a/sqlmesh/core/engine_adapter/shared.py b/sqlmesh/core/engine_adapter/shared.py
index ba2b6651..fb668bb0 100644
--- a/sqlmesh/core/engine_adapter/shared.py
+++ b/sqlmesh/core/engine_adapter/shared.py
@@ -216,10 +216,15 @@ class EngineRunMode(Enum):
 
 
 class InsertOverwriteStrategy(Enum):
+    # First, issue a DELETE to clear the data range. Then, issue an INSERT query to insert the new data
     DELETE_INSERT = 1
+    # Issue a single INSERT OVERWRITE query to replace a data range.
     INSERT_OVERWRITE = 2
+    # Issue a single INSERT INTO... REPLACE WHERE query
     # Note: Replace where on Databricks requires that `spark.sql.sources.partitionOverwriteMode` be set to `static`
     REPLACE_WHERE = 3
+    # Issue a single INSERT query to replace a data range. The assumption is that the query engine will transparently match partition bounds
+    # and replace data rather than append to it. Trino is an example of this when `hive.insert-existing-partitions-behavior=OVERWRITE` is configured
     INTO_IS_OVERWRITE = 4
 
     @property
diff --git a/sqlmesh/utils/aws.py b/sqlmesh/utils/aws.py
new file mode 100644
index 00000000..ed9c4f72
--- /dev/null
+++ b/sqlmesh/utils/aws.py
@@ -0,0 +1,39 @@
+import typing as t
+from urllib.parse import urlparse
+
+from sqlmesh.utils.errors import SQLMeshError
+
+
+def validate_s3_uri(
+    value: str, base: bool = False, error_type: t.Type[Exception] = SQLMeshError
+) -> str:
+    if not value.startswith("s3://"):
+        raise error_type(f"Location '{value}' must be a s3:// URI")
+
+    if base and not value.endswith("/"):
+        value = value + "/"
+
+    # To avoid HIVE_METASTORE_ERROR: S3 resource path length must be less than or equal to 700.
+    if len(value) > 700:
+        raise error_type(f"Location '{value}' cannot be more than 700 characters")
+
+    return value
+
+
+def parse_s3_uri(s3_uri: str) -> t.Tuple[str, str]:
+    """
+    Given a s3:// URI, parse it into a pair of (bucket, key)
+
+    Note that this could be any URI, including a file key, so we dont add a trailing / unlike validate_s3_base_uri
+    """
+    validate_s3_uri(s3_uri)
+
+    parsed_uri = urlparse(s3_uri)
+
+    bucket = parsed_uri.netloc
+    key = parsed_uri.path
+
+    if key:
+        key = key[1:]  # trim off leading /
+
+    return bucket, key


Test patch:
diff --git a/tests/core/engine_adapter/integration/__init__.py b/tests/core/engine_adapter/integration/__init__.py
index f6e9fb84..31bb0ce9 100644
--- a/tests/core/engine_adapter/integration/__init__.py
+++ b/tests/core/engine_adapter/integration/__init__.py
@@ -13,8 +13,10 @@ from sqlmesh import Config, Context, EngineAdapter
 from sqlmesh.core.config import load_config_from_paths
 from sqlmesh.core.config.connection import AthenaConnectionConfig
 from sqlmesh.core.dialect import normalize_model_name
+import sqlmesh.core.dialect as d
 from sqlmesh.core.engine_adapter import SparkEngineAdapter, TrinoEngineAdapter, AthenaEngineAdapter
 from sqlmesh.core.engine_adapter.shared import DataObject
+from sqlmesh.core.model.definition import SqlModel, load_sql_based_model
 from sqlmesh.utils import random_id
 from sqlmesh.utils.date import to_ds
 from sqlmesh.utils.pydantic import PydanticModel
@@ -68,12 +70,14 @@ class TestContext:
         self,
         test_type: str,
         engine_adapter: EngineAdapter,
+        mark: str,
         gateway: str,
         is_remote: bool = False,
         columns_to_types: t.Optional[t.Dict[str, t.Union[str, exp.DataType]]] = None,
     ):
         self.test_type = test_type
         self.engine_adapter = engine_adapter
+        self.mark = mark
         self.gateway = gateway
         self._columns_to_types = columns_to_types
         self.test_id = random_id(short=True)
@@ -148,8 +152,17 @@ class TestContext:
                 self.engine_adapter.get_catalog_type(self.engine_adapter.default_catalog) != "hive"
             )
 
+        if self.dialect == "athena":
+            return "hive" not in self.mark
+
         return True
 
+    @property
+    def default_table_format(self) -> t.Optional[str]:
+        if self.dialect in {"athena", "trino"} and "_" in self.mark:
+            return self.mark.split("_", 1)[-1]  # take eg 'athena_iceberg' and return 'iceberg'
+        return None
+
     def add_test_suffix(self, value: str) -> str:
         return f"{value}_{self.test_id}"
 
@@ -213,7 +226,7 @@ class TestContext:
             return {k: exp.Literal.string(v) if isinstance(v, str) else v for k, v in props.items()}
         return {}
 
-    def schema(self, schema_name: str, catalog_name: t.Optional[str] = None) -> str:
+    def schema(self, schema_name: str = TEST_SCHEMA, catalog_name: t.Optional[str] = None) -> str:
         return exp.table_name(
             normalize_model_name(
                 self.add_test_suffix(
@@ -498,3 +511,12 @@ class TestContext:
             self.engine_adapter.drop_schema(
                 schema_name=schema_name, ignore_if_not_exists=True, cascade=True
             )
+
+    def upsert_sql_model(self, model_definition: str) -> t.Tuple[Context, SqlModel]:
+        if not self._context:
+            self._context = self.create_context()
+
+        model = load_sql_based_model(expressions=d.parse(model_definition))
+        assert isinstance(model, SqlModel)
+        self._context.upsert_model(model)
+        return self._context, model
diff --git a/tests/core/engine_adapter/integration/conftest.py b/tests/core/engine_adapter/integration/conftest.py
index 085800ab..b8ed7ec2 100644
--- a/tests/core/engine_adapter/integration/conftest.py
+++ b/tests/core/engine_adapter/integration/conftest.py
@@ -4,12 +4,14 @@ import typing as t
 import pytest
 import pathlib
 import os
+from pytest import FixtureRequest
 
 from sqlmesh import Config, EngineAdapter
 from sqlmesh.core.config.connection import AthenaConnectionConfig
+from sqlmesh.core.engine_adapter import AthenaEngineAdapter
 from sqlmesh.core.config import load_config_from_paths
 
-from tests.core.engine_adapter.integration import TestContext
+from tests.core.engine_adapter.integration import TestContext, TEST_SCHEMA
 
 
 @pytest.fixture(scope="session")
@@ -34,7 +36,11 @@ def config() -> Config:
 
 @pytest.fixture
 def engine_adapter(
-    mark_gateway: t.Tuple[str, str], config, testrun_uid, run_count
+    mark_gateway: t.Tuple[str, str],
+    config,
+    testrun_uid: str,
+    run_count: int,
+    request: pytest.FixtureRequest,
 ) -> EngineAdapter:
     mark, gateway = mark_gateway
 
@@ -44,23 +50,25 @@ def engine_adapter(
 
     connection_config = config.gateways[gateway].connection
 
-    if mark == "athena":
-        connection_config = t.cast(AthenaConnectionConfig, connection_config)
+    engine_adapter = connection_config.create_engine_adapter()
+
+    if "athena" in mark:
+        assert isinstance(connection_config, AthenaConnectionConfig)
+        assert isinstance(engine_adapter, AthenaEngineAdapter)
+
         # S3 files need to go into a unique location for each test run
         # This is because DROP TABLE on a Hive table just drops the table from the metastore
         # The files still exist in S3, so if you CREATE TABLE to the same location, the old data shows back up
         # This is a problem for any tests like `test_init_project` that use a consistent schema like `sqlmesh_example` between runs
         # Note that the `testrun_uid` fixture comes from the xdist plugin
-        testrun_path = f"testrun_{testrun_uid}"
-        # TODO: remove the compexity of run_count and just use boto3 to clear out the S3 test location prior to running tests
-        if current_location := connection_config.s3_warehouse_location:
-            if testrun_path not in current_location:
-                # only add it if its not already there (since this setup code gets called multiple times in a full test run)
-                connection_config.s3_warehouse_location = os.path.join(
-                    current_location, testrun_path, str(run_count)
-                )
+        if connection_config.s3_warehouse_location:
+            engine_adapter.s3_warehouse_location = os.path.join(
+                connection_config.s3_warehouse_location,
+                f"testrun_{testrun_uid}",
+                str(run_count),
+                request.node.originalname,
+            )
 
-    engine_adapter = connection_config.create_engine_adapter()
     # Trino: If we batch up the requests then when running locally we get a table not found error after creating the
     # table and then immediately after trying to insert rows into it. There seems to be a delay between when the
     # metastore is made aware of the table and when it responds that it exists. I'm hoping this is not an issue
@@ -75,15 +83,27 @@ def engine_adapter(
 
 
 @pytest.fixture
-def ctx(request, engine_adapter, test_type, mark_gateway):
-    _, gateway = mark_gateway
+def ctx(
+    request: FixtureRequest,
+    engine_adapter: EngineAdapter,
+    test_type: str,
+    mark_gateway: t.Tuple[str, str],
+) -> t.Iterable[TestContext]:
+    mark, gateway = mark_gateway
     is_remote = request.node.get_closest_marker("remote") is not None
-    return TestContext(test_type, engine_adapter, gateway, is_remote=is_remote)
 
+    ctx = TestContext(test_type, engine_adapter, mark, gateway, is_remote=is_remote)
+    ctx.init()
+
+    yield ctx
 
-@pytest.fixture(autouse=True)
-def cleanup(ctx: TestContext):
-    yield  # run test
+    ctx.cleanup()
 
-    if ctx:
-        ctx.cleanup()
+
+@pytest.fixture
+def schema(ctx: TestContext) -> str:
+    schema_name = ctx.schema(TEST_SCHEMA)
+    ctx.engine_adapter.create_schema(
+        schema_name
+    )  # note: gets cleaned up when the TestContext fixture gets cleaned up
+    return schema_name
diff --git a/tests/core/engine_adapter/integration/test_integration.py b/tests/core/engine_adapter/integration/test_integration.py
index 81698c91..8ef7f035 100644
--- a/tests/core/engine_adapter/integration/test_integration.py
+++ b/tests/core/engine_adapter/integration/test_integration.py
@@ -210,7 +210,16 @@ def test_type(request):
             ],
         ),
         pytest.param(
-            "athena",
+            "athena_hive",
+            marks=[
+                pytest.mark.engine,
+                pytest.mark.remote,
+                pytest.mark.athena,
+                pytest.mark.xdist_group("engine_integration_athena"),
+            ],
+        ),
+        pytest.param(
+            "athena_iceberg",
             marks=[
                 pytest.mark.engine,
                 pytest.mark.remote,
@@ -221,7 +230,8 @@ def test_type(request):
     ]
 )
 def mark_gateway(request) -> t.Tuple[str, str]:
-    return request.param, f"inttest_{request.param}"
+    gateway = "athena" if "athena" in request.param else request.param
+    return request.param, f"inttest_{gateway}"
 
 
 @pytest.fixture
@@ -338,7 +348,9 @@ def test_temp_table(ctx: TestContext):
     )
     table = ctx.table("example")
 
-    with ctx.engine_adapter.temp_table(ctx.input_data(input_data), table.sql()) as table_name:
+    with ctx.engine_adapter.temp_table(
+        ctx.input_data(input_data), table.sql(), table_format=ctx.default_table_format
+    ) as table_name:
         results = ctx.get_metadata_results()
         assert len(results.views) == 0
         assert len(results.tables) == 1
@@ -358,6 +370,7 @@ def test_create_table(ctx: TestContext):
         {"id": exp.DataType.build("int")},
         table_description="test table description",
         column_descriptions={"id": "test id column description"},
+        table_format=ctx.default_table_format,
     )
     results = ctx.get_metadata_results()
     assert len(results.tables) == 1
@@ -387,6 +400,7 @@ def test_ctas(ctx: TestContext):
         ctx.input_data(input_data),
         table_description="test table description",
         column_descriptions={"id": "test id column description"},
+        table_format=ctx.default_table_format,
     )
     results = ctx.get_metadata_results()
     assert len(results.views) == 0
@@ -553,7 +567,9 @@ def test_replace_query(ctx: TestContext):
             {"id": 3, "ds": "2022-01-03"},
         ]
     )
-    ctx.engine_adapter.create_table(table, ctx.columns_to_types)
+    ctx.engine_adapter.create_table(
+        table, ctx.columns_to_types, table_format=ctx.default_table_format
+    )
     ctx.engine_adapter.replace_query(
         table,
         ctx.input_data(input_data),
@@ -562,6 +578,7 @@ def test_replace_query(ctx: TestContext):
         # exist prior to evaluation but when running these tests that isn't the case. As a result we just pass in
         # columns_to_types for these two engines so we can still test inference on the other ones
         columns_to_types=ctx.columns_to_types if ctx.dialect in ["spark", "databricks"] else None,
+        table_format=ctx.default_table_format,
     )
     results = ctx.get_metadata_results()
     assert len(results.views) == 0
@@ -585,6 +602,7 @@ def test_replace_query(ctx: TestContext):
             columns_to_types=(
                 ctx.columns_to_types if ctx.dialect in ["spark", "databricks"] else None
             ),
+            table_format=ctx.default_table_format,
         )
         results = ctx.get_metadata_results()
         assert len(results.views) == 0
@@ -606,7 +624,9 @@ def test_replace_query_batched(ctx: TestContext):
             {"id": 3, "ds": "2022-01-03"},
         ]
     )
-    ctx.engine_adapter.create_table(table, ctx.columns_to_types)
+    ctx.engine_adapter.create_table(
+        table, ctx.columns_to_types, table_format=ctx.default_table_format
+    )
     ctx.engine_adapter.replace_query(
         table,
         ctx.input_data(input_data),
@@ -615,6 +635,7 @@ def test_replace_query_batched(ctx: TestContext):
         # exist prior to evaluation but when running these tests that isn't the case. As a result we just pass in
         # columns_to_types for these two engines so we can still test inference on the other ones
         columns_to_types=ctx.columns_to_types if ctx.dialect in ["spark", "databricks"] else None,
+        table_format=ctx.default_table_format,
     )
     results = ctx.get_metadata_results()
     assert len(results.views) == 0
@@ -638,6 +659,7 @@ def test_replace_query_batched(ctx: TestContext):
             columns_to_types=(
                 ctx.columns_to_types if ctx.dialect in ["spark", "databricks"] else None
             ),
+            table_format=ctx.default_table_format,
         )
         results = ctx.get_metadata_results()
         assert len(results.views) == 0
@@ -650,7 +672,9 @@ def test_replace_query_batched(ctx: TestContext):
 def test_insert_append(ctx: TestContext):
     ctx.init()
     table = ctx.table("test_table")
-    ctx.engine_adapter.create_table(table, ctx.columns_to_types)
+    ctx.engine_adapter.create_table(
+        table, ctx.columns_to_types, table_format=ctx.default_table_format
+    )
     # Initial Load
     input_data = pd.DataFrame(
         [
@@ -705,6 +729,7 @@ def test_insert_overwrite_by_time_partition(ctx: TestContext):
         ctx.columns_to_types,
         partitioned_by=partitioned_by,
         partition_interval_unit="DAY",
+        table_format=ctx.default_table_format,
     )
     input_data = pd.DataFrame(
         [
@@ -730,7 +755,7 @@ def test_insert_overwrite_by_time_partition(ctx: TestContext):
     assert results.non_temp_tables[0] == table.name
     ctx.compare_with_current(table, input_data.iloc[1:])
 
-    if test_type == "df":
+    if ctx.test_type == "df":
         overwrite_data = pd.DataFrame(
             [
                 {"id": 10, ctx.time_column: "2022-01-03"},
@@ -798,7 +823,7 @@ def test_merge(ctx: TestContext):
     assert results.non_temp_tables[0] == table.name
     ctx.compare_with_current(table, input_data)
 
-    if test_type == "df":
+    if ctx.test_type == "df":
         merge_data = pd.DataFrame(
             [
                 {"id": 2, "ds": "2022-01-10"},
@@ -832,6 +857,10 @@ def test_merge(ctx: TestContext):
 
 
 def test_scd_type_2_by_time(ctx: TestContext):
+    # Athena only supports the operations required for SCD models on Iceberg tables
+    if ctx.mark == "athena_hive":
+        pytest.skip("SCD Type 2 is only supported on Athena / Iceberg")
+
     time_type = exp.DataType.build("timestamp")
 
     ctx.columns_to_types = {
@@ -847,10 +876,9 @@ def test_scd_type_2_by_time(ctx: TestContext):
         k: v for k, v in ctx.columns_to_types.items() if k not in ("valid_from", "valid_to")
     }
 
-    # Athena only supports the operations required for SCD models on Iceberg tables
-    table_format = "iceberg" if ctx.dialect == "athena" else None
-
-    ctx.engine_adapter.create_table(table, ctx.columns_to_types, table_format=table_format)
+    ctx.engine_adapter.create_table(
+        table, ctx.columns_to_types, table_format=ctx.default_table_format
+    )
     input_data = pd.DataFrame(
         [
             {"id": 1, "name": "a", "updated_at": "2022-01-01 00:00:00"},
@@ -868,7 +896,7 @@ def test_scd_type_2_by_time(ctx: TestContext):
         execution_time="2023-01-01 00:00:00",
         updated_at_as_valid_from=False,
         columns_to_types=input_schema,
-        table_format=table_format,
+        table_format=ctx.default_table_format,
         truncate=True,
     )
     results = ctx.get_metadata_results()
@@ -930,7 +958,7 @@ def test_scd_type_2_by_time(ctx: TestContext):
         execution_time="2023-01-05 00:00:00",
         updated_at_as_valid_from=False,
         columns_to_types=input_schema,
-        table_format=table_format,
+        table_format=ctx.default_table_format,
         truncate=False,
     )
     results = ctx.get_metadata_results()
@@ -983,6 +1011,10 @@ def test_scd_type_2_by_time(ctx: TestContext):
 
 
 def test_scd_type_2_by_column(ctx: TestContext):
+    # Athena only supports the operations required for SCD models on Iceberg tables
+    if ctx.mark == "athena_hive":
+        pytest.skip("SCD Type 2 is only supported on Athena / Iceberg")
+
     time_type = exp.DataType.build("timestamp")
 
     ctx.columns_to_types = {
@@ -998,10 +1030,9 @@ def test_scd_type_2_by_column(ctx: TestContext):
         k: v for k, v in ctx.columns_to_types.items() if k not in ("valid_from", "valid_to")
     }
 
-    # Athena only supports the operations required for SCD models on Iceberg tables
-    table_format = "iceberg" if ctx.dialect == "athena" else None
-
-    ctx.engine_adapter.create_table(table, ctx.columns_to_types, table_format=table_format)
+    ctx.engine_adapter.create_table(
+        table, ctx.columns_to_types, table_format=ctx.default_table_format
+    )
     input_data = pd.DataFrame(
         [
             {"id": 1, "name": "a", "status": "active"},
@@ -1164,6 +1195,7 @@ def test_get_data_objects(ctx: TestContext):
         {"id": exp.DataType.build("int")},
         table_description="test table description",
         column_descriptions={"id": "test id column description"},
+        table_format=ctx.default_table_format,
     )
     ctx.engine_adapter.create_view(
         view,
@@ -1236,10 +1268,9 @@ def test_truncate_table(ctx: TestContext):
     ctx.init()
     table = ctx.table("test_table")
 
-    # Athena only supports TRUNCATE (DELETE FROM <table>) on Iceberg tables
-    table_format = "iceberg" if ctx.dialect == "athena" else None
-
-    ctx.engine_adapter.create_table(table, ctx.columns_to_types, table_format=table_format)
+    ctx.engine_adapter.create_table(
+        table, ctx.columns_to_types, table_format=ctx.default_table_format
+    )
     input_data = pd.DataFrame(
         [
             {"id": 1, "ds": "2022-01-01"},
@@ -1280,10 +1311,15 @@ def test_transaction(ctx: TestContext):
     ctx.compare_with_current(table, input_data)
 
 
-def test_sushi(mark_gateway: t.Tuple[str, str], ctx: TestContext):
+def test_sushi(ctx: TestContext):
     if ctx.test_type != "query":
         pytest.skip("Sushi end-to-end tests only need to run for query")
 
+    if ctx.mark == "athena_hive":
+        pytest.skip(
+            "Sushi end-to-end tests only need to run once for Athena because sushi needs a hybrid of both Hive and Iceberg"
+        )
+
     config = load_config_from_paths(
         Config,
         project_paths=[
@@ -1291,12 +1327,11 @@ def test_sushi(mark_gateway: t.Tuple[str, str], ctx: TestContext):
         ],
         personal_paths=[pathlib.Path("~/.sqlmesh/config.yaml").expanduser()],
     )
-    mark, gateway = mark_gateway
 
-    if mark == "athena":
+    if ctx.dialect == "athena":
         # Ensure that this test is using the same s3_warehouse_location as TestContext (which includes the testrun_id)
         config.gateways[
-            gateway
+            ctx.gateway
         ].connection.s3_warehouse_location = ctx.engine_adapter.s3_warehouse_location
 
     # clear cache from prior runs
@@ -1306,7 +1341,7 @@ def test_sushi(mark_gateway: t.Tuple[str, str], ctx: TestContext):
 
         shutil.rmtree(cache_dir)
 
-    context = Context(paths="./examples/sushi", config=config, gateway=gateway)
+    context = Context(paths="./examples/sushi", config=config, gateway=ctx.gateway)
 
     # clean up any leftover schemas from previous runs (requires context)
     for schema in [
@@ -1369,7 +1404,7 @@ def test_sushi(mark_gateway: t.Tuple[str, str], ctx: TestContext):
             context._models.update({model_key: model_ch_cols_to_types})
 
         # create raw schema and view
-        if gateway == "inttest_clickhouse_cluster":
+        if ctx.gateway == "inttest_clickhouse_cluster":
             context.engine_adapter.execute("CREATE DATABASE IF NOT EXISTS raw ON CLUSTER cluster1;")
             context.engine_adapter.execute(
                 "DROP VIEW IF EXISTS raw.demographics ON CLUSTER cluster1;"
@@ -1378,11 +1413,12 @@ def test_sushi(mark_gateway: t.Tuple[str, str], ctx: TestContext):
                 "CREATE VIEW raw.demographics ON CLUSTER cluster1 AS SELECT 1 AS customer_id, '00000' AS zip;"
             )
 
-    # Athena needs models that get mutated after creation to be using Iceberg
     if ctx.dialect == "athena":
         for model_name in {"customer_revenue_lifetime"}:
             model_key = next(k for k in context._models if model_name in k)
-            model = context._models[model_key].copy(update={"table_format": "iceberg"})
+            model = context._models[model_key].copy(
+                update={"table_format": ctx.default_table_format}
+            )
             context._models.update({model_key: model})
 
     plan: Plan = context.plan(
@@ -1629,7 +1665,7 @@ def test_sushi(mark_gateway: t.Tuple[str, str], ctx: TestContext):
         validate_comments("sushi", is_physical_layer=False)
 
 
-def test_init_project(ctx: TestContext, mark_gateway: t.Tuple[str, str], tmp_path: pathlib.Path):
+def test_init_project(ctx: TestContext, tmp_path: pathlib.Path):
     if ctx.test_type != "query":
         pytest.skip("Init example project end-to-end tests only need to run for query")
 
@@ -1667,17 +1703,25 @@ def test_init_project(ctx: TestContext, mark_gateway: t.Tuple[str, str], tmp_pat
     if config.model_defaults.dialect != ctx.dialect:
         config.model_defaults = config.model_defaults.copy(update={"dialect": ctx.dialect})
 
-    mark, gateway = mark_gateway
-
-    if mark == "athena":
+    if ctx.dialect == "athena":
         # Ensure that this test is using the same s3_warehouse_location as TestContext (which includes the testrun_id)
         config.gateways[
-            gateway
+            ctx.gateway
         ].connection.s3_warehouse_location = ctx.engine_adapter.s3_warehouse_location
 
-    context = Context(paths=tmp_path, config=config, gateway=gateway)
+    context = Context(paths=tmp_path, config=config, gateway=ctx.gateway)
     ctx.engine_adapter = context.engine_adapter
 
+    if ctx.default_table_format:
+        # if the default table format is explicitly set, ensure its being used
+        replacement_models = {}
+        for model_key, model in context._models.items():
+            if not model.table_format:
+                replacement_models[model_key] = model.copy(
+                    update={"table_format": ctx.default_table_format}
+                )
+        context._models.update(replacement_models)
+
     # clean up any leftover schemas from previous runs (requires context)
     for schema in [
         state_schema,
diff --git a/tests/core/engine_adapter/integration/test_integration_athena.py b/tests/core/engine_adapter/integration/test_integration_athena.py
new file mode 100644
index 00000000..bd1ed6fb
--- /dev/null
+++ b/tests/core/engine_adapter/integration/test_integration_athena.py
@@ -0,0 +1,251 @@
+import typing as t
+import pytest
+import pandas as pd
+import datetime
+from sqlmesh.core.engine_adapter import AthenaEngineAdapter
+from sqlmesh.utils.aws import parse_s3_uri
+from tests.core.engine_adapter.integration import TestContext
+from sqlglot import exp
+
+pytestmark = [pytest.mark.docker, pytest.mark.engine, pytest.mark.athena]
+
+
+@pytest.fixture
+def mark_gateway() -> t.Tuple[str, str]:
+    return "athena", "inttest_athena"
+
+
+@pytest.fixture
+def test_type() -> str:
+    return "query"
+
+
+@pytest.fixture
+def s3(engine_adapter: AthenaEngineAdapter) -> t.Any:
+    return engine_adapter._s3_client
+
+
+def s3_list_objects(s3: t.Any, location: str, **list_objects_kwargs: t.Any) -> t.List[str]:
+    bucket, prefix = parse_s3_uri(location)
+    lst = []
+    for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=prefix):
+        lst.extend([o["Key"] for o in page.get("Contents", [])])
+    return lst
+
+
+def test_clear_partition_data(ctx: TestContext, engine_adapter: AthenaEngineAdapter, s3: t.Any):
+    base_uri = engine_adapter.s3_warehouse_location_or_raise
+    assert len(s3_list_objects(s3, base_uri)) == 0
+
+    src_table = ctx.table("src_table")
+    test_table = ctx.table("test_table")
+
+    base_data = pd.DataFrame(
+        [
+            {"id": 1, "ts": datetime.datetime(2023, 1, 1, 12, 13, 14)},
+            {"id": 2, "ts": datetime.datetime(2023, 1, 2, 8, 10, 0)},
+            {"id": 3, "ts": datetime.datetime(2023, 1, 3, 16, 5, 14)},
+        ]
+    )
+
+    engine_adapter.ctas(
+        table_name=src_table,
+        query_or_df=base_data,
+    )
+
+    sqlmesh_context, model = ctx.upsert_sql_model(
+        f"""
+        MODEL (
+            name {test_table},
+            kind INCREMENTAL_BY_TIME_RANGE (
+                time_column ds
+            ),
+            start '2023-01-01'
+        );
+
+        SELECT
+            id, ts, (ts::date)::varchar as ds
+        FROM {src_table}
+        WHERE ts BETWEEN @start_dt AND @end_dt
+        """
+    )
+
+    plan = sqlmesh_context.plan(no_prompts=True, auto_apply=True)
+    assert len(plan.snapshots) == 1
+    test_table_snapshot = list(plan.snapshots.values())[0]
+
+    files_before = s3_list_objects(s3, base_uri)
+    assert len(files_before) > 0
+
+    # src_table should have no partitions
+    with pytest.raises(Exception, match=r".*TABLE_NOT_FOUND.*\$partitions"):
+        engine_adapter._list_partitions(src_table)
+
+    # test_table physical snapshot table should have 3 partitions
+    test_table_physical_name = exp.to_table(test_table_snapshot.table_name())
+    partitions = engine_adapter._list_partitions(test_table_physical_name, where=None)
+    assert len(partitions) == 3
+    assert [p[0] for p in partitions] == [["2023-01-01"], ["2023-01-02"], ["2023-01-03"]]
+
+    assert engine_adapter.fetchone(f"select count(*) from {test_table}")[0] == 3  # type: ignore
+
+    # clear a partition
+    assert model.time_column
+    engine_adapter._clear_partition_data(
+        table=test_table_physical_name,
+        where=exp.Between(
+            this=model.time_column.column,
+            low=exp.Literal.string("2023-01-01"),
+            high=exp.Literal.string("2023-01-01"),
+        ),
+    )
+    partitions = engine_adapter._list_partitions(test_table_physical_name, where=None)
+    assert len(partitions) == 2
+    assert [p[0] for p in partitions] == [["2023-01-02"], ["2023-01-03"]]
+
+    # test that only S3 data for that partition was affected
+    files_after = s3_list_objects(s3, base_uri)
+    assert len(files_after) == len(files_before) - 1
+    assert len([f for f in files_before if "ds=2023-01-01" in f]) == 1
+    assert len([f for f in files_after if "ds=2023-01-01" in f]) == 0
+
+    assert engine_adapter.fetchone(f"select count(*) from {test_table}")[0] == 2  # type: ignore
+
+
+def test_clear_partition_data_multiple_columns(
+    ctx: TestContext, engine_adapter: AthenaEngineAdapter, s3: t.Any
+):
+    base_uri = engine_adapter.s3_warehouse_location_or_raise
+
+    src_table = ctx.table("src_table")
+    test_table = ctx.table("test_table")
+
+    base_data = pd.DataFrame(
+        [
+            {"id": 1, "ts": datetime.datetime(2023, 1, 1, 12, 13, 14), "system": "dev"},
+            {"id": 2, "ts": datetime.datetime(2023, 1, 1, 8, 13, 14), "system": "prod"},
+            {"id": 3, "ts": datetime.datetime(2023, 1, 2, 11, 10, 0), "system": "dev"},
+            {"id": 4, "ts": datetime.datetime(2023, 1, 2, 8, 10, 0), "system": "dev"},
+            {"id": 5, "ts": datetime.datetime(2023, 1, 3, 16, 5, 14), "system": "dev"},
+            {"id": 6, "ts": datetime.datetime(2023, 1, 3, 16, 5, 14), "system": "prod"},
+        ]
+    )
+
+    engine_adapter.ctas(
+        table_name=src_table,
+        query_or_df=base_data,
+    )
+
+    sqlmesh_context, model = ctx.upsert_sql_model(
+        f"""
+        MODEL (
+            name {test_table},
+            kind INCREMENTAL_BY_TIME_RANGE (
+                time_column ds
+            ),
+            partitioned_by (ds, system),
+            start '2023-01-01'
+        );
+
+        SELECT
+            id, ts, (ts::date)::varchar as ds, system
+        FROM {src_table}
+        WHERE ts BETWEEN @start_dt AND @end_dt
+        """
+    )
+
+    plan = sqlmesh_context.plan(no_prompts=True, auto_apply=True)
+    assert len(plan.snapshots) == 1
+    test_table_snapshot = list(plan.snapshots.values())[0]
+    test_table_physical_name = exp.to_table(test_table_snapshot.table_name())
+
+    partitions = engine_adapter._list_partitions(test_table_physical_name, where=None)
+    assert len(partitions) == 5
+    assert [p[0] for p in partitions] == [
+        ["2023-01-01", "dev"],
+        ["2023-01-01", "prod"],
+        ["2023-01-02", "dev"],
+        ["2023-01-03", "dev"],
+        ["2023-01-03", "prod"],
+    ]
+
+    files_before = s3_list_objects(s3, base_uri)
+    assert len(files_before) > 0
+
+    assert engine_adapter.fetchone(f"select count(*) from {test_table}")[0] == 6  # type: ignore
+
+    # this should clear 2 partitions, ["2023-01-01", "dev"] and ["2023-01-01", "prod"]
+    assert model.time_column
+    engine_adapter._clear_partition_data(
+        table=test_table_physical_name,
+        where=exp.Between(
+            this=model.time_column.column,
+            low=exp.Literal.string("2023-01-01"),
+            high=exp.Literal.string("2023-01-01"),
+        ),
+    )
+
+    partitions = engine_adapter._list_partitions(test_table_physical_name, where=None)
+    assert len(partitions) == 3
+    assert [p[0] for p in partitions] == [
+        ["2023-01-02", "dev"],
+        ["2023-01-03", "dev"],
+        ["2023-01-03", "prod"],
+    ]
+
+    files_after = s3_list_objects(s3, base_uri)
+    assert len(files_after) == len(files_before) - 2
+
+    def _match_partition(location_list: t.List[str], match: str):
+        return any(match in location for location in location_list)
+
+    assert _match_partition(files_before, "ds=2023-01-01/system=dev")
+    assert _match_partition(files_before, "ds=2023-01-01/system=prod")
+    assert not _match_partition(files_after, "ds=2023-01-01/system=dev")
+    assert not _match_partition(files_after, "ds=2023-01-01/system=prod")
+
+    assert engine_adapter.fetchone(f"select count(*) from {test_table}")[0] == 4  # type: ignore
+
+
+def test_hive_truncate_table(ctx: TestContext, engine_adapter: AthenaEngineAdapter, s3: t.Any):
+    base_uri = engine_adapter.s3_warehouse_location_or_raise
+
+    table_1 = ctx.table("table_one")
+    table_2 = ctx.table("table_two")
+
+    base_data = pd.DataFrame(
+        [
+            {"id": 1, "ts": datetime.datetime(2023, 1, 1, 12, 13, 14), "system": "dev"},
+            {"id": 2, "ts": datetime.datetime(2023, 1, 1, 8, 13, 14), "system": "prod"},
+            {"id": 3, "ts": datetime.datetime(2023, 1, 2, 11, 10, 0), "system": "dev"},
+            {"id": 4, "ts": datetime.datetime(2023, 1, 2, 8, 10, 0), "system": "dev"},
+            {"id": 5, "ts": datetime.datetime(2023, 1, 3, 16, 5, 14), "system": "dev"},
+            {"id": 6, "ts": datetime.datetime(2023, 1, 3, 16, 5, 14), "system": "prod"},
+        ]
+    )
+
+    assert len(s3_list_objects(s3, base_uri)) == 0
+
+    engine_adapter.ctas(table_name=table_1, query_or_df=base_data)
+
+    engine_adapter.ctas(table_name=table_2, query_or_df=base_data)
+
+    all_files = s3_list_objects(s3, base_uri)
+    assert len(all_files) > 0
+
+    table_1_location = engine_adapter._query_table_s3_location(table_1)
+    table_2_location = engine_adapter._query_table_s3_location(table_2)
+
+    table_1_files = s3_list_objects(s3, table_1_location)
+    table_2_files = s3_list_objects(s3, table_2_location)
+
+    assert len(table_1_files) < len(all_files)
+    assert len(table_2_files) < len(all_files)
+    assert len(table_1_files) + len(table_2_files) == len(all_files)
+
+    assert engine_adapter.fetchone(f"select count(*) from {table_1}")[0] == 6  # type: ignore
+    engine_adapter._truncate_table(table_1)
+    assert len(s3_list_objects(s3, table_1_location)) == 0
+    assert len(s3_list_objects(s3, table_2_location)) == len(table_2_files)
+
+    assert engine_adapter.fetchone(f"select count(*) from {table_1}")[0] == 0  # type: ignore
diff --git a/tests/core/engine_adapter/test_athena.py b/tests/core/engine_adapter/test_athena.py
index 8ba1266b..fc1b148b 100644
--- a/tests/core/engine_adapter/test_athena.py
+++ b/tests/core/engine_adapter/test_athena.py
@@ -1,5 +1,6 @@
 import typing as t
 import pytest
+from unittest.mock import Mock
 from pytest_mock import MockerFixture
 import pandas as pd
 
@@ -87,17 +88,6 @@ def test_table_location(
         assert "s3_base_location" not in table_properties
 
 
-def test_table_location_length(adapter: AthenaEngineAdapter) -> None:
-    with pytest.raises(SQLMeshError, match=r".*cannot be more than 700 characters"):
-        long_path = "foo/bar/" * 100
-        assert len(long_path) > 700
-
-        adapter._table_location(
-            table_properties={"s3_base_location": exp.Literal.string(f"s3://{long_path}")},
-            table=exp.to_table("foo"),
-        )
-
-
 def test_create_schema(adapter: AthenaEngineAdapter) -> None:
     adapter.create_schema("test")
 
@@ -307,10 +297,65 @@ def test_columns(adapter: AthenaEngineAdapter, mocker: MockerFixture):
     )
 
 
-def test_truncate_table(adapter: AthenaEngineAdapter):
+def test_truncate_table_iceberg(adapter: AthenaEngineAdapter, mocker: MockerFixture):
+    mocker.patch.object(
+        adapter,
+        "_query_table_type",
+        return_value="iceberg",
+    )
+    mocker.patch.multiple(
+        adapter, _clear_partition_data=mocker.DEFAULT, _clear_s3_location=mocker.DEFAULT
+    )
+    adapter._truncate_table(exp.to_table("foo.bar"))
+
+    assert to_sql_calls(adapter) == ['DELETE FROM "foo"."bar" WHERE TRUE']
+    t.cast(Mock, adapter._clear_partition_data).assert_not_called()
+    t.cast(Mock, adapter._clear_s3_location).assert_not_called()
+
+
+def test_truncate_table_hive(adapter: AthenaEngineAdapter, mocker: MockerFixture):
+    mocker.patch.object(
+        adapter,
+        "_query_table_type",
+        return_value="hive",
+    )
+    mocker.patch.object(
+        adapter,
+        "_is_hive_partitioned_table",
+        return_value=False,
+    )
+    mocker.patch.object(adapter, "_query_table_s3_location", return_value="s3://foo/bar")
+    mocker.patch.multiple(
+        adapter, _clear_partition_data=mocker.DEFAULT, _clear_s3_location=mocker.DEFAULT
+    )
+
+    adapter._truncate_table(exp.to_table("foo.bar"))
+
+    assert to_sql_calls(adapter) == []
+    t.cast(Mock, adapter._clear_partition_data).assert_not_called()
+    t.cast(Mock, adapter._clear_s3_location).assert_called_with("s3://foo/bar")
+
+
+def test_truncate_table_hive_partitioned(adapter: AthenaEngineAdapter, mocker: MockerFixture):
+    mocker.patch.object(
+        adapter,
+        "_query_table_type",
+        return_value="hive",
+    )
+    mocker.patch.object(
+        adapter,
+        "_is_hive_partitioned_table",
+        return_value=True,
+    )
+    mocker.patch.object(adapter, "_clear_partition_data")
+    mocker.patch.object(adapter, "_clear_s3_location")
     adapter._truncate_table(exp.to_table("foo.bar"))
 
-    assert to_sql_calls(adapter) == ['DELETE FROM "foo"."bar"']
+    assert to_sql_calls(adapter) == []
+    t.cast(Mock, adapter._clear_partition_data).assert_called_with(
+        exp.to_table("foo.bar"), exp.true()
+    )
+    t.cast(Mock, adapter._clear_s3_location).assert_not_called()
 
 
 def test_create_state_table(adapter: AthenaEngineAdapter):
diff --git a/tests/utils/test_aws.py b/tests/utils/test_aws.py
new file mode 100644
index 00000000..905cc00d
--- /dev/null
+++ b/tests/utils/test_aws.py
@@ -0,0 +1,34 @@
+import pytest
+from sqlmesh.utils.errors import SQLMeshError, ConfigError
+from sqlmesh.utils.aws import validate_s3_uri, parse_s3_uri
+
+
+def test_validate_s3_uri():
+    with pytest.raises(SQLMeshError, match=r".*must be a s3://.*"):
+        validate_s3_uri("hdfs://foo/bar")
+
+    with pytest.raises(ConfigError, match=r".*must be a s3://.*"):
+        validate_s3_uri("hdfs://foo/bar", error_type=ConfigError)
+
+    with pytest.raises(SQLMeshError, match=r".*must be a s3://.*"):
+        validate_s3_uri("/foo/bar")
+
+    with pytest.raises(SQLMeshError, match=r".*cannot be more than 700 characters"):
+        long_path = "foo/bar/" * 100
+        assert len(long_path) > 700
+        validate_s3_uri(f"s3://{long_path}")
+
+    assert validate_s3_uri("s3://foo/bar/") == "s3://foo/bar/"
+    assert validate_s3_uri("s3://foo/bar/baz") == "s3://foo/bar/baz"
+    assert validate_s3_uri("s3://foo/bar/baz", base=True) == "s3://foo/bar/baz/"
+
+
+def test_parse_s3_uri():
+    with pytest.raises(SQLMeshError, match=r".*must be a s3://.*"):
+        parse_s3_uri("hdfs://foo/bar")
+
+    assert parse_s3_uri("s3://foo") == ("foo", "")
+    assert parse_s3_uri("s3://foo/") == ("foo", "")
+    assert parse_s3_uri("s3://foo/bar") == ("foo", "bar")
+    assert parse_s3_uri("s3://foo/bar/") == ("foo", "bar/")
+    assert parse_s3_uri("s3://foo/bar/baz/bing.txt") == ("foo", "bar/baz/bing.txt")