Support extension schema contributors with IndexElement

core/src/main/kotlin/com/alecstrong/sql/psi/core/SqlFileBase.kt

@@ -7,6 +7,7 @@ import com.alecstrong.sql.psi.core.psi.SchemaContributorIndex
 import com.alecstrong.sql.psi.core.psi.SqlCreateTableStmt
 import com.alecstrong.sql.psi.core.psi.SqlCreateTriggerStmt
 import com.alecstrong.sql.psi.core.psi.SqlCreateViewStmt
+import com.alecstrong.sql.psi.core.psi.SqlExtensionStmt
 import com.alecstrong.sql.psi.core.psi.SqlStmtList
 import com.alecstrong.sql.psi.core.psi.TableElement
 import com.intellij.extapi.psi.PsiFileBase
@@ -133,8 +134,16 @@ abstract class SqlFileBase(viewProvider: FileViewProvider, language: Language) :
       ?.forEach { block(it) }
   }
 
-  private fun contributors() =
-    sqlStmtList?.stmtList?.mapNotNull { it.firstChild as? SchemaContributor }
+  private fun contributors(): List<SchemaContributor>? {
+    // Preserve schemaContributor statement ordering and add any schemaContributors that are first
+    // child of SqlExtensionStmt.
+    return sqlStmtList?.stmtList?.mapNotNull { stmt ->
+      stmt.firstChild as? SchemaContributor
+        ?: (stmt.firstChild as? SqlExtensionStmt)?.let { sqlExt ->
+          sqlExt.firstChild as? SchemaContributor
+        }
+    }
+  }
 
   /**
    * Optional files which can be used for extra Schema Contributors that are unindexed. The files

core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/IndexElement.kt

@@ -0,0 +1,12 @@
+package com.alecstrong.sql.psi.core.psi
+
+import com.intellij.psi.StubBasedPsiElement
+
+/**
+ * IndexElement represents a named index element implemented by CreateIndex, DropIndex, AlterIndex.
+ * The indexElementName can be overridden to return the new name of an AlterIndex element
+ */
+interface IndexElement :
+  SqlCompositeElement, SchemaContributor, StubBasedPsiElement<SchemaContributorStub> {
+  fun indexElementName(): String = name()
+}

core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/SchemaContributor.kt

@@ -47,4 +47,16 @@ class Schema {
   @Suppress("UNCHECKED_CAST")
   fun <Value : SchemaContributor> values(type: KClass<Value>) =
     map[type]?.values as Collection<Value>? ?: emptyList()
+
+  fun replace(
+    type: KClass<out SchemaContributor>,
+    name: String,
+    newName: String,
+    value: SchemaContributor,
+  ) {
+    @Suppress("UNCHECKED_CAST")
+    val typedMap = map.getOrPut(type) { linkedMapOf() } as MutableMap<String, SchemaContributor>
+    typedMap.remove(name)
+    typedMap[newName] = value
+  }
 }

core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/SqlCompositeElementImpl.kt

@@ -39,7 +39,7 @@ open class SqlCompositeElementImpl(node: ASTNode) :
     }
 }
 
-internal abstract class SqlSchemaContributorImpl<
+abstract class SqlSchemaContributorImpl<
   SchemaType : SchemaContributor,
   ElementType : SqlSchemaContributorElementType<SchemaType>,
 >(stub: SchemaContributorStub?, nodeType: IElementType?, node: ASTNode?) :

core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/CreateIndexMixin.kt

@@ -2,6 +2,7 @@ package com.alecstrong.sql.psi.core.psi.mixins
 
 import com.alecstrong.sql.psi.core.SqlAnnotationHolder
 import com.alecstrong.sql.psi.core.SqlSchemaContributorElementType
+import com.alecstrong.sql.psi.core.psi.IndexElement
 import com.alecstrong.sql.psi.core.psi.QueryElement.QueryResult
 import com.alecstrong.sql.psi.core.psi.Schema
 import com.alecstrong.sql.psi.core.psi.SchemaContributorStub
@@ -13,13 +14,11 @@ import com.intellij.lang.ASTNode
 import com.intellij.psi.PsiElement
 import com.intellij.psi.tree.IElementType
 
