diff --git a/marimo/_sql/sql.py b/marimo/_sql/sql.py index 37dffad8021..b2eeaeb7dd3 100644 --- a/marimo/_sql/sql.py +++ b/marimo/_sql/sql.py @@ -12,7 +12,11 @@ from marimo._sql.engines.sqlalchemy import SQLAlchemyEngine from marimo._sql.engines.types import QueryEngine from marimo._sql.get_engines import SUPPORTED_ENGINES -from marimo._sql.utils import raise_df_import_error +from marimo._sql.utils import ( + extract_explain_content, + is_explain_query, + raise_df_import_error, +) from marimo._types.ids import VariableName from marimo._utils.narwhals_utils import can_narwhalify_lazyframe @@ -103,9 +107,14 @@ def sql( raise_df_import_error("polars[pyarrow]") if output: + from marimo._plugins.stateless.plain_text import plain_text from marimo._plugins.ui._impl import table - if can_narwhalify_lazyframe(df): + if isinstance(sql_engine, DuckDBEngine) and is_explain_query(query): + # For EXPLAIN queries in DuckDB, display plain output to preserve box drawings + text_output = extract_explain_content(df) + t = plain_text(text_output) + elif can_narwhalify_lazyframe(df): # For pl.LazyFrame and DuckDBRelation, we only show the first few rows # to avoid loading all the data into memory. # Also preload the first page of data without user confirmation. diff --git a/marimo/_sql/utils.py b/marimo/_sql/utils.py index e1499497c20..213b97184ea 100644 --- a/marimo/_sql/utils.py +++ b/marimo/_sql/utils.py @@ -176,3 +176,44 @@ def sql_type_to_data_type(type_str: str) -> DataType: return "string" else: return "string" + + +def is_explain_query(query: str) -> bool: + """Check if a SQL query is an EXPLAIN query.""" + return query.lstrip().lower().startswith("explain ") + + +def extract_explain_content(df: Any) -> str: + """Extract all content from a DataFrame for EXPLAIN queries. + + Args: + df: DataFrame (pandas or polars). If not pandas / polars, return repr(df). + + Returns: + String containing content of dataframe + """ + try: + if DependencyManager.polars.has(): + import polars as pl + + if isinstance(df, pl.LazyFrame): + df = df.collect() + if isinstance(df, pl.DataFrame): + # Display full strings without truncation + with pl.Config(fmt_str_lengths=1000): + return str(df) + + if DependencyManager.pandas.has(): + import pandas as pd + + if isinstance(df, pd.DataFrame): + # Preserve newlines in the data + all_values = df.values.flatten().tolist() + return "\n".join(str(val) for val in all_values) + + # Fallback to repr for other types + return repr(df) + + except Exception as e: + LOGGER.debug("Failed to extract explain content: %s", e) + return repr(df) diff --git a/tests/_sql/test_sql.py b/tests/_sql/test_sql.py index eeae3fc061d..2861384714f 100644 --- a/tests/_sql/test_sql.py +++ b/tests/_sql/test_sql.py @@ -13,6 +13,7 @@ from marimo._sql.engines.ibis import IbisEngine from marimo._sql.engines.sqlalchemy import SQLAlchemyEngine from marimo._sql.sql import _query_includes_limit, sql +from marimo._sql.utils import extract_explain_content, is_explain_query if TYPE_CHECKING: from collections.abc import Generator @@ -360,3 +361,155 @@ def test_sql_with_ibis_expression_result(): with patch.object(IbisEngine, "sql_output_format", return_value="native"): result = sql("SELECT * FROM test", engine=duckdb_backend) assert isinstance(result, Expr) + + +class TestExplainQueries: + def test_is_explain_query(self): + """Test is_explain_query function.""" + # Test valid EXPLAIN queries + assert is_explain_query("EXPLAIN SELECT 1") + assert is_explain_query("explain SELECT 1") + assert is_explain_query(" EXPLAIN SELECT 1") + assert is_explain_query("\tEXPLAIN SELECT 1") + assert is_explain_query("\nEXPLAIN SELECT 1") + assert is_explain_query("EXPLAIN (FORMAT JSON) SELECT 1") + assert is_explain_query("EXPLAIN ANALYZE SELECT 1") + assert is_explain_query("EXPLAIN QUERY PLAN SELECT 1") + + # Test non-EXPLAIN queries + assert not is_explain_query("SELECT 1") + assert not is_explain_query("INSERT INTO t VALUES (1)") + assert not is_explain_query("UPDATE t SET col = 1") + assert not is_explain_query("DELETE FROM t") + assert not is_explain_query("CREATE TABLE t (id INT)") + assert not is_explain_query("") + assert not is_explain_query(" ") + + # Test edge cases + assert not is_explain_query("EXPLAINED") # Not exactly "explain" + assert not is_explain_query("EXPLAINING") # Not exactly "explain" + assert not is_explain_query( + "-- EXPLAIN SELECT 1" + ) # Comment, not actual query + + @pytest.fixture + def explain_df_data(self) -> dict[str, list[str]]: + return { + "explain_key": ["physical_plan", "logical_plan"], + "explain_value": [ + "┌─────────────────────────────────────┐\n│ PROJECTION │\n└─────────────────────────────────────┘", + "┌─────────────────────────────────────┐\n│ SELECTION │\n└─────────────────────────────────────┘", + ], + } + + @pytest.mark.skipif(not HAS_POLARS, reason="Polars not installed") + def test_extract_explain_content_polars( + self, explain_df_data: dict[str, list[str]] + ): + """Test extract_explain_content with polars DataFrames.""" + import polars as pl + + # Test with regular DataFrame + df = pl.DataFrame(explain_df_data) + + expected_rendering = """shape: (2, 2) +┌───────────────┬───────────────────────────────────────────┐ +│ explain_key ┆ explain_value │ +│ --- ┆ --- │ +│ str ┆ str │ +╞═══════════════╪═══════════════════════════════════════════╡ +│ physical_plan ┆ ┌─────────────────────────────────────┐ │ +│ ┆ │ PROJECTION │ │ +│ ┆ └─────────────────────────────────────┘ │ +│ logical_plan ┆ ┌─────────────────────────────────────┐ │ +│ ┆ │ SELECTION │ │ +│ ┆ └─────────────────────────────────────┘ │ +└───────────────┴───────────────────────────────────────────┘""" + + result = extract_explain_content(df) + assert result == expected_rendering + + # Test with LazyFrame + lazy_df = df.lazy() + result = extract_explain_content(lazy_df) + assert result == expected_rendering + + @pytest.mark.skipif(not HAS_PANDAS, reason="Pandas not installed") + def test_extract_explain_content_pandas( + self, explain_df_data: dict[str, list[str]] + ): + """Test extract_explain_content with pandas DataFrames.""" + import pandas as pd + + df = pd.DataFrame(explain_df_data) + + result = extract_explain_content(df) + assert ( + result + == """physical_plan +┌─────────────────────────────────────┐ +│ PROJECTION │ +└─────────────────────────────────────┘ +logical_plan +┌─────────────────────────────────────┐ +│ SELECTION │ +└─────────────────────────────────────┘""" + ) + + def test_extract_explain_content_fallback(self): + """Test extract_explain_content fallback for non-DataFrame objects.""" + # Test with non-DataFrame object + result = extract_explain_content("not a dataframe") + assert isinstance(result, str) + assert "not a dataframe" in result + + # Test with None + result = extract_explain_content(None) + assert isinstance(result, str) + assert "None" in result + + @pytest.mark.skipif(not HAS_POLARS, reason="Polars not installed") + def test_extract_explain_content_error_handling(self): + """Test extract_explain_content error handling.""" + + # Create a mock DataFrame that will raise an error + class MockDataFrame: + def __init__(self): + pass + + def __str__(self): + raise RuntimeError("Test error") + + mock_df = MockDataFrame() + result = extract_explain_content(mock_df) + assert isinstance(result, str) + assert "MockDataFrame" in result # Should fallback to repr + + @patch("marimo._sql.sql.replace") + @pytest.mark.skipif( + not HAS_POLARS or not HAS_DUCKDB, reason="polars and duckdb required" + ) + def test_sql_explain_query_display(self, mock_replace): + """Test that EXPLAIN queries are displayed as plain text.""" + import duckdb + import polars as pl + + # Create a test table + duckdb.sql( + "CREATE OR REPLACE TABLE test_explain AS SELECT * FROM range(5)" + ) + + # Test EXPLAIN query + result = sql("EXPLAIN SELECT * FROM test_explain") + assert isinstance(result, pl.DataFrame) + + # Should call replace with plain_text + mock_replace.assert_called_once() + call_args = mock_replace.call_args[0][0] + + # The call should be a plain_text object + assert hasattr(call_args, "text") + assert isinstance(call_args.text, str) + + # Clean up + duckdb.sql("DROP TABLE test_explain")