Skip to content

Fix: Enforce wrangler directive config in wrangler service and pipeline transform #1039

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
riyaa14 wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

Fix: Enforce wrangler directive config in wrangler service and pipeline transform #1039

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions ...er-service/src/main/java/io/cdap/wrangler/service/directive/AbstractDirectiveHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import io.cdap.wrangler.api.RecipeParser;
import io.cdap.wrangler.api.Row;
import io.cdap.wrangler.api.TransientStore;
import io.cdap.wrangler.dataset.workspace.ConfigStore;
import io.cdap.wrangler.executor.RecipePipelineExecutor;
import io.cdap.wrangler.parser.ConfigDirectiveContext;
import io.cdap.wrangler.parser.GrammarBasedParser;
Expand Down Expand Up @@ -128,15 +129,16 @@ protected <E extends Exception> List<Row> executeDirectives(
String recipe = migrator.migrate();

// Parse and call grammar visitor
DirectiveConfig config = getDirectiveConfig();
try {
GrammarWalker walker = new GrammarWalker(new RecipeCompiler(), new ConfigDirectiveContext(DirectiveConfig.EMPTY));
GrammarWalker walker = new GrammarWalker(new RecipeCompiler(), new ConfigDirectiveContext(config));
walker.walk(recipe, grammarVisitor);
} catch (CompileException e) {
throw new BadRequestException(e.getMessage(), e);
}

RecipeParser parser = new GrammarBasedParser(namespace, recipe, composite,
new ConfigDirectiveContext(DirectiveConfig.EMPTY));
new ConfigDirectiveContext(config));
try (RecipePipelineExecutor executor = new RecipePipelineExecutor(parser,
new ServicePipelineContext(
namespace, ExecutorContext.Environment.SERVICE,
Expand Down Expand Up @@ -281,4 +283,13 @@ private String getColumnDisplayType(Schema schema) {
type = type.substring(0, 1).toUpperCase() + type.substring(1).toLowerCase();
return type;
}

protected DirectiveConfig getDirectiveConfig() {
try {
return ConfigStore.get(getContext()).getConfig();
} catch (IOException | RuntimeException e) {
LOG.warn("Failed to fetch directive exclusions from database. Using empty configuration.", e);

@vsethi09 vsethi09 Jul 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If there was a failure reading from Database, IOException would be thrown. You are ignoring the exception and returning default configuration, which is incorrect.

@riyaa14 riyaa14 Jul 24, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, my bad, changed - throwing IOException for this method in case of error rather than handling locally.

return DirectiveConfig.EMPTY;
}
}
}
15 changes: 3 additions & 12 deletions wrangler-service/src/main/java/io/cdap/wrangler/service/directive/DirectivesHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -868,10 +868,7 @@ public void usage(HttpServiceRequest request, HttpServiceResponder responder,
respond(request, responder, namespace, ns -> {
// CDAP-15397 - reload must be called before it can be safely used
composite.reload(namespace);
DirectiveConfig config = TransactionRunners.run(getContext(), context -> {
ConfigStore store = ConfigStore.get(context);
return store.getConfig();
});
DirectiveConfig config = ConfigStore.get(getContext()).getConfig();
Map<String, List<String>> aliases = config.getReverseAlias();
List<DirectiveUsage> values = new ArrayList<>();

Expand Down Expand Up @@ -1000,10 +997,7 @@ public void uploadConfig(HttpServiceRequest request, HttpServiceResponder respon
if (config == null) {
throw new BadRequestException("Config is empty. Please check if the request is sent as HTTP POST body.");
}
TransactionRunners.run(getContext(), context -> {
ConfigStore configStore = ConfigStore.get(context);
configStore.updateConfig(config);
});
ConfigStore.get(getContext()).updateConfig(config);
return new ServiceResponse<Void>("Successfully updated configuration.");
});
}
Expand All @@ -1018,10 +1012,7 @@ public void uploadConfig(HttpServiceRequest request, HttpServiceResponder respon
@Path("config")
@TransactionPolicy(value = TransactionControl.EXPLICIT)
public void getConfig(HttpServiceRequest request, HttpServiceResponder responder) {
respond(request, responder, () -> TransactionRunners.run(getContext(), context -> {
ConfigStore configStore = ConfigStore.get(context);
return new ServiceResponse<>(configStore.getConfig());
}));
respond(request, responder, () -> new ServiceResponse<>(ConfigStore.get(getContext()).getConfig()));
}

/**
Expand Down
9 changes: 8 additions & 1 deletion ...gler-service/src/main/java/io/cdap/wrangler/service/directive/RemoteDirectiveRequest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
package io.cdap.wrangler.service.directive;

import io.cdap.cdap.api.data.schema.Schema;
import io.cdap.wrangler.api.DirectiveConfig;
import io.cdap.wrangler.parser.DirectiveClass;

import java.util.HashMap;
Expand All @@ -31,14 +32,16 @@ public class RemoteDirectiveRequest {
private final String pluginNameSpace;
private final byte[] data;
private final Schema inputSchema;
private final DirectiveConfig directiveConfig;

RemoteDirectiveRequest(String recipe, Map<String, DirectiveClass> systemDirectives,
String pluginNameSpace, byte[] data, Schema inputSchema) {
String pluginNameSpace, byte[] data, Schema inputSchema, DirectiveConfig directiveConfig) {
this.recipe = recipe;
this.systemDirectives = new HashMap<>(systemDirectives);
this.pluginNameSpace = pluginNameSpace;
this.data = data;
this.inputSchema = inputSchema;
this.directiveConfig = directiveConfig;
}

public String getRecipe() {
Expand All @@ -60,4 +63,8 @@ public String getPluginNameSpace() {
public Schema getInputSchema() {
return inputSchema;
}

public DirectiveConfig getDirectiveConfig() {
return directiveConfig;
}
}
4 changes: 3 additions & 1 deletion wrangler-service/src/main/java/io/cdap/wrangler/service/directive/RemoteExecutionTask.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,9 @@ public void run(RunnableTaskContext runnableTaskContext) throws Exception {
// Collect directives.
try (UserDirectiveRegistry userDirectiveRegistry = new UserDirectiveRegistry(systemAppContext)) {
List<Directive> directives = new ArrayList<>();
GrammarWalker walker = new GrammarWalker(new RecipeCompiler(), new ConfigDirectiveContext(DirectiveConfig.EMPTY));
DirectiveConfig config = directiveRequest.getDirectiveConfig() != null ?
directiveRequest.getDirectiveConfig() : DirectiveConfig.EMPTY;
GrammarWalker walker = new GrammarWalker(new RecipeCompiler(), new ConfigDirectiveContext(config));
walker.walk(directiveRequest.getRecipe(), (command, tokenGroup) -> {
DirectiveInfo info;
DirectiveClass directiveClass = systemDirectives.get(command);
Expand Down
10 changes: 7 additions & 3 deletions wrangler-service/src/main/java/io/cdap/wrangler/service/directive/WorkspaceHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -469,10 +469,12 @@ public void specification(HttpServiceRequest request, HttpServiceResponder respo

// check if the rows are empty before going to create a record schema, it will result in a 400 if empty fields
// are passed to a record type schema
DirectiveConfig directiveConfig = getDirectiveConfig();
Map<String, String> properties = ImmutableMap.of("directives", String.join("\n", directives),
"field", "*",
"precondition", "false",
"workspaceId", workspaceId);
"workspaceId", workspaceId,
"directiveConfig", GSON.toJson(directiveConfig));

Set<StageSpec> srcSpecs = getSourceSpecs(detail, directives);

Expand Down Expand Up @@ -656,7 +658,8 @@ private <E extends Exception> List<Row> executeRemotely(String namespace, List<S
Map<String, DirectiveClass> systemDirectives = new HashMap<>();

// Gather system directives and call additional visitor.
GrammarWalker walker = new GrammarWalker(new RecipeCompiler(), new ConfigDirectiveContext(DirectiveConfig.EMPTY));
DirectiveConfig config = getDirectiveConfig();
GrammarWalker walker = new GrammarWalker(new RecipeCompiler(), new ConfigDirectiveContext(config));
AtomicBoolean hasDirectives = new AtomicBoolean();
walker.walk(recipe, (command, tokenGroup) -> {
DirectiveInfo info = SystemDirectiveRegistry.INSTANCE.get(command);
Expand All @@ -674,7 +677,8 @@ private <E extends Exception> List<Row> executeRemotely(String namespace, List<S

RemoteDirectiveRequest directiveRequest = new RemoteDirectiveRequest(recipe, systemDirectives,
namespace, detail.getSampleAsBytes(),
TRANSIENT_STORE.get(INPUT_SCHEMA));
TRANSIENT_STORE.get(INPUT_SCHEMA),
config);
RunnableTaskRequest runnableTaskRequest = RunnableTaskRequest.getBuilder(RemoteExecutionTask.class.getName())
.withParam(GSON.toJson(directiveRequest))
.withNamespace(namespace)
Expand Down
24 changes: 24 additions & 0 deletions wrangler-storage/src/main/java/io/cdap/wrangler/dataset/workspace/ConfigStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@
import io.cdap.cdap.spi.data.table.field.Field;
import io.cdap.cdap.spi.data.table.field.FieldType;
import io.cdap.cdap.spi.data.table.field.Fields;
import io.cdap.cdap.spi.data.transaction.TransactionRunner;
import io.cdap.cdap.spi.data.transaction.TransactionRunners;
import io.cdap.wrangler.api.DirectiveConfig;

import java.io.IOException;
Expand Down Expand Up @@ -55,9 +57,16 @@ public class ConfigStore {
.withPrimaryKeys(KEY_COL)
.build();
private final StructuredTable table;
private final TransactionRunner transactionRunner;

public ConfigStore(StructuredTable table) {
this.table = table;
this.transactionRunner = null;
}

public ConfigStore(TransactionRunner transactionRunner) {

@vsethi09 vsethi09 Jul 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a need for separate constructor?

@riyaa14 riyaa14 Jul 24, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is just to make sure adding TransactionRunners is done in backward compatible way. I didn't want to change the public get method and constructor which exposed by ConfigStore currently. Thus, added new ones rather than updating the existing get method and constructor.

But rethinking this, maybe we can make this change and go forward with updating the current methods because ConfigStore is in wrangler-storage module, which is an internally used module and thus changing the constructor and methods shouldn't be a breaking change externally. If any custom plugin uses wrangler, it should be using wrangler-api module. Any insights of yours would be helpful.

this.table = null;
this.transactionRunner = transactionRunner;
}

public static ConfigStore get(StructuredTableContext context) {
Expand All @@ -70,14 +79,29 @@ public static ConfigStore get(StructuredTableContext context) {
}
}

public static ConfigStore get(TransactionRunner transactionRunner) {

@vsethi09 vsethi09 Jul 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a need for a separate overload?

@riyaa14 riyaa14 Jul 24, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same reason as mentioned in #1039 (comment)

return new ConfigStore(transactionRunner);
}

public void updateConfig(DirectiveConfig config) throws IOException {
if (transactionRunner != null) {
TransactionRunners.run(transactionRunner, context -> {
get(context).updateConfig(config);
});
return;
}
List<Field<?>> fields = new ArrayList<>(2);
fields.add(keyField);
fields.add(Fields.stringField(VAL_COL, GSON.toJson(config)));
table.upsert(fields);
}

public DirectiveConfig getConfig() throws IOException {
if (transactionRunner != null) {
return TransactionRunners.run(transactionRunner, context -> {
return get(context).getConfig();
});
}
Optional<StructuredRow> row = table.read(Collections.singletonList(keyField));
String configStr = row.map(r -> r.getString(VAL_COL)).orElse("{}");
return GSON.fromJson(configStr, DirectiveConfig.class);
Expand Down
Loading
Loading