1 files changed, 75 insertions, 12 deletions

diff --git a/SQLiteStudio3/coreSQLiteStudio/selectresolver.cpp b/SQLiteStudio3/coreSQLiteStudio/selectresolver.cpp
index 1b7e33d..64da5b2 100644
--- a/SQLiteStudio3/coreSQLiteStudio/selectresolver.cpp
+++ b/SQLiteStudio3/coreSQLiteStudio/selectresolver.cpp
@@ -47,12 +47,14 @@ QList<QList<SelectResolver::Column>> SelectResolver::resolveColumns()
 QList<SelectResolver::Column> SelectResolver::resolve(SqliteSelect::Core *selectCore)
 {
     errors.clear();
+    extractCte(selectCore);
     return resolveCore(selectCore);
 }
 
 QList<QList<SelectResolver::Column>> SelectResolver::resolve(SqliteSelect *select)
 {
     errors.clear();
+    extractCte(select);
     QList<QList<SelectResolver::Column> > results;
     for (SqliteSelect::Core* core : select->coreSelects)
     {
@@ -66,12 +68,14 @@ QList<QList<SelectResolver::Column>> SelectResolver::resolve(SqliteSelect *selec
 QList<SelectResolver::Column> SelectResolver::resolveAvailableColumns(SqliteSelect::Core *selectCore)
 {
     errors.clear();
+    extractCte(selectCore);
     return resolveAvailableCoreColumns(selectCore);
 }
 
 QList<QList<SelectResolver::Column> > SelectResolver::resolveAvailableColumns(SqliteSelect *select)
 {
     errors.clear();
+    extractCte(select);
     QList<QList<SelectResolver::Column> > results;
     for (SqliteSelect::Core* core : select->coreSelects)
         results << resolveAvailableCoreColumns(core);
@@ -82,6 +86,7 @@ QList<QList<SelectResolver::Column> > SelectResolver::resolveAvailableColumns(Sq
 QSet<SelectResolver::Table> SelectResolver::resolveTables(SqliteSelect::Core *selectCore)
 {
     QSet<Table> tables;
+    extractCte(selectCore);
     QList<Column> columns = resolveAvailableColumns(selectCore);
     for (const Column& col : columns)
     {
@@ -97,6 +102,7 @@ QSet<SelectResolver::Table> SelectResolver::resolveTables(SqliteSelect::Core *se
 QList<QSet<SelectResolver::Table> > SelectResolver::resolveTables(SqliteSelect *select)
 {
     QList<QSet<Table> > results;
+    extractCte(select);
     QList<QList<Column> > columnLists = resolveAvailableColumns(select);
     for (const QList<Column>& columns : columnLists)
     {
@@ -118,6 +124,7 @@ QList<QSet<SelectResolver::Table> > SelectResolver::resolveTables(SqliteSelect *
 QList<SelectResolver::Column> SelectResolver::translateToColumns(SqliteSelect* select, const TokenList& columnTokens)
 {
     errors.clear();
+    extractCte(select);
     QList<SelectResolver::Column> results;
     for (const TokenPtr& token : columnTokens)
         results << translateTokenToColumn(select, token);
@@ -175,7 +182,7 @@ QList<SelectResolver::Column> SelectResolver::resolveAvailableCoreColumns(Sqlite
 
     SqliteSelect* select = dynamic_cast<SqliteSelect*>(selectCore->parentStatement());
     if (select && select->with)
-        markCteColumns();
+        markCteColumns(&columns);
 
     return columns;
 }
@@ -237,9 +244,9 @@ void SelectResolver::markCompoundColumns()
     markCurrentColumnsWithFlag(FROM_COMPOUND_SELECT);
 }
 
-void SelectResolver::markCteColumns()
+void SelectResolver::markCteColumns(QList<Column>* columnList)
 {
-    markCurrentColumnsWithFlag(FROM_CTE_SELECT);
+    markCurrentColumnsWithFlag(FROM_CTE_SELECT, columnList);
 }
 
 void SelectResolver::markGroupedColumns()
@@ -283,9 +290,9 @@ void SelectResolver::fixColumnNames()
     }
 }
 
-void SelectResolver::markCurrentColumnsWithFlag(SelectResolver::Flag flag)
+void SelectResolver::markCurrentColumnsWithFlag(SelectResolver::Flag flag, QList<Column>* columnList)
 {
-    QMutableListIterator<Column> it(currentCoreResults);
+    QMutableListIterator<Column> it(columnList ? *columnList : currentCoreResults);
     while (it.hasNext())
         it.next().flags |= flag;
 }
@@ -528,6 +535,24 @@ TokenList SelectResolver::getResColTokensWithoutAlias(SqliteSelect::Core::Result
     return allTokens;
 }
 
+void SelectResolver::extractCte(SqliteSelect *select)
+{
+    cteList.clear();
+    if (!select->with)
+        return;
+
+    for (SqliteWith::CommonTableExpression* cte : select->with->cteList)
+        cteList[cte->table] = cte;
+}
+
+void SelectResolver::extractCte(SqliteSelect::Core *core)
+{
+    if (!core->parentStatement())
+        return;
+
+    extractCte(dynamic_cast<SqliteSelect*>(core->parentStatement()));
+}
+
 QList<SelectResolver::Column> SelectResolver::resolveJoinSource(SqliteSelect::Core::JoinSource *joinSrc)
 {
     QList<SelectResolver::Column> columnSources;
@@ -549,28 +574,56 @@ QList<SelectResolver::Column> SelectResolver::resolveSingleSource(SqliteSelect::
     if (joinSrc->joinSource)
         return resolveJoinSource(joinSrc->joinSource);
 
+    if (!joinSrc->funcName.isNull())
+        return resolveTableFunctionColumns(joinSrc);
+
     if (isView(joinSrc->database, joinSrc->table))
         return resolveView(joinSrc->database, joinSrc->table, joinSrc->alias);
 
     QList<Column> columnSources;
     QStringList columns = getTableColumns(joinSrc->database, joinSrc->table, joinSrc->alias);
     Column column;
+    column.type = Column::COLUMN;
+    column.table = joinSrc->table;;
+    column.database = joinSrc->database;
+    column.originalDatabase = resolveDatabase(joinSrc->database);
+    if (!joinSrc->alias.isNull())
+        column.tableAlias = joinSrc->alias;
+
     for (const QString& columnName : columns)
     {
-        column.type = Column::COLUMN;
         column.column = columnName;
-        column.table = joinSrc->table;;
-        column.database = joinSrc->database;
-        column.originalDatabase = resolveDatabase(joinSrc->database);
-        if (!joinSrc->alias.isNull())
-            column.tableAlias = joinSrc->alias;
-
         columnSources << column;
     }
 
     return columnSources;
 }
 
+QList<SelectResolver::Column> SelectResolver::resolveTableFunctionColumns(SqliteSelect::Core::SingleSource *joinSrc)
+{
+    static_qstring(columnSqlTpl, "SELECT * FROM %1 LIMIT 0");
+    SqlQueryPtr result = db->exec(columnSqlTpl.arg(joinSrc->detokenize()));
+    if (result->isError())
+        errors << result->getErrorText();
+
+    QStringList columnNames = result->getColumnNames();
+
+    QList<Column> columnSources;
+    Column column;
+    column.type = Column::OTHER;
+    column.database = joinSrc->database;
+    column.originalDatabase = resolveDatabase(joinSrc->database);
+    if (!joinSrc->alias.isNull())
+        column.tableAlias = joinSrc->alias;
+
+    for (const QString& columnName : columnNames)
+    {
+        column.column = columnName;
+        columnSources << column;
+    }
+    return columnSources;
+}
+
 QList<SelectResolver::Column> SelectResolver::resolveSingleSourceSubSelect(SqliteSelect::Core::SingleSource *joinSrc)
 {
     QList<Column> columnSources = resolveSubSelect(joinSrc->select);
@@ -652,7 +705,17 @@ QStringList SelectResolver::getTableColumns(const QString &database, const QStri
     dbTable.tableAlias = alias;
 
     if (tableColumnsCache.contains(dbTable))
+    {
         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);