fleet-python 0.2.18__tar.gz → 0.2.20__tar.gz

{fleet_python-0.2.18 → fleet_python-0.2.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fleet-python
-Version: 0.2.18
+Version: 0.2.20
 Summary: Python SDK for Fleet environments
 Author-email: Fleet AI <nic@fleet.so>
 License: Apache-2.0

{fleet_python-0.2.18 → fleet_python-0.2.20}/fleet/resources/mcp.py

@@ -1,5 +1,4 @@
 from typing import Dict
-import aiohttp
 import json
 
 
@@ -24,6 +23,7 @@ class MCPResource:
         }
 
     async def list_tools(self):
+        import aiohttp
         """
         Make an async request to list available tools from the MCP endpoint.
         

{fleet_python-0.2.18 → fleet_python-0.2.20}/fleet/resources/sqlite.py

@@ -18,7 +18,7 @@ from ..verifiers.db import IgnoreConfig, _get_row_identifier, _format_row_for_er
 
 
 class SyncDatabaseSnapshot:
-    """Async database snapshot that fetches data through API and stores locally for diffing."""
+    """Lazy database snapshot that fetches data on-demand through API."""
     
     def __init__(self, resource: "SQLiteResource", name: str | None = None):
         self.resource = resource
@@ -26,11 +26,12 @@ class SyncDatabaseSnapshot:
         self.created_at = datetime.utcnow()
         self._data: dict[str, list[dict[str, Any]]] = {}
         self._schemas: dict[str, list[str]] = {}
-        self._fetched = False
+        self._table_names: list[str] | None = None
+        self._fetched_tables: set[str] = set()
         
-    def _ensure_fetched(self):
-        """Fetch all data from remote database if not already fetched."""
-        if self._fetched:
+    def _ensure_tables_list(self):
+        """Fetch just the list of table names if not already fetched."""
+        if self._table_names is not None:
             return
             
         # Get all tables
@@ -39,34 +40,37 @@ class SyncDatabaseSnapshot:
         )
         
         if not tables_response.rows:
-            self._fetched = True
+            self._table_names = []
             return
             
-        table_names = [row[0] for row in tables_response.rows]
-        
-        # Fetch data from each table
-        for table in table_names:
-            # Get table schema
-            schema_response = self.resource.query(f"PRAGMA table_info({table})")
-            if schema_response.rows:
-                self._schemas[table] = [row[1] for row in schema_response.rows]  # Column names
-            
-            # Get all data
-            data_response = self.resource.query(f"SELECT * FROM {table}")
-            if data_response.rows and data_response.columns:
-                self._data[table] = [
-                    dict(zip(data_response.columns, row))
-                    for row in data_response.rows
-                ]
-            else:
-                self._data[table] = []
-                
-        self._fetched = True
+        self._table_names = [row[0] for row in tables_response.rows]
+        
+    def _ensure_table_data(self, table: str):
+        """Fetch data for a specific table on demand."""
+        if table in self._fetched_tables:
+            return
+            
+        # Get table schema
+        schema_response = self.resource.query(f"PRAGMA table_info({table})")
+        if schema_response.rows:
+            self._schemas[table] = [row[1] for row in schema_response.rows]  # Column names
+        
+        # Get all data for this table
+        data_response = self.resource.query(f"SELECT * FROM {table}")
+        if data_response.rows and data_response.columns:
+            self._data[table] = [
+                dict(zip(data_response.columns, row))
+                for row in data_response.rows
+            ]
+        else:
+            self._data[table] = []
+            
+        self._fetched_tables.add(table)
         
     def tables(self) -> list[str]:
         """Get list of all tables in the snapshot."""
-        self._ensure_fetched()
-        return list(self._data.keys())
+        self._ensure_tables_list()
+        return list(self._table_names) if self._table_names else []
         
     def table(self, table_name: str) -> "SyncSnapshotQueryBuilder":
         """Create a query builder for snapshot data."""
@@ -78,13 +82,12 @@ class SyncDatabaseSnapshot:
         ignore_config: IgnoreConfig | None = None,
     ) -> "SyncSnapshotDiff":
         """Compare this snapshot with another."""
-        self._ensure_fetched()
-        other._ensure_fetched()
+        # No need to fetch all data upfront - diff will fetch on demand
         return SyncSnapshotDiff(self, other, ignore_config)
 
 
 class SyncSnapshotQueryBuilder:
-    """Query builder that works on local snapshot data."""
+    """Query builder that works on snapshot data - can use targeted queries when possible."""
     
     def __init__(self, snapshot: SyncDatabaseSnapshot, table: str):
         self._snapshot = snapshot
@@ -94,10 +97,63 @@ class SyncSnapshotQueryBuilder:
         self._limit: int | None = None
         self._order_by: str | None = None
         self._order_desc: bool = False
+        self._use_targeted_query = True  # Try to use targeted queries when possible
+        
+    def _can_use_targeted_query(self) -> bool:
+        """Check if we can use a targeted query instead of loading all data."""
+        # We can use targeted query if:
+        # 1. We have simple equality conditions
+        # 2. No complex operations like joins
+        # 3. The query is selective (has conditions)
+        if not self._conditions:
+            return False
+        for col, op, val in self._conditions:
+            if op not in ["=", "IS", "IS NOT"]:
+                return False
+        return True
+        
+    def _execute_targeted_query(self) -> list[dict[str, Any]]:
+        """Execute a targeted query directly instead of loading all data."""
+        # Build WHERE clause
+        where_parts = []
+        for col, op, val in self._conditions:
+            if op == "=" and val is None:
+                where_parts.append(f"{col} IS NULL")
+            elif op == "IS":
+                where_parts.append(f"{col} IS NULL")
+            elif op == "IS NOT":
+                where_parts.append(f"{col} IS NOT NULL")
+            elif op == "=":
+                if isinstance(val, str):
+                    escaped_val = val.replace("'", "''")
+                    where_parts.append(f"{col} = '{escaped_val}'")
+                else:
+                    where_parts.append(f"{col} = '{val}'")
+        
+        where_clause = " AND ".join(where_parts)
+        
+        # Build full query
+        cols = ", ".join(self._select_cols)
+        query = f"SELECT {cols} FROM {self._table} WHERE {where_clause}"
+        
+        if self._order_by:
+            query += f" ORDER BY {self._order_by}"
+        if self._limit is not None:
+            query += f" LIMIT {self._limit}"
+            
+        # Execute query
+        response = self._snapshot.resource.query(query)
+        if response.rows and response.columns:
+            return [dict(zip(response.columns, row)) for row in response.rows]
+        return []
         
     def _get_data(self) -> list[dict[str, Any]]:
-        """Get table data from snapshot."""
-        self._snapshot._ensure_fetched()
+        """Get table data - use targeted query if possible, otherwise load all data."""
+        if self._use_targeted_query and self._can_use_targeted_query():
+            return self._execute_targeted_query()
+        
+        # Fall back to loading all data
+        self._snapshot._ensure_table_data(self._table)
         return self._snapshot._data.get(self._table, [])
         
     def eq(self, column: str, value: Any) -> "SyncSnapshotQueryBuilder":
@@ -121,6 +177,11 @@ class SyncSnapshotQueryBuilder:
         return rows[0] if rows else None
         
     def all(self) -> list[dict[str, Any]]:
+        # If we can use targeted query, _get_data already applies filters
+        if self._use_targeted_query and self._can_use_targeted_query():
+            return self._get_data()
+            
+        # Otherwise, get all data and apply filters manually
         data = self._get_data()
         
         # Apply filters
@@ -188,6 +249,7 @@ class SyncSnapshotDiff:
         self.after = after
         self.ignore_config = ignore_config or IgnoreConfig()
         self._cached: dict[str, Any] | None = None
+        self._targeted_mode = False  # Flag to use targeted queries
         
     def _get_primary_key_columns(self, table: str) -> list[str]:
         """Get primary key columns for a table."""
@@ -228,6 +290,10 @@ class SyncSnapshotDiff:
             # Get primary key columns
             pk_columns = self._get_primary_key_columns(tbl)
             
+            # Ensure data is fetched for this table
+            self.before._ensure_table_data(tbl)
+            self.after._ensure_table_data(tbl)
+            
             # Get data from both snapshots
             before_data = self.before._data.get(tbl, [])
             after_data = self.after._data.get(tbl, [])
@@ -303,10 +369,348 @@ class SyncSnapshotDiff:
         self._cached = diff
         return diff
         
-    def expect_only(self, allowed_changes: list[dict[str, Any]]):
-        """Ensure only specified changes occurred."""
+    def _can_use_targeted_queries(self, allowed_changes: list[dict[str, Any]]) -> bool:
+        """Check if we can use targeted queries for optimization."""
+        # We can use targeted queries if all allowed changes specify table and pk
+        for change in allowed_changes:
+            if "table" not in change or "pk" not in change:
+                return False
+        return True
+        
+    def _expect_only_targeted(self, allowed_changes: list[dict[str, Any]]):
+        """Optimized version that only queries specific rows mentioned in allowed_changes."""
+        import concurrent.futures
+        from threading import Lock
+        
+        # Group allowed changes by table
+        changes_by_table: dict[str, list[dict[str, Any]]] = {}
+        for change in allowed_changes:
+            table = change["table"]
+            if table not in changes_by_table:
+                changes_by_table[table] = []
+            changes_by_table[table].append(change)
+        
+        errors = []
+        errors_lock = Lock()
+        
+        # Function to check a single row
+        def check_row(table: str, pk: Any, table_changes: list[dict[str, Any]], pk_columns: list[str]):
+            try:
+                # Build WHERE clause for this PK
+                where_sql = self._build_pk_where_clause(pk_columns, pk)
+                
+                # Query before snapshot
+                before_query = f"SELECT * FROM {table} WHERE {where_sql}"
+                before_response = self.before.resource.query(before_query)
+                before_row = dict(zip(before_response.columns, before_response.rows[0])) if before_response.rows else None
+                
+                # Query after snapshot  
+                after_response = self.after.resource.query(before_query)
+                after_row = dict(zip(after_response.columns, after_response.rows[0])) if after_response.rows else None
+                
+                # Check changes for this row
+                if before_row and after_row:
+                    # Modified row - check fields
+                    for field in set(before_row.keys()) | set(after_row.keys()):
+                        if self.ignore_config.should_ignore_field(table, field):
+                            continue
+                        before_val = before_row.get(field)
+                        after_val = after_row.get(field)
+                        if not _values_equivalent(before_val, after_val):
+                            # Check if this change is allowed
+                            if not self._is_field_change_allowed(table_changes, pk, field, after_val):
+                                error_msg = (
+                                    f"Unexpected change in table '{table}', "
+                                    f"row {pk}, field '{field}': "
+                                    f"{repr(before_val)} -> {repr(after_val)}"
+                                )
+                                with errors_lock:
+                                    errors.append(AssertionError(error_msg))
+                                return  # Stop checking this row
+                elif not before_row and after_row:
+                    # Added row
+                    if not self._is_row_change_allowed(table_changes, pk, "__added__"):
+                        error_msg = f"Unexpected row added in table '{table}': {pk}"
+                        with errors_lock:
+                            errors.append(AssertionError(error_msg))
+                elif before_row and not after_row:
+                    # Removed row
+                    if not self._is_row_change_allowed(table_changes, pk, "__removed__"):
+                        error_msg = f"Unexpected row removed from table '{table}': {pk}"
+                        with errors_lock:
+                            errors.append(AssertionError(error_msg))
+            except Exception as e:
+                with errors_lock:
+                    errors.append(e)
+        
+        # Prepare all row checks
+        row_checks = []
+        for table, table_changes in changes_by_table.items():
+            if self.ignore_config.should_ignore_table(table):
+                continue
+                
+            # Get primary key columns once per table
+            pk_columns = self._get_primary_key_columns(table)
+            
+            # Extract unique PKs to check
+            pks_to_check = {change["pk"] for change in table_changes}
+            
+            for pk in pks_to_check:
+                row_checks.append((table, pk, table_changes, pk_columns))
+        
+        # Execute row checks in parallel
+        if row_checks:
+            with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+                futures = [
+                    executor.submit(check_row, table, pk, table_changes, pk_columns)
+                    for table, pk, table_changes, pk_columns in row_checks
+                ]
+                concurrent.futures.wait(futures)
+        
+        # Check for errors from row checks
+        if errors:
+            raise errors[0]
+        
+        # Now check tables not mentioned in allowed_changes to ensure no changes
+        all_tables = set(self.before.tables()) | set(self.after.tables())
+        tables_to_verify = []
+        
+        for table in all_tables:
+            if table not in changes_by_table and not self.ignore_config.should_ignore_table(table):
+                tables_to_verify.append(table)
+        
+        # Function to verify no changes in a table
+        def verify_no_changes(table: str):
+            try:
+                # For tables with no allowed changes, just check row counts
+                before_count_response = self.before.resource.query(f"SELECT COUNT(*) FROM {table}")
+                before_count = before_count_response.rows[0][0] if before_count_response.rows else 0
+                
+                after_count_response = self.after.resource.query(f"SELECT COUNT(*) FROM {table}")
+                after_count = after_count_response.rows[0][0] if after_count_response.rows else 0
+                
+                if before_count != after_count:
+                    error_msg = (
+                        f"Unexpected change in table '{table}': "
+                        f"row count changed from {before_count} to {after_count}"
+                    )
+                    with errors_lock:
+                        errors.append(AssertionError(error_msg))
+            except Exception as e:
+                with errors_lock:
+                    errors.append(e)
+        
+        # Execute table verification in parallel
+        if tables_to_verify:
+            with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+                futures = [
+                    executor.submit(verify_no_changes, table)
+                    for table in tables_to_verify
+                ]
+                concurrent.futures.wait(futures)
+        
+        # Final error check
+        if errors:
+            raise errors[0]
+        
+        return self
+        
+    def _build_pk_where_clause(self, pk_columns: list[str], pk_value: Any) -> str:
+        """Build WHERE clause for primary key lookup."""
+        # Escape single quotes in values to prevent SQL injection
+        def escape_value(val: Any) -> str:
+            if val is None:
+                return "NULL"
+            elif isinstance(val, str):
+                escaped = str(val).replace("'", "''")
+                return f"'{escaped}'"
+            else:
+                return f"'{val}'"
+                
+        if len(pk_columns) == 1:
+            return f"{pk_columns[0]} = {escape_value(pk_value)}"
+        else:
+            # Composite key
+            if isinstance(pk_value, tuple):
+                conditions = [f"{col} = {escape_value(val)}" for col, val in zip(pk_columns, pk_value)]
+                return " AND ".join(conditions)
+            else:
+                # Shouldn't happen if data is consistent
+                return f"{pk_columns[0]} = {escape_value(pk_value)}"
+                
+    def _is_field_change_allowed(self, table_changes: list[dict[str, Any]], pk: Any, field: str, after_val: Any) -> bool:
+        """Check if a specific field change is allowed."""
+        for change in table_changes:
+            if (str(change.get("pk")) == str(pk) and 
+                change.get("field") == field and
+                _values_equivalent(change.get("after"), after_val)):
+                return True
+        return False
+        
+    def _is_row_change_allowed(self, table_changes: list[dict[str, Any]], pk: Any, change_type: str) -> bool:
+        """Check if a row addition/deletion is allowed."""
+        for change in table_changes:
+            if str(change.get("pk")) == str(pk) and change.get("after") == change_type:
+                return True
+        return False
+        
+    def _expect_no_changes(self):
+        """Efficiently verify that no changes occurred between snapshots using row counts."""
+        try:
+            import concurrent.futures
+            from threading import Lock
+            
+            # Get all tables from both snapshots
+            before_tables = set(self.before.tables())
+            after_tables = set(self.after.tables())
+            
+            # Check for added/removed tables (excluding ignored ones)
+            added_tables = after_tables - before_tables
+            removed_tables = before_tables - after_tables
+            
+            for table in added_tables:
+                if not self.ignore_config.should_ignore_table(table):
+                    raise AssertionError(f"Unexpected table added: {table}")
+                    
+            for table in removed_tables:
+                if not self.ignore_config.should_ignore_table(table):
+                    raise AssertionError(f"Unexpected table removed: {table}")
+            
+            # Prepare tables to check
+            tables_to_check = []
+            all_tables = before_tables | after_tables
+            for table in all_tables:
+                if not self.ignore_config.should_ignore_table(table):
+                    tables_to_check.append(table)
+            
+            # If no tables to check, we're done
+            if not tables_to_check:
+                return self
+            
+            # Use ThreadPoolExecutor to parallelize count queries
+            # We use threads instead of processes since the queries are I/O bound
+            errors = []
+            errors_lock = Lock()
+            tables_needing_verification = []
+            verification_lock = Lock()
+            
+            def check_table_counts(table: str):
+                """Check row counts for a single table."""
+                try:
+                    # Get row counts from both snapshots
+                    before_count = 0
+                    after_count = 0
+                    
+                    if table in before_tables:
+                        before_count_response = self.before.resource.query(f"SELECT COUNT(*) FROM {table}")
+                        before_count = before_count_response.rows[0][0] if before_count_response.rows else 0
+                        
+                    if table in after_tables:
+                        after_count_response = self.after.resource.query(f"SELECT COUNT(*) FROM {table}")
+                        after_count = after_count_response.rows[0][0] if after_count_response.rows else 0
+                    
+                    if before_count != after_count:
+                        error_msg = (
+                            f"Unexpected change in table '{table}': "
+                            f"row count changed from {before_count} to {after_count}"
+                        )
+                        with errors_lock:
+                            errors.append(AssertionError(error_msg))
+                    elif before_count > 0 and before_count <= 1000:
+                        # Mark for detailed verification
+                        with verification_lock:
+                            tables_needing_verification.append(table)
+                            
+                except Exception as e:
+                    with errors_lock:
+                        errors.append(e)
+            
+            # Execute count checks in parallel
+            with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+                futures = [executor.submit(check_table_counts, table) for table in tables_to_check]
+                concurrent.futures.wait(futures)
+            
+            # Check if any errors occurred during count checking
+            if errors:
+                # Raise the first error
+                raise errors[0]
+            
+            # Now verify small tables for data changes (also in parallel)
+            if tables_needing_verification:
+                verification_errors = []
+                
+                def verify_table(table: str):
+                    """Verify a single table's data hasn't changed."""
+                    try:
+                        self._verify_table_unchanged(table)
+                    except AssertionError as e:
+                        with errors_lock:
+                            verification_errors.append(e)
+                
+                with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
+                    futures = [executor.submit(verify_table, table) for table in tables_needing_verification]
+                    concurrent.futures.wait(futures)
+                
+                # Check if any errors occurred during verification
+                if verification_errors:
+                    raise verification_errors[0]
+                    
+            return self
+            
+        except AssertionError:
+            # Re-raise assertion errors (these are expected failures)
+            raise
+        except Exception as e:
+            # If the optimized check fails for other reasons, fall back to full diff
+            print(f"Warning: Optimized no-changes check failed: {e}")
+            print("Falling back to full diff...")
+            return self._expect_only_fallback([])
+            
+    def _verify_table_unchanged(self, table: str):
+        """Verify that a table's data hasn't changed (for small tables)."""
+        # Get primary key columns
+        pk_columns = self._get_primary_key_columns(table)
+        
+        # Get sorted data from both snapshots
+        order_by = ", ".join(pk_columns) if pk_columns else "rowid"
+        
+        before_response = self.before.resource.query(f"SELECT * FROM {table} ORDER BY {order_by}")
+        after_response = self.after.resource.query(f"SELECT * FROM {table} ORDER BY {order_by}")
+        
+        # Quick check: if column counts differ, there's a schema change
+        if before_response.columns != after_response.columns:
+            raise AssertionError(f"Schema changed in table '{table}'")
+            
+        # Compare row by row
+        if len(before_response.rows) != len(after_response.rows):
+            raise AssertionError(
+                f"Row count mismatch in table '{table}': "
+                f"{len(before_response.rows)} vs {len(after_response.rows)}"
+            )
+            
+        for i, (before_row, after_row) in enumerate(zip(before_response.rows, after_response.rows)):
+            before_dict = dict(zip(before_response.columns, before_row))
+            after_dict = dict(zip(after_response.columns, after_row))
+            
+            # Compare fields, ignoring those in ignore config
+            for field in before_response.columns:
+                if self.ignore_config.should_ignore_field(table, field):
+                    continue
+                    
+                if not _values_equivalent(before_dict.get(field), after_dict.get(field)):
+                    pk_val = before_dict.get(pk_columns[0]) if pk_columns else i
+                    raise AssertionError(
+                        f"Unexpected change in table '{table}', row {pk_val}, "
+                        f"field '{field}': {repr(before_dict.get(field))} -> {repr(after_dict.get(field))}"
+                    )
+                    
+    def _expect_only_fallback(self, allowed_changes: list[dict[str, Any]]):
+        """Fallback to full diff collection when optimized methods fail."""
         diff = self._collect()
+        return self._validate_diff_against_allowed_changes(diff, allowed_changes)
         
+    def _validate_diff_against_allowed_changes(self, diff: dict[str, Any], allowed_changes: list[dict[str, Any]]):
+        """Validate a collected diff against allowed changes."""
         def _is_change_allowed(
             table: str, row_id: Any, field: str | None, after_value: Any
         ) -> bool:
@@ -419,6 +823,20 @@ class SyncSnapshotDiff:
             raise AssertionError("\n".join(error_lines))
             
         return self
+        
+    def expect_only(self, allowed_changes: list[dict[str, Any]]):
+        """Ensure only specified changes occurred."""
+        # Special case: empty allowed_changes means no changes should have occurred
+        if not allowed_changes:
+            return self._expect_no_changes()
+            
+        # For expect_only, we can optimize by only checking the specific rows mentioned
+        if self._can_use_targeted_queries(allowed_changes):
+            return self._expect_only_targeted(allowed_changes)
+        
+        # Fall back to full diff for complex cases
+        diff = self._collect()
+        return self._validate_diff_against_allowed_changes(diff, allowed_changes)
 
 
 class SyncQueryBuilder:
@@ -669,9 +1087,8 @@ class SQLiteResource(Resource):
 
     def snapshot(self, name: str | None = None) -> SyncDatabaseSnapshot:
         """Create a snapshot of the current database state."""
-        snapshot = SyncDatabaseSnapshot(self, name)
-        snapshot._ensure_fetched()
-        return snapshot
+        # No longer fetch all data upfront - let it be lazy
+        return SyncDatabaseSnapshot(self, name)
 
     def diff(
         self,

{fleet_python-0.2.18 → fleet_python-0.2.20}/fleet_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fleet-python
-Version: 0.2.18
+Version: 0.2.20
 Summary: Python SDK for Fleet environments
 Author-email: Fleet AI <nic@fleet.so>
 License: Apache-2.0

{fleet_python-0.2.18 → fleet_python-0.2.20}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fleet-python"
-version = "0.2.18"
+version = "0.2.20"
 description = "Python SDK for Fleet environments"
 authors = [
     {name = "Fleet AI", email = "nic@fleet.so"},