server/mssql: refine and generalize mssql OUTPUT clause and temporary table syntax AST

PR-URL: https://github.com/hasura/graphql-engine-mono/pull/2931
GitOrigin-RevId: e6aef494a096ca119597dc92bcaf9f16e644a72a

server/src-lib/Hasura/Backends/MSSQL/FromIr.hs

@@ -1078,7 +1078,7 @@ fromInsert IR.AnnInsert {..} =
       insertColumnNames = maybe [] (map fst) $ listToMaybe insertRows
       insertValues = map (Values . map snd) insertRows
       primaryKeyColumns = map OutputColumn $ _mssqlPrimaryKeyColumns _aiExtraInsertData
-   in Insert _aiTableName insertColumnNames (InsertOutput primaryKeyColumns) insertValues
+   in Insert _aiTableName insertColumnNames (Output Inserted primaryKeyColumns) insertValues
 
 -- | Normalize a row by adding missing columns with 'DEFAULT' value and sort by column name to make sure
 -- all rows are consistent in column values and order.
@@ -1123,6 +1123,7 @@ fromDelete (IR.AnnDel tableName (permFilter, whereClause) _ allColumns) = do
     ( do
         permissionsFilter <- fromGBoolExp permFilter
         whereExpression <- fromGBoolExp whereClause
+        let columnNames = map (ColumnName . unName . IR.pgiName) allColumns
         pure
           Delete
             { deleteTable =
@@ -1130,16 +1131,16 @@ fromDelete (IR.AnnDel tableName (permFilter, whereClause) _ allColumns) = do
                   { aliasedAlias = entityAliasText tableAlias,
                     aliasedThing = tableName
                   },
-              deleteColumns = map (ColumnName . unName . IR.pgiName) allColumns,
+              deleteOutput = Output Deleted (map OutputColumn columnNames),
+              deleteTempTable = TempTable tempTableNameDeleted columnNames,
               deleteWhere = Where [permissionsFilter, whereExpression]
             }
     )
     tableAlias
 
