fleet-python 0.2.21__py3-none-any.whl → 0.2.23__py3-none-any.whl

fleet/resources/sqlite.py

@@ -14,11 +14,11 @@ if TYPE_CHECKING:
 
 
 # Import types from verifiers module
-from ..verifiers.db import IgnoreConfig, _get_row_identifier, _format_row_for_error, _values_equivalent
+from fleet.verifiers.db import IgnoreConfig, _get_row_identifier, _format_row_for_error, _values_equivalent
 
 
 class SyncDatabaseSnapshot:
-    """Lazy database snapshot that fetches data on-demand through API."""
+    """Async database snapshot that fetches data through API and stores locally for diffing."""
     
     def __init__(self, resource: "SQLiteResource", name: str | None = None):
         self.resource = resource
@@ -26,12 +26,11 @@ class SyncDatabaseSnapshot:
         self.created_at = datetime.utcnow()
         self._data: dict[str, list[dict[str, Any]]] = {}
         self._schemas: dict[str, list[str]] = {}
-        self._table_names: list[str] | None = None
-        self._fetched_tables: set[str] = set()
+        self._fetched = False
         
-    def _ensure_tables_list(self):
-        """Fetch just the list of table names if not already fetched."""
-        if self._table_names is not None:
+    def _ensure_fetched(self):
+        """Fetch all data from remote database if not already fetched."""
+        if self._fetched:
             return
             
         # Get all tables
@@ -40,37 +39,34 @@ class SyncDatabaseSnapshot:
         )
         
         if not tables_response.rows:
-            self._table_names = []
+            self._fetched = True
             return
             
-        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)
+        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
         
     def tables(self) -> list[str]:
         """Get list of all tables in the snapshot."""
-        self._ensure_tables_list()
-        return list(self._table_names) if self._table_names else []
+        self._ensure_fetched()
+        return list(self._data.keys())
         
     def table(self, table_name: str) -> "SyncSnapshotQueryBuilder":
         """Create a query builder for snapshot data."""
@@ -82,12 +78,13 @@ class SyncDatabaseSnapshot:
         ignore_config: IgnoreConfig | None = None,
     ) -> "SyncSnapshotDiff":
         """Compare this snapshot with another."""
-        # No need to fetch all data upfront - diff will fetch on demand
+        self._ensure_fetched()
+        other._ensure_fetched()
         return SyncSnapshotDiff(self, other, ignore_config)
 
 
 class SyncSnapshotQueryBuilder:
-    """Query builder that works on snapshot data - can use targeted queries when possible."""
+    """Query builder that works on local snapshot data."""
     
     def __init__(self, snapshot: SyncDatabaseSnapshot, table: str):
         self._snapshot = snapshot
@@ -97,63 +94,10 @@ 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 - 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)
+        """Get table data from snapshot."""
+        self._snapshot._ensure_fetched()
         return self._snapshot._data.get(self._table, [])
         
     def eq(self, column: str, value: Any) -> "SyncSnapshotQueryBuilder":
@@ -177,11 +121,6 @@ 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
@@ -249,7 +188,6 @@ 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."""
@@ -290,10 +228,6 @@ 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, [])
@@ -369,348 +303,10 @@ class SyncSnapshotDiff:
         self._cached = diff
         return diff
         
-    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."""
+    def expect_only(self, allowed_changes: list[dict[str, Any]]):
+        """Ensure only specified changes occurred."""
         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:
@@ -823,20 +419,6 @@ 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:
@@ -1087,22 +669,23 @@ class SQLiteResource(Resource):
 
     def snapshot(self, name: str | None = None) -> SyncDatabaseSnapshot:
         """Create a snapshot of the current database state."""
-        # No longer fetch all data upfront - let it be lazy
-        return SyncDatabaseSnapshot(self, name)
+        snapshot = SyncDatabaseSnapshot(self, name)
+        snapshot._ensure_fetched()
+        return snapshot
 
     def diff(
         self,
         other: "SQLiteResource",
         ignore_config: IgnoreConfig | None = None,
     ) -> SyncSnapshotDiff:
-        """Compare this database with another SQLiteResource.
+        """Compare this database with another AsyncSQLiteResource.
         
         Args:
-            other: Another SQLiteResource to compare against
+            other: Another AsyncSQLiteResource to compare against
             ignore_config: Optional configuration for ignoring specific tables/fields
             
         Returns:
-            SyncSnapshotDiff: Object containing the differences between the two databases
+            AsyncSnapshotDiff: Object containing the differences between the two databases
         """
         # Create snapshots of both databases
         before_snapshot = self.snapshot(name=f"before_{datetime.utcnow().isoformat()}")

fleet/tasks.py

@@ -10,7 +10,7 @@ from uuid import UUID
 from pydantic import BaseModel, Field, validator
 
 # Import the shared VerifierFunction type that works for both async and sync
-from ..types import VerifierFunction
+from fleet.types import VerifierFunction
 
 
 class Task(BaseModel):
@@ -41,4 +41,4 @@ class Task(BaseModel):
             datetime: lambda v: v.isoformat(),
         }
         # Allow arbitrary types for the verifier field
-        arbitrary_types_allowed = True 
+        arbitrary_types_allowed = True