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

fix: environment variable loading and improve error handling in DB retrieval (#4168)

Browse source 
* test: rewrite and enable variable loading test

* refactor: simplify environment variable storage logic

* refactor: simplify parameter loading logic from db

* fix: handle additional error case when loading variables
This commit is contained in:
Ítalo Johnny 2024-10-16 08:53:40 -03:00 • committed by GitHub
commit 1a72aa71b6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 60 additions and 78 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -109,38 +109,28 @@ def update_params_with_load_from_db_fields(
*,
fallback_to_env_vars=False,
):
# For each field in load_from_db_fields, we will check if it's in the params
# and if it is, we will get the value from the custom_component.keys(name)
# and update the params with the value
for field in load_from_db_fields:
if field in params:
try:
key = None
try:
key = custom_component.variables(params[field], field)
except ValueError as e:
# check if "User id is not set" is in the error message, this is an internal bug
if "User id is not set" in str(e):
raise
logger.debug(str(e))
if fallback_to_env_vars and key is None:
key = os.getenv(params[field])
if key is None:
msg = f"Environment variable {params[field]} is not set."
logger.error(msg)
else:
logger.info(f"Using environment variable {params[field]} for {field}")
if key is None:
logger.warning(f"Could not get value for {field}. Setting it to None.")
if field not in params:
continue
params[field] = key
except TypeError:
try:
key = custom_component.variables(params[field], field)
except ValueError as e:
if any(reason in str(e) for reason in ["User id is not set", "variable not found."]):
raise
logger.debug(str(e))
key = None
except Exception: # noqa: BLE001
logger.exception(f"Failed to get value for {field} from custom component. Setting it to None.")
params[field] = None
if fallback_to_env_vars and key is None:
key = os.getenv(params[field])
if key:
logger.info(f"Using environment variable {params[field]} for {field}")
else:
logger.error(f"Environment variable {params[field]} is not set.")
params[field] = key if key is not None else None
if key is None:
logger.warning(f"Could not get value for {field}. Setting it to None.")
return params

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

@ -27,46 +27,31 @@ class DatabaseVariableService(VariableService, Service):
self.settings_service = settings_service
def initialize_user_variables(self, user_id: UUID | str, session: Session = Depends(get_session)):
# Check for environment variables that should be stored in the database
should_or_should_not = "Should" if self.settings_service.settings.store_environment_variables else "Should not"
logger.info(f"{should_or_should_not} store environment variables in the database.")
if self.settings_service.settings.store_environment_variables:
for var in self.settings_service.settings.variables_to_get_from_environment:
if var in os.environ:
logger.debug(f"Creating {var} variable from environment.")
if found_variable := session.exec(
select(Variable).where(Variable.user_id == user_id, Variable.name == var)
).first():
# Update it
value = os.environ[var]
if isinstance(value, str):
value = value.strip()
# If the secret_key changes the stored value could be invalid
# so we need to re-encrypt it
encrypted = auth_utils.encrypt_api_key(value, settings_service=self.settings_service)
found_variable.value = encrypted
session.add(found_variable)
session.commit()
else:
# Create it
try:
value = os.environ[var]
if isinstance(value, str):
value = value.strip()
self.create_variable(
user_id=user_id,
name=var,
value=value,
default_fields=[],
_type=CREDENTIAL_TYPE,
session=session,
)
except Exception: # noqa: BLE001
logger.exception(f"Error creating {var} variable")
else:
if not self.settings_service.settings.store_environment_variables:
logger.info("Skipping environment variable storage.")
return
logger.info("Storing environment variables in the database.")
for var_name in self.settings_service.settings.variables_to_get_from_environment:
if var_name in os.environ and os.environ[var_name].strip():
value = os.environ[var_name].strip()
query = select(Variable).where(Variable.user_id == user_id, Variable.name == var_name)
existing = session.exec(query).first()
try:
if existing:
self.update_variable(user_id, var_name, value, session)
else:
self.create_variable(
user_id=user_id,
name=var_name,
value=value,
default_fields=[],
_type=CREDENTIAL_TYPE,
session=session,
)
logger.info(f"Processed {var_name} variable from environment.")
except Exception as e: # noqa: BLE001
logger.exception(f"Error processing {var_name} variable: {e!s}")
def get_variable(
self,