11from __future__ import annotations
22
3- from typing import TYPE_CHECKING , Any , Optional , Union
3+ import warnings
4+ from typing import TYPE_CHECKING , Any , Literal , Optional , Union , overload
45from urllib .parse import urlencode
56
67from pydantic import BaseModel
@@ -72,14 +73,44 @@ class InfrahubBranchManager(InfraHubBranchManagerBase):
7273 def __init__ (self , client : InfrahubClient ):
7374 self .client = client
7475
76+ @overload
7577 async def create (
7678 self ,
7779 branch_name : str ,
7880 sync_with_git : bool = True ,
7981 description : str = "" ,
80- background_execution : bool = False ,
81- ) -> BranchData :
82+ wait_until_completion : Literal [True ] = True ,
83+ background_execution : Optional [bool ] = False ,
84+ ) -> BranchData : ...
85+
86+ @overload
87+ async def create (
88+ self ,
89+ branch_name : str ,
90+ sync_with_git : bool = True ,
91+ description : str = "" ,
92+ wait_until_completion : Literal [False ] = False ,
93+ background_execution : Optional [bool ] = False ,
94+ ) -> str : ...
95+
96+ async def create (
97+ self ,
98+ branch_name : str ,
99+ sync_with_git : bool = True ,
100+ description : str = "" ,
101+ wait_until_completion : bool = True ,
102+ background_execution : Optional [bool ] = False ,
103+ ) -> Union [BranchData , str ]:
104+ if background_execution is not None :
105+ warnings .warn (
106+ "`background_execution` is deprecated, please use `wait_until_completion` instead." ,
107+ DeprecationWarning ,
108+ stacklevel = 1 ,
109+ )
110+
111+ background_execution = background_execution or not wait_until_completion
82112 input_data = {
113+ # Should be switched to `wait_until_completion` once `background_execution` is removed server side.
83114 "background_execution" : background_execution ,
84115 "data" : {
85116 "name" : branch_name ,
@@ -91,6 +122,10 @@ async def create(
91122 query = Mutation (mutation = "BranchCreate" , input_data = input_data , query = MUTATION_QUERY_DATA )
92123 response = await self .client .execute_graphql (query = query .render (), tracker = "mutation-branch-create" )
93124
125+ # Make sure server version is recent enough to support background execution, as previously
126+ # using background_execution=True had no effect.
127+ if background_execution and "task" in response ["BranchCreate" ]:
128+ return BranchData (** response ["BranchCreate" ]["task" ]["id" ])
94129 return BranchData (** response ["BranchCreate" ]["object" ])
95130
96131 async def delete (self , branch_name : str ) -> bool :
@@ -209,14 +244,44 @@ def get(self, branch_name: str) -> BranchData:
209244 raise BranchNotFoundError (identifier = branch_name )
210245 return BranchData (** data ["Branch" ][0 ])
211246
247+ @overload
248+ def create (
249+ self ,
250+ branch_name : str ,
251+ sync_with_git : bool = True ,
252+ description : str = "" ,
253+ wait_until_completion : Literal [True ] = True ,
254+ background_execution : Optional [bool ] = False ,
255+ ) -> BranchData : ...
256+
257+ @overload
258+ def create (
259+ self ,
260+ branch_name : str ,
261+ sync_with_git : bool = True ,
262+ description : str = "" ,
263+ wait_until_completion : Literal [False ] = False ,
264+ background_execution : Optional [bool ] = False ,
265+ ) -> str : ...
266+
212267 def create (
213268 self ,
214269 branch_name : str ,
215270 sync_with_git : bool = True ,
216271 description : str = "" ,
217- background_execution : bool = False ,
218- ) -> BranchData :
272+ wait_until_completion : bool = True ,
273+ background_execution : Optional [bool ] = False ,
274+ ) -> Union [BranchData , str ]:
275+ if background_execution is not None :
276+ warnings .warn (
277+ "`background_execution` is deprecated, please use `wait_until_completion` instead." ,
278+ DeprecationWarning ,
279+ stacklevel = 1 ,
280+ )
281+
282+ background_execution = background_execution or not wait_until_completion
219283 input_data = {
284+ # Should be switched to `wait_until_completion` once `background_execution` is removed server side.
220285 "background_execution" : background_execution ,
221286 "data" : {
222287 "name" : branch_name ,
@@ -228,6 +293,10 @@ def create(
228293 query = Mutation (mutation = "BranchCreate" , input_data = input_data , query = MUTATION_QUERY_DATA )
229294 response = self .client .execute_graphql (query = query .render (), tracker = "mutation-branch-create" )
230295
296+ # Make sure server version is recent enough to support background execution, as previously
297+ # using background_execution=True had no effect.
298+ if background_execution and "task" in response ["BranchCreate" ]:
299+ return BranchData (** response ["BranchCreate" ]["task" ]["id" ])
231300 return BranchData (** response ["BranchCreate" ]["object" ])
232301
233302 def delete (self , branch_name : str ) -> bool :
0 commit comments