vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

No need to pass aws client around (#2672)

Browse Source
This commit is contained in:
Asher Foa
2025-06-13 12:41:22 -04:00
committed by GitHub
parent ff3b6b16ec
commit 59f11af271
1 changed files with 17 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
skyvern/forge/sdk/workflow/context_manager.py
View File

@@ -63,7 +63,7 @@ class WorkflowRunContext:
], ],
) -> Self: ) -> Self:
# key is label name # key is label name
workflow_run_context = cls() workflow_run_context = cls(aws_client=aws_client)
for parameter, run_parameter in workflow_parameter_tuples: for parameter, run_parameter in workflow_parameter_tuples:
if parameter.workflow_parameter_type == WorkflowParameterType.CREDENTIAL_ID: if parameter.workflow_parameter_type == WorkflowParameterType.CREDENTIAL_ID:
await workflow_run_context.register_secret_workflow_parameter_value( await workflow_run_context.register_secret_workflow_parameter_value(
@@ -87,22 +87,22 @@ class WorkflowRunContext:
for secrete_parameter in secret_parameters: for secrete_parameter in secret_parameters:
if isinstance(secrete_parameter, AWSSecretParameter): if isinstance(secrete_parameter, AWSSecretParameter):
await workflow_run_context.register_aws_secret_parameter_value(aws_client, secrete_parameter) await workflow_run_context.register_aws_secret_parameter_value(secrete_parameter)
elif isinstance(secrete_parameter, CredentialParameter): elif isinstance(secrete_parameter, CredentialParameter):
await workflow_run_context.register_credential_parameter_value(secrete_parameter, organization) await workflow_run_context.register_credential_parameter_value(secrete_parameter, organization)
elif isinstance(secrete_parameter, OnePasswordCredentialParameter): elif isinstance(secrete_parameter, OnePasswordCredentialParameter):
await workflow_run_context.register_onepassword_credential_parameter_value(secrete_parameter) await workflow_run_context.register_onepassword_credential_parameter_value(secrete_parameter)
elif isinstance(secrete_parameter, BitwardenLoginCredentialParameter): elif isinstance(secrete_parameter, BitwardenLoginCredentialParameter):
await workflow_run_context.register_bitwarden_login_credential_parameter_value( await workflow_run_context.register_bitwarden_login_credential_parameter_value(
aws_client, secrete_parameter, organization secrete_parameter, organization
) )
elif isinstance(secrete_parameter, BitwardenCreditCardDataParameter): elif isinstance(secrete_parameter, BitwardenCreditCardDataParameter):
await workflow_run_context.register_bitwarden_credit_card_data_parameter_value( await workflow_run_context.register_bitwarden_credit_card_data_parameter_value(
aws_client, secrete_parameter, organization secrete_parameter, organization
) )
elif isinstance(secrete_parameter, BitwardenSensitiveInformationParameter): elif isinstance(secrete_parameter, BitwardenSensitiveInformationParameter):
await workflow_run_context.register_bitwarden_sensitive_information_parameter_value( await workflow_run_context.register_bitwarden_sensitive_information_parameter_value(
aws_client, secrete_parameter, organization secrete_parameter, organization
) )
for context_parameter in context_parameters: for context_parameter in context_parameters:
@@ -113,11 +113,12 @@ class WorkflowRunContext:
return workflow_run_context return workflow_run_context
def __init__(self) -> None: def __init__(self, aws_client: AsyncAWSClient) -> None:
self.blocks_metadata: dict[str, BlockMetadata] = {} self.blocks_metadata: dict[str, BlockMetadata] = {}
self.parameters: dict[str, PARAMETER_TYPE] = {} self.parameters: dict[str, PARAMETER_TYPE] = {}
self.values: dict[str, Any] = {} self.values: dict[str, Any] = {}
self.secrets: dict[str, Any] = {} self.secrets: dict[str, Any] = {}
self._aws_client = aws_client
def get_parameter(self, key: str) -> Parameter: def get_parameter(self, key: str) -> Parameter:
return self.parameters[key] return self.parameters[key]
@@ -372,13 +373,12 @@ class WorkflowRunContext:
async def register_aws_secret_parameter_value( async def register_aws_secret_parameter_value(
self, self,
aws_client: AsyncAWSClient,
parameter: AWSSecretParameter, parameter: AWSSecretParameter,
) -> None: ) -> None:
# If the parameter is an AWS secret, fetch the secret value and store it in the secrets dict # If the parameter is an AWS secret, fetch the secret value and store it in the secrets dict
# The value of the parameter will be the random secret id with format `secret_<uuid>`. # The value of the parameter will be the random secret id with format `secret_<uuid>`.
# We'll replace the random secret id with the actual secret value when we need to use it. # We'll replace the random secret id with the actual secret value when we need to use it.
secret_value = await aws_client.get_secret(parameter.aws_key) secret_value = await self._aws_client.get_secret(parameter.aws_key)
if secret_value is not None: if secret_value is not None:
random_secret_id = self.generate_random_secret_id() random_secret_id = self.generate_random_secret_id()
self.secrets[random_secret_id] = secret_value self.secrets[random_secret_id] = secret_value
@@ -427,19 +427,18 @@ class WorkflowRunContext:
async def register_bitwarden_login_credential_parameter_value( async def register_bitwarden_login_credential_parameter_value(
self, self,
aws_client: AsyncAWSClient,
parameter: BitwardenLoginCredentialParameter, parameter: BitwardenLoginCredentialParameter,
organization: Organization, organization: Organization,
) -> None: ) -> None:
try: try:
# Get the Bitwarden login credentials from AWS secrets # Get the Bitwarden login credentials from AWS secrets
client_id = settings.BITWARDEN_CLIENT_ID or await aws_client.get_secret( client_id = settings.BITWARDEN_CLIENT_ID or await self._aws_client.get_secret(
parameter.bitwarden_client_id_aws_secret_key parameter.bitwarden_client_id_aws_secret_key
) )
client_secret = settings.BITWARDEN_CLIENT_SECRET or await aws_client.get_secret( client_secret = settings.BITWARDEN_CLIENT_SECRET or await self._aws_client.get_secret(
parameter.bitwarden_client_secret_aws_secret_key parameter.bitwarden_client_secret_aws_secret_key
) )
master_password = settings.BITWARDEN_MASTER_PASSWORD or await aws_client.get_secret( master_password = settings.BITWARDEN_MASTER_PASSWORD or await self._aws_client.get_secret(
parameter.bitwarden_master_password_aws_secret_key parameter.bitwarden_master_password_aws_secret_key
) )
except Exception as e: except Exception as e:
@@ -531,19 +530,18 @@ class WorkflowRunContext:
async def register_bitwarden_sensitive_information_parameter_value( async def register_bitwarden_sensitive_information_parameter_value(
self, self,
aws_client: AsyncAWSClient,
parameter: BitwardenSensitiveInformationParameter, parameter: BitwardenSensitiveInformationParameter,
organization: Organization, organization: Organization,
) -> None: ) -> None:
try: try:
# Get the Bitwarden login credentials from AWS secrets # Get the Bitwarden login credentials from AWS secrets
client_id = settings.BITWARDEN_CLIENT_ID or await aws_client.get_secret( client_id = settings.BITWARDEN_CLIENT_ID or await self._aws_client.get_secret(
parameter.bitwarden_client_id_aws_secret_key parameter.bitwarden_client_id_aws_secret_key
) )
client_secret = settings.BITWARDEN_CLIENT_SECRET or await aws_client.get_secret( client_secret = settings.BITWARDEN_CLIENT_SECRET or await self._aws_client.get_secret(
parameter.bitwarden_client_secret_aws_secret_key parameter.bitwarden_client_secret_aws_secret_key
) )
master_password = settings.BITWARDEN_MASTER_PASSWORD or await aws_client.get_secret( master_password = settings.BITWARDEN_MASTER_PASSWORD or await self._aws_client.get_secret(
parameter.bitwarden_master_password_aws_secret_key parameter.bitwarden_master_password_aws_secret_key
) )
except Exception as e: except Exception as e:
@@ -599,19 +597,18 @@ class WorkflowRunContext:
async def register_bitwarden_credit_card_data_parameter_value( async def register_bitwarden_credit_card_data_parameter_value(
self, self,
aws_client: AsyncAWSClient,
parameter: BitwardenCreditCardDataParameter, parameter: BitwardenCreditCardDataParameter,
organization: Organization, organization: Organization,
) -> None: ) -> None:
try: try:
# Get the Bitwarden login credentials from AWS secrets # Get the Bitwarden login credentials from AWS secrets
client_id = settings.BITWARDEN_CLIENT_ID or await aws_client.get_secret( client_id = settings.BITWARDEN_CLIENT_ID or await self._aws_client.get_secret(
parameter.bitwarden_client_id_aws_secret_key parameter.bitwarden_client_id_aws_secret_key
) )
client_secret = settings.BITWARDEN_CLIENT_SECRET or await aws_client.get_secret( client_secret = settings.BITWARDEN_CLIENT_SECRET or await self._aws_client.get_secret(
parameter.bitwarden_client_secret_aws_secret_key parameter.bitwarden_client_secret_aws_secret_key
) )
master_password = settings.BITWARDEN_MASTER_PASSWORD or await aws_client.get_secret( master_password = settings.BITWARDEN_MASTER_PASSWORD or await self._aws_client.get_secret(
parameter.bitwarden_master_password_aws_secret_key parameter.bitwarden_master_password_aws_secret_key
) )
except Exception as e: except Exception as e: