Repository: AlecStrong/sql-psi
Pull request #368 (2022-04-07 22:25:29)

Pass the projection columns to the group by/having statements
Closes https://github.com/AlecStrong/sql-psi/issues/192


Patch:
diff --git a/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/QueryElement.kt b/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/QueryElement.kt
index d518f6f..e967baa 100644
--- a/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/QueryElement.kt
+++ b/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/QueryElement.kt
@@ -35,7 +35,7 @@ interface QueryElement : PsiElement {
     constructor(column: PsiElement) : this(columns = listOf(QueryColumn(column)))
 
     override fun toString(): String {
-      return "${table?.name} : [${columns.joinToString { it.element.text }}]"
+      return "${table?.name} : [${columns.joinToString { it.element.text }} [${synthesizedColumns.joinToString { it.acceptableValues.toString() }}]]"
     }
   }
 
diff --git a/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/ColumnAliasMixin.kt b/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/ColumnAliasMixin.kt
index 80524a3..dd5cf66 100644
--- a/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/ColumnAliasMixin.kt
+++ b/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/ColumnAliasMixin.kt
@@ -19,7 +19,7 @@ internal abstract class ColumnAliasMixin(
   override fun source(): PsiElement {
     parent.let {
       return when (it) {
-        is ResultColumnMixin -> it.expr!!
+        is ResultColumnMixin -> it
 
         is SqlCteTableName -> {
           val index = it.columnAliasList.indexOf(this)
diff --git a/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/SelectStmtMixin.kt b/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/SelectStmtMixin.kt
index 89d216d..e628e53 100644
--- a/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/SelectStmtMixin.kt
+++ b/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/SelectStmtMixin.kt
@@ -20,6 +20,7 @@ import com.alecstrong.sql.psi.core.psi.SqlTypes
 import com.alecstrong.sql.psi.core.psi.asColumns
 import com.intellij.lang.ASTNode
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiNamedElement
 import com.intellij.psi.util.PsiTreeUtil
 
 internal abstract class SelectStmtMixin(
@@ -27,6 +28,12 @@ internal abstract class SelectStmtMixin(
 ) : SqlCompositeElementImpl(node),
   SqlSelectStmt,
   FromQuery {
+  /**
+   * During some resolution steps we don't care about the parent's projection and can safely ignore
+   * it to avoid recursing too far.
+   */
+  private var ignoreParentProjection = false
+
   private val queryExposed = ModifiableFileLazy {
     if (valuesExpressionList.isNotEmpty()) {
       return@ModifiableFileLazy listOf(QueryResult(null, valuesExpressionList.first().exprList.asColumns()))
@@ -47,7 +54,27 @@ internal abstract class SelectStmtMixin(
   }
 
   override fun queryAvailable(child: PsiElement): Collection<QueryResult> {
-    if (child in exprList || child in resultColumnList) {
+    if (child in exprList) {
+      val available = fromQuery().map { it.copy(adjacent = true) } +
+        super.queryAvailable(this).map { it.copy(adjacent = false) }
+      if (ignoreParentProjection) return available
+
+      val projection = (parent as CompoundSelectStmtMixin).queryExposed().map { selectStmt ->
+        selectStmt.copy(
+          adjacent = false,
+          columns = selectStmt.columns.filter { projectionColumn ->
+            // Avoid including any projection columns that would create name collisions.
+            (projectionColumn.element as? PsiNamedElement)?.name !in
+              available.flatMap {
+                it.columns.mapNotNull { (it.element as? PsiNamedElement)?.name }
+              }
+          }
+        )
+      }
+
+      return available + projection
+    }
+    if (child in resultColumnList) {
       return fromQuery().map { it.copy(adjacent = true) } +
         super.queryAvailable(this).map { it.copy(adjacent = false) }
     }
@@ -83,11 +110,16 @@ internal abstract class SelectStmtMixin(
   }
 
   private fun SqlColumnName.isSameAs(other: SqlColumnName): Boolean {
-    if (this == other) return true
-    val thisRef = reference?.resolve() ?: return false
-    if (thisRef == other) return true
-    val otherRef = other.reference?.resolve() ?: return false
-    return thisRef == otherRef
+    try {
+      ignoreParentProjection = true
+      if (this == other) return true
+      val thisRef = reference?.resolve() ?: return false
+      if (thisRef == other) return true
+      val otherRef = other.reference?.resolve() ?: return false
+      return thisRef == otherRef
+    } finally {
+      ignoreParentProjection = false
+    }
   }
 
   override fun annotate(annotationHolder: SqlAnnotationHolder) {
diff --git a/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/TableOrSubqueryMixin.kt b/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/TableOrSubqueryMixin.kt
index 851f494..4593aff 100644
--- a/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/TableOrSubqueryMixin.kt
+++ b/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/TableOrSubqueryMixin.kt
@@ -17,7 +17,7 @@ internal abstract class TableOrSubqueryMixin(
         return@lazy emptyList<QueryResult>()
       }
       tableAlias?.let { alias ->
-        return@lazy listOf(QueryResult(alias, result.flatMap { it.columns }))
+        return@lazy listOf(QueryResult(alias, result.flatMap { it.columns }, result.flatMap { it.synthesizedColumns }))
       }
       return@lazy result
     }
diff --git a/core/src/main/kotlin/com/alecstrong/sql/psi/core/sql.bnf b/core/src/main/kotlin/com/alecstrong/sql/psi/core/sql.bnf
index ab21b4d..b138938 100644
--- a/core/src/main/kotlin/com/alecstrong/sql/psi/core/sql.bnf
+++ b/core/src/main/kotlin/com/alecstrong/sql/psi/core/sql.bnf
@@ -57,7 +57,7 @@ stmt_list ::= ( stmt ';' ) * {
 }
 stmt ::= [ EXPLAIN [ QUERY PLAN ] ] ( alter_table_stmt | analyze_stmt | attach_stmt | begin_stmt | commit_stmt | create_index_stmt | create_table_stmt | create_trigger_stmt | create_view_stmt |
                                           create_virtual_table_stmt | delete_stmt_limited | detach_stmt | drop_index_stmt | drop_table_stmt | drop_trigger_stmt | drop_view_stmt | insert_stmt |
-                                          pragma_stmt | reindex_stmt | release_stmt | rollback_stmt | savepoint_stmt | compound_select_stmt | update_stmt_limited | vacuum_stmt
+                                          pragma_stmt | reindex_stmt | release_stmt | rollback_stmt | savepoint_stmt | compound_select_stmt | update_stmt_limited | vacuum_stmt | extension_stmt
                                         ) {
   recoverWhile=sql_stmt_recovery
   pin = 2
@@ -116,7 +116,8 @@ column_def ::= [ javadoc ] column_name column_type ( column_constraint ) * {
 }
 column_type ::= type_name
 type_name ::= identifier [ '(' signed_number ')' | '(' signed_number ',' signed_number ')' ]
-column_constraint ::= [ CONSTRAINT identifier ] ( PRIMARY KEY [ ASC | DESC ] conflict_clause [ AUTOINCREMENT ] | NOT NULL conflict_clause | UNIQUE conflict_clause | check_constraint | default_constraint | COLLATE collation_name | foreign_key_clause )
+column_constraint ::= [ CONSTRAINT identifier ] ( PRIMARY KEY [ ASC | DESC ] conflict_clause [ AUTOINCREMENT ] | NOT NULL conflict_clause | UNIQUE conflict_clause | check_constraint | default_constraint | COLLATE collation_name | foreign_key_clause | generated_clause )
+generated_clause ::= [ GENERATED ] ALWAYS AS '(' expr ')'
 check_constraint ::= CHECK '(' expr ')'
 default_constraint ::= DEFAULT ( signed_number | literal_value | '(' expr ')' )
 signed_number ::= [ '+' | '-' ] numeric_literal
@@ -225,6 +226,7 @@ expr ::= ( raise_expr
 
 other_expr ::= extension_expr
 extension_expr ::= NULL
+extension_stmt ::= NULL
 literal_expr ::= literal_value
 bind_expr ::= bind_parameter
 column_expr ::= [ database_name '.' table_name '.' | table_name '.' ] column_name
@@ -255,7 +257,7 @@ binary_and_expr ::= expr AND expr {
 binary_or_expr ::= expr OR expr {
   implements = "com.alecstrong.sql.psi.core.psi.SqlBinaryExpr"
 }
-function_expr ::= function_name '(' [ [ DISTINCT ] expr ( ( ',' | FROM | FOR ) expr ) * | '*' ] ')'
+function_expr ::= function_name '(' [ [ DISTINCT ] expr ( ( ',' | FROM | FOR | SEPARATOR ) expr ) * | '*' ] ')'
 paren_expr ::= '(' expr ')' {
   pin = 1
 }


Test patch:
diff --git a/test-fixtures/src/main/resources/fixtures/alias-in-having/Test.s b/test-fixtures/src/main/resources/fixtures/alias-in-having/Test.s
new file mode 100644
index 0000000..85b68b6
--- /dev/null
+++ b/test-fixtures/src/main/resources/fixtures/alias-in-having/Test.s
@@ -0,0 +1,12 @@
+-- Table definition edited down to the relevant columns:
+CREATE TABLE `models` (
+  `model_id` TEXT NOT NULL AUTOINCREMENT,
+  `model_descriptor_id` TEXT NOT NULL,
+  `build_date` TEXT NOT NULL,
+  PRIMARY KEY (`model_id`)
+);
+
+SELECT max(build_date) AS build_date_seconds
+FROM models
+GROUP BY model_descriptor_id
+HAVING build_date_seconds > 0 LIMIT 1;
\ No newline at end of file
diff --git a/test-fixtures/src/main/resources/fixtures/rowid-orderby/Test.s b/test-fixtures/src/main/resources/fixtures/rowid-orderby/Test.s
new file mode 100644
index 0000000..90b412a
--- /dev/null
+++ b/test-fixtures/src/main/resources/fixtures/rowid-orderby/Test.s
@@ -0,0 +1,14 @@
+CREATE TABLE Test1(id TEXT);
+CREATE TABLE Test2(id TEXT);
+
+SELECT t1.*
+FROM Test1 t1
+WHERE t1.rowid IS NOT NULL
+ORDER BY t1.rowid ASC
+LIMIT 1;
+
+SELECT t1.*, t2.* FROM Test1 t1
+JOIN Test2 t2 ON t1.id = t2.id
+WHERE t1.id = 'something good'
+ORDER BY t1.rowid ASC
+LIMIT 1;
\ No newline at end of file