added flag to do consistent read if needed / possible for DB (#1539)

Author: ilongin

src/datachain/data_storage/db_engine.py

@@ -77,6 +77,7 @@ def execute(
         query,
         cursor: Any | None = None,
         conn: Any | None = None,
+        consistent_read: bool = False,
     ) -> Iterator[tuple[Any, ...]]: ...
 
     def get_table(self, name: str) -> "Table":

src/datachain/data_storage/sqlite.py

@@ -230,6 +230,7 @@ def execute(
         query,
         cursor: sqlite3.Cursor | None = None,
         conn=None,
+        consistent_read: bool = False,
     ) -> sqlite3.Cursor:
         if self.is_closed:
             # Reconnect in case of being closed previously.

src/datachain/data_storage/warehouse.py

@@ -237,9 +237,9 @@ def query_count(self, query: sa.Select) -> int:
         count_query = sa.select(sa.func.count(1)).select_from(query.subquery())
         return next(self.db.execute(count_query))[0]
 
-    def table_rows_count(self, table) -> int:
+    def table_rows_count(self, table, consistent_read: bool = False) -> int:
         count_query = sa.select(sa.func.count(1)).select_from(table)
-        return next(self.db.execute(count_query))[0]
+        return next(self.db.execute(count_query, consistent_read=consistent_read))[0]
 
     def dataset_select_paginated(
         self,
@@ -494,7 +494,7 @@ def dataset_stats(
         if size_columns:
             expressions = (*expressions, sa.func.sum(sum(size_columns)))
         query = sa.select(*expressions)
-        ((nrows, *rest),) = self.db.execute(query)
+        ((nrows, *rest),) = self.db.execute(query, consistent_read=True)
         return nrows, rest[0] if rest else 0
 
     @abstractmethod

src/datachain/query/dataset.py

@@ -1509,6 +1509,8 @@ def as_iterable(self, **kwargs) -> Iterator[ResultIter]:
         try:
             query = self.apply_steps().select()
             selected_columns = [c.name for c in query.selected_columns]
+            if "consistent_read" not in kwargs:
+                kwargs["consistent_read"] = True
             yield ResultIter(
                 self.catalog.warehouse.dataset_rows_select(query, **kwargs),
                 selected_columns,