Fix 4441 PostgreSql Dialect auxilary CTE (#4493)

* add PostgreSql Dialect SqlInsertStmtMixin

Add mixin to override tablesAvailable

For insert statement check the parent auxillary with clause

* Add fixture

Auxillary with clause using insert populated by select

* Include support for delete stmt with mixin

Delete statements can reference CTE

* Add fixture

referencing cte from auxilary delete statement

* add support for update statement

reference cte from update statement

* add fixture

for update statement

---------

Co-authored-by: griffio <griffio@users.noreply.github.com>

dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/grammar/PostgreSql.bnf

@@ -203,6 +203,7 @@ with_clause_auxiliary_stmt ::= {compound_select_stmt} | delete_stmt_limited | in
 }
 
 delete_stmt_limited ::= [ {with_clause} ] DELETE FROM {qualified_table_name} [ WHERE <<expr '-1'>> ] [ [ ORDER BY {ordering_term} ( COMMA {ordering_term} ) * ] LIMIT <<expr '-1'>> [ ( OFFSET | COMMA ) <<expr '-1'>> ] ] [ returning_clause ] {
+  mixin =  "app.cash.sqldelight.dialects.postgresql.grammar.mixins.SqlDeleteStmtLimitedMixin"
   extends = "com.alecstrong.sql.psi.core.psi.impl.SqlDeleteStmtLimitedImpl"
   implements = "com.alecstrong.sql.psi.core.psi.SqlDeleteStmtLimited"
   pin = 2
@@ -233,7 +234,7 @@ insert_stmt ::= [ {with_clause} ]
     [ LP {column_name} ( COMMA {column_name} ) * RP ] {insert_stmt_values}
     [ ON CONFLICT ( [conflict_target] DO NOTHING | conflict_target conflict_update ) ]
     [ returning_clause ] {
-  extends = "com.alecstrong.sql.psi.core.psi.impl.SqlInsertStmtImpl"
+  mixin = "app.cash.sqldelight.dialects.postgresql.grammar.mixins.SqlInsertStmtMixin"
   implements = "com.alecstrong.sql.psi.core.psi.SqlInsertStmt"
   override = true
   pin = 5

dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/grammar/mixins/SqlDeleteStmtLimitedMixin.kt

@@ -0,0 +1,21 @@
+package app.cash.sqldelight.dialects.postgresql.grammar.mixins
+
+import app.cash.sqldelight.dialects.postgresql.grammar.psi.PostgreSqlDeleteStmtLimited
+import com.alecstrong.sql.psi.core.psi.LazyQuery
+import com.alecstrong.sql.psi.core.psi.SqlWithClause
+import com.alecstrong.sql.psi.core.psi.SqlWithClauseAuxiliaryStmt
+import com.alecstrong.sql.psi.core.psi.impl.SqlDeleteStmtLimitedImpl
+import com.intellij.lang.ASTNode
+import com.intellij.psi.PsiElement
+
+internal abstract class SqlDeleteStmtLimitedMixin(
+  node: ASTNode,
+) : SqlDeleteStmtLimitedImpl(node), PostgreSqlDeleteStmtLimited {
+
+  override fun tablesAvailable(child: PsiElement): Collection<LazyQuery> {
+    val tablesAvailable = super.tablesAvailable(child)
+    val withClauseAuxiliaryStmts = parent as? SqlWithClauseAuxiliaryStmt ?: return tablesAvailable
+    val withClause = withClauseAuxiliaryStmts.parent as SqlWithClause
+    return tablesAvailable + withClause.tablesExposed()
+  }
+}

dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/grammar/mixins/SqlInsertStmtMixin.kt

@@ -0,0 +1,21 @@
+package app.cash.sqldelight.dialects.postgresql.grammar.mixins
+
+import app.cash.sqldelight.dialects.postgresql.grammar.psi.PostgreSqlInsertStmt
+import com.alecstrong.sql.psi.core.psi.LazyQuery
+import com.alecstrong.sql.psi.core.psi.SqlWithClause
+import com.alecstrong.sql.psi.core.psi.SqlWithClauseAuxiliaryStmt
+import com.alecstrong.sql.psi.core.psi.impl.SqlInsertStmtImpl
+import com.intellij.lang.ASTNode
+import com.intellij.psi.PsiElement
+
+internal abstract class SqlInsertStmtMixin(
+  node: ASTNode,
+) : SqlInsertStmtImpl(node), PostgreSqlInsertStmt {
+
+  override fun tablesAvailable(child: PsiElement): Collection<LazyQuery> {
+    val tablesAvailable = super.tablesAvailable(child)
+    val withClauseAuxiliaryStmts = parent as? SqlWithClauseAuxiliaryStmt ?: return tablesAvailable
+    val withClause = withClauseAuxiliaryStmts.parent as SqlWithClause
+    return tablesAvailable + withClause.tablesExposed()
+  }
+}

dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/grammar/mixins/UpdateStmtLimitedMixin.kt

@@ -2,8 +2,11 @@ package app.cash.sqldelight.dialects.postgresql.grammar.mixins
 
 import app.cash.sqldelight.dialects.postgresql.grammar.psi.PostgreSqlUpdateStmtLimited
 import com.alecstrong.sql.psi.core.psi.FromQuery
+import com.alecstrong.sql.psi.core.psi.LazyQuery
 import com.alecstrong.sql.psi.core.psi.QueryElement
 import com.alecstrong.sql.psi.core.psi.SqlJoinClause
+import com.alecstrong.sql.psi.core.psi.SqlWithClause
+import com.alecstrong.sql.psi.core.psi.SqlWithClauseAuxiliaryStmt
 import com.alecstrong.sql.psi.core.psi.impl.SqlUpdateStmtLimitedImpl
 import com.intellij.lang.ASTNode
 import com.intellij.psi.PsiElement
@@ -28,5 +31,12 @@ internal abstract class UpdateStmtLimitedMixin(
     return emptyList()
   }
 
+  override fun tablesAvailable(child: PsiElement): Collection<LazyQuery> {
+    val tablesAvailable = super.tablesAvailable(child)
+    val withClauseAuxiliaryStmts = parent as? SqlWithClauseAuxiliaryStmt ?: return tablesAvailable
+    val withClause = withClauseAuxiliaryStmts.parent as SqlWithClause
+    return tablesAvailable + withClause.tablesExposed()
+  }
+
   private val joinClause: SqlJoinClause? get() = findChildByClass(SqlJoinClause::class.java)
 }

dialects/postgresql/src/testFixtures/resources/fixtures_postgresql/cte-multiple-auxiliary-statements/Sample.s

@@ -20,3 +20,32 @@ INSERT
 INTO sample
 SELECT test_id0, 100
 FROM test_ids0;
+
+WITH insert_test_sample AS (
+   INSERT INTO test
+   VALUES(1, 'Foo')
+   RETURNING id, name
+),
+insert_sample AS (
+    INSERT INTO sample (id, val)
+    SELECT id, 31 FROM insert_test_sample
+)
+SELECT * FROM insert_test_sample;
+
+WITH deleted_test_sample AS (
+  DELETE FROM test
+  RETURNING id
+), deleted_sample AS (
+  DELETE FROM sample WHERE id IN (SELECT id FROM deleted_test_sample)
+  RETURNING id
+)
+SELECT * FROM deleted_test_sample;
+
+WITH updated_test_sample AS (
+  UPDATE test SET name = 'Bar'
+  RETURNING id
+), updated_sample AS (
+  UPDATE sample SET val = 42 WHERE id IN (SELECT id FROM updated_test_sample)
+  RETURNING id
+)
+SELECT * FROM updated_sample;