feat: improve handling of s3 table storage

moto/s3/models.py

@@ -1646,6 +1646,10 @@ def default_retention(self) -> str:
         return now.strftime("%Y-%m-%dT%H:%M:%SZ")
 
 
+class FakeTableStorageBucket(FakeBucket):
+    ...
+
+
 class S3Backend(BaseBackend, CloudWatchMetricProvider):
     """
     Custom S3 endpoints are supported, if you are using a S3-compatible storage solution like Ceph.
@@ -1711,6 +1715,7 @@ def test_my_custom_endpoint():
     def __init__(self, region_name: str, account_id: str):
         super().__init__(region_name, account_id)
         self.buckets: Dict[str, FakeBucket] = {}
+        self.table_buckets: Dict[str, FakeTableStorageBucket] = {}
         self.tagger = TaggingService()
         self._pagination_tokens: Dict[str, str] = {}
 
@@ -1719,7 +1724,7 @@ def reset(self) -> None:
         # Ensure that these TemporaryFile-objects are closed, and leave no filehandles open
         #
         # First, check all known buckets/keys
-        for bucket in self.buckets.values():
+        for bucket in itertools.chain(self.buckets.values(), self.table_buckets.values()):
             for key in bucket.keys.values():  # type: ignore
                 if isinstance(key, FakeKey):
                     key.dispose()
@@ -1876,13 +1881,27 @@ def create_bucket(self, bucket_name: str, region_name: str) -> FakeBucket:
 
         return new_bucket
 
+    def create_table_storage_bucket(self, bucket_name: str, region_name: str) -> FakeTableStorageBucket:
+        if bucket_name in s3_backends.bucket_accounts.keys():
+            raise BucketAlreadyExists(bucket=bucket_name)
+        new_bucket = FakeTableStorageBucket(
+            name=bucket_name, account_id=self.account_id, region_name=region_name
+        )
+
+        self.table_buckets[bucket_name] = new_bucket
+        return new_bucket
+
+
     def list_buckets(self) -> List[FakeBucket]:
         return list(self.buckets.values())
 
     def get_bucket(self, bucket_name: str) -> FakeBucket:
         if bucket_name in self.buckets:
             return self.buckets[bucket_name]
 
+        if bucket_name in self.table_buckets:
+            return self.table_buckets[bucket_name]
+
         if bucket_name in s3_backends.bucket_accounts:
             if not s3_allow_crossdomain_access():
                 raise AccessDeniedByLock
@@ -1903,6 +1922,19 @@ def delete_bucket(self, bucket_name: str) -> Optional[FakeBucket]:
             s3_backends.bucket_accounts.pop(bucket_name, None)
             return self.buckets.pop(bucket_name)
 
+    def delete_table_storage_bucket(self, bucket_name: str) -> Optional[FakeBucket]:
+        bucket = self.get_bucket(bucket_name)
+        assert isinstance(bucket, FakeTableStorageBucket)
+        # table storage buckets can be deleted while not empty
+        if bucket.keys:
+            for key in bucket.keys.values():  # type: ignore
+                if isinstance(key, FakeKey):
+                    key.dispose()
+            for part in bucket.multiparts.values():
+                part.dispose()
+        s3_backends.bucket_accounts.pop(bucket_name, None)
+        return self.table_buckets.pop(bucket_name)
+
     def get_bucket_accelerate_configuration(self, bucket_name: str) -> Optional[str]:
         bucket = self.get_bucket(bucket_name)
         return bucket.accelerate_configuration

moto/s3tables/models.py

@@ -81,6 +81,7 @@ def __init__(
         self.modified_by: Optional[str] = None
         self.namespace = namespace
         self.table_bucket_arn = table_bucket_arn
+        self.region_name = table_bucket_arn.split(":")[3]
         self.managed_by_service = managed_by_service
         self.metadata_location: Optional[str] = None
         self._bucket = self._create_underlying_bucket()
@@ -107,12 +108,8 @@ def _create_underlying_bucket(self) -> FakeBucket:
         hash.update(self.name.encode())
         bucket_name = f"{str(uuid4())[:21]}{hash.hexdigest()[:34]}--table-s3"
 
-        # TODO: properly integrate s3 table buckets into s3
-        # bucket = s3_backends[self.account_id][self.partition].create_bucket(
-        #     bucket_name, region_name="us-east-1"
-        # )
-        bucket = s3_backends["__s3tables__"][self.partition].create_bucket(
-            bucket_name, region_name="us-east-1"
+        bucket = s3_backends[self.account_id][self.partition].create_table_storage_bucket(
+            bucket_name, region_name=self.region_name
         )
         return bucket
 
@@ -404,12 +401,15 @@ def delete_table(
         self, table_bucket_arn: str, namespace: str, name: str, version_token: str
     ) -> None:
         bucket = self.table_buckets.get(table_bucket_arn)
-        if bucket and namespace in bucket.namespaces:
-            ns = bucket.namespaces[namespace]
-            if name in ns.tables:
-                ns.tables.pop(name)
-                return
-        raise TableDoesNotExist()
+        if not bucket or namespace not in bucket.namespaces or (name not in bucket.namespaces[namespace].tables):
+            raise TableDoesNotExist()
+
+        ns = bucket.namespaces[namespace]
+        table = ns.tables[name]
+        from moto.s3.models import s3_backends
+        s3_backends[self.account_id][self.partition].delete_table_storage_bucket(table._bucket.name)
+
+        ns.tables.pop(name)
 
     def update_table_metadata_location(
         self,

tests/test_s3tables/test_s3tables.py

@@ -215,6 +215,27 @@ def test_delete_table():
     assert len(resp["tables"]) == 0
 
 
+@mock_aws
+def test_delete_table_deletes_underlying_table_storage():
+    client = boto3.client("s3tables", region_name="us-east-2")
+    s3 = boto3.client("s3", region_name="us-east-2")
+    arn = client.create_table_bucket(name="foo")["arn"]
+    client.create_namespace(tableBucketARN=arn, namespace=["bar"])
+    client.create_table(
+        tableBucketARN=arn, namespace="bar", name="baz", format="ICEBERG"
+    )
+    warehouse = client.get_table(
+        tableBucketARN=arn, namespace="bar", name="baz"
+    )["warehouseLocation"]
+
+    bucket_name = warehouse.replace("s3://", "")
+    s3.head_bucket(Bucket=bucket_name)
+
+    client.delete_table(tableBucketARN=arn, namespace="bar", name="baz")
+    with pytest.raises(s3.exceptions.ClientError):
+        s3.head_bucket(Bucket=bucket_name)
+
+
 @mock_aws
 def test_get_table_metadata_location():
     client = boto3.client("s3tables", region_name="us-east-2")