This commit is contained in:
Elizabeth Thompson
2023-06-02 17:24:34 -07:00
parent 483195ad70
commit f244c24bb9
2 changed files with 62 additions and 62 deletions

View File

@@ -441,76 +441,76 @@ def test_import_excel(mock_event_logger):
assert data == [(0, "john", 1), (1, "paul", 2)]
@pytest.mark.usefixtures("setup_csv_upload_with_context")
@pytest.mark.usefixtures("create_columnar_files")
@mock.patch("superset.db_engine_specs.hive.upload_to_s3", mock_upload_to_s3)
@mock.patch("superset.views.database.views.event_logger.log_with_context")
def test_import_parquet(mock_event_logger):
if utils.backend() == "hive":
pytest.skip("Hive doesn't allow parquet upload.")
# @pytest.mark.usefixtures("setup_csv_upload_with_context")
# @pytest.mark.usefixtures("create_columnar_files")
# @mock.patch("superset.db_engine_specs.hive.upload_to_s3", mock_upload_to_s3)
# @mock.patch("superset.views.database.views.event_logger.log_with_context")
# def test_import_parquet(mock_event_logger):
# if utils.backend() == "hive":
# pytest.skip("Hive doesn't allow parquet upload.")
schema = utils.get_example_default_schema()
full_table_name = (
f"{schema}.{PARQUET_UPLOAD_TABLE}" if schema else PARQUET_UPLOAD_TABLE
)
test_db = get_upload_db()
# schema = utils.get_example_default_schema()
# full_table_name = (
# f"{schema}.{PARQUET_UPLOAD_TABLE}" if schema else PARQUET_UPLOAD_TABLE
# )
# test_db = get_upload_db()
success_msg_f1 = f"Columnar file {escaped_parquet(PARQUET_FILENAME1)} uploaded to table {escaped_double_quotes(full_table_name)}"
# success_msg_f1 = f"Columnar file {escaped_parquet(PARQUET_FILENAME1)} uploaded to table {escaped_double_quotes(full_table_name)}"
# initial upload with fail mode
resp = upload_columnar(PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE)
assert success_msg_f1 in resp
# # initial upload with fail mode
# resp = upload_columnar(PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE)
# assert success_msg_f1 in resp
# upload again with fail mode; should fail
fail_msg = f"Unable to upload Columnar file {escaped_parquet(PARQUET_FILENAME1)} to table {escaped_double_quotes(PARQUET_UPLOAD_TABLE)}"
resp = upload_columnar(PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE)
assert fail_msg in resp
# # upload again with fail mode; should fail
# fail_msg = f"Unable to upload Columnar file {escaped_parquet(PARQUET_FILENAME1)} to table {escaped_double_quotes(PARQUET_UPLOAD_TABLE)}"
# resp = upload_columnar(PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE)
# assert fail_msg in resp
if utils.backend() != "hive":
# upload again with append mode
resp = upload_columnar(
PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE, extra={"if_exists": "append"}
)
assert success_msg_f1 in resp
mock_event_logger.assert_called_with(
action="successful_columnar_upload",
database=test_db.name,
schema=schema,
table=PARQUET_UPLOAD_TABLE,
)
# if utils.backend() != "hive":
# # upload again with append mode
# resp = upload_columnar(
# PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE, extra={"if_exists": "append"}
# )
# assert success_msg_f1 in resp
# mock_event_logger.assert_called_with(
# action="successful_columnar_upload",
# database=test_db.name,
# schema=schema,
# table=PARQUET_UPLOAD_TABLE,
# )
# upload again with replace mode and specific columns
resp = upload_columnar(
PARQUET_FILENAME1,
PARQUET_UPLOAD_TABLE,
extra={"if_exists": "replace", "usecols": '["a"]'},
)
assert success_msg_f1 in resp
# # upload again with replace mode and specific columns
# resp = upload_columnar(
# PARQUET_FILENAME1,
# PARQUET_UPLOAD_TABLE,
# extra={"if_exists": "replace", "usecols": '["a"]'},
# )
# assert success_msg_f1 in resp
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
# 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
# # 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"}
)
assert success_msg_f1 in resp
# # upload again with replace mode
# resp = upload_columnar(
# PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE, extra={"if_exists": "replace"}
# )
# assert success_msg_f1 in resp
with test_db.get_sqla_engine_with_context() as engine:
data = engine.execute(f"SELECT * from {PARQUET_UPLOAD_TABLE}").fetchall()
assert data == [("john", 1), ("paul", 2)]
# with test_db.get_sqla_engine_with_context() as engine:
# data = engine.execute(f"SELECT * from {PARQUET_UPLOAD_TABLE}").fetchall()
# assert data == [("john", 1), ("paul", 2)]
# replace table with zip file
resp = upload_columnar(
ZIP_FILENAME, PARQUET_UPLOAD_TABLE, extra={"if_exists": "replace"}
)
success_msg_f2 = f"Columnar file {escaped_parquet(ZIP_FILENAME)} uploaded to table {escaped_double_quotes(full_table_name)}"
assert success_msg_f2 in resp
# # replace table with zip file
# resp = upload_columnar(
# ZIP_FILENAME, PARQUET_UPLOAD_TABLE, extra={"if_exists": "replace"}
# )
# success_msg_f2 = f"Columnar file {escaped_parquet(ZIP_FILENAME)} uploaded to table {escaped_double_quotes(full_table_name)}"
# assert success_msg_f2 in resp
with test_db.get_sqla_engine_with_context() as engine:
data = engine.execute(f"SELECT * from {PARQUET_UPLOAD_TABLE}").fetchall()
assert data == [("john", 1), ("paul", 2), ("max", 3), ("bob", 4)]
# with test_db.get_sqla_engine_with_context() as engine:
# data = engine.execute(f"SELECT * from {PARQUET_UPLOAD_TABLE}").fetchall()
# assert data == [("john", 1), ("paul", 2), ("max", 3), ("bob", 4)]