bito-code-review[bot] commented on code in PR #43805:
URL: https://github.com/apache/superset/pull/43805#discussion_r4085565481
##########
superset/utils/s3.py:
##########
@@ -15,69 +15,101 @@
# specific language governing permissions and limitations
# under the License.
"""
-Minimal S3 helpers for uploading export artifacts and minting pre-signed URLs.
+The AWS S3 export storage backend.
-Credentials and region come from the standard boto3 resolution chain (env vars,
-shared config, instance role). Operators can override client construction via
-the ``EXCEL_EXPORT_S3_CLIENT_KWARGS`` config (e.g. ``region_name`` or an
-``endpoint_url`` for S3-compatible stores such as MinIO/LocalStack).
+Set ``EXPORT_STORAGE["backend"] = S3ExportStorage()`` in
+``superset_config.py`` when the export bucket is an S3 (or S3-compatible)
+bucket. Credentials and region come from the standard boto3 resolution chain
+(env vars, shared config, instance role); client construction can be overridden
+via the constructor (e.g. ``region_name``, or an ``endpoint_url`` for
+S3-compatible stores such as MinIO/LocalStack):
+
+ EXPORT_STORAGE["backend"] = S3ExportStorage(
+ client_kwargs={"endpoint_url": "http://minio:9000"}
+ )
"""
from __future__ import annotations
import logging
+from collections.abc import Iterator
from typing import Any
-from flask import current_app
+from superset.utils.export_storage import ExportDownload
logger = logging.getLogger(__name__)
+DOWNLOAD_CHUNK_BYTES = 8 * 1024 * 1024
-def _get_s3_client() -> Any:
- """Build an S3 client using operator-provided client kwargs (if any)."""
- # boto3 is imported lazily so that importing this module (which happens at
- # app startup via the dashboard API) does not require boto3 to be
installed.
- # The dependency is only needed when an export actually runs; if it is
- # missing, surface an actionable install hint rather than a bare
ImportError.
- try:
- import boto3 # pylint: disable=import-outside-toplevel
- except ImportError as ex:
- raise ImportError(
- "boto3 is required for dashboard Excel export but is not
installed. "
- "Install it with `pip install apache-superset[excel-export]`."
- ) from ex
-
- client_kwargs: dict[str, Any] = current_app.config.get(
- "EXCEL_EXPORT_S3_CLIENT_KWARGS", {}
- )
- return boto3.client("s3", **client_kwargs)
+class S3ExportStorage:
+ """Store export artifacts in AWS S3 via boto3.
+
+ Implements ``superset.utils.export_storage.ExportStorage``.
-def upload_file_to_s3(local_path: str, bucket: str, key: str) -> None:
+ :param client_kwargs: Extra kwargs passed to ``boto3.client("s3", ...)``
"""
- Upload a local file to S3.
- ``boto3``'s ``upload_file`` automatically uses a managed multipart transfer
- for large files, so no manual chunking is required.
+ def __init__(self, client_kwargs: dict[str, Any] | None = None) -> None:
+ self._client_kwargs = client_kwargs or {}
- :param local_path: Path to the file on local disk
- :param bucket: Destination S3 bucket
- :param key: Destination S3 object key
- """
- _get_s3_client().upload_file(local_path, bucket, key)
+ def _client(self) -> Any:
+ # boto3 is imported lazily so that importing this module (which happens
+ # at config-load time if EXPORT_STORAGE references this class)
+ # does not require boto3 to be installed. The dependency is only needed
+ # when an export actually runs; if it is missing, surface an actionable
+ # install hint rather than a bare ImportError.
+ try:
+ import boto3 # pylint: disable=import-outside-toplevel
+ except ImportError as ex:
+ raise ImportError(
+ "boto3 is required for S3ExportStorage but is not installed. "
+ "Install it with `pip install apache-superset[excel-export]`."
+ ) from ex
+ return boto3.client("s3", **self._client_kwargs)
+ def upload_file(self, local_path: str, bucket: str, key: str) -> None:
+ """
+ Upload a local file to S3.
-def generate_presigned_url(bucket: str, key: str, expires_in: int) -> str:
- """
- Generate a time-limited pre-signed URL for downloading an S3 object.
+ ``boto3``'s ``upload_file`` automatically uses a managed multipart
+ transfer for large files, so no manual chunking is required.
- :param bucket: The S3 bucket
- :param key: The S3 object key
- :param expires_in: URL lifetime in seconds
- :returns: A pre-signed ``get_object`` URL
- """
- return _get_s3_client().generate_presigned_url(
- "get_object",
- Params={"Bucket": bucket, "Key": key},
- ExpiresIn=expires_in,
- )
+ :param local_path: Path to the file on local disk
+ :param bucket: Destination S3 bucket
+ :param key: Destination S3 object key
+ """
+ self._client().upload_file(local_path, bucket, key)
+
+ def download(self, bucket: str, key: str) -> ExportDownload:
+ """
+ An S3 object as ``(size, chunks)``, existence checked eagerly.
+
+ :param bucket: The S3 bucket
+ :param key: The S3 object key
+ :raises FileNotFoundError: when the object does not exist
+ """
+ # Build the client first so a missing boto3 surfaces the install
+ # hint from _client() instead of a bare error on this import.
+ client = self._client()
+ import botocore.exceptions # pylint: disable=import-outside-toplevel
+
+ try:
+ # Metadata only: the body is opened when the stream is consumed, so
+ # a HEAD (which never iterates it) leaves no connection behind.
+ head = client.head_object(Bucket=bucket, Key=key)
+ except botocore.exceptions.ClientError as ex:
+ if ex.response.get("Error", {}).get("Code") in ("NoSuchKey",
"404"):
+ raise FileNotFoundError(f"s3://{bucket}/{key}") from ex
+ raise
+
+ def chunks() -> Iterator[bytes]:
+ body = client.get_object(Bucket=bucket, Key=key)["Body"]
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Uncaught GET error in stream</b></div>
<div id="fix">
`client.get_object(...)` now runs inside the `chunks()` generator, outside
the `try/except botocore.exceptions.ClientError` that still only wraps
`head_object`. If the object is removed between the HEAD and the GET (or the
GET fails transiently), the `NoSuchKey`/`404` no longer becomes
`FileNotFoundError`, so `download_xlsx`'s `except FileNotFoundError`
(api.py:2073) can't turn it into a 410 — it propagates as a 500 mid-stream.
Wrap the GET in the same ClientError→FileNotFoundError handling.
</div>
</div>
<small><i>Code Review Run #d60a4e</i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a
href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
