Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: S3 feature_ check #1

Open
wants to merge 26 commits into
base: master
Choose a base branch
from
Open
Changes from all commits
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
91fd1ac
Added S3 as features
deepak-kumar-007 Mar 18, 2024
952cf2f
Updated the all required changes till what we did on 2.1.0
deepak-kumar-007 Mar 18, 2024
30e2040
corrected the required changes
deepak-kumar-007 Mar 18, 2024
bacf6ac
new changes
deepak-kumar-007 Mar 18, 2024
8fa4f61
Added formatter
deepak-kumar-007 Mar 18, 2024
18ebd0f
added license
deepak-kumar-007 Mar 18, 2024
7716627
added pytest for schemas_test
deepak-kumar-007 Mar 18, 2024
2f8edb0
Custom width test case
deepak-kumar-007 Mar 18, 2024
999bb8b
updated schemas validation
deepak-kumar-007 Mar 18, 2024
6332d64
pylint changes
deepak-kumar-007 Mar 18, 2024
6899d43
formatted changes
deepak-kumar-007 Mar 18, 2024
c771db3
added pre-commit changes
deepak-kumar-007 Mar 18, 2024
ee2a769
pylint error resolved
deepak-kumar-007 Mar 18, 2024
50eac36
snake case for s3 file
deepak-kumar-007 Mar 19, 2024
e7dc570
Testing integration
deepak-kumar-007 Mar 19, 2024
b4dfaf3
reverted changes
deepak-kumar-007 Mar 19, 2024
da19ce1
test integration
deepak-kumar-007 Mar 19, 2024
6419493
Testing api_tests
deepak-kumar-007 Mar 19, 2024
f11a454
test 2 api_test
deepak-kumar-007 Mar 19, 2024
540575b
Test3 api_test
deepak-kumar-007 Mar 19, 2024
2b6da35
test4 api_test
deepak-kumar-007 Mar 19, 2024
8f85f37
Test5 api_test
deepak-kumar-007 Mar 19, 2024
929d756
test7 api-test
deepak-kumar-007 Mar 19, 2024
6019ae6
test8 api_test
deepak-kumar-007 Mar 19, 2024
b83cf42
test 9 api_test
deepak-kumar-007 Mar 19, 2024
6afe4b7
test command report
deepak-kumar-007 Mar 19, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docker/docker-frontend.sh
Original file line number Diff line number Diff line change
@@ -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
7 changes: 7 additions & 0 deletions docker/pythonpath_dev/superset_config.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion superset/commands/base.py
Original file line number Diff line number Diff line change
@@ -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]:
"""
25 changes: 20 additions & 5 deletions superset/commands/report/execute.py
Original file line number Diff line number Diff line change
@@ -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:
Original file line number Diff line number Diff line change
@@ -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")
6 changes: 6 additions & 0 deletions superset/reports/api.py
Original file line number Diff line number Diff line change
@@ -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",
7 changes: 5 additions & 2 deletions superset/reports/models.py
Original file line number Diff line number Diff line change
@@ -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:
13 changes: 10 additions & 3 deletions superset/reports/notifications/__init__.py
Original file line number Diff line number Diff line change
@@ -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"
)
13 changes: 12 additions & 1 deletion superset/reports/notifications/base.py
Original file line number Diff line number Diff line change
@@ -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()
145 changes: 145 additions & 0 deletions superset/reports/notifications/s3.py
Original file line number Diff line number Diff line change
@@ -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
Loading