bito-code-review[bot] commented on code in PR #37547: URL: https://github.com/apache/superset/pull/37547#discussion_r2740839027
########## tests/integration_tests/databases/commands/upload_mysql_pk_test.py: ########## @@ -0,0 +1,236 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Tests for CSV upload with MySQL databases requiring primary keys. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, Mock, patch + +import pandas as pd +import pytest + +from superset.commands.database.uploaders.base import UploadCommand +from superset.commands.database.uploaders.csv_reader import CSVReader +from superset.db_engine_specs.mysql import MySQLEngineSpec +from superset.models.core import Database +from superset.sql.parse import Table + + +def test_mysql_requires_primary_key_detection(): + """Test that _requires_primary_key correctly detects MySQL setting""" + # Create a mock engine + mock_engine = Mock() + mock_conn = Mock() + mock_result = Mock() + + # Test case: sql_require_primary_key is ON (1) + mock_result.scalar.return_value = 1 + mock_conn.execute.return_value = mock_result + mock_conn.__enter__ = Mock(return_value=mock_conn) + mock_conn.__exit__ = Mock(return_value=False) + mock_engine.connect.return_value = mock_conn + + assert MySQLEngineSpec._requires_primary_key(mock_engine) is True + + # Test case: sql_require_primary_key is OFF (0) + mock_result.scalar.return_value = 0 + assert MySQLEngineSpec._requires_primary_key(mock_engine) is False + + # Test case: Query fails (e.g., older MySQL version without the setting) + mock_conn.execute.side_effect = Exception("Unknown system variable") + assert MySQLEngineSpec._requires_primary_key(mock_engine) is False + + +def test_mysql_df_to_sql_adds_primary_key(): + """Test that df_to_sql adds a primary key when required""" + # Create test data + df = pd.DataFrame({ + "name": ["Alice", "Bob", "Charlie"], + "age": [30, 25, 35], + "city": ["NYC", "LA", "SF"] + }) + + # Create mock database and table + mock_database = Mock(spec=Database) + mock_table = Table(table="test_table", schema=None) + + # Mock engine with sql_require_primary_key enabled + mock_engine = Mock() + mock_conn = Mock() + mock_result = Mock() + mock_result.scalar.return_value = 1 # Primary key required + mock_conn.execute.return_value = mock_result + mock_conn.commit = Mock() + mock_conn.__enter__ = Mock(return_value=mock_conn) + mock_conn.__exit__ = Mock(return_value=False) + mock_engine.connect.return_value = mock_conn + + # Mock engine.begin() for transaction context + mock_transaction = Mock() + mock_transaction.__enter__ = Mock(return_value=mock_conn) + mock_transaction.__exit__ = Mock(return_value=False) + mock_engine.begin.return_value = mock_transaction + + mock_engine.dialect.supports_multivalues_insert = True + + # Track to_sql calls + to_sql_called = False + captured_df = None + captured_kwargs = None + + def capture_to_sql(con, **kwargs): + nonlocal to_sql_called, captured_df, captured_kwargs + to_sql_called = True + captured_df = kwargs.get("name") # Just capture what we can + captured_kwargs = kwargs + + # Mock df.to_sql to capture the call + with patch.object(pd.DataFrame, "to_sql", side_effect=capture_to_sql): + with patch.object( + MySQLEngineSpec, + "get_engine", + return_value=mock_engine.__enter__() + ): + to_sql_kwargs = { + "if_exists": "fail", + "index": False, + "chunksize": 1000, + } + + MySQLEngineSpec.df_to_sql( + mock_database, + mock_table, + df, + to_sql_kwargs + ) + + # Verify to_sql was called + assert to_sql_called + assert "test_table" in str(captured_kwargs) + + # Verify ALTER TABLE was called to add PRIMARY KEY + execute_calls = [call for call in mock_conn.execute.call_args_list] + # First call is checking sql_require_primary_key + # Second call should be ALTER TABLE + assert len(execute_calls) >= 2 + alter_call = execute_calls[1][0][0] + assert "ALTER TABLE" in alter_call + assert "AUTO_INCREMENT" in alter_call + assert "PRIMARY KEY" in alter_call + + +def test_mysql_df_to_sql_skips_primary_key_when_not_required(): + """Test that df_to_sql doesn't add primary key when not required""" + df = pd.DataFrame({ + "name": ["Alice", "Bob"], + "age": [30, 25] + }) + + mock_database = Mock(spec=Database) + mock_table = Table(table="test_table", schema=None) + + # Mock engine with sql_require_primary_key disabled + mock_engine = Mock() + mock_conn = Mock() + mock_result = Mock() + mock_result.scalar.return_value = 0 # Primary key NOT required + mock_conn.execute.return_value = mock_result + mock_conn.__enter__ = Mock(return_value=mock_conn) + mock_conn.__exit__ = Mock(return_value=False) + mock_engine.connect.return_value = mock_conn + mock_engine.dialect.supports_multivalues_insert = True + + # Use parent's df_to_sql + with patch.object( + MySQLEngineSpec, + "get_engine", + return_value=mock_engine.__enter__() + ): + with patch("superset.db_engine_specs.base.BaseEngineSpec.df_to_sql") as mock_parent: + to_sql_kwargs = { + "if_exists": "fail", + "index": False, + } + + MySQLEngineSpec.df_to_sql( + mock_database, + mock_table, + df, + to_sql_kwargs + ) + + # Verify parent's df_to_sql was called (normal path) + mock_parent.assert_called_once() + # Verify no ALTER TABLE was executed + # Only one call should be to check sql_require_primary_key + assert mock_conn.execute.call_count == 1 + + +def test_mysql_df_to_sql_handles_column_name_conflicts(): + """Test that primary key column name is adjusted if it conflicts""" + df = pd.DataFrame({ + "__superset_upload_id__": ["existing_data"], + "name": ["Alice"] + }) + + mock_database = Mock(spec=Database) + mock_table = Table(table="test_table", schema=None) + + mock_engine = Mock() + mock_conn = Mock() + mock_result = Mock() + mock_result.scalar.return_value = 1 # Primary key required + mock_conn.execute.return_value = mock_result + mock_conn.commit = Mock() + mock_conn.__enter__ = Mock(return_value=mock_conn) + mock_conn.__exit__ = Mock(return_value=False) + mock_engine.connect.return_value = mock_conn + mock_engine.dialect.supports_multivalues_insert = True + + captured_df_columns = None + + def capture_to_sql(con, **kwargs): + # Can't directly access df from kwargs in mock, but we can check ALTER + pass + + with patch.object(pd.DataFrame, "to_sql", side_effect=capture_to_sql): + with patch.object( + MySQLEngineSpec, + "get_engine", + return_value=mock_engine.__enter__() + ): + to_sql_kwargs = { + "if_exists": "fail", + "index": False, + } + + MySQLEngineSpec.df_to_sql( + mock_database, + mock_table, + df, + to_sql_kwargs + ) + + # Verify ALTER TABLE uses a different column name (prefixed with _) + execute_calls = [call for call in mock_conn.execute.call_args_list] + if len(execute_calls) >= 2: + alter_call = execute_calls[1][0][0] + # The column name should be modified to avoid conflict + assert "__superset_upload_id__" in alter_call or "___superset_upload_id__" in alter_call Review Comment: <div> <div id="suggestion"> <div id="issue"><b>Test assertion too permissive</b></div> <div id="fix"> The assertion allows either the original or prefixed column name, but the code always prefixes when there's a conflict, so the test should verify the exact adjusted name. </div> <details> <summary> <b>Code suggestion</b> </summary> <blockquote>Check the AI-generated fix before applying</blockquote> <div id="code"> ``` - assert "___superset_upload_id__" in alter_call + assert "___superset_upload_id__" in alter_call and "__superset_upload_id__" not in alter_call ``` </div> </details> </div> <small><i>Code Review Run #528e4b</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 ########## superset/db_engine_specs/mysql.py: ########## @@ -401,3 +407,122 @@ def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> bool: return False return True + + @classmethod + def _requires_primary_key(cls, engine: Any) -> bool: + """ + Check if the MySQL database requires a primary key for table creation. + + :param engine: SQLAlchemy engine + :return: True if primary key is required + """ + try: + with engine.connect() as conn: + result = conn.execute( + "SELECT @@session.sql_require_primary_key" + ).scalar() + return bool(result) + except Exception: # pylint: disable=broad-except + # If we can't determine the setting, assume it's not required + # to maintain backward compatibility + return False + + @classmethod + def df_to_sql( + cls, + database: Database, + table: Table, + df: pd.DataFrame, + to_sql_kwargs: dict[str, Any], + ) -> None: + """ + Upload data from a Pandas DataFrame to a MySQL database. + + Automatically adds a primary key column when the database requires it + (e.g., sql_require_primary_key = ON) and the table is being created. + + :param database: The database to upload the data to + :param table: The table to upload the data to + :param df: The dataframe with data to be uploaded + :param to_sql_kwargs: The kwargs to be passed to pandas.DataFrame.to_sql + """ + with cls.get_engine( + database, + catalog=table.catalog, + schema=table.schema, + ) as engine: + # Check if we need to add a primary key + if_exists = to_sql_kwargs.get("if_exists", "fail") + needs_primary_key = ( + if_exists == "fail" # Only for new table creation + and not to_sql_kwargs.get("index", False) # No index column exists Review Comment: <div> <div id="suggestion"> <div id="issue"><b>Unclosed parenthesis in conditional expression</b></div> <div id="fix"> Line 456 has an unclosed opening parenthesis `(` in the `needs_primary_key` assignment. The condition started on line 456 is missing a closing parenthesis before the assignment completes on line 458. </div> <details> <summary> <b>Code suggestion</b> </summary> <blockquote>Check the AI-generated fix before applying</blockquote> <div id="code"> ````suggestion needs_primary_key = ( if_exists == "fail" # Only for new table creation and not to_sql_kwargs.get("index", False) # No index column exists ) ```` </div> </details> </div> <small><i>Code Review Run #528e4b</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]
