Tighten branch time helpers to require Timestamp inputs (#7608)

* Tighten branch query helpers to Timestamp inputs and is related to issue #63. API tightening on get_branches_and_times_to_query[_global], with tests adjusted accordingly.

* refactor: require Timestamp inputs for branch range filters. This captures that both get_query_filter_range and get_query_filter_relationships_range now enforce Timestamp arguments only.

* refactor: make Branch.get_query_filter_relationships Timestamp-only

* Altered spacing and added towncrier notes as well

Author: iddocohen

backend/infrahub/core/branch/models.py

@@ -206,7 +206,7 @@ def get_branches_in_scope(self) -> list[str]:
 
         return [default_branch, self.name]
 
-    def get_branches_and_times_to_query(self, at: Optional[Union[Timestamp, str]] = None) -> dict[frozenset, str]:
+    def get_branches_and_times_to_query(self, at: Optional[Timestamp] = None) -> dict[frozenset, str]:
         """Return all the names of the branches that are constituing this branch with the associated times excluding the global branch"""
 
         at = Timestamp(at)
@@ -227,7 +227,7 @@ def get_branches_and_times_to_query(self, at: Optional[Union[Timestamp, str]] =
 
     def get_branches_and_times_to_query_global(
         self,
-        at: Optional[Union[Timestamp, str]] = None,
+        at: Optional[Timestamp] = None,
         is_isolated: bool = True,
     ) -> dict[frozenset, str]:
         """Return all the names of the branches that are constituting this branch with the associated times."""
@@ -297,7 +297,7 @@ async def delete(self, db: InfrahubDatabase) -> None:
         await super().delete(db=db)
 
     def get_query_filter_relationships(
-        self, rel_labels: list, at: Optional[Union[Timestamp, str]] = None, include_outside_parentheses: bool = False
+        self, rel_labels: list, at: Optional[Timestamp] = None, include_outside_parentheses: bool = False
     ) -> tuple[list, dict]:
         """
         Generate a CYPHER Query filter based on a list of relationships to query a part of the graph at a specific time and on a specific branch.
@@ -360,7 +360,7 @@ def get_query_filter_path(
             params[f"{pp}time1"] = at_str
             return filter_str, params
 
-        branches_times = self.get_branches_and_times_to_query_global(at=at_str, is_isolated=is_isolated)
+        branches_times = self.get_branches_and_times_to_query_global(at=at, is_isolated=is_isolated)
 
         for idx, (branch_name, time_to_query) in enumerate(branches_times.items()):
             params[f"{pp}branch{idx}"] = list(branch_name)
@@ -382,8 +382,8 @@ def get_query_filter_path(
     def get_query_filter_relationships_range(
         self,
         rel_labels: list,
-        start_time: Union[Timestamp, str],
-        end_time: Union[Timestamp, str],
+        start_time: Timestamp,
+        end_time: Timestamp,
         include_outside_parentheses: bool = False,
         include_global: bool = False,
     ) -> tuple[list, dict]:
@@ -463,9 +463,7 @@ def get_query_filter_relationships_diff(
 
         return filters, params
 
-    def get_query_filter_range(
-        self, rel_label: list, start_time: Union[Timestamp, str], end_time: Union[Timestamp, str]
-    ) -> tuple[list, dict]:
+    def get_query_filter_range(self, rel_label: list, start_time: Timestamp, end_time: Timestamp) -> tuple[list, dict]:
         """
         Generate a CYPHER Query filter to query a range of values in the graph between start_time and end_time."""
 

backend/infrahub/core/query/relationship.py

@@ -871,7 +871,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs) -> None:  # noqa: ARG
         self.params["branch"] = self.branch.name
 
         rels_filter, rels_params = self.branch.get_query_filter_relationships(
-            rel_labels=["r1", "r2"], at=self.at.to_string(), include_outside_parentheses=True
+            rel_labels=["r1", "r2"], at=self.at, include_outside_parentheses=True
         )
 
         self.params.update(rels_params)
@@ -954,7 +954,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs) -> None:  # noqa: ARG
         self.params["at"] = self.at.to_string()
 
         rels_filter, rels_params = self.branch.get_query_filter_relationships(
-            rel_labels=["r1", "r2"], at=self.at.to_string(), include_outside_parentheses=True
+            rel_labels=["r1", "r2"], at=self.at, include_outside_parentheses=True
         )
         self.params.update(rels_params)
 

backend/tests/unit/core/test_branch.py

@@ -120,7 +120,7 @@ async def test_get_query_filter_relationships_main(db: InfrahubDatabase, base_da
     default_branch = await registry.get_branch(branch="main", db=db)
 
     filters, params = default_branch.get_query_filter_relationships(
-        rel_labels=["r1", "r2"], at=Timestamp().to_string(), include_outside_parentheses=False
+        rel_labels=["r1", "r2"], at=Timestamp(), include_outside_parentheses=False
     )
 
     expected_filters = [
@@ -139,7 +139,7 @@ async def test_get_query_filter_relationships_branch1(db: InfrahubDatabase, base
     branch1 = await registry.get_branch(branch="branch1", db=db)
 
     filters, params = branch1.get_query_filter_relationships(
-        rel_labels=["r1", "r2"], at=Timestamp().to_string(), include_outside_parentheses=False
+        rel_labels=["r1", "r2"], at=Timestamp(), include_outside_parentheses=False
     )
 
     assert isinstance(filters, list)
@@ -157,7 +157,7 @@ async def test_get_branches_and_times_to_query_main(db: InfrahubDatabase, base_d
     assert Timestamp(results[frozenset(["main"])]) > now
 
     t1 = Timestamp("2s")
-    results = main_branch.get_branches_and_times_to_query(at=t1.to_string())
+    results = main_branch.get_branches_and_times_to_query(at=t1)
     assert results[frozenset(["main"])] == t1.to_string()
 
 
@@ -186,7 +186,7 @@ async def test_get_branches_and_times_to_query_global_main(db: InfrahubDatabase,
     assert Timestamp(results[frozenset((GLOBAL_BRANCH_NAME, "main"))]) > now
 
     t1 = Timestamp("2s")
-    results = main_branch.get_branches_and_times_to_query_global(at=t1.to_string())
+    results = main_branch.get_branches_and_times_to_query_global(at=t1)
     assert results[frozenset((GLOBAL_BRANCH_NAME, "main"))] == t1.to_string()
 
 
@@ -201,7 +201,7 @@ async def test_get_branches_and_times_to_query_global_branch1(db: InfrahubDataba
     assert results[frozenset((GLOBAL_BRANCH_NAME, "main"))] == base_dataset_02["time_m45"]
 
     t1 = Timestamp("2s")
-    results = branch1.get_branches_and_times_to_query_global(at=t1.to_string())
+    results = branch1.get_branches_and_times_to_query_global(at=t1)
     assert results[frozenset((GLOBAL_BRANCH_NAME, "branch1"))] == t1.to_string()
     assert results[frozenset((GLOBAL_BRANCH_NAME, "main"))] == base_dataset_02["time_m45"]
 

changelog/63.changed.md

@@ -0,0 +1 @@
+This refactors timestamp handling across query and relationship functions to accept only Timestamp types, updates all dependent code and tests accordingly, refreshes development environment files, and includes no user-facing changes.