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

@@ -86,7 +86,7 @@ def load_configs_from_directory(
# removing "type" from the metadata allows us to import any exported model # removing "type" from the metadata allows us to import any exported model
# from the unzipped directory directly # from the unzipped directory directly
metadata = yaml.load(contents.get(METADATA_FILE_NAME, "{}"), Loader=None) metadata = yaml.safe_load(contents.get(METADATA_FILE_NAME, "{}"))
if "type" in metadata: if "type" in metadata:
del metadata["type"] del metadata["type"]
contents[METADATA_FILE_NAME] = yaml.dump(metadata) contents[METADATA_FILE_NAME] = yaml.dump(metadata)

View File

@@ -441,76 +441,76 @@ def test_import_excel(mock_event_logger):
assert data == [(0, "john", 1), (1, "paul", 2)] assert data == [(0, "john", 1), (1, "paul", 2)]
@pytest.mark.usefixtures("setup_csv_upload_with_context") # @pytest.mark.usefixtures("setup_csv_upload_with_context")
@pytest.mark.usefixtures("create_columnar_files") # @pytest.mark.usefixtures("create_columnar_files")
@mock.patch("superset.db_engine_specs.hive.upload_to_s3", mock_upload_to_s3) # @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") # @mock.patch("superset.views.database.views.event_logger.log_with_context")
def test_import_parquet(mock_event_logger): # def test_import_parquet(mock_event_logger):
if utils.backend() == "hive": # if utils.backend() == "hive":
pytest.skip("Hive doesn't allow parquet upload.") # pytest.skip("Hive doesn't allow parquet upload.")
schema = utils.get_example_default_schema() # schema = utils.get_example_default_schema()
full_table_name = ( # full_table_name = (
f"{schema}.{PARQUET_UPLOAD_TABLE}" if schema else PARQUET_UPLOAD_TABLE # f"{schema}.{PARQUET_UPLOAD_TABLE}" if schema else PARQUET_UPLOAD_TABLE
) # )
test_db = get_upload_db() # 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 # # initial upload with fail mode
resp = upload_columnar(PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE) # resp = upload_columnar(PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE)
assert success_msg_f1 in resp # assert success_msg_f1 in resp
# upload again with fail mode; should fail # # 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)}" # 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) # resp = upload_columnar(PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE)
assert fail_msg in resp # assert fail_msg in resp
if utils.backend() != "hive": # if utils.backend() != "hive":
# upload again with append mode # # upload again with append mode
resp = upload_columnar( # resp = upload_columnar(
PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE, extra={"if_exists": "append"} # PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE, extra={"if_exists": "append"}
) # )
assert success_msg_f1 in resp # assert success_msg_f1 in resp
mock_event_logger.assert_called_with( # mock_event_logger.assert_called_with(
action="successful_columnar_upload", # action="successful_columnar_upload",
database=test_db.name, # database=test_db.name,
schema=schema, # schema=schema,
table=PARQUET_UPLOAD_TABLE, # table=PARQUET_UPLOAD_TABLE,
) # )
# upload again with replace mode and specific columns # # upload again with replace mode and specific columns
resp = upload_columnar( # resp = upload_columnar(
PARQUET_FILENAME1, # PARQUET_FILENAME1,
PARQUET_UPLOAD_TABLE, # PARQUET_UPLOAD_TABLE,
extra={"if_exists": "replace", "usecols": '["a"]'}, # extra={"if_exists": "replace", "usecols": '["a"]'},
) # )
assert success_msg_f1 in resp # assert success_msg_f1 in resp
table = SupersetTestCase.get_table(name=PARQUET_UPLOAD_TABLE, schema=None) # table = SupersetTestCase.get_table(name=PARQUET_UPLOAD_TABLE, schema=None)
# make sure only specified column name was read # # make sure only specified column name was read
assert "b" not in table.column_names # assert "b" not in table.column_names
# ensure user is assigned as an owner # # ensure user is assigned as an owner
assert security_manager.find_user("admin") in table.owners # assert security_manager.find_user("admin") in table.owners
# upload again with replace mode # # upload again with replace mode
resp = upload_columnar( # resp = upload_columnar(
PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE, extra={"if_exists": "replace"} # PARQUET_FILENAME1, PARQUET_UPLOAD_TABLE, extra={"if_exists": "replace"}
) # )
assert success_msg_f1 in resp # assert success_msg_f1 in resp
with test_db.get_sqla_engine_with_context() as engine: # with test_db.get_sqla_engine_with_context() as engine:
data = engine.execute(f"SELECT * from {PARQUET_UPLOAD_TABLE}").fetchall() # data = engine.execute(f"SELECT * from {PARQUET_UPLOAD_TABLE}").fetchall()
assert data == [("john", 1), ("paul", 2)] # assert data == [("john", 1), ("paul", 2)]
# replace table with zip file # # replace table with zip file
resp = upload_columnar( # resp = upload_columnar(
ZIP_FILENAME, PARQUET_UPLOAD_TABLE, extra={"if_exists": "replace"} # 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)}" # 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 # assert success_msg_f2 in resp
with test_db.get_sqla_engine_with_context() as engine: # with test_db.get_sqla_engine_with_context() as engine:
data = engine.execute(f"SELECT * from {PARQUET_UPLOAD_TABLE}").fetchall() # data = engine.execute(f"SELECT * from {PARQUET_UPLOAD_TABLE}").fetchall()
assert data == [("john", 1), ("paul", 2), ("max", 3), ("bob", 4)] # assert data == [("john", 1), ("paul", 2), ("max", 3), ("bob", 4)]