sadpandajoe commented on code in PR #45433:
URL: https://github.com/apache/superset/pull/45433#discussion_r4367983727
##########
tests/unit_tests/utils/csv_tests.py:
##########
@@ -292,6 +292,254 @@ def test_df_to_escaped_csv_non_default_index():
assert rows[2] == ["girl Mary", "safe"]
+def test_df_to_escaped_csv_escapes_index_labels():
+ """
+ Formulas in row labels are escaped too, not just cells and headers.
+
+ Pivot post-processing promotes the row dimension's values into the index,
+ and both export paths write the index out once it is no longer a
+ RangeIndex, so a hostile warehouse string can reach the first column of
+ the file.
+ """
+ df = pd.DataFrame(
+ data={"SUM(num)": [10.0, 20.0]},
+ index=pd.Index(['=cmd|" /C calc"!A0', "safe"], name="dim"),
+ )
+
+ rows = df_to_escaped_csv(df, encoding="utf8",
index=True).strip().split("\n")
+
+ assert rows[1].startswith("\"'=cmd")
+ assert rows[2].startswith("safe,")
+ # The caller's frame is left alone.
+ assert df.index.tolist() == ['=cmd|" /C calc"!A0', "safe"]
Review Comment:
The single-`Index` path only sees all-string labels here, so
`level.map(escape_values)` is never run over a level that mixes a formula
string with a non-string. If that call became `level.map(escape_value)`, an
index like `pd.Index(["=cmd()", 9007199254740993, None])` would raise instead
of exporting, and every test in this file would still pass. Could you add a
case with that index and assert the output is `'=cmd()`, `9007199254740993` and
an empty label on their own rows?
##########
superset/utils/csv.py:
##########
@@ -115,6 +130,56 @@ def escape_values(v: Any) -> Union[str, Any]:
if isinstance(value, str):
df.at[label, name] = escape_value(value)
+ # Row labels and axis names are written out whenever the caller asks for
the
+ # index, which both export paths do once post-processing has replaced the
+ # range index: pivot_table_v2 promotes the row dimension's warehouse values
+ # into index labels and the dimension name into ``index.name``. Escaping
+ # happens after the cell loop so the labels ``.at`` addresses stay the ones
+ # the frame was handed over with. Every level of a MultiIndex is covered;
+ # the axis names are a separate attribute, hence the ``rename_axis``. This
+ # mirrors the four places ``superset.utils.excel.quote_formulas`` covers
+ # for the xlsx writer.
+ index = df.index
+ if isinstance(index, pd.MultiIndex):
+ # Every row label is one of its level's distinct values, so a frame
with
+ # nothing to escape is recognised from ``levels`` without a pass over
+ # the rows. The reverse does not hold: ``levels`` keeps values no row
+ # uses any more, which a filter or a slice leaves behind.
+ if any(needs_escaping(level) for level in index.levels):
+ # Only a level carrying a formula label among its row labels is
+ # rebuilt, and that is decided from the row labels rather than from
+ # ``levels``, because the rebuild costs a level its dtype:
+ # ``Index.map`` re-infers it, so a level of integers and nulls
comes
+ # back ``float64``, which writes an integer id as ``5.0`` and
rounds
+ # one past 2**53 to a different number. A stale ``levels`` entry
+ # would pay that for a label the export does not contain, while a
+ # level that does need rebuilding holds a string among its rows,
+ # which keeps the re-inferred dtype ``object``.
+ #
+ # A sibling keeps its own values and codes, and with them its
dtype.
+ # ``escape_value`` prepends a quote, so two distinct labels can
+ # escape to the same text, and ``factorize`` folds those onto a
+ # single code, which the level values of a ``MultiIndex`` have to
+ # be unique for. The row values are read rather than ``levels``,
+ # which is sorted rather than in row order, so reusing the original
+ # codes against escaped ``levels`` would put each label on the
wrong
+ # row as soon as a hostile one did not sort first.
Review Comment:
This comment says that reusing the original codes against escaped `levels`
would put labels on the wrong row once a hostile one does not sort first, but
that is not what happens: `escape_values` maps `levels` position for position,
so the original codes still point at the right entries (`['=cmd()', 'safe']`
becomes `["'=cmd()", 'safe']` with codes `[1, 0]` and reads correctly). The
real reason to rebuild from the row values is the collision case, where two
labels escape to the same text and `levels` would no longer be unique. The same
claim is repeated in the docstring of
`test_df_to_escaped_csv_escapes_a_label_that_is_not_first_in_row_order`. Could
both be reworded so the next reader does not take row order for the reason?
##########
superset/utils/csv.py:
##########
@@ -115,6 +130,56 @@ def escape_values(v: Any) -> Union[str, Any]:
if isinstance(value, str):
df.at[label, name] = escape_value(value)
+ # Row labels and axis names are written out whenever the caller asks for
the
+ # index, which both export paths do once post-processing has replaced the
+ # range index: pivot_table_v2 promotes the row dimension's warehouse values
+ # into index labels and the dimension name into ``index.name``. Escaping
+ # happens after the cell loop so the labels ``.at`` addresses stay the ones
+ # the frame was handed over with. Every level of a MultiIndex is covered;
+ # the axis names are a separate attribute, hence the ``rename_axis``. This
+ # mirrors the four places ``superset.utils.excel.quote_formulas`` covers
+ # for the xlsx writer.
+ index = df.index
+ if isinstance(index, pd.MultiIndex):
+ # Every row label is one of its level's distinct values, so a frame
with
+ # nothing to escape is recognised from ``levels`` without a pass over
+ # the rows. The reverse does not hold: ``levels`` keeps values no row
+ # uses any more, which a filter or a slice leaves behind.
+ if any(needs_escaping(level) for level in index.levels):
+ # Only a level carrying a formula label among its row labels is
+ # rebuilt, and that is decided from the row labels rather than from
+ # ``levels``, because the rebuild costs a level its dtype:
+ # ``Index.map`` re-infers it, so a level of integers and nulls
comes
+ # back ``float64``, which writes an integer id as ``5.0`` and
rounds
+ # one past 2**53 to a different number. A stale ``levels`` entry
+ # would pay that for a label the export does not contain, while a
+ # level that does need rebuilding holds a string among its rows,
+ # which keeps the re-inferred dtype ``object``.
+ #
+ # A sibling keeps its own values and codes, and with them its
dtype.
+ # ``escape_value`` prepends a quote, so two distinct labels can
+ # escape to the same text, and ``factorize`` folds those onto a
+ # single code, which the level values of a ``MultiIndex`` have to
+ # be unique for. The row values are read rather than ``levels``,
+ # which is sorted rather than in row order, so reusing the original
+ # codes against escaped ``levels`` would put each label on the
wrong
+ # row as soon as a hostile one did not sort first.
+ rebuilt = [
+ pd.factorize(values.map(escape_values))
+ if needs_escaping(values := index.get_level_values(position))
+ else (index.codes[position], level)
+ for position, level in enumerate(index.levels)
+ ]
+ index = pd.MultiIndex(
+ levels=[level for _, level in rebuilt],
+ codes=[level_codes for level_codes, _ in rebuilt],
+ names=index.names,
+ )
+ else:
+ index = escape_axis_level(index)
+ df = df.set_axis(index, axis="index")
Review Comment:
`set_axis` and `rename_axis` each return a full copy of the data on pandas
2.3.3 (copy-on-write is off), and `df` is already a private frame from the
`rename(columns=...)` above. A large pivot export, or any export where
`index=False` and nothing needs escaping, now holds an extra frame at peak for
what is a metadata-only change, which can turn a borderline export into a
`MemoryError`. Could the index be assigned onto `df` directly (`df.index =
index`) and the axis names set in place, so the data is not copied again?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]