Relational Transform Tests

wrangler-core/src/test/java/io/cdap/MockEngine.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright © 2023 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.cdap;
+
+import io.cdap.cdap.etl.api.relational.Capability;
+import io.cdap.cdap.etl.api.relational.Engine;
+import io.cdap.cdap.etl.api.relational.ExpressionFactory;
+import io.cdap.cdap.etl.api.relational.StringExpressionFactoryType;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+
+public class MockEngine implements Engine {
+
+    @Override
+    public Set<Capability> getCapabilities() {
+        return Collections.singleton(StringExpressionFactoryType.SQL);
+    }
+
+    @Override
+    public List<ExpressionFactory<?>> getExpressionFactories() {
+    return Collections.singletonList(new MockExpressionFactory());
+    }
+}

wrangler-core/src/test/java/io/cdap/MockExpression.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright © 2023 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.cdap;
+
+import io.cdap.cdap.etl.api.relational.Expression;
+
+public class MockExpression implements Expression {
+
+    private final String expression;
+
+
+    public MockExpression(String expression) {
+        this.expression = expression;
+    }
+
+    public boolean isValid() {
+        return true;
+    }
+
+    public String getExpression() {
+        return this.expression;
+    }
+
+    @Override
+    public String getValidationError() {
+        return null;
+    }
+}

wrangler-core/src/test/java/io/cdap/MockExpressionFactory.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright © 2023 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.cdap;
+
+import io.cdap.cdap.etl.api.relational.Capability;
+import io.cdap.cdap.etl.api.relational.CoreExpressionCapabilities;
+import io.cdap.cdap.etl.api.relational.Expression;
+import io.cdap.cdap.etl.api.relational.ExpressionFactory;
+import io.cdap.cdap.etl.api.relational.ExpressionFactoryType;
+import io.cdap.cdap.etl.api.relational.StringExpressionFactoryType;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+
+public class MockExpressionFactory implements ExpressionFactory<String> {
+
+    @Override
+public ExpressionFactoryType<String> getType() {
+    return StringExpressionFactoryType.SQL;
+}
+    private static final Set<Capability> CAPABILITIES = Collections.unmodifiableSet(
+            new HashSet<Capability>() {{
+                add(StringExpressionFactoryType.SQL);
+                add(CoreExpressionCapabilities.CAN_GET_QUALIFIED_DATASET_NAME);
+                add(CoreExpressionCapabilities.CAN_GET_QUALIFIED_COLUMN_NAME);
+                add(CoreExpressionCapabilities.CAN_SET_DATASET_ALIAS);
+            }});
+
+    @Override
+    public Set<Capability> getCapabilities() {
+        return CAPABILITIES;
+    }
+
+    @Override
+    public Expression compile(String expression) {
+    return new MockExpression(expression);
+    }
+}

wrangler-core/src/test/java/io/cdap/MockRelation.java

