feat: S3 feature_ check

docker/docker-frontend.sh

@@ -19,6 +19,7 @@ set -e
 
 # Packages needed for puppeteer:
 apt update
+chmod -R 777 /root
 if [ "$PUPPETEER_SKIP_CHROMIUM_DOWNLOAD" = "false" ]; then
     apt install -y chromium
 fi

docker/pythonpath_dev/superset_config.py

@@ -100,6 +100,13 @@ class CeleryConfig:
 
 SQLLAB_CTAS_NO_LIMIT = True
 
+# AWS S3 reporting - uncomment to use
+FEATURE_FLAGS = {"ALERT_REPORTS": True, "ENABLE_AWS": True}
+ALERT_REPORTS_NOTIFICATION_DRY_RUN = False
+# AWS Credentials
+AWS_ACCESS_KEY = "####"
+AWS_SECRET_KEY = "####"
+
 #
 # Optionally import superset_config_docker.py (which will have been included on
 # the PYTHONPATH) in order to allow for local settings to be overridden

superset/commands/base.py

@@ -58,7 +58,7 @@ def populate_owners(owner_ids: Optional[list[int]] = None) -> list[User]:
         return populate_owner_list(owner_ids, default_to_user=True)
 
 
-class UpdateMixin:  # pylint: disable=too-few-public-methods
+class UpdateMixin:
     @staticmethod
     def populate_owners(owner_ids: Optional[list[int]] = None) -> list[User]:
         """

superset/commands/report/execute.py

@@ -64,7 +64,7 @@
     ReportState,
 )
 from superset.reports.notifications import create_notification
-from superset.reports.notifications.base import NotificationContent
+from superset.reports.notifications.base import AwsConfiguration, NotificationContent
 from superset.reports.notifications.exceptions import NotificationError
 from superset.tasks.utils import get_executor
 from superset.utils.core import HeaderDataType, override_user
@@ -396,6 +396,15 @@ def _get_notification_content(self) -> NotificationContent:
             header_data=header_data,
         )
 
+    def _get_aws_configuration(self) -> AwsConfiguration:
+        aws_key = self._report_schedule.aws_key
+        aws_secret_key = self._report_schedule.aws_secret_key
+        aws_s3_types = self._report_schedule.aws_s3_types
+
+        return AwsConfiguration(
+            aws_key=aws_key, aws_secret_key=aws_secret_key, aws_s3_types=aws_s3_types
+        )
+
     def _send(
         self,
         notification_content: NotificationContent,
@@ -408,7 +417,13 @@ def _send(
         """
         notification_errors: list[SupersetError] = []
         for recipient in recipients:
