Repository: AlecStrong/sql-psi
Pull request #417 (2022-08-11 20:59:41)

Support functions in index
Fixes https://github.com/cashapp/sqldelight/issues/3412

I think it is okay to not support the indexing of columns used as parameter, isn't it?


Patch:
diff --git a/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/CreateTableMixin.kt b/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/CreateTableMixin.kt
index 49b5120..4d4c8cc 100644
--- a/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/CreateTableMixin.kt
+++ b/core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/CreateTableMixin.kt
@@ -7,6 +7,7 @@ import com.alecstrong.sql.psi.core.psi.QueryElement.QueryResult
 import com.alecstrong.sql.psi.core.psi.QueryElement.SynthesizedColumn
 import com.alecstrong.sql.psi.core.psi.Schema
 import com.alecstrong.sql.psi.core.psi.SchemaContributorStub
+import com.alecstrong.sql.psi.core.psi.SqlColumnExpr
 import com.alecstrong.sql.psi.core.psi.SqlColumnName
 import com.alecstrong.sql.psi.core.psi.SqlCompositeElement
 import com.alecstrong.sql.psi.core.psi.SqlCompoundSelectStmt
@@ -91,7 +92,10 @@ internal abstract class CreateTableMixin private constructor(
   private fun primaryKey(): List<String> {
     val compositeKey = tableConstraintList.firstOrNull { it.hasPrimaryKey() }
     if (compositeKey != null) {
-      return compositeKey.indexedColumnList.mapNotNull { it.columnName?.name }
+      return compositeKey.indexedColumnList.mapNotNull {
+        val expr = it.expr
+        if (expr is SqlColumnExpr) expr.columnName.name else null
+      }
     }
 
     return columnDefList.filter { it.columnConstraintList.any { it.hasPrimaryKey() } }
@@ -101,7 +105,12 @@ internal abstract class CreateTableMixin private constructor(
 
   private fun isCollectivelyUnique(columns: List<SqlColumnName>): Boolean {
     tableConstraintList.filter { it.hasPrimaryKey() || it.isUnique() }
-      .map { it.indexedColumnList.mapNotNull { it.columnName?.name } }
+      .map {
+        it.indexedColumnList.mapNotNull {
+          val expr = it.expr
+          if (expr is SqlColumnExpr) expr.columnName.name else null
+        }
+      }
       .plus(
         listOf(
           columnDefList.filter {
@@ -117,7 +126,10 @@ internal abstract class CreateTableMixin private constructor(
     containingFile.schema<SqlCreateIndexStmt>(this)
       .filter { it.isUnique() && it.indexedColumnList.all { it.collationName == null } }
       .forEach {
-        val indexedColumns = it.indexedColumnList.mapNotNull { it.columnName?.name }
+        val indexedColumns = it.indexedColumnList.mapNotNull {
+          val expr = it.expr
+          if (expr is SqlColumnExpr) expr.columnName.name else null
+        }
         if (columns.map { it.name }.containsAll(indexedColumns) &&
           columns.size == indexedColumns.size
         ) {
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 708359b..7be85af 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
@@ -106,7 +106,7 @@ create_index_stmt ::= CREATE [ UNIQUE ] INDEX [ IF NOT EXISTS ] [ database_name
   stubClass = "com.alecstrong.sql.psi.core.psi.SchemaContributorStub"
   pin = 6
 }
-indexed_column ::= [ column_name | expr ] [ COLLATE collation_name ] [ ASC | DESC ]
+indexed_column ::= expr [ COLLATE collation_name ] [ ASC | DESC ]
 create_table_stmt ::= CREATE [ TEMP | TEMPORARY ] TABLE [ IF NOT EXISTS ] [ database_name '.' ] table_name ( '(' column_def ( ',' column_def ) * ( ',' table_constraint ) * ')' [ table_options ] | AS compound_select_stmt ) {
   mixin = "com.alecstrong.sql.psi.core.psi.mixins.CreateTableMixin"
   implements = "com.alecstrong.sql.psi.core.psi.TableElement"


Test patch:
diff --git a/test-fixtures/src/main/resources/fixtures/create-index-validation-failures/Test.s b/test-fixtures/src/main/resources/fixtures/create-index-validation-failures/Test.s
index 4e294b4..98b44a7 100644
--- a/test-fixtures/src/main/resources/fixtures/create-index-validation-failures/Test.s
+++ b/test-fixtures/src/main/resources/fixtures/create-index-validation-failures/Test.s
@@ -13,4 +13,8 @@ WHERE column1 IS NOT NULL;
 
 -- valid.
 CREATE INDEX index_3
-ON test (_id COLLATE some_collation_name ASC);
\ No newline at end of file
+ON test (_id COLLATE some_collation_name ASC);
+
+-- valid.
+CREATE INDEX index_4
+ON test (coalesce(_id) COLLATE some_collation_name ASC);