-internal abstract class CreateIndexMixin(
-  stub: SchemaContributorStub?,
-  nodeType: IElementType?,
-  node: ASTNode?,
-) :
-  SqlSchemaContributorImpl<SqlCreateIndexStmt, CreateIndexElementType>(stub, nodeType, node),
-  SqlCreateIndexStmt {
+internal abstract class CreateIndexMixin
+private constructor(stub: SchemaContributorStub?, nodeType: IElementType?, node: ASTNode?) :
+  SqlSchemaContributorImpl<IndexElement, CreateIndexElementType>(stub, nodeType, node),
+  SqlCreateIndexStmt,
+  IndexElement {
   constructor(node: ASTNode) : this(null, null, node)
 
   constructor(stub: SchemaContributorStub, nodeType: IElementType) : this(stub, nodeType, null)
@@ -32,7 +31,7 @@ internal abstract class CreateIndexMixin(
   }
 
   override fun modifySchema(schema: Schema) {
-    schema.put<SqlCreateIndexStmt>(this)
+    schema.put<IndexElement>(this)
   }
 
   override fun queryAvailable(child: PsiElement): Collection<QueryResult> {
@@ -47,8 +46,8 @@ internal abstract class CreateIndexMixin(
   override fun annotate(annotationHolder: SqlAnnotationHolder) {
     if (
       node.findChildByType(SqlTypes.EXISTS) == null &&
-        containingFile.schema<SqlCreateIndexStmt>(this).any {
-          it != this && it.indexName.textMatches(indexName)
+        containingFile.schema<IndexElement>(this).any {
+          it != this && it.indexElementName() == indexName.text
         }
     ) {
       annotationHolder.createErrorAnnotation(indexName, "Duplicate index name ${indexName.text}")
@@ -58,7 +57,7 @@ internal abstract class CreateIndexMixin(
 }
 
 open class CreateIndexElementType(name: String) :
-  SqlSchemaContributorElementType<SqlCreateIndexStmt>(name, SqlCreateIndexStmt::class.java) {
+  SqlSchemaContributorElementType<IndexElement>(name, IndexElement::class.java) {
   override fun nameType() = SqlTypes.INDEX_NAME
 
   override fun createPsi(stub: SchemaContributorStub) = SqlCreateIndexStmtImpl(stub, this)

core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/CreateTableMixin.kt

@@ -2,6 +2,7 @@ package com.alecstrong.sql.psi.core.psi.mixins
 
 import com.alecstrong.sql.psi.core.SqlAnnotationHolder
 import com.alecstrong.sql.psi.core.SqlSchemaContributorElementType
+import com.alecstrong.sql.psi.core.psi.IndexElement
 import com.alecstrong.sql.psi.core.psi.LazyQuery
 import com.alecstrong.sql.psi.core.psi.QueryElement.QueryResult
 import com.alecstrong.sql.psi.core.psi.QueryElement.SynthesizedColumn
@@ -120,7 +121,8 @@ private constructor(stub: SchemaContributorStub?, nodeType: IElementType?, node:
 
     // Check in file for `CREATE UNIQUE INDEX` statement that matches the given columns.
     containingFile
-      .schema<SqlCreateIndexStmt>() // sqlStmtElement is null to include all statements
+      .schema<IndexElement>()
+      .filterIsInstance<SqlCreateIndexStmt>()
       .filter { it.isUnique() && it.indexedColumnList.all { it.collationName == null } }
       .forEach {
         val indexedColumns =

core/src/main/kotlin/com/alecstrong/sql/psi/core/psi/mixins/DropIndexMixin.kt

@@ -2,9 +2,9 @@ package com.alecstrong.sql.psi.core.psi.mixins
 
 import com.alecstrong.sql.psi.core.SqlAnnotationHolder
 import com.alecstrong.sql.psi.core.SqlSchemaContributorElementType
+import com.alecstrong.sql.psi.core.psi.IndexElement
 import com.alecstrong.sql.psi.core.psi.Schema
 import com.alecstrong.sql.psi.core.psi.SchemaContributorStub
-import com.alecstrong.sql.psi.core.psi.SqlCreateIndexStmt
 import com.alecstrong.sql.psi.core.psi.SqlDropIndexStmt
 import com.alecstrong.sql.psi.core.psi.SqlSchemaContributorImpl
 import com.alecstrong.sql.psi.core.psi.SqlTypes
@@ -14,7 +14,7 @@ import com.intellij.psi.tree.IElementType
 
 internal abstract class DropIndexMixin
 private constructor(stub: SchemaContributorStub?, nodeType: IElementType?, node: ASTNode?) :
-  SqlSchemaContributorImpl<SqlCreateIndexStmt, DropIndexElementType>(stub, nodeType, node),
+  SqlSchemaContributorImpl<IndexElement, DropIndexElementType>(stub, nodeType, node),
   SqlDropIndexStmt {
   constructor(node: ASTNode) : this(null, null, node)
 
@@ -28,15 +28,15 @@ private constructor(stub: SchemaContributorStub?, nodeType: IElementType?, node:
   }
 
   override fun modifySchema(schema: Schema) {
-    schema.forType<SqlCreateIndexStmt>().remove(name())
+    schema.forType<IndexElement>().remove(name())
   }
 
   override fun annotate(annotationHolder: SqlAnnotationHolder) {
     indexName?.let { indexName ->
       if (
         node.findChildByType(SqlTypes.EXISTS) == null &&
-          containingFile.schema<SqlCreateIndexStmt>(this).none {
-            it != this && it.indexName.textMatches(indexName)
+          containingFile.schema<IndexElement>(this).none {
+            it != this && it.indexElementName() == indexName.text
           }
       ) {
         annotationHolder.createErrorAnnotation(
@@ -51,8 +51,8 @@ private constructor(stub: SchemaContributorStub?, nodeType: IElementType?, node:
 }
 
 internal class DropIndexElementType(name: String) :
-  SqlSchemaContributorElementType<SqlCreateIndexStmt>(name, SqlCreateIndexStmt::class.java) {
-  override fun nameType() = SqlTypes.TABLE_NAME
+  SqlSchemaContributorElementType<IndexElement>(name, IndexElement::class.java) {
+  override fun nameType() = SqlTypes.INDEX_NAME
 
   override fun createPsi(stub: SchemaContributorStub) = SqlDropIndexStmtImpl(stub, this)
 }

core/src/main/kotlin/com/alecstrong/sql/psi/core/sql.bnf

@@ -103,7 +103,7 @@ release_stmt ::= RELEASE [ SAVEPOINT ] savepoint_name {
 }
 create_index_stmt ::= CREATE [ UNIQUE ] INDEX [ IF NOT EXISTS ] [ database_name DOT ] index_name ON table_name LP indexed_column ( COMMA indexed_column ) * RP [ WHERE expr ] {
   mixin = "com.alecstrong.sql.psi.core.psi.mixins.CreateIndexMixin"
-  implements = "com.alecstrong.sql.psi.core.psi.SchemaContributor"
+  implements = "com.alecstrong.sql.psi.core.psi.IndexElement"
   elementTypeClass = "com.alecstrong.sql.psi.core.psi.mixins.CreateIndexElementType"
   stubClass = "com.alecstrong.sql.psi.core.psi.SchemaContributorStub"
   pin = 6