Skip to content

Commit

Permalink
Preserve timestamps when encrypting and redacting template
Browse files Browse the repository at this point in the history 
Signed-off-by: Daniel Widdis <[email protected]>
  • Loading branch information
@dbwiddis
dbwiddis committed Mar 5, 2024
1 parent 234a791 commit 7f5cd10
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 26 deletions.
31 changes: 15 additions & 16 deletions src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ public CreateWorkflowTransportAction(
protected void doExecute(Task task, WorkflowRequest request, ActionListener<WorkflowResponse> listener) {

User user = getUserContext(client);
Template templateWithUserAndTimestamps = new Template(
Template templateWithUser = new Template(
request.getTemplate().name(),
request.getTemplate().description(),
request.getTemplate().useCase(),
Expand All @@ -105,17 +105,17 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
request.getTemplate().workflows(),
request.getTemplate().getUiMetadata(),
user,
null,
null,
null
request.getTemplate().createdTime(),
request.getTemplate().lastUpdatedTime(),
request.getTemplate().lastProvisionedTime()
);

String[] validateAll = { "all" };
if (Arrays.equals(request.getValidation(), validateAll)) {
try {
validateWorkflows(templateWithUserAndTimestamps);
validateWorkflows(templateWithUser);
} catch (Exception e) {
String errorMessage = "Workflow validation failed for template " + templateWithUserAndTimestamps.name();
String errorMessage = "Workflow validation failed for template " + templateWithUser.name();
logger.error(errorMessage, e);
listener.onFailure(
e instanceof FlowFrameworkException ? e : new FlowFrameworkException(errorMessage, ExceptionsHelper.status(e))
Expand Down Expand Up @@ -147,7 +147,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
} else {
// Create new global context and state index entries
flowFrameworkIndicesHandler.putTemplateToGlobalContext(
templateWithUserAndTimestamps,
templateWithUser,
ActionListener.wrap(globalContextResponse -> {
flowFrameworkIndicesHandler.putInitialStateToWorkflowState(
globalContextResponse.getId(),
Expand Down Expand Up @@ -238,16 +238,15 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<Work
context.restore();
if (getResponse.isExists()) {
Template existingTemplate = Template.parse(getResponse.getSourceAsString());
Template newTemplate = request.getTemplate();
// Update existing entry, full document replacement
Template template = new Template.Builder().name(newTemplate.name())
.description(newTemplate.description())
.useCase(newTemplate.useCase())
.templateVersion(newTemplate.templateVersion())
.compatibilityVersion(newTemplate.compatibilityVersion())
.workflows(newTemplate.workflows())
.uiMetadata(newTemplate.getUiMetadata())
.user(newTemplate.getUser()) // Should we care about old user here?
Template template = new Template.Builder().name(templateWithUser.name())
.description(templateWithUser.description())
.useCase(templateWithUser.useCase())
.templateVersion(templateWithUser.templateVersion())
.compatibilityVersion(templateWithUser.compatibilityVersion())
.workflows(templateWithUser.workflows())
.uiMetadata(templateWithUser.getUiMetadata())
.user(templateWithUser.getUser()) // Should we care about old user here?
.createdTime(existingTemplate.createdTime())
.lastUpdatedTime(Instant.now())
.lastProvisionedTime(existingTemplate.lastProvisionedTime())
Expand Down
18 changes: 8 additions & 10 deletions src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -126,8 +126,6 @@ public Template decryptTemplateCredentials(Template template) {
* @return template with encrypted credentials
*/
private Template processTemplateCredentials(Template template, Function<String, String> cipherFunction) {
Template.Builder processedTemplateBuilder = new Template.Builder();

Map<String, Workflow> processedWorkflows = new HashMap<>();
for (Map.Entry<String, Workflow> entry : template.workflows().entrySet()) {

Expand Down Expand Up @@ -161,17 +159,18 @@ private Template processTemplateCredentials(Template template, Function<String,
processedWorkflows.put(entry.getKey(), new Workflow(entry.getValue().userParams(), processedNodes, entry.getValue().edges()));
}

Template processedTemplate = processedTemplateBuilder.name(template.name())
return new Template.Builder().name(template.name())
.description(template.description())
.useCase(template.useCase())
.templateVersion(template.templateVersion())
.compatibilityVersion(template.compatibilityVersion())
.workflows(processedWorkflows)
.uiMetadata(template.getUiMetadata())
.user(template.getUser())
.createdTime(template.createdTime())
.lastUpdatedTime(template.lastUpdatedTime())
.lastProvisionedTime(template.lastProvisionedTime())
.build();

return processedTemplate;
}

/**
Expand Down Expand Up @@ -217,8 +216,6 @@ String decrypt(final String encryptedCredential) {
* @return the redacted template
*/
public Template redactTemplateCredentials(Template template) {
Template.Builder redactedTemplateBuilder = new Template.Builder();

Map<String, Workflow> processedWorkflows = new HashMap<>();
for (Map.Entry<String, Workflow> entry : template.workflows().entrySet()) {

Expand All @@ -241,17 +238,18 @@ public Template redactTemplateCredentials(Template template) {
processedWorkflows.put(entry.getKey(), new Workflow(entry.getValue().userParams(), processedNodes, entry.getValue().edges()));
}

Template processedTemplate = redactedTemplateBuilder.name(template.name())
return new Template.Builder().name(template.name())
.description(template.description())
.useCase(template.useCase())
.templateVersion(template.templateVersion())
.compatibilityVersion(template.compatibilityVersion())
.workflows(processedWorkflows)
.uiMetadata(template.getUiMetadata())
.user(template.getUser())
.createdTime(template.createdTime())
.lastUpdatedTime(template.lastUpdatedTime())
.lastProvisionedTime(template.lastProvisionedTime())
.build();

return processedTemplate;
}

/**
Expand Down

0 comments on commit 7f5cd10

Please sign in to comment.