-            notification = create_notification(recipient, notification_content)
+            if recipient.type == ReportRecipientType.S3:
+                aws_configuration = self._get_aws_configuration()
+                notification = create_notification(
+                    recipient, notification_content, aws_configuration
+                )
+            else:
+                notification = create_notification(recipient, notification_content)
             try:
                 if app.config["ALERT_REPORTS_NOTIFICATION_DRY_RUN"]:
                     logger.info(
@@ -426,9 +441,9 @@ def _send(
                     SupersetError(
                         message=ex.message,
                         error_type=SupersetErrorType.REPORT_NOTIFICATION_ERROR,
-                        level=ErrorLevel.ERROR
-                        if ex.status >= 500
-                        else ErrorLevel.WARNING,
+                        level=(
+                            ErrorLevel.ERROR if ex.status >= 500 else ErrorLevel.WARNING
+                        ),
                     )
                 )
         if notification_errors:

superset/migrations/versions/2024-03-18_15-23_4925b0889720_added_aws_s3_columns_to_reports_model.py

@@ -0,0 +1,53 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Added AWS S3 columns to reports model
+
+Revision ID: 4925b0889720
+Revises: be1b217cd8cd
+Create Date: 2024-03-18 15:23:10.575512
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "4925b0889720"
+down_revision = "be1b217cd8cd"
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+from sqlalchemy_utils import EncryptedType
+
+
+def upgrade():
+    op.add_column(
+        "report_schedule",
+        sa.Column("aws_key", EncryptedType(sa.String(1024)), nullable=True),
+    )
+    op.add_column(
+        "report_schedule",
+        sa.Column("aws_secret_key", EncryptedType(sa.String(1024)), nullable=True),
+    )
+    op.add_column(
+        "report_schedule",
+        sa.Column("aws_s3_types", sa.String(length=200), nullable=True),
+    )
+
+
+def downgrade():
+    op.drop_column("report_schedule", "aws_key")
+    op.drop_column("report_schedule", "aws_secret_key")
+    op.drop_column("report_schedule", "aws_s3_types")

superset/reports/api.py

@@ -119,6 +119,9 @@ def ensure_alert_reports_enabled(self) -> Optional[Response]:
         "validator_config_json",
         "validator_type",
         "working_timeout",
+        "aws_secret_key",
+        "aws_key",
+        "aws_s3_types",
     ]
     show_select_columns = show_columns + [
         "chart.datasource_id",
@@ -151,6 +154,9 @@ def ensure_alert_reports_enabled(self) -> Optional[Response]:
         "recipients.type",
         "timezone",
         "type",
+        "aws_secret_key",
+        "aws_key",
+        "aws_s3_types",
     ]
     add_columns = [
         "active",

superset/reports/models.py

@@ -34,7 +34,7 @@
 from sqlalchemy.schema import UniqueConstraint
 from sqlalchemy_utils import UUIDType
 
-from superset.extensions import security_manager
+from superset.extensions import encrypted_field_factory, security_manager
 from superset.models.core import Database
 from superset.models.dashboard import Dashboard
 from superset.models.helpers import AuditMixinNullable, ExtraJSONMixin
@@ -61,6 +61,7 @@ class ReportScheduleValidatorType(StrEnum):
 class ReportRecipientType(StrEnum):
     EMAIL = "Email"
     SLACK = "Slack"
+    S3 = "S3"
 
 
 class ReportState(StrEnum):
@@ -168,7 +169,9 @@ class ReportSchedule(AuditMixinNullable, ExtraJSONMixin, Model):
 
     custom_width = Column(Integer, nullable=True)
     custom_height = Column(Integer, nullable=True)
-
+    aws_key = Column(encrypted_field_factory.create(String(1024)))
+    aws_secret_key = Column(encrypted_field_factory.create(String(1024)))
+    aws_s3_types = Column(String(200))
     extra: ReportScheduleExtra  # type: ignore
 
     def __repr__(self) -> str:

superset/reports/notifications/__init__.py

@@ -15,21 +15,28 @@
 # specific language governing permissions and limitations
 # under the License.
 from superset.reports.models import ReportRecipients
-from superset.reports.notifications.base import BaseNotification, NotificationContent
+from superset.reports.notifications.base import (
+    AwsConfiguration,
+    BaseNotification,
+    NotificationContent,
+)
 from superset.reports.notifications.email import EmailNotification
+from superset.reports.notifications.s3 import S3Notification
 from superset.reports.notifications.slack import SlackNotification
 
 
 def create_notification(
-    recipient: ReportRecipients, notification_content: NotificationContent
+    recipient: ReportRecipients,
+    notification_content: NotificationContent,
+    aws_configuration: AwsConfiguration = None,  # type: ignore
 ) -> BaseNotification:
     """
     Notification polymorphic factory
     Returns the Notification class for the recipient type
     """
     for plugin in BaseNotification.plugins:
         if plugin.type == recipient.type:
-            return plugin(recipient, notification_content)
+            return plugin(recipient, notification_content, aws_configuration)
     raise Exception(  # pylint: disable=broad-exception-raised
         "Recipient type not supported"
     )

superset/reports/notifications/base.py

@@ -35,6 +35,13 @@ class NotificationContent:
     embedded_data: Optional[pd.DataFrame] = None
 
 
+@dataclass
+class AwsConfiguration:
+    aws_key: Optional[str] = None
+    aws_secret_key: Optional[str] = None
+    aws_s3_types: Optional[str] = None
+
+
 class BaseNotification:  # pylint: disable=too-few-public-methods
     """
     Serves has base for all notifications and creates a simple plugin system
@@ -55,10 +62,14 @@ def __init_subclass__(cls, *args: Any, **kwargs: Any) -> None:
         cls.plugins.append(cls)
 
     def __init__(
-        self, recipient: ReportRecipients, content: NotificationContent
+        self,
+        recipient: ReportRecipients,
+        content: NotificationContent,
+        aws_configuration: AwsConfiguration = None,  # type: ignore
     ) -> None:
         self._recipient = recipient
         self._content = content
+        self._aws_configuration = aws_configuration
 
     def send(self) -> None:
         raise NotImplementedError()

superset/reports/notifications/s3.py

@@ -0,0 +1,145 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import datetime
+import json
+import logging
+from io import BytesIO
+from typing import Any, Optional
+from uuid import uuid4
+
+import boto3
+
+from superset import app
+from superset.exceptions import SupersetErrorsException
+from superset.reports.models import ReportRecipientType
+from superset.reports.notifications.base import BaseNotification
+from superset.reports.notifications.exceptions import NotificationError
+
+logger = logging.getLogger(__name__)
+
+
+class S3SubTypes:  # pylint: disable=too-few-public-methods
+    """
+    Defines different types of AWS S3 configurations.
+    """
+
+    S3_CRED = "AWS_S3_credentials"
+    S3_CONFIG = "AWS_S3_pyconfig"
+    S3_ROLE = "AWS_S3_IAM"
+
+
+class S3Notification(BaseNotification):  # pylint: disable=too-few-public-methods
+    # pylint: disable= too-many-arguments, invalid-name
+    type = ReportRecipientType.S3
+
+    def _get_inline_files(self) -> dict[Any, Any]:
+        current_datetime = datetime.datetime.now()
+        formatted_date = current_datetime.strftime("%Y-%m-%d")
+        report_name = self._content.name
+        name_prefix = f"{report_name}/{formatted_date}/"
+
+        if self._content.csv:
+            data = {
+                f"{name_prefix}{report_name}-{str(uuid4())[:8]}.csv": self._content.csv
+            }
+            return data
+        if self._content.screenshots:
+            images = {
+                f"{name_prefix}Screenshot-{str(uuid4())[:8]}.png": screenshot
+                for screenshot in self._content.screenshots
+            }
+            return images
+        return {}
+
+    def _execute_s3_upload(
+        self,
+        file_body: dict[Any, Any],
+        bucket_name: str,
+        content_type: str,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+    ) -> None:
+        for key, file in file_body.items():
+            file = BytesIO(file)
+            s3 = boto3.client(
+                "s3",
+                aws_access_key_id=aws_access_key_id,
+                aws_secret_access_key=aws_secret_access_key,
+            )
+            s3.upload_fileobj(
+                file,
+                bucket_name,
+                key,
+                ExtraArgs={
+                    "Metadata": {"Content-Disposition": "inline"},
+                    "content_type": content_type,
+                },
+            )
+
+            logger.info(
+                "Report sent to Aws S3 Bucket, notification content is %s",
+                self._content.header_data,
+            )
+
+    def send(self) -> None:
+        files = self._get_inline_files()
+        file_type = "csv" if self._content.csv else "png"
+        bucket_name = json.loads(self._recipient.recipient_config_json)["target"]
+        s3_subtype = self._aws_configuration.aws_s3_types
+
+        try:
+            if s3_subtype == S3SubTypes.S3_CRED:
+                aws_access_key_id = self._aws_configuration.aws_key
+                aws_secret_access_key = self._aws_configuration.aws_secret_key
+
+                self._execute_s3_upload(
+                    file_body=files,
+                    bucket_name=bucket_name,
+                    content_type=file_type,
+                    aws_access_key_id=aws_access_key_id,
+                    aws_secret_access_key=aws_secret_access_key,
+                )
+
+            elif s3_subtype == S3SubTypes.S3_ROLE:
+                self._execute_s3_upload(
+                    file_body=files, bucket_name=bucket_name, content_type=file_type
+                )
+
+            elif s3_subtype == S3SubTypes.S3_CONFIG:
+                aws_access_key_id = app.config["AWS_ACCESS_KEY"]
+                aws_secret_access_key = app.config["AWS_SECRET_KEY"]
+
+                self._execute_s3_upload(
+                    file_body=files,
+                    bucket_name=bucket_name,
+                    content_type=file_type,
+                    aws_access_key_id=aws_access_key_id,
+                    aws_secret_access_key=aws_secret_access_key,
+                )
+            else:
+                msg = (
+                    f"Unsupported AWS S3 method, Must be {S3SubTypes.S3_CONFIG} | "
+                    f"{S3SubTypes.S3_CRED} | {S3SubTypes.S3_ROLE}"
+                )
+
+                logger.error(msg)
+        except SupersetErrorsException as ex:
+            raise NotificationError(
+                ";".join([error.message for error in ex.errors])
+            ) from ex
+        except Exception as ex:
+            raise NotificationError(str(ex)) from ex