Skip to content

Commit

Permalink
Added flag for removing the index on the report CSV
Browse files Browse the repository at this point in the history
  • Loading branch information
SkinnyPigeon committed Jun 23, 2024
1 parent fe3ba12 commit 8e4dad8
Show file tree
Hide file tree
Showing 4 changed files with 88 additions and 0 deletions.
4 changes: 4 additions & 0 deletions superset/commands/report/execute.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

import pandas as pd
from celery.exceptions import SoftTimeLimitExceeded
from flask import current_app

from superset import app, db, security_manager
from superset.commands.base import BaseCommand
Expand All @@ -44,6 +45,7 @@
ReportScheduleUnexpectedError,
ReportScheduleWorkingTimeoutError,
)
from superset.commands.report.utils import remove_post_processed
from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
from superset.daos.report import (
REPORT_SCHEDULE_ERROR_NOTIFICATION_MARKER,
Expand Down Expand Up @@ -251,6 +253,8 @@ def _get_pdf(self) -> bytes:

def _get_csv_data(self) -> bytes:
url = self._get_url(result_format=ChartDataResultFormat.CSV)
if not current_app.config["CSV_INDEX"]:
url = remove_post_processed(url)

Check warning on line 257 in superset/commands/report/execute.py

View check run for this annotation

Codecov / codecov/patch

superset/commands/report/execute.py#L257

Added line #L257 was not covered by tests
_, username = get_executor(
executor_types=app.config["ALERT_REPORTS_EXECUTE_AS"],
model=self._report_schedule,
Expand Down
33 changes: 33 additions & 0 deletions superset/commands/report/utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
# 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.


def remove_post_processed(url: str) -> str:
"""Remove the type=post_processed parameter from the URL query string.
Args:
url (str): The URL to process.
Returns:
str: The URL with the type=post_processed parameter removed."""
if "?" not in url:
return url
base_url, query_string = url.split("?", 1)
params = query_string.split("&")
filtered_params = [param for param in params if param != "type=post_processed"]
filtered_query_string = "&".join(filtered_params)
filtered_url = f"{base_url}?{filtered_query_string}"
return filtered_url
3 changes: 3 additions & 0 deletions superset/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -826,6 +826,9 @@ class D3TimeFormat(TypedDict, total=False):
# note: index option should not be overridden
CSV_EXPORT = {"encoding": "utf-8"}

# Include the CSV index when sending reports
CSV_INDEX = True

# Excel Options: key/value pairs that will be passed as argument to DataFrame.to_excel
# method.
# note: index option should not be overridden
Expand Down
48 changes: 48 additions & 0 deletions tests/unit_tests/commands/report/test_report_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
# 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.

from superset.commands.report.utils import remove_post_processed


def test_remove_post_processed():
url = "https://superset.com/?param1=value1&type=post_processed&param2=value2"
expected = "https://superset.com/?param1=value1&param2=value2"
assert remove_post_processed(url) == expected


def test_retain_other_parameters():
url = "https://superset.com/?param1=value1&param2=value2"
expected = "https://superset.com/?param1=value1&param2=value2"
assert remove_post_processed(url) == expected


def test_no_post_processed_present():
url = "https://superset.com/?param1=value1&param2=value2"
expected = "https://superset.com/?param1=value1&param2=value2"
assert remove_post_processed(url) == expected


def test_empty_query_string():
url = "https://superset.com/?"
expected = "https://superset.com/?"
assert remove_post_processed(url) == expected


def test_no_query_string():
url = "https://superset.com"
expected = "https://superset.com"
assert remove_post_processed(url) == expected

0 comments on commit 8e4dad8

Please sign in to comment.