Merge pull request #176 from poissoncorp/RDBC-685

RDBC-685 Python counters - session #2

Author: ml054

ravendb/documents/session/in_memory_document_session_operations.py

@@ -459,7 +459,7 @@ def __init__(self, store: "DocumentStore", key: uuid.UUID, options: SessionOptio
         self._included_documents_by_id = CaseInsensitiveDict()
         self._documents_by_entity: DocumentsByEntityHolder = DocumentsByEntityHolder()
 
-        self._counters_by_doc_id: Dict[str, List[Dict[str, int]]] = {}
+        self._counters_by_doc_id: Dict[str, List[Dict[str, int]]] = CaseInsensitiveDict()
         self._time_series_by_doc_id: Dict[str, Dict[str, List[TimeSeriesRangeResult]]] = {}
 
         self._deleted_entities: Union[
@@ -1244,38 +1244,44 @@ def __include(key):
                     continue
                 IncludesUtil.include(result, include, __include)
 
+    def register_counters_ids_list(
+        self, result_counters: dict, keys: List[str], counters_to_include: List[str], got_all: bool
+    ) -> None:
+        if self.no_tracking:
+            return
+
+        if not result_counters:
+            if got_all:
+                for key in keys:
+                    self.__set_got_all_counters_for_document(key)
+
+                return
+        else:
+            self._register_counters_internal(result_counters, None, False, got_all)
+
+        self.__register_missing_counters_for_keys(keys, counters_to_include)
+
     def register_counters(
         self,
         result_counters: dict,
-        counters_to_include: Union[Dict[str, List[str]], List[str]],
-        keys: List[str] = None,
-        got_all: bool = None,
-    ):
-        keys_case = keys is not None and got_all is not None
+        counters_to_include: Dict[str, List[str]],
+    ) -> None:
         if self.no_tracking:
             return
+
         if not result_counters:
-            if keys_case:
-                if got_all:
-                    for key in keys:
-                        self.__set_got_all_counters_for_document(key)
-                    return
-            else:
-                self.__set_got_all_in_cache_if_needed(counters_to_include)
+            self.__set_got_all_in_cache_if_needed(counters_to_include)
         else:
-            self._register_counters_internal(
-                result_counters,
-                {} if keys_case else counters_to_include,
-                not keys_case,
-                got_all if keys_case else False,
-            )
+            self._register_counters_internal(result_counters, counters_to_include, True, False)
 
-        self.__register_missing_counters_for_keys(
-            keys, counters_to_include
-        ) if keys_case else self.__register_missing_counters(counters_to_include)
+        self.__register_missing_counters(counters_to_include)
 
     def _register_counters_internal(
-        self, result_counters: dict, counters_to_include: Dict[str, List[str]], from_query_result: bool, got_all: bool
+        self,
+        result_counters: dict,
+        counters_to_include: Optional[Dict[str, List[str]]],
+        from_query_result: bool,
+        got_all: bool,
     ):
         for key, result_counters in result_counters.items():
             if not result_counters:
@@ -1361,8 +1367,8 @@ def __register_missing_counters_for_keys(self, keys: List[str], counters_to_incl
 
         for counter in counters_to_include:
             for key in keys:
-                cache = self._counters_by_doc_id.get(key)
-                if not cache:
+                cache = self._counters_by_doc_id.get(key, None)
+                if cache is None:
                     cache = [False, CaseInsensitiveDict()]
                     self._counters_by_doc_id[key] = cache
 

ravendb/documents/session/loaders/include.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import datetime
-from typing import Set, Tuple, Dict, Union, Optional
+from typing import Set, Tuple, Dict, Union, Optional, List
 from ravendb.documents.conventions import DocumentConventions
 from ravendb.data.timeseries import TimeSeriesRange
 from ravendb.tools.utils import CaseInsensitiveDict, CaseInsensitiveSet
@@ -43,24 +43,21 @@ def _include_compare_exchange_value(self, path: str) -> None:
 
     def _include_counter_with_alias(self, path: str, *names: str) -> None:
         self._with_alias()
-        self._include_counters(path, *names)
+        self._include_counters(path, list(names) if isinstance(names, tuple) else [names])
 
     def _include_documents(self, path: str):
         if not self._documents_to_include:
             self._documents_to_include = set()
         self._documents_to_include.add(path)
 
-    def _include_counters(self, path: str, *names: str):
+    def _include_counters(self, path: str, names: List[str]):
         if not names:
             raise ValueError("Expected at least one name")
 
         self._assert_not_all_and_add_new_entry_if_needed(path)
         for name in names:
             if name.isspace():
                 raise ValueError("Counters(names): 'names' should not contain null or whitespace elements")
-            if isinstance(name, list):
-                self._include_counters(path, *name)
-                continue
             self._counters_to_include_by_source_path.get(path)[1].add(name)
 
     def _include_all_counters_with_alias(self, path: str):
@@ -72,14 +69,14 @@ def _include_all_counters(self, source_path: str):
             self._counters_to_include_by_source_path = CaseInsensitiveDict()
         val = self._counters_to_include_by_source_path.get(source_path)
         if val is not None and val[1] is not None:
-            raise ValueError("You cannot use all_counters() after using counters(*names)")
+            raise RuntimeError("You cannot use all_counters() after using counters(*names)")
         self._counters_to_include_by_source_path[source_path] = (True, None)
 
     def _assert_not_all_and_add_new_entry_if_needed(self, path: str):
         if self._counters_to_include_by_source_path is not None:
             val = self._counters_to_include_by_source_path.get(path, None)
             if val is not None and val[0]:
-                raise ValueError("You cannot use counters(*names) after using all_counters()")
+                raise RuntimeError("You cannot use counters(*names) after using all_counters()")
 
         if self._counters_to_include_by_source_path is None:
             self._counters_to_include_by_source_path = CaseInsensitiveDict()
@@ -116,8 +113,13 @@ def include_documents(self, path: str) -> IncludeBuilderBase:
         self._include_documents(path)
         return self
 
-    def include_counter(self, *names: str) -> IncludeBuilderBase:
-        self._include_counters("", *names)
+    def include_counter(self, name: str) -> IncludeBuilderBase:
+        self._include_counters("", [name])
+        return self
+
+    def include_counters(self, *names: str) -> IncludeBuilderBase:
+        for name in names:
+            self.include_counter(name)
         return self
 
     def include_all_counters(self) -> IncludeBuilderBase:
@@ -140,8 +142,12 @@ def include_compare_exchange_value(self, path: str) -> IncludeBuilderBase:
 
 
 class QueryIncludeBuilder(IncludeBuilderBase):
+    def include_counter(self, name: str, path: Optional[str] = ""):
+        self._include_counter_with_alias(path, name)
+        return self
+
     def include_counters(self, *names: str, path: Optional[str] = ""):
-        self._include_counter_with_alias(path, *names)
+        self._include_counter_with_alias(path, *names if isinstance(names, tuple) else names)
         return self
 
     def include_all_counters(self, path: Optional[str] = ""):

ravendb/documents/session/operations/load_operation.py

@@ -181,7 +181,7 @@ def set_result(self, result: GetDocumentsResult) -> None:
         self._session.register_includes(result.includes)
 
         if self._include_all_counters or self._counters_to_include:
-            self._session.register_counters(
+            self._session.register_counters_ids_list(
                 result.counter_includes, self._keys, self._counters_to_include, self._include_all_counters
             )
 

ravendb/documents/session/operations/query.py

@@ -130,10 +130,10 @@ def __complete_internal(
                 query_result.results, query_result.includes, query_result.included_paths
             )
 
-            # todo: counters and timeseries
             if query_result.counter_includes is not None:
-                raise NotImplementedError()
+                self.__session.register_counters(query_result.counter_includes, query_result.included_counter_names)
 
+            # todo: timeseries
             if query_result.time_series_includes is not None:
                 raise NotImplementedError()