1 files changed, 51 insertions, 27 deletions

diff --git a/SQLiteStudio3/coreSQLiteStudio/selectresolver.cpp b/SQLiteStudio3/coreSQLiteStudio/selectresolver.cpp
index 64da5b2..3db029d 100644
--- a/SQLiteStudio3/coreSQLiteStudio/selectresolver.cpp
+++ b/SQLiteStudio3/coreSQLiteStudio/selectresolver.cpp
@@ -168,9 +168,6 @@ QList<SelectResolver::Column> SelectResolver::resolveCore(SqliteSelect::Core* se
     if (select && select->coreSelects.size() > 1)
         markCompoundColumns();
 
-    if (select && select->with)
-        markCteColumns();
-
     return currentCoreResults;
 }
 
@@ -180,10 +177,6 @@ QList<SelectResolver::Column> SelectResolver::resolveAvailableCoreColumns(Sqlite
     if (selectCore->from)
         columns = resolveJoinSource(selectCore->from);
 
-    SqliteSelect* select = dynamic_cast<SqliteSelect*>(selectCore->parentStatement());
-    if (select && select->with)
-        markCteColumns(&columns);
-
     return columns;
 }
 
@@ -244,11 +237,6 @@ void SelectResolver::markCompoundColumns()
     markCurrentColumnsWithFlag(FROM_COMPOUND_SELECT);
 }
 
-void SelectResolver::markCteColumns(QList<Column>* columnList)
-{
-    markCurrentColumnsWithFlag(FROM_CTE_SELECT, columnList);
-}
-
 void SelectResolver::markGroupedColumns()
 {
     markCurrentColumnsWithFlag(FROM_GROUPED_SELECT);
@@ -361,16 +349,16 @@ void SelectResolver::resolveStar(SqliteSelect::Core::ResultColumn *resCol)
 
 void SelectResolver::resolveExpr(SqliteSelect::Core::ResultColumn *resCol)
 {
-    SelectResolver::Column column;
-    column.alias = resCol->alias;
-    column.originalColumn = resCol;
-    column.column = getResColTokensWithoutAlias(resCol).detokenize().trimmed();
-    column.displayName = !resCol->alias.isNull() ? column.alias : column.column;
-
     SqliteExpr* expr = resCol->expr;
     if (expr->mode != SqliteExpr::Mode::ID)
     {
         // Not a simple column, but some expression
+        SelectResolver::Column column;
+        column.alias = resCol->alias;
+        column.originalColumn = resCol;
+        column.column = getResColTokensWithoutAlias(resCol).detokenize().trimmed();
+        column.displayName = !resCol->alias.isNull() ? column.alias : column.column;
+
         column.type = Column::OTHER;
         currentCoreResults << column;
 
@@ -580,6 +568,9 @@ QList<SelectResolver::Column> SelectResolver::resolveSingleSource(SqliteSelect::
     if (isView(joinSrc->database, joinSrc->table))
         return resolveView(joinSrc->database, joinSrc->table, joinSrc->alias);
 
+    if (joinSrc->database.isNull() && cteList.contains(joinSrc->table))
+        return resolveCteColumns(joinSrc);
+
     QList<Column> columnSources;
     QStringList columns = getTableColumns(joinSrc->database, joinSrc->table, joinSrc->alias);
     Column column;
@@ -599,6 +590,47 @@ QList<SelectResolver::Column> SelectResolver::resolveSingleSource(SqliteSelect::
     return columnSources;
 }
 
+QList<SelectResolver::Column> SelectResolver::resolveCteColumns(SqliteSelect::Core::SingleSource* joinSrc)
+{
+    SqliteWith::CommonTableExpression* cte = cteList[joinSrc->table];
+
+    Column column;
+    column.type = Column::COLUMN;
+    column.flags |= FROM_CTE_SELECT;
+    column.tableAlias = cte->table;
+
+    QList<Column> columnSources;
+
+    static_qstring(cteSelectTpl, "WITH %1 SELECT * FROM %2");
+    QString selectQuery = cte->detokenize();
+    QString query = cteSelectTpl.arg(selectQuery, cte->table);
+    QList<AliasedColumn> queryColumns = db->columnsForQuery(query);
+    if (queryColumns.isEmpty())
+    {
+        qWarning() << "Could not detect query columns. Probably due to db error:" << db->getErrorText();
+        return columnSources;
+    }
+
+    for (const AliasedColumn& queryColumn : queryColumns)
+    {
+        if (!queryColumn.getDatabase().isNull())
+            column.database = resolveDatabase(queryColumn.getDatabase());
+        else
+            column.database = queryColumn.getDatabase();
+
+        column.table = queryColumn.getTable();
+
+        // From CTE perspective, however the column is received as "result column name" from SQLite API
+        // is what we report back to user of the CTE as available column. No matter if it's actual alias,
+        // or simply name of a column.
+        column.column = queryColumn.getAlias();
+        column.displayName = queryColumn.getAlias();
+        columnSources << column;
+    }
+
+    return columnSources;
+}
+
 QList<SelectResolver::Column> SelectResolver::resolveTableFunctionColumns(SqliteSelect::Core::SingleSource *joinSrc)
 {
     static_qstring(columnSqlTpl, "SELECT * FROM %1 LIMIT 0");
@@ -708,14 +740,6 @@ QStringList SelectResolver::getTableColumns(const QString &database, const QStri
     {
         return tableColumnsCache.value(dbTable);
     }
-    else if (database.isNull() && cteList.contains(table))
-    {
-        QStringList columns;
-        for (SqliteIndexedColumn* idxCol : cteList[table]->indexedColumns)
-            columns << idxCol->name;
-
-        return columns;
-    }
     else
     {
         QStringList columns = schemaResolver->getTableColumns(database, table);
@@ -760,7 +784,7 @@ bool SelectResolver::parseOriginalQuery()
     if (originalQueryParsed)
         return true;
 
-    Parser parser(db->getDialect());
+    Parser parser;
     if (!parser.parse(query) || parser.getQueries().isEmpty())
     {
         qWarning() << "Could not parse query in SelectResolver:" << query;