Add configurable MAX_ROWS to cap fetched data per table scan

benaiad · benaiad · commit e5bc861a17e2 · 2026-03-24T04:15:16.000-04:00
Replace MAX_PAGES with MAX_ROWS (default 1000) as the pagination cap.
When no SQL LIMIT is specified, each table scan fetches up to max_rows
before stopping. The current page is always completed, so actual row
count may slightly exceed the cap.

SQL LIMIT is still respected when DataFusion pushes it down (simple
queries). For JOINs where LIMIT can't be pushed, max_rows ensures
both sides fetch enough data for meaningful results.

Configurable via:
  --max-rows 5000          (CLI flag, one-off override)
  SQLIZE_MAX_ROWS=5000     (env var, persistent override)
  Default: 1000
diff --git a/crates/sqlize-core/src/datafusion/exec.rs b/crates/sqlize-core/src/datafusion/exec.rs
@@ -20,10 +20,6 @@ use crate::exec::pagination;
 
 use super::arrow_convert::json_response_to_batch;
 
-/// Maximum number of HTTP requests (pages) per table scan.
-/// At ~30 rows per page (GitHub default), this yields ~1500 rows.
-const MAX_PAGES: usize = 50;
-
 /// A custom DataFusion `ExecutionPlan` that fetches data from a REST API.
 /// Returns a lazy stream that fetches one page per `poll_next()`.
 #[derive(Debug)]
