From c53775ddf90006660ef39bc8dc2cf27ad2062ab2 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 24 Nov 2023 11:19:17 -0300 Subject: [PATCH] Fix user_id modification in Component class --- .../langflow/interface/custom/component.py | 6 ++++++ .../interface/custom/custom_component.py | 16 ++++++++++++++-- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/interface/custom/component.py b/src/backend/langflow/interface/custom/component.py index 98ccf8db3..bc4b4e605 100644 --- a/src/backend/langflow/interface/custom/component.py +++ b/src/backend/langflow/interface/custom/component.py @@ -4,6 +4,7 @@ from typing import Any, ClassVar, Optional from cachetools import TTLCache, cachedmethod from fastapi import HTTPException + from langflow.interface.custom.code_parser import CodeParser from langflow.utils import validate @@ -29,6 +30,11 @@ class Component: for key, value in data.items(): setattr(self, key, value) + def __setattr__(self, key, value): + if key == "user_id" and hasattr(self, key): + raise AttributeError("Modification of user_id is not allowed") + super().__setattr__(key, value) + @cachedmethod(cache=operator.attrgetter("cache")) def get_code_tree(self, code: str): parser = CodeParser(code) diff --git a/src/backend/langflow/interface/custom/custom_component.py b/src/backend/langflow/interface/custom/custom_component.py index 032ed1476..7f1722241 100644 --- a/src/backend/langflow/interface/custom/custom_component.py +++ b/src/backend/langflow/interface/custom/custom_component.py @@ -5,7 +5,6 @@ from uuid import UUID import yaml from cachetools import TTLCache, cachedmethod from fastapi import HTTPException - from langflow.field_typing.constants import CUSTOM_COMPONENT_SUPPORTED_TYPES from langflow.interface.custom.component import Component from langflow.interface.custom.directory_reader import DirectoryReader @@ -15,7 +14,7 @@ from langflow.interface.custom.utils import ( ) from langflow.services.database.models.flow import Flow from langflow.services.database.utils import session_getter -from langflow.services.deps import get_db_service +from langflow.services.deps import get_credential_service, get_db_service from langflow.utils import validate @@ -184,6 +183,19 @@ class CustomComponent(Component): return super().build_template_config(attributes) + @property + def keys(self): + def get_credential(name: str): + if not self.user_id: + raise ValueError(f"User id is not set for {self.__class__.__name__}") + credential_service = get_credential_service() # Get service instance + # Retrieve and decrypt the credential by name for the current user + db_service = get_db_service() + with session_getter(db_service) as session: + return credential_service.get_credential(user_id=self.user_id, name=name, session=session) + + return get_credential + @property def get_function(self): return validate.create_function(self.code, self.function_entrypoint_name)