Repository: apache/superset
File: tests/unit_tests/commands/importers/v1/examples_test.py (Python)

Code before:
    config = {"table_name": "my_table", "sql": original_sql}

    # Should not raise, should keep original SQL
    transpile_virtual_dataset_sql(config, 1)
    assert config["sql"] == original_sql


@patch("superset.commands.importers.v1.examples.db")
@patch("superset.commands.importers.v1.examples.transpile_to_dialect")
def test_transpile_virtual_dataset_sql_complex_query(mock_transpile, mock_db):
    """Test transpilation of a more complex SQL query."""
    mock_database = MagicMock()
    mock_database.db_engine_spec.engine = "duckdb"
    mock_db.session.query.return_value.get.return_value = mock_database

    original_sql = """
        SELECT
            DATE_TRUNC('month', created_at) as month,
            COUNT(*) as count
        FROM orders
        WHERE status = 'completed'
        GROUP BY 1
    """
    transpiled_sql = """
        SELECT
            DATE_TRUNC('month', created_at) AS month,
            COUNT(*) AS count
        FROM orders
        WHERE status = 'completed'
        GROUP BY 1
    """
    mock_transpile.return_value = transpiled_sql

    config = {"table_name": "monthly_orders", "sql": original_sql}
    transpile_virtual_dataset_sql(config, 1)

    assert config["sql"] == transpiled_sql
    mock_transpile.assert_called_once_with(original_sql, "duckdb")


Review comment (suggestion):
This one is easy. Let's add tests transpiling to a bigger variety of DBs -- at least something that doesn't have `DATE_TRUNC`, like Clickhouse (I think?).

Code after:
        "table_name": "my_table",
        "sql": original_sql,
        "source_db_engine": "postgresql",
    }
    transpile_virtual_dataset_sql(config, 1)

    assert config["sql"] == original_sql


@patch("superset.commands.importers.v1.examples.db")
@patch("superset.commands.importers.v1.examples.transpile_to_dialect")
def test_transpile_virtual_dataset_sql_error_fallback(mock_transpile, mock_db):
    """Test graceful fallback when transpilation fails."""
    from superset.exceptions import QueryClauseValidationException

    mock_database = MagicMock()
    mock_database.db_engine_spec.engine = "mysql"
    mock_db.session.query.return_value.get.return_value = mock_database

    mock_transpile.side_effect = QueryClauseValidationException("Parse error")

    original_sql = "SELECT SOME_POSTGRES_SPECIFIC_FUNCTION() FROM foo"
    config = {
        "table_name": "my_table",
        "sql": original_sql,
        "source_db_engine": "postgresql",
    }

    # Should not raise, should keep original SQL
    transpile_virtual_dataset_sql(config, 1)
    assert config["sql"] == original_sql


@patch("superset.commands.importers.v1.examples.db")
@patch("superset.commands.importers.v1.examples.transpile_to_dialect")
def test_transpile_virtual_dataset_sql_postgres_to_duckdb(mock_transpile, mock_db):
    """Test transpilation from PostgreSQL to DuckDB."""
    mock_database = MagicMock()
    mock_database.db_engine_spec.engine = "duckdb"
    mock_db.session.query.return_value.get.return_value = mock_database

    original_sql = """
        SELECT DATE_TRUNC('month', created_at) AS month, COUNT(*) AS cnt
        FROM orders WHERE status = 'completed' GROUP BY 1
    """
    transpiled_sql = """
        SELECT DATE_TRUNC('month', created_at) AS month, COUNT(*) AS cnt
        FROM orders WHERE status = 'completed' GROUP BY 1
    """
    mock_transpile.return_value = transpiled_sql