Support JSON table functions (#3090)

* Support the json table function names

* Json Integration test

.gitignore

@@ -1,5 +1,8 @@
 .gradle
-/.idea
+**/.idea
+sqldelight-gradle-plugin/**/gradle
+sqldelight-gradle-plugin/**/gradlew
+sqldelight-gradle-plugin/**/gradlew.bat
 *.iml
 build
 local.properties

dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/PostgreSqlTypeResolver.kt

@@ -125,7 +125,7 @@ class PostgreSqlTypeResolver(private val parentResolver: TypeResolver) : TypeRes
     val columnDef = intermediateType.column ?: return intermediateType
     val tableDef = columnDef.parent as? SqlCreateTableStmt ?: return intermediateType
     tableDef.tableConstraintList.forEach {
-      if (columnDef.columnName.name in it.indexedColumnList.map { it.columnName.name }) {
+      if (columnDef.columnName.name in it.indexedColumnList.mapNotNull { it.columnName?.name }) {
         return intermediateType.asNonNullable()
       }
     }

dialects/sqlite-3-35/src/main/kotlin/app/cash/sqldelight/dialects/sqlite_3_35/grammar/mixins/AlterTableDropColumnMixin.kt

@@ -65,7 +65,7 @@ internal abstract class AlterTableDropColumnMixin(
         containingFile
           .schema(SqlCreateIndexStmt::class, this)
           .find { index ->
-            index.indexedColumnList.any { it.columnName.textMatches(columnName) }
+            index.indexedColumnList.any { it.columnName?.textMatches(columnName) == true }
           }
           ?.let { indexForColumnToDrop ->
             annotationHolder.createErrorAnnotation(

dialects/sqlite/json-module/build.gradle

@@ -1,11 +1,25 @@
 plugins {
   alias(deps.plugins.kotlin.jvm)
+  alias(deps.plugins.grammarKitComposer)
   alias(deps.plugins.publish)
   alias(deps.plugins.dokka)
 }
 
+grammarKit {
+  intellijRelease.set(deps.versions.idea)
+}
+
 dependencies {
   compileOnly project(':sqldelight-compiler:dialect')
+  compileOnly deps.intellij.lang
+
+  testImplementation project(':dialects:sqlite-3-18')
+  testImplementation deps.intellij.core
+  testImplementation deps.intellij.lang
+  testImplementation deps.junit
+  testImplementation deps.truth
+  testImplementation project(':sqldelight-compiler:dialect')
+  testImplementation deps.sqlPsiTestFixtures
 }
 
 apply from: "$rootDir/gradle/gradle-mvn-push.gradle"

dialects/sqlite/json-module/src/main/kotlin/app/cash/sqldelight/dialects/sqlite/json/module/JsonFunctions.kt

@@ -4,11 +4,17 @@ import app.cash.sqldelight.dialect.api.IntermediateType
 import app.cash.sqldelight.dialect.api.PrimitiveType
 import app.cash.sqldelight.dialect.api.SqlDelightModule
 import app.cash.sqldelight.dialect.api.TypeResolver
+import app.cash.sqldelight.dialects.sqlite.json.module.grammar.JsonParserUtil
 import com.alecstrong.sql.psi.core.psi.SqlFunctionExpr
 
 class JsonModule : SqlDelightModule {
   override fun typeResolver(parentResolver: TypeResolver): TypeResolver =
     JsonTypeResolver(parentResolver)
+
+  override fun setup() {
+    JsonParserUtil.reset()
+    JsonParserUtil.overrideSqlParser()
+  }
 }
 
 private class JsonTypeResolver(private val parentResolver: TypeResolver) :

dialects/sqlite/json-module/src/main/kotlin/app/cash/sqldelight/dialects/sqlite/json/module/grammar/json.bnf

@@ -0,0 +1,27 @@
+{
+  // Specify the parent parser.
+  overrides="com.alecstrong.sql.psi.core.SqlParser"
+  elementTypeClass = "com.alecstrong.sql.psi.core.SqlElementType"
+
+  implements="com.alecstrong.sql.psi.core.psi.SqlCompositeElement"
+  extends="com.alecstrong.sql.psi.core.psi.SqlCompositeElementImpl"
+  psiClassPrefix = "SqliteJson"
+}
+overrides ::= table_or_subquery
+
+table_or_subquery ::= ( json_function_name '(' <<expr '-1'>> ( ',' <<expr '-1'>> ) * ')'
+                      | [ {database_name} '.' ] {table_name} [ [ 'AS' ] {table_alias} ] [ 'INDEXED' 'BY' {index_name} | 'NOT' 'INDEXED' ]
+                      | '(' ( {table_or_subquery} ( ',' {table_or_subquery} ) * | {join_clause} ) ')'
+                      | '(' {compound_select_stmt} ')' [ [ 'AS' ] {table_alias} ] ) {
+  mixin = "app.cash.sqldelight.dialects.sqlite.json.module.grammar.mixins.TableOrSubqueryMixin"
+  implements = "com.alecstrong.sql.psi.core.psi.SqlTableOrSubquery"
+  override = true
+}
+
+json_function_name ::= 'json_each' | 'json_tree' {
+  mixin = "app.cash.sqldelight.dialects.sqlite.json.module.grammar.mixins.JsonFunctionNameMixin"
+  implements = [
+    "com.alecstrong.sql.psi.core.psi.NamedElement";
+    "com.alecstrong.sql.psi.core.psi.SqlCompositeElement"
+  ]
+}

dialects/sqlite/json-module/src/main/kotlin/app/cash/sqldelight/dialects/sqlite/json/module/grammar/mixins/TableOrSubqueryMixin.kt

@@ -0,0 +1,52 @@
+package app.cash.sqldelight.dialects.sqlite.json.module.grammar.mixins
+
+import app.cash.sqldelight.dialect.api.ExposableType
+import app.cash.sqldelight.dialect.api.IntermediateType
+import app.cash.sqldelight.dialect.api.PrimitiveType
+import app.cash.sqldelight.dialects.sqlite.json.module.grammar.JsonParser
+import app.cash.sqldelight.dialects.sqlite.json.module.grammar.psi.SqliteJsonTableOrSubquery
+import com.alecstrong.sql.psi.core.ModifiableFileLazy
+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.SqlExpr
+import com.alecstrong.sql.psi.core.psi.SqlJoinClause
+import com.alecstrong.sql.psi.core.psi.SqlNamedElementImpl
+import com.alecstrong.sql.psi.core.psi.SqlTableName
+import com.alecstrong.sql.psi.core.psi.impl.SqlTableOrSubqueryImpl
+import com.intellij.lang.ASTNode
+import com.intellij.lang.PsiBuilder
+import com.intellij.psi.PsiElement
+
+internal abstract class TableOrSubqueryMixin(node: ASTNode?) : SqlTableOrSubqueryImpl(node), SqliteJsonTableOrSubquery {
+  private val queryExposed = ModifiableFileLazy lazy@{
+    if (jsonFunctionName != null) {
+      return@lazy listOf(
+        QueryResult(
+          table = jsonFunctionName!!,
+          columns = emptyList(),
+          synthesizedColumns = listOf(
+            SynthesizedColumn(jsonFunctionName!!, acceptableValues = listOf("key", "value", "type", "atom", "id", "parent", "fullkey", "path", "json", "root"))
+          )
+        )
+      )
+    }
+    super.queryExposed()
+  }
+
+  override fun queryExposed() = queryExposed.forFile(containingFile)
+
+  override fun queryAvailable(child: PsiElement): Collection<QueryResult> {
+    if (child is SqlExpr) {
+      val parent = parent as SqlJoinClause
+      return parent.tableOrSubqueryList.takeWhile { it != this }.flatMap { it.queryExposed() }
+    }
+    return super.queryAvailable(child)
+  }
+}
+
+internal abstract class JsonFunctionNameMixin(node: ASTNode) : SqlNamedElementImpl(node), SqlTableName, ExposableType {
+  override fun getId(): PsiElement? = null
+  override fun getString(): PsiElement? = null
+  override val parseRule: (PsiBuilder, Int) -> Boolean = JsonParser::json_function_name_real
+  override fun type() = IntermediateType(PrimitiveType.TEXT)
+}

dialects/sqlite/json-module/src/test/fixtures/json_table_functions/Test.s

@@ -0,0 +1,27 @@
+CREATE TABLE user(name TEXT, phone TEXT);
+
+SELECT DISTINCT user.name
+  FROM user, json_each(user.phone)
+ WHERE json_each.value LIKE '704-%';
+
+SELECT name FROM user WHERE phone LIKE '704-%'
+UNION
+SELECT user.name
+  FROM user, json_each(user.phone)
+ WHERE json_valid(user.phone)
+   AND json_each.value LIKE '704-%';
+
+CREATE TABLE big(json TEXT);
+
+SELECT big.rowid, fullkey, value
+  FROM big, json_tree(big.json)
+ WHERE json_tree.type NOT IN ('object','array');
+
+SELECT big.rowid, fullkey, atom
+  FROM big, json_tree(big.json)
+ WHERE atom IS NOT NULL;
+
+SELECT DISTINCT json_extract(big.json,'$.id')
+  FROM big, json_tree(big.json, '$.partlist')
+ WHERE json_tree.key='uuid'
+   AND json_tree.value='6fa5181e-5721-11e5-a04e-57f3d7b32808';

dialects/sqlite/json-module/src/test/kotlin/app/cash/sqldelight/dialects/sqlite/json/module/JsonModuleTest.kt

@@ -0,0 +1,28 @@
+package app.cash.sqldelight.dialects.sqlite.json.module
+
+import app.cash.sqldelight.dialects.sqlite_3_18.SqliteDialect
+import com.alecstrong.sql.psi.test.fixtures.FixturesTest
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+import org.junit.runners.Parameterized.Parameters
+import java.io.File
+
+@RunWith(Parameterized::class)
+class JsonModuleTest(name: String, fixtureRoot: File) : FixturesTest(name, fixtureRoot) {
+  override fun setupDialect() {
+    SqliteDialect().setup()
+    JsonModule().setup()
+  }
+
+  companion object {
+    private val fixtures = arrayOf("src/test/fixtures")
+
+    @Suppress("unused") // Used by Parameterized JUnit runner reflectively.
+    @Parameters(name = "{0}")
+    @JvmStatic fun parameters() = fixtures.flatMap { fixtureFolder ->
+      File(fixtureFolder).listFiles()!!
+        .filter { it.isDirectory }
+        .map { arrayOf(it.name, it) }
+    } + ansiFixtures
+  }
+}

sqldelight-compiler/dialect/src/main/kotlin/app/cash/sqldelight/dialect/api/ExposableType.kt

@@ -0,0 +1,7 @@
+package app.cash.sqldelight.dialect.api
+
+import com.alecstrong.sql.psi.core.psi.SqlAnnotatedElement
+
+interface ExposableType : SqlAnnotatedElement {
+  fun type(): IntermediateType
+}

sqldelight-compiler/dialect/src/main/kotlin/app/cash/sqldelight/dialect/api/NoOp.kt

@@ -1,6 +1,6 @@
 package app.cash.sqldelight.dialect.api
 
-class NoOp : SqlGeneratorStrategy {
+internal class NoOp : SqlGeneratorStrategy {
 
   override fun tableNameChanged(oldName: String, newName: String): String {
     return ""

sqldelight-compiler/dialect/src/main/kotlin/app/cash/sqldelight/dialect/api/SqlDelightModule.kt

@@ -2,4 +2,5 @@ package app.cash.sqldelight.dialect.api
 
 interface SqlDelightModule {
   fun typeResolver(parentResolver: TypeResolver): TypeResolver
+  fun setup() {}
 }

sqldelight-compiler/dialect/src/main/kotlin/app/cash/sqldelight/dialect/api/TypeResolver.kt

@@ -1,6 +1,5 @@
 package app.cash.sqldelight.dialect.api
 
-import com.alecstrong.sql.psi.core.psi.SqlBindExpr
 import com.alecstrong.sql.psi.core.psi.SqlExpr
 import com.alecstrong.sql.psi.core.psi.SqlFunctionExpr
 import com.alecstrong.sql.psi.core.psi.SqlStmt
@@ -14,8 +13,6 @@ interface TypeResolver {
    */
   fun resolvedType(expr: SqlExpr): IntermediateType
 
-  fun argumentType(bindArg: SqlBindExpr): IntermediateType
-
   /**
    * In the context of [parent], @return the type [argument] (which is a child expression) should have.
    */

sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/compiler/model/BindableQuery.kt

@@ -20,6 +20,7 @@ import app.cash.sqldelight.core.lang.acceptsTableInterface
 import app.cash.sqldelight.core.lang.psi.ColumnTypeMixin.ValueTypeDialectType
 import app.cash.sqldelight.core.lang.psi.StmtIdentifierMixin
 import app.cash.sqldelight.core.lang.types.typeResolver
+import app.cash.sqldelight.core.lang.util.argumentType
 import app.cash.sqldelight.core.lang.util.childOfType
 import app.cash.sqldelight.core.lang.util.columns
 import app.cash.sqldelight.core.lang.util.findChildrenOfType

sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/ParserUtil.kt

@@ -5,9 +5,11 @@ import app.cash.sqldelight.core.SqldelightParserUtil
 import app.cash.sqldelight.core.compiler.model.SqlDelightPragmaName
 import app.cash.sqldelight.core.lang.psi.FunctionExprMixin
 import app.cash.sqldelight.dialect.api.SqlDelightDialect
+import app.cash.sqldelight.dialect.api.SqlDelightModule
 import com.alecstrong.sql.psi.core.SqlParserUtil
 import com.alecstrong.sql.psi.core.psi.SqlTypes
 import com.intellij.openapi.project.Project
+import java.util.ServiceLoader
 
 internal class ParserUtil {
   private var dialect: Class<out SqlDelightDialect>? = null
@@ -19,6 +21,9 @@ internal class ParserUtil {
       SqldelightParserUtil.reset()
 
       newDialect.setup()
+      ServiceLoader.load(SqlDelightModule::class.java, newDialect::class.java.classLoader).forEach {
+        it.setup()
+      }
       SqldelightParserUtil.overrideSqlParser()
       dialect = newDialect::class.java
 

sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/SqlDelightQueriesFile.kt

@@ -23,6 +23,7 @@ import app.cash.sqldelight.core.compiler.model.NamedMutator.Insert
 import app.cash.sqldelight.core.compiler.model.NamedMutator.Update
 import app.cash.sqldelight.core.compiler.model.NamedQuery
 import app.cash.sqldelight.core.lang.psi.StmtIdentifierMixin
+import app.cash.sqldelight.core.lang.util.argumentType
 import app.cash.sqldelight.core.psi.SqlDelightStmtList
 import com.alecstrong.sql.psi.core.SqlAnnotationHolder
 import com.alecstrong.sql.psi.core.psi.SqlAnnotatedElement

sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/util/ExprUtil.kt

@@ -75,10 +75,6 @@ internal class AnsiSqlTypeResolver : TypeResolver {
     return functionExpr.typeReturned()
   }
 
-  override fun argumentType(bindArg: SqlBindExpr): IntermediateType {
-    return bindArg.inferredType().copy(bindArg = bindArg)
-  }
-
   override fun definitionType(typeName: SqlTypeName) =
     throw UnsupportedOperationException("ANSI SQL is not supported for being used as a dialect.")
 
@@ -167,6 +163,10 @@ internal class AnsiSqlTypeResolver : TypeResolver {
   }
 }
 
+internal fun TypeResolver.argumentType(bindArg: SqlBindExpr): IntermediateType {
+  return bindArg.inferredType().copy(bindArg = bindArg)
+}
+
 private fun SqlExpr.type(): IntermediateType {
   return typeResolver.resolvedType(this)
 }

sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/util/SelectStmtUtil.kt

@@ -57,7 +57,7 @@ internal fun PsiElement.referencedTables(
             .findChildOfType<SqlCompoundSelectStmt>()?.tablesObserved().orEmpty()
         }
       }
-      else -> reference!!.resolve()!!.referencedTables()
+      else -> reference?.resolve()?.referencedTables().orEmpty()
     }
   }
   else -> throw IllegalStateException("Cannot get reference table for psi type ${this.javaClass}")

sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/util/TreeUtil.kt

@@ -20,6 +20,7 @@ import app.cash.sqldelight.core.lang.SqlDelightQueriesFile
 import app.cash.sqldelight.core.lang.acceptsTableInterface
 import app.cash.sqldelight.core.lang.psi.ColumnTypeMixin
 import app.cash.sqldelight.core.lang.psi.InsertStmtValuesMixin
+import app.cash.sqldelight.dialect.api.ExposableType
 import app.cash.sqldelight.dialect.api.IntermediateType
 import app.cash.sqldelight.dialect.api.PrimitiveType
 import app.cash.sqldelight.dialect.api.PrimitiveType.INTEGER
@@ -53,6 +54,7 @@ internal inline fun <reified R : PsiElement> PsiElement.parentOfType(): R {
 }
 
 internal fun PsiElement.type(): IntermediateType = when (this) {
+  is ExposableType -> type()
   is SqlTypeName -> sqFile().typeResolver.definitionType(this)
   is AliasElement -> source().type().copy(name = name)
   is ColumnDefMixin -> (columnType as ColumnTypeMixin).type()

sqldelight-compiler/src/test/kotlin/app/cash/sqldelight/core/BindArgsTest.kt

@@ -1,6 +1,7 @@
 package app.cash.sqldelight.core
 
 import app.cash.sqldelight.core.lang.types.typeResolver
+import app.cash.sqldelight.core.lang.util.argumentType
 import app.cash.sqldelight.core.lang.util.findChildrenOfType
 import app.cash.sqldelight.core.lang.util.isArrayParameter
 import app.cash.sqldelight.dialect.api.PrimitiveType

sqldelight-gradle-plugin/src/test/integration-sqlite-json/build.gradle

@@ -0,0 +1,28 @@
+buildscript {
+  apply from: "${projectDir.absolutePath}/../buildscript.gradle"
+}
+
+apply plugin: 'org.jetbrains.kotlin.jvm'
+apply plugin: 'app.cash.sqldelight'
+
+repositories {
+  maven {
+    url "file://${projectDir.absolutePath}/../../../../build/localMaven"
+  }
+  mavenCentral()
+}
+
+sqldelight {
+  QueryWrapper {
+    packageName = "app.cash.sqldelight.integration"
+    dialect("app.cash.sqldelight:sqlite-3-18-dialect:${app.cash.sqldelight.VersionKt.VERSION}")
+    module("app.cash.sqldelight:sqlite-json-module:${app.cash.sqldelight.VersionKt.VERSION}")
+  }
+}
+
+dependencies {
+  implementation deps.sqliteJdbc
+  implementation "app.cash.sqldelight:sqlite-driver:${app.cash.sqldelight.VersionKt.VERSION}"
+  implementation deps.truth
+  implementation("com.squareup.moshi:moshi:1.13.0")
+}

sqldelight-gradle-plugin/src/test/integration-sqlite-json/settings.gradle

@@ -0,0 +1 @@
+apply from: "../settings.gradle"

sqldelight-gradle-plugin/src/test/integration-sqlite-json/src/main/sqldelight/app/cash/sqldelight/integration/JsonTable.sq

@@ -0,0 +1,21 @@
+CREATE TABLE user(
+  name TEXT NOT NULL PRIMARY KEY,
+  phone TEXT
+);
+
+insertUser:
+INSERT INTO user
+VALUES (?, ?);
+
+byAreaCode:
+SELECT DISTINCT user.name
+  FROM user, json_each(user.phone)
+ WHERE json_each.value LIKE :areaCode || '-%';
+
+byAreaCode2:
+SELECT name FROM user WHERE phone LIKE :areaCode || '-%'
+UNION
+SELECT user.name
+  FROM user, json_each(user.phone)
+ WHERE json_valid(user.phone)
+   AND json_each.value LIKE :areaCode || '-%';