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

add exception to Credentials vars in session_id fields

Browse source
This commit is contained in:
italojohnny 2024-06-09 16:01:42 -03:00
commit 565df927dc
3 changed files with 16 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

5
src/backend/base/langflow/custom/custom_component/custom_component.py
View file

@ -378,13 +378,14 @@ class CustomComponent(Component):
The variable for the current user with the specified name. The variable for the current user with the specified name.
""" """
def get_variable(name: str): def get_variable(name: str, field: str):
if hasattr(self, "_user_id") and not self._user_id: if hasattr(self, "_user_id") and not self._user_id:
raise ValueError(f"User id is not set for {self.__class__.__name__}") raise ValueError(f"User id is not set for {self.__class__.__name__}")
variable_service = get_variable_service() # Get service instance variable_service = get_variable_service() # Get service instance
# Retrieve and decrypt the variable by name for the current user # Retrieve and decrypt the variable by name for the current user
with session_scope() as session: with session_scope() as session:
return variable_service.get_variable(user_id=self._user_id or "", name=name, session=session) user_id = self._user_id or ""
return variable_service.get_variable(user_id=user_id, name=name, field=field, session=session)
return get_variable return get_variable

8
src/backend/base/langflow/interface/initialize/loading.py
View file

@ -71,7 +71,7 @@ def update_params_with_load_from_db_fields(
try: try:
key = None key = None
try: try:
key = custom_component.variables(params[field]) key = custom_component.variables(params[field], field)
except ValueError as e: except ValueError as e:
# check if "User id is not set" is in the error message # check if "User id is not set" is in the error message
if "User id is not set" in str(e) and not fallback_to_env_vars: if "User id is not set" in str(e) and not fallback_to_env_vars:
@ -86,8 +86,10 @@ def update_params_with_load_from_db_fields(
if key is None: if key is None:
logger.warning(f"Could not get value for {field}. Setting it to None.") logger.warning(f"Could not get value for {field}. Setting it to None.")
if field != "session_id": params[field] = key
params[field] = key
except TypeError as exc:
raise exc
except Exception as exc: except Exception as exc:
logger.error(f"Failed to get value for {field} from custom component. Setting it to None. Error: {exc}") logger.error(f"Failed to get value for {field} from custom component. Setting it to None. Error: {exc}")

8
src/backend/base/langflow/services/variable/service.py
View file

@ -54,11 +54,19 @@ class VariableService(Service):
self, self,
user_id: Union[UUID, str], user_id: Union[UUID, str],
name: str, name: str,
field: str,
session: Session = Depends(get_session), session: Session = Depends(get_session),
) -> str: ) -> str:
# we get the credential from the database # we get the credential from the database
# credential = session.query(Variable).filter(Variable.user_id == user_id, Variable.name == name).first() # credential = session.query(Variable).filter(Variable.user_id == user_id, Variable.name == name).first()
variable = session.exec(select(Variable).where(Variable.user_id == user_id, Variable.name == name)).first() variable = session.exec(select(Variable).where(Variable.user_id == user_id, Variable.name == name)).first()
if variable.type == "Credential" and field == "session_id":
raise TypeError(
f"variable {name} of type 'Credential' cannot be used in a Session ID field "
"because its purpose is to prevent the exposure of values."
)
# we decrypt the value # we decrypt the value
if not variable or not variable.value: if not variable or not variable.value:
raise ValueError(f"{name} variable not found.") raise ValueError(f"{name} variable not found.")