@@ -33,7 +29,7 @@ pub struct ApiTableExec {
     projected_schema: SchemaRef,
     params: HashMap<String, String>,
     projection: Option<Vec<usize>>,
-    limit: Option<usize>,
+    max_rows: usize,
     auth: AuthConfig,
     client: reqwest::Client,
     properties: PlanProperties,
@@ -45,7 +41,7 @@ impl ApiTableExec {
         full_schema: SchemaRef,
         params: HashMap<String, String>,
         projection: Option<Vec<usize>>,
-        limit: Option<usize>,
+        max_rows: usize,
         auth: AuthConfig,
         client: reqwest::Client,
     ) -> Self {
@@ -73,7 +69,7 @@ impl ApiTableExec {
             projected_schema,
             params,
             projection,
-            limit,
+            max_rows,
             auth,
             client,
             properties,
@@ -129,7 +125,7 @@ impl ExecutionPlan for ApiTableExec {
         let projected_schema = self.projected_schema.clone();
         let params = self.params.clone();
         let projection = self.projection.clone();
-        let limit = self.limit;
+        let max_rows = self.max_rows;
         let auth = self.auth.clone();
         let client = self.client.clone();
 
@@ -150,8 +146,7 @@ impl ExecutionPlan for ApiTableExec {
             next_url: Option<String>,
             is_first_page: bool,
             total_rows: usize,
-            pages_fetched: usize,
-            row_limit: Option<usize>,
+            max_rows: usize,
             table: VirtualTable,
             full_schema: SchemaRef,
             params: HashMap<String, String>,
@@ -165,8 +160,7 @@ impl ExecutionPlan for ApiTableExec {
             next_url: Some(first_url),
             is_first_page: true,
             total_rows: 0,
-            pages_fetched: 0,
-            row_limit: limit,
+            max_rows,
             table,
             full_schema,
             params,
@@ -179,21 +173,13 @@ impl ExecutionPlan for ApiTableExec {
         let stream = futures::stream::unfold(initial_state, |mut state| async move {
             let url = state.next_url.take()?;
 
-            if state.pages_fetched >= MAX_PAGES {
-                tracing::warn!(
-                    table = %state.table.name,
-                    pages = state.pages_fetched,
-                    "reached MAX_PAGES limit, stopping pagination"
-                );
+            // Stop if we've already fetched enough rows.
+            // We check >= rather than > so that the current page that pushed us
+            // over the limit is still returned (the caller's LIMIT will trim).
+            if state.total_rows >= state.max_rows {
                 return None;
             }
 
-            if let Some(limit) = state.row_limit {
-                if state.total_rows >= limit {
-                    return None;
-                }
-            }
-
             let (body, headers) = match fetch_page(
                 &state.client,
                 &state.auth,
@@ -230,7 +216,6 @@ impl ExecutionPlan for ApiTableExec {
             };
 
             state.total_rows += batch.num_rows();
-            state.pages_fetched += 1;
 
             // Determine next page
             let ctx = pagination::PageContext {
diff --git a/crates/sqlize-core/src/datafusion/mod.rs b/crates/sqlize-core/src/datafusion/mod.rs
@@ -15,40 +15,44 @@ use crate::exec::AuthConfig;
 use self::arrow_convert::batches_to_result_set;
 use self::schema::ApiSchemaProvider;
 
+/// Default maximum rows fetched per table scan when no SQL LIMIT is specified.
+/// Overridable via `SQLIZE_MAX_ROWS` env var or `--max-rows` CLI flag.
+pub const DEFAULT_MAX_ROWS: usize = 1000;
+
 /// The main entry point for executing SQL against REST APIs via DataFusion.
 ///
 /// Wraps a DataFusion `SessionContext` with registered API table providers.
 /// Supports multiple specs (schemas) for federated queries.
 pub struct SqlizeContext {
     ctx: SessionContext,
+    max_rows: usize,
 }
 
 impl Default for SqlizeContext {
     fn default() -> Self {
-        Self::new()
+        Self::new(DEFAULT_MAX_ROWS)
     }
 }
 
 impl SqlizeContext {
-    pub fn new() -> Self {
+    pub fn new(max_rows: usize) -> Self {
         let config = SessionConfig::new()
             .with_information_schema(false)
             .with_default_catalog_and_schema("sqlize", "default");
         let ctx = SessionContext::new_with_config(config);
-        Self { ctx }
+        Self { ctx, max_rows }
     }
 
     /// Register a spec's tables under the given schema name.
-    /// If `is_default` is true, sets this schema as the default so bare table
-    /// names resolve to it.
     pub fn register_spec(
         &self,
         schema_name: &str,
         catalog: &Catalog,
         auth: AuthConfig,
         client: reqwest::Client,
     ) -> Result<()> {
-        let schema_provider = Arc::new(ApiSchemaProvider::new(catalog, auth, client));
+        let schema_provider =
+            Arc::new(ApiSchemaProvider::new(catalog, auth, client, self.max_rows));
 
         let df_catalog = self
             .ctx
diff --git a/crates/sqlize-core/src/datafusion/provider.rs b/crates/sqlize-core/src/datafusion/provider.rs
@@ -24,6 +24,7 @@ pub struct ApiTableProvider {
     schema: SchemaRef,
     auth: AuthConfig,
     client: reqwest::Client,
+    max_rows: usize,
 }
 
 impl fmt::Debug for ApiTableProvider {
@@ -35,13 +36,19 @@ impl fmt::Debug for ApiTableProvider {
 }
 
 impl ApiTableProvider {
-    pub fn new(table: VirtualTable, auth: AuthConfig, client: reqwest::Client) -> Self {
+    pub fn new(
+        table: VirtualTable,
+        auth: AuthConfig,
+        client: reqwest::Client,
+        max_rows: usize,
+    ) -> Self {
         let schema = virtual_table_to_schema(&table);
         Self {
             table,
             schema,
             auth,
             client,
+            max_rows,
         }
     }
 }
@@ -113,12 +120,15 @@ impl TableProvider for ApiTableProvider {
             }
         }
 
+        // Use DataFusion's pushed-down limit if available, otherwise use max_rows.
+        let effective_limit = limit.unwrap_or(self.max_rows);
+
         let exec = ApiTableExec::new(
             self.table.clone(),
             self.schema.clone(),
             params,
             projection.cloned(),
-            limit,
+            effective_limit,
             self.auth.clone(),
             self.client.clone(),
         );
@@ -136,11 +146,7 @@ fn classify_filter(table: &VirtualTable, expr: &Expr) -> TableProviderFilterPush
                     ColumnRole::PathParam | ColumnRole::QueryParam => {
                         TableProviderFilterPushDown::Exact
                     }
-                    ColumnRole::QueryParamAndResponse => {
-                        // DataFusion should also evaluate this locally since the
-                        // API may return items that don't match (e.g., paginated results)
-                        TableProviderFilterPushDown::Inexact
-                    }
+                    ColumnRole::QueryParamAndResponse => TableProviderFilterPushDown::Inexact,
                     _ => TableProviderFilterPushDown::Unsupported,
                 };
             }
diff --git a/crates/sqlize-core/src/datafusion/schema.rs b/crates/sqlize-core/src/datafusion/schema.rs
@@ -18,10 +18,16 @@ pub struct ApiSchemaProvider {
     tables: HashMap<String, VirtualTable>,
     auth: AuthConfig,
     client: reqwest::Client,
+    max_rows: usize,
 }
 
 impl ApiSchemaProvider {
-    pub fn new(catalog: &Catalog, auth: AuthConfig, client: reqwest::Client) -> Self {
+    pub fn new(
+        catalog: &Catalog,
+        auth: AuthConfig,
+        client: reqwest::Client,
+        max_rows: usize,
+    ) -> Self {
         let tables: HashMap<String, VirtualTable> = catalog
             .tables()
             .map(|t| (t.name.as_str().to_owned(), t.clone()))
@@ -30,6 +36,7 @@ impl ApiSchemaProvider {
             tables,
             auth,
             client,
+            max_rows,
         }
     }
 }
@@ -52,8 +59,12 @@ impl SchemaProvider for ApiSchemaProvider {
     ) -> datafusion::common::Result<Option<Arc<dyn TableProvider>>> {
         match self.tables.get(name) {
             Some(vt) => {
-                let provider =
-                    ApiTableProvider::new(vt.clone(), self.auth.clone(), self.client.clone());
+                let provider = ApiTableProvider::new(
+                    vt.clone(),
+                    self.auth.clone(),
+                    self.client.clone(),
+                    self.max_rows,
+                );
                 Ok(Some(Arc::new(provider)))
             }
             None => Ok(None),
diff --git a/crates/sqlize/src/main.rs b/crates/sqlize/src/main.rs
@@ -8,7 +8,7 @@ use clap::{Parser, Subcommand};
 
 use rmcp::ServiceExt;
 use sqlize_core::catalog::Catalog;
-use sqlize_core::datafusion::SqlizeContext;
+use sqlize_core::datafusion::{DEFAULT_MAX_ROWS, SqlizeContext};
 use sqlize_core::exec::AuthConfig;
 use sqlize_core::spec::SpecInfo;
 
@@ -29,6 +29,10 @@ struct Cli {
     /// Output format for query results
     #[arg(short, long, default_value = "table", global = true)]
     format: repl::OutputFormat,
+
+    /// Max rows fetched per table scan when no SQL LIMIT is specified
+    #[arg(long, global = true)]
+    max_rows: Option<usize>,
 }
 
 #[derive(Subcommand)]
@@ -172,7 +176,16 @@ async fn main() -> anyhow::Result<()> {
     }
 
     let client = sqlize_core::exec::Client::new();
-    let sqlize_ctx = Arc::new(SqlizeContext::new());
+
+    // Resolve max_rows: CLI flag > env var > default
+    let max_rows = cli.max_rows.unwrap_or_else(|| {
+        std::env::var("SQLIZE_MAX_ROWS")
+            .ok()
+            .and_then(|s| s.parse().ok())
+            .unwrap_or(DEFAULT_MAX_ROWS)
+    });
+
+    let sqlize_ctx = Arc::new(SqlizeContext::new(max_rows));
     let is_single_spec = specs.len() == 1;
 
     // Register each spec as a named schema