@@ -0,0 +1,81 @@
+/*
+ * Copyright © 2023 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.cdap;
+
+import io.cdap.cdap.etl.api.relational.Expression;
+import io.cdap.cdap.etl.api.relational.Relation;
+import java.util.Map;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+
+public class MockRelation implements Relation {
+    String column;
+
+    Expression expression;
+
+    public MockRelation(@Nullable String column, @Nullable Expression expression) {
+        this.column = column;
+        this.expression = expression;
+    }
+    @Override
+    public boolean isValid() {
+        return true;
+    }
+
+    @Override
+    public String getValidationError() {
+        return null;
+    }
+
+    @Override
+    public Relation setColumn(String column, Expression value) {
+        return new MockRelation(column, value);
+    }
+
+    @Override
+    public Relation dropColumn(String column) {
+        return new MockRelation(column, null);
+    }
+
+    @Override
+    public Relation select(Map<String, Expression> columnExpMap) {
+        String columns = columnExpMap.entrySet()
+                .stream()
+                .map(column -> column.getKey())
+                .collect(Collectors.joining(","));
+
+        String expressions = columnExpMap.entrySet()
+                .stream()
+                .map(column -> ((MockExpression) column.getValue()).getExpression())
+                .collect(Collectors.joining(","));
+
+    return new MockRelation(columns, new MockExpression(expressions));
+    }
+
+    @Override
+    public Relation filter(Expression filter) {
+        return new MockRelation(null, filter);
+    }
+
+    public String getColumn() {
+        return this.column;
+    }
+
+    public Expression getExpression() {
+        return this.expression;
+    }
+}

wrangler-core/src/test/java/io/cdap/MockRelationalTransformContext.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright © 2023 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.cdap;
+
+import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.cdap.api.feature.FeatureFlagsProvider;
+import io.cdap.cdap.etl.api.engine.sql.StandardSQLCapabilities;
+import io.cdap.cdap.etl.api.relational.Capability;
+import io.cdap.cdap.etl.api.relational.Engine;
+import io.cdap.cdap.etl.api.relational.Relation;
+import io.cdap.cdap.etl.api.relational.RelationalTranformContext;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+
+public class MockRelationalTransformContext implements RelationalTranformContext {
+    private Engine engine;
+    private Map<String, Relation> inputMap;
+    private Map<String, Schema> inputSchemas;
+    private Schema outputSchema;
+    private  FeatureFlagsProvider featureFlagsProvider;
+    private Relation outputRelation;
+
+    public MockRelationalTransformContext(Engine engine,
+                                          Map<String, Relation> inputMap,
+                                          Map<String, Schema> inputSchemas,
+                                          Schema outputSchema,
+                                          FeatureFlagsProvider featureFlagsProvider) {
+        this.engine = engine;
+        this.inputMap = inputMap;
+        this.inputSchemas = inputSchemas;
+        this.outputSchema = outputSchema;
+        this.featureFlagsProvider = featureFlagsProvider;
+    }
+
+    @Override
+    public Engine getEngine() {
+        return engine;
+    }
+
+    @Override
+    public Relation getInputRelation(String inputStage) {
+        return inputMap.get(inputStage);
+    }
+
+    @Override
+    public Set<String> getInputRelationNames() {
+        return Collections.unmodifiableSet(inputMap.keySet());
+    }
+
+    @Override
+    public Schema getInputSchema(String inputStage) {
+        return inputSchemas.get(inputStage);
+    }
+
+    @Override
+    public Schema getOutputSchema() {
+        return outputSchema;
+    }
+
+    @Override
+    public void setOutputRelation(Relation outputRelation) {
+        this.outputRelation = outputRelation;
+    }
+
+    @Override
+    public void setOutputRelation(String portName, Relation outputDataSet) {
+        throw new UnsupportedOperationException("Only single output is supported");
+    }
+
+    @Override
+    public boolean isFeatureEnabled(String name) {
+        return this.featureFlagsProvider.isFeatureEnabled(name);
+    }
+
+
+    public Collection<Capability> getDefaultLanguageCapabilityList() {
+        return Collections.singleton(StandardSQLCapabilities.POSTGRES);
+    }
+}

wrangler-core/src/test/java/io/cdap/RelationalDirectiveTest.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright © 2023 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.cdap;
+
+import io.cdap.cdap.etl.api.relational.Relation;
+import io.cdap.cdap.etl.api.relational.RelationalTranformContext;
+import io.cdap.wrangler.api.Directive;
+import io.cdap.wrangler.api.DirectiveLoadException;
+import io.cdap.wrangler.api.DirectiveParseException;
+import io.cdap.wrangler.api.RecipeException;
+import io.cdap.wrangler.parser.GrammarBasedParser;
+import io.cdap.wrangler.parser.MigrateToV2;
+import io.cdap.wrangler.registry.DirectiveRegistry;
+import io.cdap.wrangler.registry.SystemDirectiveRegistry;
+import java.util.List;
+
+public class RelationalDirectiveTest {
+
+    public static Relation runTransform(String[] recipe,
+                                        RelationalTranformContext relationalTranformContext,
+                                        Relation relation)
+            throws DirectiveParseException, RecipeException {
+        DirectiveRegistry registry;
+        registry = SystemDirectiveRegistry.INSTANCE;
+        try {
+            registry.reload("default");
+        } catch (DirectiveLoadException e) {
+            throw new RuntimeException(e);
+        }
+
+        GrammarBasedParser parser = new GrammarBasedParser("default",
+                new MigrateToV2(recipe).migrate(), registry);
+        List<Directive> directives = parser.parse();
+        for (Directive directive : directives) {
+            relation = directive.transform(relationalTranformContext, relation);
+        }
+        return relation;
+    }
+}