Skip to content

Commit

Permalink
fix: Assign an owner when creating a dataset from a csv, excel or tab…
Browse files Browse the repository at this point in the history
…ular (apache#17986)

* Assign an owner when creating a dataset from a csv, excel or columnar

* Added some unit tests

* Code cleanup

* Removed blank line

* Attempt to fix a broken test

* Attempt # 2 to fix a broken test

* Attempt # 3 to fix a broken test

* Attempt # 4 to fix a broken test

* Attempt # 5 to fix a broken test

* Attempt # 6 to fix a broken test

* Broken test fixed, code cleanup
  • Loading branch information
cccs-joel authored and bwang221 committed Feb 10, 2022
1 parent 9eb12c6 commit 992afaa
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 10 deletions.
6 changes: 3 additions & 3 deletions superset/views/database/views.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,7 @@ def form_post(self, form: CsvToDatabaseForm) -> Response:
sqla_table = SqlaTable(table_name=csv_table.table)
sqla_table.database = expore_database
sqla_table.database_id = database.id
sqla_table.user_id = g.user.get_id()
sqla_table.owners = [g.user]
sqla_table.schema = csv_table.schema
sqla_table.fetch_metadata()
db.session.add(sqla_table)
Expand Down Expand Up @@ -369,7 +369,7 @@ def form_post(self, form: ExcelToDatabaseForm) -> Response:
sqla_table = SqlaTable(table_name=excel_table.table)
sqla_table.database = expore_database
sqla_table.database_id = database.id
sqla_table.user_id = g.user.get_id()
sqla_table.owners = [g.user]
sqla_table.schema = excel_table.schema
sqla_table.fetch_metadata()
db.session.add(sqla_table)
Expand Down Expand Up @@ -521,7 +521,7 @@ def form_post( # pylint: disable=too-many-locals
sqla_table = SqlaTable(table_name=columnar_table.table)
sqla_table.database = expore_database
sqla_table.database_id = database.id
sqla_table.user_id = g.user.get_id()
sqla_table.owners = [g.user]
sqla_table.schema = columnar_table.schema
sqla_table.fetch_metadata()
db.session.add(sqla_table)
Expand Down
29 changes: 22 additions & 7 deletions tests/integration_tests/csv_upload_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@

import superset.utils.database
from superset.sql_parse import Table
from superset import security_manager
from tests.integration_tests.conftest import ADMIN_SCHEMA_NAME
from tests.integration_tests.test_app import app # isort:skip
from superset import db
Expand Down Expand Up @@ -142,15 +143,19 @@ def upload_csv(filename: str, table_name: str, extra: Optional[Dict[str, str]] =
def upload_excel(
filename: str, table_name: str, extra: Optional[Dict[str, str]] = None
):
excel_upload_db_id = get_upload_db().id
schema = utils.get_example_default_schema()
form_data = {
"excel_file": open(filename, "rb"),
"name": table_name,
"con": get_upload_db().id,
"con": excel_upload_db_id,
"sheet_name": "Sheet1",
"if_exists": "fail",
"index_label": "test_label",
"mangle_dupe_cols": False,
}
if schema:
form_data["schema"] = schema
if extra:
form_data.update(extra)
return get_resp(test_client, "/exceltodatabaseview/form", data=form_data)
Expand Down Expand Up @@ -346,6 +351,9 @@ def test_import_csv(mock_event_logger):
# make sure the new column name is reflected in the table metadata
assert "d" in table.column_names

# ensure user is assigned as an owner
assert security_manager.find_user("admin") in table.owners

# null values are set
upload_csv(
CSV_FILENAME2,
Expand All @@ -372,22 +380,26 @@ def test_import_excel(mock_event_logger):
if utils.backend() == "hive":
pytest.skip("Hive doesn't excel upload.")

schema = utils.get_example_default_schema()
full_table_name = f"{schema}.{EXCEL_UPLOAD_TABLE}" if schema else EXCEL_UPLOAD_TABLE
test_db = get_upload_db()

success_msg = (
f'Excel file "{EXCEL_FILENAME}" uploaded to table "{EXCEL_UPLOAD_TABLE}"'
)
success_msg = f'Excel file "{EXCEL_FILENAME}" uploaded to table "{full_table_name}"'

# initial upload with fail mode
resp = upload_excel(EXCEL_FILENAME, EXCEL_UPLOAD_TABLE)
assert success_msg in resp
mock_event_logger.assert_called_with(
action="successful_excel_upload",
database=test_db.name,
schema=None,
schema=schema,
table=EXCEL_UPLOAD_TABLE,
)

# ensure user is assigned as an owner
table = SupersetTestCase.get_table(name=EXCEL_UPLOAD_TABLE)
assert security_manager.find_user("admin") in table.owners

# upload again with fail mode; should fail
fail_msg = f'Unable to upload Excel file "{EXCEL_FILENAME}" to table "{EXCEL_UPLOAD_TABLE}"'
resp = upload_excel(EXCEL_FILENAME, EXCEL_UPLOAD_TABLE)
Expand All @@ -408,7 +420,7 @@ def test_import_excel(mock_event_logger):
mock_event_logger.assert_called_with(
action="successful_excel_upload",
database=test_db.name,
schema=None,
schema=schema,
table=EXCEL_UPLOAD_TABLE,
)

Expand Down Expand Up @@ -467,10 +479,13 @@ def test_import_parquet(mock_event_logger):
)
assert success_msg_f1 in resp

# make sure only specified column name was read
table = SupersetTestCase.get_table(name=PARQUET_UPLOAD_TABLE, schema=None)
# make sure only specified column name was read
assert "b" not in table.column_names

# ensure user is assigned as an owner
assert security_manager.find_user("admin") in table.owners

# upload again with replace mode
resp = upload_columnar(
PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE, extra={"if_exists": "replace"}
Expand Down

0 comments on commit 992afaa

Please sign in to comment.