TEIID-2392 ensuring consistent column names with a pushdown common table

expression

Author: shawkins

api/src/main/java/org/teiid/language/visitor/SQLStringVisitor.java

@@ -58,6 +58,7 @@ public class SQLStringVisitor extends AbstractLanguageVisitor {
 
     protected StringBuilder buffer = new StringBuilder();
     private boolean appendedSourceComment;
+    private boolean shortNameOnly = false;
 
     /**
      * Gets the name of a group or element from the RuntimeMetadata
@@ -238,7 +239,7 @@ protected String replaceElementName(String group, String element) {
     }
 
     public void visit(ColumnReference obj) {
-        buffer.append(getElementName(obj, true));
+        buffer.append(getElementName(obj, !shortNameOnly));
     }
 
 	private String getElementName(ColumnReference obj, boolean qualify) {
@@ -482,14 +483,9 @@ public void visit(Insert obj) {
 		append(obj.getTable());
 		buffer.append(Tokens.SPACE).append(Tokens.LPAREN);
 
-		int elementCount = obj.getColumns().size();
-		for (int i = 0; i < elementCount; i++) {
-			buffer.append(getElementName(obj.getColumns().get(i), false));
-			if (i < elementCount - 1) {
-				buffer.append(Tokens.COMMA);
-				buffer.append(Tokens.SPACE);
-			}
-		}
+		this.shortNameOnly = true;
+		append(obj.getColumns());
+		this.shortNameOnly = false;
 
 		buffer.append(Tokens.RPAREN);
         buffer.append(Tokens.SPACE);
@@ -925,7 +921,9 @@ public void visit(WithItem obj) {
     	buffer.append(Tokens.SPACE);
     	if (obj.getColumns() != null) {
     		buffer.append(Tokens.LPAREN);
+    		shortNameOnly = true;
     		append(obj.getColumns());
+    		shortNameOnly = false;
     		buffer.append(Tokens.RPAREN);
     		buffer.append(Tokens.SPACE);
     	}

connectors/translator-jdbc/src/test/java/org/teiid/translator/jdbc/sqlserver/TestSqlServerConversionVisitor.java

@@ -203,7 +203,7 @@ public void testDateFunctions() throws Exception {
 
     @Test public void testWith() throws Exception {
         String input = "with x as (select intkey from bqt1.smalla) select intkey from x limit 100"; //$NON-NLS-1$
-        String output = "WITH x AS (SELECT SmallA.IntKey FROM SmallA) SELECT TOP 100 g_0.intkey AS c_0 FROM x g_0"; //$NON-NLS-1$
+        String output = "WITH x (IntKey) AS (SELECT SmallA.IntKey FROM SmallA) SELECT TOP 100 g_0.intkey AS c_0 FROM x g_0"; //$NON-NLS-1$
 
 		CommandBuilder commandBuilder = new CommandBuilder(RealMetadataFactory.exampleBQTCached());
         Command obj = commandBuilder.getCommand(input, true, true);

engine/src/main/java/org/teiid/dqp/internal/datamgr/LanguageBridgeFactory.java

@@ -35,13 +35,13 @@
 import org.teiid.core.types.ArrayImpl;
 import org.teiid.core.types.DataTypeManager;
 import org.teiid.language.*;
-import org.teiid.language.DerivedColumn;
-import org.teiid.language.Select;
-import org.teiid.language.WindowSpecification;
 import org.teiid.language.Argument.Direction;
 import org.teiid.language.Comparison.Operator;
 import org.teiid.language.SortSpecification.Ordering;
 import org.teiid.language.SubqueryComparison.Quantifier;
+import org.teiid.language.DerivedColumn;
+import org.teiid.language.Select;
+import org.teiid.language.WindowSpecification;
 import org.teiid.metadata.Procedure;
 import org.teiid.metadata.ProcedureParameter;
 import org.teiid.query.QueryPlugin;
@@ -261,6 +261,7 @@ public With translate(List<WithQueryCommand> with) {
 		        for (ElementSymbol es: withQueryCommand.getColumns()) {
 		            translatedElements.add(translate(es));
 		        }
+				item.setColumns(translatedElements);
 			}
 			item.setSubquery(translate(withQueryCommand.getCommand()));
 			items.add(item);

engine/src/main/java/org/teiid/query/optimizer/relational/RelationalPlanner.java

@@ -241,6 +241,7 @@ public ProcessorPlan optimize(
 				if (queryCommand == null || CriteriaCapabilityValidatorVisitor.validateCommandPushdown(modelID, metadata, capFinder, aNode) == null) {
 					supportsWithPushdown = false;
 				} else {
+					//TODO: do this prior to alias generation, or rerun alias generation here
 					queryCommand.setWith(pushDownWith);
 				}
         	} else {

engine/src/main/java/org/teiid/query/sql/visitor/SQLStringVisitor.java

@@ -83,6 +83,7 @@ protected DDLStringVisitor append(Object o) {
     private static final String END_HINT = "*/"; //$NON-NLS-1$
     private static final char ID_ESCAPE_CHAR = '\"';
 	protected StringBuilder parts = new StringBuilder();
+	private boolean shortNameOnly = false;
 
     /**
      * Helper to quickly get the parser string for an object using the visitor.
@@ -722,7 +723,9 @@ public void visit(WithQueryCommand obj) {
     	append(SPACE);
     	if (obj.getColumns() != null && !obj.getColumns().isEmpty()) {
     		append(Tokens.LPAREN);
+    		shortNameOnly = true;
     		registerNodes(obj.getColumns(), 0);
+    		shortNameOnly = false;
     		append(Tokens.RPAREN);
     		append(SPACE);
     	}
@@ -1274,7 +1277,7 @@ static String escapeStringValue( String str,
     }
 
     public void visit( ElementSymbol obj ) {
-        if (obj.getDisplayMode().equals(ElementSymbol.DisplayMode.SHORT_OUTPUT_NAME)) {
+        if (obj.getDisplayMode().equals(ElementSymbol.DisplayMode.SHORT_OUTPUT_NAME) || shortNameOnly) {
             outputShortName(obj);
             return;
         }

engine/src/test/java/org/teiid/query/processor/TestWithClauseProcessing.java

@@ -77,6 +77,28 @@ public class TestWithClauseProcessing {
 	    helpProcess(plan, dataManager, expected);
 	}
 
+	@Test public void testWithPushdown1() throws TeiidException {
+		FakeCapabilitiesFinder capFinder = new FakeCapabilitiesFinder();
+        BasicSourceCapabilities caps = TestOptimizer.getTypicalCapabilities();
+        caps.setCapabilitySupport(Capability.COMMON_TABLE_EXPRESSIONS, true);
+        caps.setCapabilitySupport(Capability.QUERY_FROM_JOIN_SELFJOIN, true);
+        caps.setCapabilitySupport(Capability.QUERY_AGGREGATES_COUNT, true);
+        capFinder.addCapabilities("pm1", caps); //$NON-NLS-1$
+       
+	    String sql = "with a as (select x, y, z from (select e1 as x, e2 as y, e3 as z from pm1.g1) v) SELECT count(a.x) from a, a z"; //$NON-NLS-1$
+	    
+	    HardcodedDataManager dataManager = new HardcodedDataManager(RealMetadataFactory.example1Cached());
+	    List<?>[] expected = new List[] { 
+		        Arrays.asList("a", 1, Boolean.FALSE),
+		    };    
+
+	    dataManager.addData("WITH a (x, y, z) AS (SELECT g_0.e1, g_0.e2, g_0.e3 FROM g1 AS g_0) SELECT COUNT(g_0.x) FROM a AS g_0, a AS g_1", expected);
+	    
+	    ProcessorPlan plan = TestOptimizer.helpPlan(sql, RealMetadataFactory.example1Cached(), null, capFinder, new String[] {"WITH a (x, y, z) AS (SELECT g_0.e1, g_0.e2, g_0.e3 FROM pm1.g1 AS g_0) SELECT COUNT(g_0.x) FROM a AS g_0, a AS g_1"}, ComparisonMode.EXACT_COMMAND_STRING);
+	    
+	    helpProcess(plan, dataManager, expected);
+	}
+	
 	@Test public void testWithPushdownWithConstants() throws TeiidException {
 		 FakeCapabilitiesFinder capFinder = new FakeCapabilitiesFinder();
         BasicSourceCapabilities caps = TestOptimizer.getTypicalCapabilities();