jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

🔧 fix(custom_component.py): change list_flows and get_flow methods to accept an optional get_session parameter for better flexibility and testability

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-07-31 17:28:29 -03:00
commit 2e7b35ddd5
1 changed files with 7 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

10
src/backend/langflow/interface/custom/custom_component.py
View file

@ -153,8 +153,9 @@ class CustomComponent(Component, extra=Extra.allow):
graph_data = process_tweaks(graph_data=graph_data, tweaks=tweaks) graph_data = process_tweaks(graph_data=graph_data, tweaks=tweaks)
return build_sorted_vertices_with_caching(graph_data) return build_sorted_vertices_with_caching(graph_data)
def list_flows(self) -> List[Flow]: def list_flows(self, *, get_session: Optional[Callable] = None) -> List[Flow]:
with session_getter() as session: get_session = get_session or session_getter
with get_session() as session:
flows = session.query(Flow).all() flows = session.query(Flow).all()
return flows return flows
@ -164,8 +165,11 @@ class CustomComponent(Component, extra=Extra.allow):
flow_name: Optional[str] = None, flow_name: Optional[str] = None,
flow_id: Optional[str] = None, flow_id: Optional[str] = None,
tweaks: Optional[dict] = None, tweaks: Optional[dict] = None,
get_session: Optional[Callable] = None,
) -> Flow: ) -> Flow:
with session_getter() as session: get_session = get_session or session_getter
with get_session() as session:
if flow_id: if flow_id:
flow = session.query(Flow).get(flow_id) flow = session.query(Flow).get(flow_id)
elif flow_name: elif flow_name: