add support for RETURNING clause (PostgreSQL, SQLite)

doc/_cli-options.md

@@ -21,7 +21,7 @@
 
  Dialect and checks:
   -dialect mysql|postgresql|sqlite|tidb            Set SQL dialect. Queries can only use its features
-  -no-check {all|<feature>{,<feature>}+}           Disable dialect feature checks (possible features: collation|join_on_subquery|create_table_as_select|on_duplicate_key|on_conflict|straight_join|lock_in_share_mode|fulltext_index|unsigned_types|autoincrement|replace_into|row_locking|default_expr|ttl|alter_column|user_defined_type)
+  -no-check {all|<feature>{,<feature>}+}           Disable dialect feature checks (possible features: collation|join_on_subquery|create_table_as_select|on_duplicate_key|on_conflict|straight_join|lock_in_share_mode|fulltext_index|unsigned_types|autoincrement|replace_into|row_locking|default_expr|ttl|alter_column|user_defined_type|returning)
   -allow-write-notnull-null                        Accept writing a nullable value into a NOT NULL column, instead of failing (MySQL, TiDB and SQLite only)
 
  Generated header:

lib/dialect.ml

@@ -28,6 +28,7 @@ type feature =
   | Ttl [@as "ttl"]
   | AlterColumn [@as "alter_column"]
   | UserDefinedType [@as "user_defined_type"]
+  | Returning [@as "returning"]
 [@@deriving show { with_path = false }, enumerate, to_string, of_string]
 
 let show_feature x = 
@@ -106,6 +107,8 @@ let get_on_duplicate_key pos = only OnDuplicateKey [MySQL; TiDB] pos
 
 let get_on_conflict pos = only OnConflict [SQLite; PostgreSQL] pos
 
+let get_returning pos = only Returning [PostgreSQL; SQLite] pos
+
 let get_straight_join pos = only StraightJoin [MySQL; TiDB] pos
 
 let get_lock_in_share_mode pos = only LockInShareMode [MySQL] pos
@@ -220,6 +223,10 @@ and analyze_column acc cols k = match cols with
     | { value = (All | AllOf _); _ } -> analyze_column acc rest k
     | { value = Expr ({ value = expr; _ }, _); _ } -> analyze_expr acc [expr] (fun acc -> analyze_column acc rest k)
 
+and analyze_returning acc returning k = match returning with
+  | None -> k acc
+  | Some { value = cols; pos } -> analyze_column (get_returning pos :: acc) cols k
+
 and analyze_source acc srcs k = match srcs with
   | [] -> k acc
   | src :: rest ->
@@ -364,7 +371,7 @@ and analyze_alter_action acc actions k = match actions with
 
 and analyze_insert_action acc ias k = match ias with
   | [] -> k acc
-  | { action; on_conflict_clause; insert_action_kind; _ } :: rest ->
+  | { action; on_conflict_clause; insert_action_kind; returning; _ } :: rest ->
     let acc = match insert_action_kind with
       | Replace_into pos -> get_replace_into pos :: acc
       | Insert_into -> acc
@@ -391,7 +398,8 @@ and analyze_insert_action acc ias k = match ias with
     analyze_action acc (fun acc ->
       let conflict_aes = List.map snd conflict_assignments in
       analyze_assignment_expr acc conflict_aes (fun acc ->
-        analyze_insert_action acc rest k))
+        analyze_returning acc returning (fun acc ->
+          analyze_insert_action acc rest k)))
 
 let analyze_schema_index idx = match idx.value.Sql.idx_kind with
   | Regular_idx -> None
@@ -414,8 +422,9 @@ let rec analyze stmt =
   | CreateIndex { ci_cols; _ } -> List.concat_map check_collated ci_cols
   | Insert insert_action ->
       analyze_insert_action acc [insert_action] List.rev
-  | Delete (_, where_opt) ->
-      analyze_expr acc (option_list where_opt) List.rev
+  | Delete (_, where_opt, returning) ->
+      analyze_expr acc (option_list where_opt) (fun acc ->
+        analyze_returning acc returning List.rev)
   | DeleteMulti (_, nested, where_opt) ->
       analyze_nested acc [nested] (fun acc ->
         analyze_expr acc (option_list where_opt) List.rev)
@@ -424,11 +433,12 @@ let rec analyze stmt =
       analyze_expr acc exprs (fun acc ->
         let stmt_features = Option.map_default analyze [] stmt_opt in
         List.rev (List.rev_append stmt_features acc))
-  | Update (_, assignments, where_opt, order, _) ->
+  | Update (_, assignments, where_opt, order, _, returning) ->
       let aes = List.map snd assignments in
       analyze_assignment_expr acc aes (fun acc ->
         let exprs = option_list where_opt @ List.map fst order in
-        analyze_expr acc exprs List.rev)
+        analyze_expr acc exprs (fun acc ->
+          analyze_returning acc returning List.rev))
   | UpdateMulti (nesteds, assignments, where_opt, order, _) ->
       analyze_nested acc nesteds (fun acc ->
         let aes = List.map snd assignments in

lib/sql.ml

@@ -753,6 +753,8 @@ and column_kind =
 
 type columns = column list [@@deriving show]
 
+type returning = columns located [@@deriving show]
+
 let source_fun_kind_to_infer = function
   | Ret t -> Ret (Source_type.to_infer_type t)
   | Agg (Self | Count | Avg | With_order _) 
@@ -826,6 +828,7 @@ type insert_action =
            | `Param of (string list option * param_id)
            | `Select of (string list option * select_full) ];
   on_conflict_clause : conflict_clause located option;
+  returning : returning option;
 } [@@deriving show {with_path=false}]
 
 type table_constraints = [ `Ignore | `Primary of string list | `Unique of string option * string list ] [@@deriving show {with_path=false}]
@@ -971,10 +974,10 @@ type stmt =
   | Rename of (table_name * table_name) list
   | CreateIndex of create_index_def
   | Insert of insert_action
-  | Delete of table_name * expr option
+  | Delete of table_name * expr option * returning option
   | DeleteMulti of table_name list * nested * expr option
   | Set of (string * expr) list * stmt option
-  | Update of table_name * assignments * expr option * order * Source_type.t param list (* where, order, limit *)
+  | Update of table_name * assignments * expr option * order * Source_type.t param list * returning option (* where, order, limit, returning *)
   | UpdateMulti of nested list * assignments * expr option * order * Source_type.t param list (* where, order, limit *)
   | Select of select_full
   | CreateRoutine of table_name * Source_type.kind collated located option * (string * Source_type.kind collated located * expr option) list (* table_name represents possibly namespaced function name *)

lib/sql_lexer.mll

@@ -146,6 +146,7 @@ let keywords =
    "rename",RENAME;
    "replace",REPLACE;
    "restrict",RESTRICT;
+   "returning",RETURNING;
    "returns", RETURNS;
    "row", ROW;
    "rows", ROWS;

lib/sql_parser.mly

@@ -35,7 +35,7 @@
        LIMIT ORDER BY DESC ASC EQUAL DELETE FROM DEFAULT OFFSET SET STRAIGHT_JOIN JOIN LIKE_OP LIKE
        EXCL TILDE NOT BETWEEN AND XOR ESCAPE USING UNION EXCEPT INTERSECT AS TO
        CONCAT_OP LEFT RIGHT FULL INNER OUTER NATURAL CROSS REPLACE IN GROUP HAVING
-       UNIQUE PRIMARY KEY FOREIGN AUTOINCREMENT ON CONFLICT DO NOTHING TEMPORARY IF EXISTS
+       UNIQUE PRIMARY KEY FOREIGN AUTOINCREMENT ON CONFLICT DO NOTHING TEMPORARY IF EXISTS RETURNING
        PRECISION SIGNED UNSIGNED ZEROFILL VARYING CHARSET NATIONAL ASCII UNICODE COLLATE BINARY CHARACTER
        DATETIME_FUNC DATE TIME TIMESTAMP ALTER RENAME ADD COLUMN CASCADE RESTRICT DROP
        GLOBAL LOCAL REFERENCES CHECK CONSTRAINT IGNORED AFTER INDEX FULLTEXT SPATIAL FIRST
@@ -141,34 +141,34 @@ statement: CREATE ioption(temporary) TABLE ioption(if_not_exists) name=table_nam
                 CreateIndex { ci_name = name; ci_table = table; ci_cols = cols; ci_kind }
               }
          | select_stmt { Select $1 }
-         | insert_action_kind=insert_cmd target=table_name names=sequence(ident)? VALUES values=commas(sequence(set_column_expr))? ss=located(conflict_clause)?
+         | insert_action_kind=insert_cmd target=table_name names=sequence(ident)? VALUES values=commas(sequence(set_column_expr))? ss=located(conflict_clause)? r=located(returning)?
               {
-                Insert { insert_action_kind; target; action=`Values (names, values); on_conflict_clause=ss; }
+                Insert { insert_action_kind; target; action=`Values (names, values); on_conflict_clause=ss; returning=r; }
               }
-         | insert_action_kind=insert_cmd target=table_name names=sequence(ident)? VALUES p=param ss=located(conflict_clause)?
+         | insert_action_kind=insert_cmd target=table_name names=sequence(ident)? VALUES p=param ss=located(conflict_clause)? r=located(returning)?
               {
-                Insert { insert_action_kind; target; action=`Param (names, p); on_conflict_clause=ss; }
+                Insert { insert_action_kind; target; action=`Param (names, p); on_conflict_clause=ss; returning=r; }
               }
-         | insert_action_kind=insert_cmd target=table_name names=sequence(ident)? select=maybe_parenth(select_stmt) ss=located(conflict_clause)?
+         | insert_action_kind=insert_cmd target=table_name names=sequence(ident)? select=maybe_parenth(select_stmt) ss=located(conflict_clause)? r=located(returning)?
               {
-                Insert { insert_action_kind; target; action=`Select (names, select); on_conflict_clause=ss; }
+                Insert { insert_action_kind; target; action=`Select (names, select); on_conflict_clause=ss; returning=r; }
               }
-         | insert_action_kind=insert_cmd target=table_name SET set=commas(set_column)? ss=located(conflict_clause)?
+         | insert_action_kind=insert_cmd target=table_name SET set=commas(set_column)? ss=located(conflict_clause)? r=located(returning)?
               {
-                Insert { insert_action_kind; target; action=`Set set; on_conflict_clause=ss; }
+                Insert { insert_action_kind; target; action=`Set set; on_conflict_clause=ss; returning=r; }
               }
-         | update_cmd table=table_name SET ss=commas(set_column) w=where? o=loption(order) lim=loption(limit)
+         | update_cmd table=table_name SET ss=commas(set_column) w=where? o=loption(order) lim=loption(limit) r=located(returning)?
               {
-                Update (table,ss,w,o,lim)
+                Update (table,ss,w,o,lim,r)
               }
          /* http://dev.mysql.com/doc/refman/5.1/en/update.html multi-table syntax */
          | update_cmd tables=commas(table_list) SET ss=commas(set_column) w=where? o=loption(order) lim=loption(limit)
               {
                 UpdateMulti (tables,ss,w,o,lim)
               }
-         | DELETE FROM table=table_name w=where?
+         | DELETE FROM table=table_name w=where? r=located(returning)?
               {
-                Delete (table,w)
+                Delete (table,w,r)
               }
          /* https://dev.mysql.com/doc/refman/5.7/en/delete.html multi-table syntax */
          | DELETE targets=commas(table_name) FROM tables=table_list w=where?
@@ -349,6 +349,7 @@ from: FROM t=table_list { t }
 where: WHERE e=expr { e }
 group: GROUP BY l=expr_list { l }
 having: HAVING e=expr { e }
+returning: RETURNING r=commas(column1) { r }
 
 column1:
        | c=located(column1_kind) { c }

lib/stmt.ml

@@ -13,7 +13,7 @@ type cardinality = [`Zero_one | `One | `Nat] [@@deriving show]
 let cardinality_to_string = show_cardinality
 
 type kind = | Select of cardinality
-            | Insert of inferred * Sql.table_name
+            | Insert of inferred * Sql.table_name * cardinality
             | Create of Sql.table_name
             | CreateIndex of string
             | Update of Sql.table_name option (** name for single-table UPDATEs *)
@@ -26,6 +26,12 @@ type kind = | Select of cardinality
             | Other
             [@@deriving show {with_path=false}]
 
+let cardinality_of_kind = function
+| Select c -> c
+| Insert (_, _, c) -> c
+| Create _ | CreateIndex _ | Update _ | Delete _ | Alter _ | Drop _
+| CreateRoutine _ | CreateType _ | DropType _ | Other -> `Nat
+
 type category = DDL | DQL | DML | DCL | TCL | OTHER [@@deriving show {with_path=false}, enum]
 
 let all_categories = List.init (max_category - min_category) (fun i -> Option.get @@ category_of_enum @@ min_category + i)

lib/syntax.ml

@@ -1590,11 +1590,33 @@ let with_constraints attrs constraints : Schema.t =
   ) attrs
 
 
+let with_returning env (returning : Sql.returning option) ((schema, params, kind) as stmt) =
+  match returning with
+  | None -> stmt
+  | Some returning ->
+    let projection = make_dynamic_select ~env returning.value in
+    schema @ List.map drop_sources (infer_schema ~not_null_keys:[] env projection),
+    params @ get_params_of_columns env projection,
+    kind
+
+let single_row_insert_cardinality (on_conflict_clause : Sql.conflict_clause Sql.located option) =
+  match on_conflict_clause with
+  | Some { value = On_conflict { action = Do_nothing; _ }; _ } -> `Zero_one
+  | Some { value = (On_conflict { action = Do_update _; _ } | On_duplicate _); _ }
+  | None -> `One
+
 let rec eval (stmt:Sql.stmt) =
   let open Stmt in
   let open Schema.Source in
   let open Attr in
   match stmt with
+  | Insert { action = (`Values (_, None) | `Set None); returning = Some { pos; _ }; _ } ->
+    (* currently not handled as generated VALUES tuple require shifting indices *)
+    failed ~at:pos "RETURNING is not supported when inserted columns are inferred"
+  | Insert { action = (`Values (_, None) | `Set None); on_conflict_clause = Some { pos; value }; _ } ->
+    (* currently not handled as generated VALUES tuple require shifting indices *)
+    let what = match value with On_conflict _ -> "ON CONFLICT" | On_duplicate _ -> "ON DUPLICATE KEY UPDATE" in
+    failed ~at:pos "%s is not supported when inserted columns are inferred" what
   | Create (name, Schema { schema; constraints; indexes }) ->
       let attrs = List.map Alter_action_attr.to_attr schema in
       let attrs = with_constraints attrs constraints in
@@ -1684,14 +1706,14 @@ let rec eval (stmt:Sql.stmt) =
       Sql.Schema.project cols (Tables.get_schema ci_table) |> ignore;
       Tables.index_add ci_table ~index_name:ci_name ~kind:ci_kind ~cols;
       [],[],CreateIndex ci_name
-  | Insert { target=table; action=`Values (names, values); on_conflict_clause; _ } ->
+  | Insert { target=table; action=`Values (names, values); on_conflict_clause; returning; _ } ->
     let expect = values_or_all table names in
     let t = Tables.get_schema table in
     let schema = List.map (fun attr -> { sources=[table]; attr }) t in
     let env = { empty_env with tables = [Tables.get table]; schema; } in
     begin match values with
-    | None -> 
-      [], [], Insert(Some (Values, expect), table)
+    | None ->
+      [], [], Insert (Some (Values, expect), table, `One)
     | Some values ->
       let vl = List.map List.length values in
       let cl = List.length expect in
@@ -1746,18 +1768,23 @@ let rec eval (stmt:Sql.stmt) =
       let p1 = List.concat_map (fun (_c, p, _t) -> p) resolved in
       let conflict_assigns = resolve_on_conflict_clause ~env table.tn on_conflict_clause in
       let params2 = params_of_assigns { env with is_update = true; } conflict_assigns in
-      [], p1 @ params2, Insert (None, table)
+      let cardinality =
+        match values with
+        | [_] -> single_row_insert_cardinality on_conflict_clause
+        | _ -> `Nat
+      in
+      with_returning env returning ([], p1 @ params2, Insert (None, table, cardinality))
     end
-  | Insert { target=table; action=`Param (names, param_id); on_conflict_clause; _ } ->
+  | Insert { target=table; action=`Param (names, param_id); on_conflict_clause; returning; _ } ->
     let schema = List.map (fun attr -> { Schema.Source.Attr.sources=[table]; attr }) (Tables.get_schema table) in
     let env = { empty_env with tables = [Tables.get table]; schema; } in
     let conflict_assigns = resolve_on_conflict_clause ~env table.tn on_conflict_clause in
     let expect = values_or_all table names in
     List.iter (fun a -> Hashtbl.add env.insert_resolved_types a.attr.name a.attr.domain ) schema;
     let params2 = params_of_assigns { env with is_update = true } conflict_assigns in
     let params = [ TupleList (param_id, Insertion expect) ] in
-    [], params @ params2, Insert (None, table)
-  | Insert { target=table; action=`Select (names, select); on_conflict_clause; _ } ->
+    with_returning env returning ([], params @ params2, Insert (None, table, `Nat))
+  | Insert { target=table; action=`Select (names, select); on_conflict_clause; returning; _ } ->
     let expect = values_or_all table names in
     let env = { empty_env with tables = [Tables.get table]; 
       schema = List.map (fun attr -> { sources=[table]; attr }) (Tables.get_schema table); 
@@ -1775,8 +1802,8 @@ let rec eval (stmt:Sql.stmt) =
     let conflict_assigns = resolve_on_conflict_clause ~env table.tn on_conflict_clause in
     List.iter2 (fun a1 a2 -> Hashtbl.add env.insert_resolved_types a2.name a1.attr.domain ) schema expect;
     let params2 = params_of_assigns { env with is_update = true } conflict_assigns in
-    [], params @ params2, Insert (None,table)
-  | Insert { target=table; action=`Set ss; on_conflict_clause; _ } ->
+    with_returning env returning ([], params @ params2, Insert (None, table, `Nat))
+  | Insert { target=table; action=`Set ss; on_conflict_clause; returning; _ } ->
     let env = { empty_env with tables = [Tables.get table]; 
       schema = List.map (fun attr -> { sources=[table]; attr }) (Tables.get_schema table);
     } in
@@ -1786,14 +1813,15 @@ let rec eval (stmt:Sql.stmt) =
     in
     let conflict_assigns = resolve_on_conflict_clause ~env table.tn on_conflict_clause in
     let params2 = params_of_assigns { env with is_update = true } conflict_assigns in
-    [], params @ params2, Insert (inferred,table)
-  | Delete (table, where) ->
+    with_returning env returning
+      ([], params @ params2, Insert (inferred, table, single_row_insert_cardinality on_conflict_clause))
+  | Delete (table, where, returning) ->
     let t = Tables.get table in
-    let p = get_params_opt { empty_env with tables=[t]; 
-      schema=List.map (fun attr -> { Schema.Source.Attr.sources=[t |> fst]; attr }) (t |> snd); 
-      set_tyvar_strict = true 
-    } where in
-    [], p, Delete [table]
+    let env = { empty_env with tables=[t];
+      schema=List.map (fun attr -> { Schema.Source.Attr.sources=[t |> fst]; attr }) (t |> snd);
+    } in
+    let p = get_params_opt { env with set_tyvar_strict = true } where in
+    with_returning env returning ([], p, Delete [table])
   | DeleteMulti (targets, tables, where) ->
     (* use dummy columns to verify targets match the provided tables  *)
     let select = ({ columns = [dummy_loc All]; from = Some tables; where; group = []; having = None }, []) in
@@ -1811,7 +1839,7 @@ let rec eval (stmt:Sql.stmt) =
     | None -> [], p, Other
     | Some stmt -> let (schema,p2,kind) = eval stmt in (schema, p @ p2, kind)
     end
-  | Update (table,ss,w,o,lim) ->
+  | Update (table,ss,w,o,lim,returning) ->
     let f, s = Tables.get table in
     let env = { empty_env with is_update = true } in
     let r = List.map (fun attr -> {Schema.Source.Attr.attr; sources=[f] }) s in
@@ -1820,7 +1848,8 @@ let rec eval (stmt:Sql.stmt) =
     let env = { env with schema = update_schema_with_aliases [] r; is_update = true } in
     let p3 = params_of_order o [] { env with tables = [(f, s)] } in
     let lim = List.map (fun p -> make_param ~id:p.id ~typ:(Source_type.to_infer_type p.typ)) lim in
-    [], params @ p3 @ (List.map (fun p -> Single (p, Meta.empty())) lim), Update (Some table)
+    with_returning { empty_env with tables = [(f, s)]; schema = r } returning
+      ([], params @ p3 @ (List.map (fun p -> Single (p, Meta.empty())) lim), Update (Some table))
   | UpdateMulti (tables,ss,w,o,lim) ->
     let env = { empty_env with is_update = true } in
     let sources = List.map (fun src -> resolve_source { env with scope = Subquery } ((`Nested src), None)) tables in
@@ -1991,7 +2020,7 @@ let common_prefix = function
 (* fill inferred sql for VALUES or SET *)
 let complete_sql kind sql =
   match kind with
-  | Stmt.Insert (Some (kind,schema), _) ->
+  | Stmt.Insert (Some (kind,schema), _, _) ->
     let (pre,each,post) = match kind with
     | Values -> "(", (fun _ -> ""), ")"
     | Assign -> "", (fun name -> name ^" = "), ""