bito-code-review[bot] commented on code in PR #44411:
URL: https://github.com/apache/superset/pull/44411#discussion_r4137759558
##########
tests/unit_tests/db_engine_specs/test_mysql.py:
##########
@@ -510,3 +510,391 @@ def test_extended_aggregation_func_median_unsupported()
-> None:
from superset.db_engine_specs.mysql import MySQLEngineSpec as spec #
noqa: N813
assert spec.get_extended_aggregation_func("MEDIAN") is None
+
+
+def test_df_to_sql_adds_primary_key_when_mysql_requires_one() -> None:
+ """
+ apache/superset#37399: uploading a CSV/Excel/columnar file creates the
+ table via ``pandas.DataFrame.to_sql``, which never declares a primary
+ key. A MySQL server configured with ``sql_require_primary_key = ON``
+ rejects such a ``CREATE TABLE`` with error 3750.
+
+ No live MySQL server is available in this environment, so an in-memory
+ SQLite engine stands in for the target database, with a SQLAlchemy event
+ listener reproducing MySQL's documented enforcement: any executed
+ ``CREATE TABLE`` lacking a primary key raises the same error 3750 seen
+ in the issue.
+ """
+ import pandas as pd
+ import sqlalchemy as sa
+ from sqlalchemy import create_engine, event
+ from sqlalchemy.exc import OperationalError
+
+ from superset.db_engine_specs.mysql import MySQLEngineSpec
+ from superset.sql.parse import Table
+
+ engine = create_engine("sqlite://")
+
+ @event.listens_for(engine, "before_cursor_execute")
+ def _enforce_sql_require_primary_key( # pylint: disable=unused-argument
+ conn: Any,
+ cursor: Any,
+ statement: str,
+ parameters: Any,
+ context: Any,
+ executemany: bool,
+ ) -> None:
+ if (
+ statement.strip().upper().startswith("CREATE TABLE")
+ and "PRIMARY KEY" not in statement.upper()
+ ):
+ raise OperationalError(
+ statement,
+ parameters,
+ Exception(
+ '(3750, "Unable to create or change a table without a '
+ "primary key, when the system variable "
+ "'sql_require_primary_key' is set.\")"
+ ),
+ )
+
+ df = pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})
+
+ with (
+ patch.object(MySQLEngineSpec, "get_engine") as mock_get_engine,
+ patch.object(MySQLEngineSpec, "_requires_primary_key",
return_value=True),
+ ):
+ mock_get_engine.return_value.__enter__.return_value = engine
+ mock_get_engine.return_value.__exit__.return_value = False
+
+ MySQLEngineSpec.df_to_sql(
+ database=Mock(),
+ table=Table(table="my_table"),
+ df=df,
+ to_sql_kwargs={"if_exists": "fail", "index": False},
+ )
+
+ with engine.connect() as conn:
+ rows = conn.execute(sa.text("SELECT a, b FROM my_table ORDER BY
a")).fetchall()
+ assert [tuple(row) for row in rows] == [(1, "x"), (2, "y"), (3, "z")]
+
+ pk = sa.inspect(engine).get_pk_constraint("my_table")
+ assert pk["constrained_columns"], (
+ "expected the created table to have a primary key so MySQL's "
+ "sql_require_primary_key would accept the CREATE TABLE"
+ )
+
+
+def test_df_to_sql_promotes_pandas_index_to_primary_key() -> None:
+ """
+ The literal scenario reported in apache/superset#37399: the "Dataframe
+ index" upload option is enabled, so pandas writes an extra ``index``
+ column -- the reporter's CREATE TABLE showed this column present but
+ never marked PRIMARY KEY. When MySQL requires one, that pandas index
+ column should be promoted to the primary key instead of adding a
+ redundant second column.
+ """
+ import pandas as pd
+ import sqlalchemy as sa
+ from sqlalchemy import create_engine
+
+ from superset.db_engine_specs.mysql import MySQLEngineSpec
+ from superset.sql.parse import Table
+
+ engine = create_engine("sqlite://")
+ df = pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})
+
+ with (
+ patch.object(MySQLEngineSpec, "get_engine") as mock_get_engine,
+ patch.object(MySQLEngineSpec, "_requires_primary_key",
return_value=True),
+ ):
+ mock_get_engine.return_value.__enter__.return_value = engine
+ mock_get_engine.return_value.__exit__.return_value = False
+
+ MySQLEngineSpec.df_to_sql(
+ database=Mock(),
+ table=Table(table="my_table"),
+ df=df,
+ to_sql_kwargs={"if_exists": "fail", "index": True},
+ )
+
+ with engine.connect() as conn:
+ rows = conn.execute(
+ sa.text('SELECT "index", a, b FROM my_table ORDER BY "index"')
+ ).fetchall()
+ assert [tuple(row) for row in rows] == [(0, 1, "x"), (1, 2, "y"), (2, 3,
"z")]
+
+ pk = sa.inspect(engine).get_pk_constraint("my_table")
+ assert pk["constrained_columns"] == ["index"], (
+ "expected the pandas index column to become the primary key, not "
+ "an extra synthesized column"
+ )
+
+
+def test_df_to_sql_disables_autoincrement_on_synthesized_primary_key() -> None:
+ """
+ pandas declares the primary key via a table-level PrimaryKeyConstraint
+ (never Column(primary_key=True)), but SQLAlchemy's MySQL DDL compiler
+ still infers AUTO_INCREMENT for a lone integer primary-key column by
+ default. The synthesized key values here are explicit (the promoted
+ index, or the 1..n range for the synthesized "id" column), not
+ DB-generated -- and pandas' default RangeIndex starts at 0, so inserting
+ 0 into an AUTO_INCREMENT column asks MySQL to generate a value instead
+ of storing 0 literally, colliding with the row whose key is 1.
+
+ No live MySQL server is available in this environment; compile the
+ exact table ``SQLTable.create()`` would hand to MySQL against
+ SQLAlchemy's MySQL dialect directly and assert AUTO_INCREMENT never
+ appears.
+ """
+ import pandas as pd
+ from sqlalchemy import create_engine
+ from sqlalchemy.dialects import mysql
+ from sqlalchemy.schema import CreateTable
+
+ from superset.db_engine_specs.mysql import MySQLEngineSpec
+ from superset.sql.parse import Table
+
+ engine = create_engine("sqlite://")
+ df = pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})
+ captured_tables: list[Any] = []
+
+ def _capture_instead_of_create(self: Any) -> None:
+ captured_tables.append(self.table)
+
+ with (
+ patch.object(MySQLEngineSpec, "get_engine") as mock_get_engine,
+ patch.object(MySQLEngineSpec, "_requires_primary_key",
return_value=True),
+ patch.object(pd.io.sql.SQLTable, "create", _capture_instead_of_create),
+ patch.object(pd.io.sql.SQLTable, "insert"),
+ ):
+ mock_get_engine.return_value.__enter__.return_value = engine
+ mock_get_engine.return_value.__exit__.return_value = False
+
+ MySQLEngineSpec.df_to_sql(
+ database=Mock(),
+ table=Table(table="my_table"),
+ df=df,
+ to_sql_kwargs={"if_exists": "fail", "index": False},
+ )
+
+ assert len(captured_tables) == 1
+ ddl = str(CreateTable(captured_tables[0]).compile(dialect=mysql.dialect()))
+ assert "AUTO_INCREMENT" not in ddl.upper()
+
+
+def test_df_to_sql_falls_back_to_synthesized_key_for_non_unique_index() ->
None:
+ """
+ apache/superset#37399: the "Dataframe index" upload option promotes the
+ DataFrame's index straight to PRIMARY KEY, but a CSV/Excel upload can
+ point that option at a column that repeats values or contains missing
+ ones -- neither of which a primary key can hold. When the index isn't
+ unique (or has NaNs), fall back to the synthesized key used for the
+ ``index=False`` path instead of letting MySQL reject the INSERT with a
+ duplicate-key or NOT-NULL error.
+ """
+ import pandas as pd
+ import sqlalchemy as sa
+ from sqlalchemy import create_engine
+
+ from superset.db_engine_specs.mysql import MySQLEngineSpec
+ from superset.sql.parse import Table
+
+ engine = create_engine("sqlite://")
+ df = pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]}, index=[0, 0, 1])
+
+ with (
+ patch.object(MySQLEngineSpec, "get_engine") as mock_get_engine,
+ patch.object(MySQLEngineSpec, "_requires_primary_key",
return_value=True),
+ ):
+ mock_get_engine.return_value.__enter__.return_value = engine
+ mock_get_engine.return_value.__exit__.return_value = False
+
+ MySQLEngineSpec.df_to_sql(
+ database=Mock(),
+ table=Table(table="my_table"),
+ df=df,
+ to_sql_kwargs={"if_exists": "fail", "index": True},
+ )
+
+ with engine.connect() as conn:
+ rows = conn.execute(
+ sa.text("SELECT id, a, b FROM my_table ORDER BY id")
+ ).fetchall()
+ assert [tuple(row) for row in rows] == [(1, 1, "x"), (2, 2, "y"), (3, 3,
"z")]
+
+ columns = {col["name"] for col in
sa.inspect(engine).get_columns("my_table")}
+ assert "index" not in columns, (
+ "the non-unique pandas index should not also be written as a redundant
column"
+ )
+
+ pk = sa.inspect(engine).get_pk_constraint("my_table")
+ assert pk["constrained_columns"] == ["id"], (
+ "expected the synthesized key to be used since the pandas index is not
unique"
+ )
+
+
+def test_df_to_sql_synthesized_key_avoids_case_insensitive_collision() -> None:
+ """
+ MySQL compares column identifiers case-insensitively, so an existing
+ "ID" column would collide with a lowercase synthesized "id" primary key
+ at the MySQL level even though Python sees them as different strings.
+ """
+ import pandas as pd
+ import sqlalchemy as sa
+ from sqlalchemy import create_engine
+
+ from superset.db_engine_specs.mysql import MySQLEngineSpec
+ from superset.sql.parse import Table
+
+ engine = create_engine("sqlite://")
+ df = pd.DataFrame({"ID": [10, 20, 30], "b": ["x", "y", "z"]})
+
+ with (
+ patch.object(MySQLEngineSpec, "get_engine") as mock_get_engine,
+ patch.object(MySQLEngineSpec, "_requires_primary_key",
return_value=True),
+ ):
+ mock_get_engine.return_value.__enter__.return_value = engine
+ mock_get_engine.return_value.__exit__.return_value = False
+
+ MySQLEngineSpec.df_to_sql(
+ database=Mock(),
+ table=Table(table="my_table"),
+ df=df,
+ to_sql_kwargs={"if_exists": "fail", "index": False},
+ )
+
+ columns = [col["name"] for col in
sa.inspect(engine).get_columns("my_table")]
+ assert "_id" in columns, (
+ "the synthesized key should be renamed to avoid the case-"
+ "insensitive collision with the existing 'ID' column"
+ )
+
+ pk = sa.inspect(engine).get_pk_constraint("my_table")
+ assert pk["constrained_columns"] == ["_id"]
+
+
+def
test_df_to_sql_promoted_index_name_collision_falls_back_to_synthesized_key() ->
(
+ None
+):
+ """
+ The pandas index is unique and NaN-free, so it would normally be
+ promoted straight to PRIMARY KEY. But its name ("ID") collides
+ case-insensitively with an existing "id" column, and pandas'
+ ``SQLTable`` only rejects exact (case-sensitive) name clashes when
+ resetting the index -- so promoting it would produce a CREATE TABLE
+ with two columns MySQL sees as the same identifier (error 1060). Fall
+ back to the synthesized key instead, the same as a non-unique index.
+ """
+ import pandas as pd
+ import sqlalchemy as sa
+ from sqlalchemy import create_engine
+
+ from superset.db_engine_specs.mysql import MySQLEngineSpec
+ from superset.sql.parse import Table
+
+ engine = create_engine("sqlite://")
+ df = pd.DataFrame({"id": [1, 2, 3], "b": ["x", "y", "z"]})
+ df.index.name = "ID"
+
+ with (
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Duplicated test setup and call</b></div>
<div id="fix">
There is duplicated code in the test file. The same pattern of mocking
MySQLEngineSpec and calling df_to_sql appears at lines 800-812, 607-619,
706-718, 563-574, and 753-764. Consider creating a shared helper or fixture to
avoid repetition and improve maintainability.
</div>
</div>
<small><i>Code Review Run #0c82c8</i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a
href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
--
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]
