Function test_query_includes_limit

tests/_sql/test_sql.py:398–430

Test _query_includes_limit function.

()

@pytest.mark.skipif(not HAS_SQLGLOT, reason="sqlglot not installed")
def test_query_includes_limit() -> None:
    """Test _query_includes_limit function."""
    # Test queries with LIMIT
    assert _query_includes_limit("SELECT * FROM table LIMIT 5")
    assert _query_includes_limit("SELECT * FROM table LIMIT\n5")
    assert _query_includes_limit("SELECT * FROM table limit\n5")
    assert _query_includes_limit(
        """
        SELECT *
        FROM table
        LIMIT 5
        """
    )

    # Test queries without LIMIT
    assert not _query_includes_limit("SELECT * FROM table")
    assert not _query_includes_limit("SELECT LIMIT_COL FROM table")
    assert not _query_includes_limit("-- LIMIT 5\nSELECT * FROM table")

    # Test invalid SQL
    assert not _query_includes_limit("NOT A VALID SQL QUERY")
    assert not _query_includes_limit("")
    assert not _query_includes_limit("INSERT INTO t VALUES (1, 2, 3)")

    # Multiple queries
    assert (
        _query_includes_limit("SELECT * FROM t; SELECT * FROM t2 LIMIT 5")
        is True
    )
    assert (
        _query_includes_limit("SELECT * FROM t LIMIT 5; SELECT * FROM t2")
        is False
    )


@patch("marimo._sql.sql.replace")