feature(paginator): adding paginator functionality (#1644)

scaliseraoul · ellipsis-dev[bot] · web-flow · commit 357c97cc4c42 · 2025-02-27T19:22:56.000+01:00
* feature(paginator): adding paginator functionality

* Update pandasai/helpers/sql_sanitizer.py

Co-authored-by: ellipsis-dev[bot] &lt;65095814+ellipsis-dev[bot]@users.noreply.github.com&gt;

* Update pandasai/query_builders/sql_parser.py

Co-authored-by: ellipsis-dev[bot] &lt;65095814+ellipsis-dev[bot]@users.noreply.github.com&gt;

* feature(paginator): docstring

* feature(paginator): format issues

---------

Co-authored-by: ellipsis-dev[bot] &lt;65095814+ellipsis-dev[bot]@users.noreply.github.com&gt;
diff --git a/pandasai/helpers/sql_sanitizer.py b/pandasai/helpers/sql_sanitizer.py
@@ -97,3 +97,26 @@ def is_sql_query_safe(query: str, dialect: str = "postgres") -> bool:
 
     except sqlglot.errors.ParseError:
         return False
+
+
+def is_sql_query(query: str) -> bool:
+    # Define SQL patterns with context to avoid standalone keyword matches
+    sql_patterns = [
+        r"\bSELECT\b.*\bFROM\b",
+        r"\bINSERT\b.*\bINTO\b",
+        r"\bUPDATE\b.*\bSET\b",
+        r"\bDELETE\b.*\bFROM\b",
+        r"\bDROP\b.*\b(TABLE|DATABASE)\b",
+        r"\bCREATE\b.*\b(DATABASE|TABLE)\b",
+        r"\bALTER\b.*\bTABLE\b",
+        r"\bJOIN\b.*\bON\b",
+        r"\bWHERE\b",
+    ]
+
+    # Combine all patterns into a single regex
+    sql_regex = re.compile("|".join(sql_patterns), re.IGNORECASE)
+
+    # If the query matches any SQL pattern, it's considered a SQL query
+    if sql_regex.search(query):
+        return True
+    return False
diff --git a/pandasai/query_builders/base_query_builder.py b/pandasai/query_builders/base_query_builder.py
@@ -2,7 +2,6 @@
 
 import sqlglot
 from sqlglot import select
-from sqlglot.expressions import false
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
 from pandasai.data_loader.semantic_layer_schema import SemanticLayerSchema, Source
diff --git a/pandasai/query_builders/paginator.py b/pandasai/query_builders/paginator.py
@@ -0,0 +1,209 @@
+import datetime
+import json
+import uuid
+from typing import List, Optional, Tuple
+
+import sqlglot
+from pydantic import BaseModel, Field, field_validator
+
+from pandasai.helpers.sql_sanitizer import is_sql_query
+
+
+class PaginationParams(BaseModel):
+    """Parameters for pagination requests"""
+
+    page: int = Field(ge=1, description="Page number, starting from 1")
+    page_size: int = Field(
+        ge=1, le=100, description="Number of items per page, maximum 100"
+    )
+    search: Optional[str] = Field(
+        None, description="Search term to filter across all fields"
+    )
+    sort_by: Optional[str] = Field(None, description="Column to sort by")
+    sort_order: Optional[str] = Field(
+        None, pattern="^(asc|desc)$", description="Sort order (asc or desc)"
+    )
+    filters: Optional[str] = Field(None, description="Filters to apply to the data")
+
+    @field_validator("search", "filters", "sort_by", "sort_order")
+    @classmethod
+    def not_sql(cls, field):
+        if is_sql_query(str(field)):
+            raise ValueError(
+                f"SQL queries are not allowed in pagination parameters: {field}"
+            )
+        return field
+
+
+class DatasetPaginator:
+    @staticmethod
+    def is_float(value: str) -> bool:
+        try:
+            # Try to cast the value to a number
+            float(value)
+            return True
+        except (ValueError, TypeError):
+            # If it fails, it's not a number
+            return False
+
+    @staticmethod
+    def is_valid_boolean(value):
+        """Check if the value is a valid boolean."""
+        return (
+            value.lower() in ["true", "false"]
+            if isinstance(value, str)
+            else isinstance(value, bool)
+        )
+
+    @staticmethod
+    def is_valid_uuid(value):
+        try:
+            uuid.UUID(value)
+            return True
+        except ValueError:
+            return False
+
+    @staticmethod
+    def is_valid_datetime(value: str) -> bool:
+        try:
+            datetime.datetime.strptime(value, "%Y-%m-%d %H:%M:%S")
+            return True
+        except ValueError:
+            return False
+
+    @staticmethod
+    def apply_pagination(
+        query: str,
+        columns: List[dict],
+        pagination: Optional[PaginationParams],
+        target_dialect: str = "postgres",
+    ) -> Tuple[str, List]:
+        """
+        Apply pagination to a SQL query.
+
+        Args:
+            query (str): The SQL query to apply pagination to
+            columns (List[dict]): A list of dictionaries containing
+                information about the columns in the result set. Each
+                dictionary should have the following structure:
+                    {
+                        "name": str,
+                        "type": str
+                    }
+                The type should be one of: "string", "number", "integer", "float",
+                "boolean", "datetime"
+            pagination (Optional[PaginationParams]): The pagination parameters
+                to apply to the query. If None, the query is returned unchanged
+            target_dialect (str): The SQL dialect to generate the query for.
+                Defaults to "postgres".
+
+        Returns:
+            Tuple[str, List]: A tuple containing the modified SQL query and a
+                list of parameters to pass to the query.
+        """
+
+        params = []
+
+        if not pagination:
+            return query, params
+
+        filtering_query = f"SELECT * FROM ({query}) AS filtered_data"
+        conditions = []
+
+        # Handle search functionality
+        if pagination.search:
+            search_conditions = []
+            for column in columns:
+                column_name = column["name"]
+                column_type = column["type"]
+
+                if column_type == "string":
+                    search_conditions.append(f"{column_name} ILIKE %s")
+                    params.append(f"%{pagination.search}%")
+
+                elif column_type == "float" and DatasetPaginator.is_float(
+                    pagination.search
+                ):
+                    search_conditions.append(f"{column_name} = %s")
+                    params.append(pagination.search)
+
+                elif (
+                    column_type in ["number", "integer"]
+                    and pagination.search.isnumeric()
+                ):
+                    search_conditions.append(f"{column_name} = %s")
+                    params.append(pagination.search)
+
+                elif column_type == "datetime" and DatasetPaginator.is_valid_datetime(
+                    pagination.search
+                ):
+                    search_conditions.append(f"{column_name} = %s")
+                    params.append(
+                        datetime.datetime.strptime(
+                            pagination.search, "%Y-%m-%d %H:%M:%S"
+                        )
+                    )
+
+                elif column_type == "boolean" and DatasetPaginator.is_valid_boolean(
+                    pagination.search
+                ):
+                    search_conditions.append(f"{column_name} = %s")
+                    params.append(pagination.search)
+
+                elif column_type == "uuid" and DatasetPaginator.is_valid_uuid(
+                    pagination.search
+                ):
+                    search_conditions.append(f"{column_name}::TEXT = %s")
+                    params.append(pagination.search)
+
+            if search_conditions:
+                conditions.append(" OR ".join(search_conditions))
+
+        # Handle filters
+        if pagination.filters:
+            try:
+                filters = (
+                    json.loads(pagination.filters)
+                    if isinstance(pagination.filters, str)
+                    else pagination.filters
+                )
+                for column, values in filters.items():
+                    if not isinstance(values, list):
+                        values = [values]
+                    placeholders = ", ".join(["%s"] * len(values))
+                    conditions.append(f"{column} IN ({placeholders})")
+                    params.extend(values)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Invalid filters format: {e}")
+
+        # Add WHERE clause if conditions exist
+        if conditions:
+            filtering_query += " WHERE " + " AND ".join(conditions)
+
+        # Handle sorting
+        if pagination.sort_by and pagination.sort_order:
+            if not any(pagination.sort_by == column["name"] for column in columns):
+                raise ValueError(
+                    f"Sort column '{pagination.sort_by}' not found in available columns"
+                )
+
+            filtering_query += (
+                f" ORDER BY {pagination.sort_by} {pagination.sort_order.upper()}"
+            )
+
+        # Handle page and page_size
+        if pagination.page and pagination.page_size:
+            filtering_query += " LIMIT %s OFFSET %s"
+            params.extend(
+                [pagination.page_size, (pagination.page - 1) * pagination.page_size]
+            )
+
+        # Replace placeholders for target dialect
+        placeholder = "___PLACEHOLDER___"
+        temp_query = filtering_query.replace("%s", placeholder)
+        transpiled_query = sqlglot.transpile(
+            temp_query, read="postgres", write=target_dialect
+        )[0]
+        final_query = transpiled_query.replace(placeholder, "%s")
+
+        return final_query, params
diff --git a/pandasai/query_builders/sql_parser.py b/pandasai/query_builders/sql_parser.py
@@ -57,11 +57,16 @@ def transform_node(node):
         return transformed.sql(pretty=True)
 
     @staticmethod
-    def transpile_sql_dialect(query, to_dialect, from_dialect=None):
+    def transpile_sql_dialect(
+        query: str, to_dialect: str, from_dialect: Optional[str] = None
+    ):
+        placeholder = "___PLACEHOLDER___"
+        query = query.replace("%s", placeholder)
         query = (
             parse_one(query, read=from_dialect) if from_dialect else parse_one(query)
         )
-        return query.sql(dialect=to_dialect, pretty=True)
+        result = query.sql(dialect=to_dialect, pretty=True)
+        return result.replace(placeholder, "%s")
 
     @staticmethod
     def extract_table_names(sql_query: str, dialect: str = "postgres") -> List[str]:
diff --git a/tests/unit_tests/helpers/test_sql_sanitizer.py b/tests/unit_tests/helpers/test_sql_sanitizer.py
@@ -1,4 +1,5 @@
 from pandasai.helpers.sql_sanitizer import (
+    is_sql_query,
     is_sql_query_safe,
     sanitize_file_name,
     sanitize_view_column_name,
@@ -94,3 +95,37 @@ def test_safe_query_with_subquery(self):
     def test_safe_query_with_query_params(self):
         query = "SELECT * FROM (SELECT * FROM heart_data) AS filtered_data LIMIT %s OFFSET %s"
         assert is_sql_query_safe(query)
+
+    def test_plain_text(self):
+        """Test with plain text input that is not a SQL query."""
+        assert not is_sql_query("Hello, how are you?")
+        assert not is_sql_query("This is just some text.")
+
+    def test_sql_queries(self):
+        """Test with typical SQL queries."""
+        assert is_sql_query("SELECT * FROM users")
+        assert is_sql_query("insert into users values ('john', 25)")
+        assert is_sql_query("delete from orders where id=10")
+        assert is_sql_query("DROP TABLE users")
+        assert is_sql_query("update products set price=100 where id=1")
+
+    def test_case_insensitivity(self):
+        """Test with queries in different cases."""
+        assert is_sql_query("select id from users")
+        assert is_sql_query("SeLeCt id FROM users")
+        assert is_sql_query("DROP table orders")
+        assert is_sql_query("cReAtE DATABASE testdb")
+
+    def test_edge_cases(self):
+        """Test with edge cases like empty strings and special characters."""
+        assert not is_sql_query("")
+        assert not is_sql_query(" ")
+        assert not is_sql_query("1234567890")
+        assert not is_sql_query("#$%^&*()")
+        assert not is_sql_query("JOIN the party")  # Not SQL context
+
+    def test_mixed_input(self):
+        """Test with mixed input containing SQL keywords in non-SQL contexts."""
+        assert not is_sql_query("Let's SELECT a movie to watch")
+        assert not is_sql_query("CREATE a new painting")
+        assert not is_sql_query("DROP by my house later")
diff --git a/tests/unit_tests/query_builders/test_paginator.py b/tests/unit_tests/query_builders/test_paginator.py