--- | Convert IR AST representing delete into MSSQL AST representing a creation of a temporary table
---   with the same schema as the deleted from table.
-toSelectIntoTempTable :: TempTableName -> IR.AnnDel 'MSSQL -> SelectIntoTempTable
-toSelectIntoTempTable tempTableName (IR.AnnDel {IR.dqp1Table = fromTable, IR.dqp1AllCols = allColumns}) = do
+-- | Create a temporary table with the same schema as the given table.
+toSelectIntoTempTable :: TempTableName -> TableName -> [IR.ColumnInfo 'MSSQL] -> SelectIntoTempTable
+toSelectIntoTempTable tempTableName fromTable allColumns = do
   SelectIntoTempTable
     { sittTempTableName = tempTableName,
       sittColumns = map columnInfoToUnifiedColumn allColumns,

server/src-lib/Hasura/Backends/MSSQL/Instances/Execute.hs

@@ -419,12 +419,14 @@ buildDeleteTx ::
   Bool ->
   Tx.TxET QErr IO EncJSON
 buildDeleteTx deleteOperation stringifyNum = do
-  let tempTableName = TSQL.TempTableName "deleted"
-      withAlias = "with_alias"
-      createTempTableQuery = toQueryFlat $ TQ.fromSelectIntoTempTable $ TSQL.toSelectIntoTempTable tempTableName deleteOperation
+  let withAlias = "with_alias"
+      createTempTableQuery =
+        toQueryFlat $
+          TQ.fromSelectIntoTempTable $
+            TSQL.toSelectIntoTempTable tempTableNameDeleted (dqp1Table deleteOperation) (dqp1AllCols deleteOperation)
   -- Create a temp table
   Tx.unitQueryE fromMSSQLTxError createTempTableQuery
-  let deleteQuery = TQ.fromDelete tempTableName <$> TSQL.fromDelete deleteOperation
+  let deleteQuery = TQ.fromDelete <$> TSQL.fromDelete deleteOperation
   deleteQueryValidated <- toQueryFlat <$> V.runValidate (runFromIr deleteQuery) `onLeft` (throw500 . tshow)
   -- Execute DELETE statement
   Tx.unitQueryE fromMSSQLTxError deleteQueryValidated
@@ -432,7 +434,7 @@ buildDeleteTx deleteOperation stringifyNum = do
   let withSelect =
         emptySelect
           { selectProjections = [StarProjection],
-            selectFrom = Just $ FromTempTable $ Aliased tempTableName "deleted_alias"
+            selectFrom = Just $ FromTempTable $ Aliased tempTableNameDeleted "deleted_alias"
           }
       finalMutationOutputSelect = mutationOutputSelect {selectWith = Just $ With $ pure $ Aliased withSelect withAlias}
       mutationOutputSelectQuery = toQueryFlat $ TQ.fromSelect finalMutationOutputSelect

server/src-lib/Hasura/Backends/MSSQL/ToQuery.hs

@@ -163,13 +163,25 @@ fromFieldName :: FieldName -> Printer
 fromFieldName (FieldName {..}) =
   fromNameText fieldNameEntity <+> "." <+> fromNameText fieldName
 
-fromOutputColumn :: OutputColumn -> Printer
-fromOutputColumn (OutputColumn columnName) =
-  "INSERTED." <+> fromNameText (columnNameText columnName)
+fromInserted :: Inserted -> Printer
+fromInserted Inserted = "INSERTED"
+
+fromDeleted :: Deleted -> Printer
+fromDeleted Deleted = "DELETED"
+
+fromOutputColumn :: Printer -> OutputColumn -> Printer
+fromOutputColumn prefix (OutputColumn columnName) =
+  prefix <+> "." <+> fromNameText (columnNameText columnName)
+
+fromOutput :: (t -> Printer) -> Output t -> Printer
+fromOutput typePrinter (Output ty outputColumns) =
+  "OUTPUT " <+> SepByPrinter ", " (map (fromOutputColumn (typePrinter ty)) outputColumns)
 
 fromInsertOutput :: InsertOutput -> Printer
-fromInsertOutput (InsertOutput outputColumns) =
-  "OUTPUT " <+> SepByPrinter ", " (map fromOutputColumn outputColumns)
+fromInsertOutput = fromOutput fromInserted
+
+fromDeleteOutput :: DeleteOutput -> Printer
+fromDeleteOutput = fromOutput fromDeleted
 
 fromValues :: Values -> Printer
 fromValues (Values values) =
@@ -209,13 +221,13 @@ fromSetIdentityInsert SetIdentityInsert {..} =
 -- Becomes:
 --
 -- > DELETE [alias] OUTPUT DELETED.[id], DELETED.[name] INTO #deleted([id], [name]) FROM [schema].[table] AS [alias] WHERE ((1) = (1))
-fromDelete :: TempTableName -> Delete -> Printer
-fromDelete tempTableName Delete {deleteTable, deleteColumns, deleteWhere} =
+fromDelete :: Delete -> Printer
+fromDelete Delete {deleteTable, deleteOutput, deleteTempTable, deleteWhere} =
   SepByPrinter
     NewlinePrinter
     [ "DELETE " <+> fromNameText (aliasedAlias deleteTable),
-      "OUTPUT " <+> SepByPrinter ", " (map ((<+>) "DELETED." . fromColumnName) deleteColumns),
-      "INTO " <+> fromTempTableName tempTableName <+> parens (SepByPrinter ", " (map fromColumnName deleteColumns)),
+      fromDeleteOutput deleteOutput,
+      "INTO " <+> fromTempTable deleteTempTable,
       "FROM " <+> fromAliased (fmap fromTableName deleteTable),
       fromWhere deleteWhere
     ]
@@ -264,6 +276,10 @@ fromSelectIntoTempTable SelectIntoTempTable {sittTempTableName, sittColumns, sit
 fromTempTableName :: TempTableName -> Printer
 fromTempTableName (TempTableName v) = QueryPrinter (fromString . T.unpack $ "#" <> v)
 
+fromTempTable :: TempTable -> Printer
+fromTempTable (TempTable table columns) =
+  fromTempTableName table <+> parens (SepByPrinter ", " (map fromColumnName columns))
+
 fromSelect :: Select -> Printer
 fromSelect Select {..} = fmap fromWith selectWith ?<+> wrapFor selectFor result
   where

server/src-lib/Hasura/Backends/MSSQL/Types/Internal.hs

@@ -11,6 +11,7 @@ module Hasura.Backends.MSSQL.Types.Internal
     Comment (..),
     Countable (..),
     Delete (..),
+    DeleteOutput,
     EntityAlias (..),
     Expression (..),
     FieldName (..),
@@ -19,7 +20,7 @@ module Hasura.Backends.MSSQL.Types.Internal
     From (..),
     FunctionName,
     Insert (..),
-    InsertOutput (..),
+    InsertOutput,
     Join (..),
     JoinAlias (..),
     JoinSource (..),
@@ -32,6 +33,9 @@ module Hasura.Backends.MSSQL.Types.Internal
     Order (..),
     OrderBy (..),
     OutputColumn (..),
+    Inserted (..),
+    Deleted (..),
+    Output (..),
     Projection (..),
     Reselect (..),
     Root (..),
@@ -40,6 +44,7 @@ module Hasura.Backends.MSSQL.Types.Internal
     Select (..),
     SetIdentityInsert (..),
     TempTableName (..),
+    TempTable (..),
     SetValue (..),
     SelectIntoTempTable (..),
     SpatialOp (..),
@@ -65,6 +70,7 @@ module Hasura.Backends.MSSQL.Types.Internal
     scalarTypeDBName,
     snakeCaseTableName,
     stringTypes,
+    tempTableNameDeleted,
   )
 where
 
@@ -157,9 +163,18 @@ emptySelect =
       selectOffset = Nothing
     }
 
-newtype OutputColumn = OutputColumn ColumnName
+newtype OutputColumn = OutputColumn {unOutputColumn :: ColumnName}
 
-newtype InsertOutput = InsertOutput [OutputColumn]
+data Inserted = Inserted
+
+data Deleted = Deleted
+
+data Output t = Output
+  { outputType :: !t,
+    outputColumns :: ![OutputColumn]
+  }
+
+type InsertOutput = Output Inserted
 
 newtype Values = Values [Expression]
 
@@ -179,9 +194,12 @@ data SetIdentityInsert = SetIdentityInsert
     setValue :: !SetValue
   }
 
+type DeleteOutput = Output Deleted
+
 data Delete = Delete
   { deleteTable :: !(Aliased TableName),
-    deleteColumns :: [ColumnName],
+    deleteOutput :: !DeleteOutput,
+    deleteTempTable :: !TempTable,
     deleteWhere :: !Where
   }
 
@@ -196,6 +214,14 @@ data SelectIntoTempTable = SelectIntoTempTable
 -- | A temporary table name is prepended by a hash-sign
 newtype TempTableName = TempTableName Text
 
+tempTableNameDeleted :: TempTableName
+tempTableNameDeleted = TempTableName "deleted"
+
+data TempTable = TempTable
+  { ttName :: !TempTableName,
+    ttColumns :: ![ColumnName]
+  }
+
 data Reselect = Reselect
   { reselectProjections :: ![Projection],
     reselectFor :: !For,