[tarantool-patches] [PATCH v5 1/1] sql: return all generated ids via IPROTO

[imeevma@tarantool.org]

According to documentation some JDBC functions have an ability to
return all ids that were generated in executed INSERT statement.
This patch gives a way to implement such functionality.

Closes #2618
---
Branch: https://github.com/tarantool/tarantool/tree/imeevma/gh-2618-return-all-generated-ids
Issue: https://github.com/tarantool/tarantool/issues/2618

 src/box/execute.c        | 24 ++++++++++++++++
 src/box/execute.h        |  1 +
 src/box/lua/net_box.c    | 21 ++++++++++++--
 src/box/sequence.c       | 29 ++++++++++++++++++++
 src/box/sql/vdbe.c       |  2 +-
 src/box/sql/vdbe.h       |  3 +-
 src/box/sql/vdbeInt.h    |  3 ++
 src/box/sql/vdbeaux.c    | 15 ++++++++--
 src/box/txn.h            | 13 +++++++++
 test/sql/iproto.result   | 71 ++++++++++++++++++++++++++++++++++++++++++++++++
 test/sql/iproto.test.lua | 15 ++++++++++
 11 files changed, 190 insertions(+), 7 deletions(-)

diff --git a/src/box/execute.c b/src/box/execute.c
index 24459b4..d9fe33f 100644
--- a/src/box/execute.c
+++ b/src/box/execute.c
@@ -42,6 +42,7 @@
 #include "schema.h"
 #include "port.h"
 #include "tuple.h"
+#include "sql/vdbeInt.h"
 
 const char *sql_type_strs[] = {
 	NULL,
@@ -639,14 +640,37 @@ err:
 		assert(port_tuple->size == 0);
 		if (iproto_reply_map_key(out, 1, IPROTO_SQL_INFO) != 0)
 			goto err;
+		struct stailq *id_list = &((struct Vdbe *)stmt)->id_list;
+		uint64_t id_count = 0;
 		int changes = sqlite3_changes(db);
 		int size = mp_sizeof_uint(SQL_INFO_ROW_COUNT) +
 			   mp_sizeof_uint(changes);
+		if (stailq_empty(id_list) == 0) {
+			struct id_entry *id_entry;
+			stailq_foreach_entry(id_entry, id_list, link) {
+				size += id_entry->id >= 0 ?
+					mp_sizeof_uint(id_entry->id) :
+					mp_sizeof_int(id_entry->id);
+				id_count++;
+			}
+			size += mp_sizeof_uint(SQL_INFO_GENERATED_IDS) +
+				mp_sizeof_array(id_count);
+		}
 		char *buf = obuf_alloc(out, size);
 		if (buf == NULL) {
 			diag_set(OutOfMemory, size, "obuf_alloc", "buf");
 			goto err;
 		}
+		if (id_count > 0) {
+			buf = mp_encode_uint(buf, SQL_INFO_GENERATED_IDS);
+			buf = mp_encode_array(buf, id_count);
+			struct id_entry *id_entry;
+			stailq_foreach_entry(id_entry, id_list, link) {
+				buf = id_entry->id >= 0 ?
+				      mp_encode_uint(buf, id_entry->id) :
+				      mp_encode_int(buf, id_entry->id);
+			}
+		}
 		buf = mp_encode_uint(buf, SQL_INFO_ROW_COUNT);
 		buf = mp_encode_uint(buf, changes);
 	}
diff --git a/src/box/execute.h b/src/box/execute.h
index f21393b..614d3d0 100644
--- a/src/box/execute.h
+++ b/src/box/execute.h
@@ -42,6 +42,7 @@ extern "C" {
 /** Keys of IPROTO_SQL_INFO map. */
 enum sql_info_key {
 	SQL_INFO_ROW_COUNT = 0,
+	SQL_INFO_GENERATED_IDS = 1,
 	sql_info_key_MAX,
 };
 
diff --git a/src/box/lua/net_box.c b/src/box/lua/net_box.c
index a928a4c..5d2ac78 100644
--- a/src/box/lua/net_box.c
+++ b/src/box/lua/net_box.c
@@ -671,11 +671,28 @@ netbox_decode_sql_info(struct lua_State *L, const char **data)
 	/* Only SQL_INFO_ROW_COUNT is available. */
 	assert(map_size == 1);
 	(void) map_size;
+	lua_newtable(L);
 	uint32_t key = mp_decode_uint(data);
+	/*
+	 * If SQL_INFO_GENERATED_IDS in data then it should be
+	 * just before SQL_INFO_ROW_COUNT.
+	 */
+	if (key == SQL_INFO_GENERATED_IDS) {
+		uint64_t count = mp_decode_array(data);
+		assert (count > 0);
+		lua_createtable(L, 0, count);
+		lua_setfield(L, -2, "generated_ids");
+		lua_getfield(L, -1, "generated_ids");
+		for (uint32_t j = 0; j < count; ++j) {
+			int64_t value = mp_decode_uint(data);
+			lua_pushinteger(L, value);
+			lua_rawseti(L, -2, j + 1);
+		}
+		lua_pop(L, 1);
+		key = mp_decode_uint(data);
+	}
 	assert(key == SQL_INFO_ROW_COUNT);
-	(void) key;
 	uint32_t row_count = mp_decode_uint(data);
-	lua_createtable(L, 0, 1);
 	lua_pushinteger(L, row_count);
 	lua_setfield(L, -2, "rowcount");
 }
diff --git a/src/box/sequence.c b/src/box/sequence.c
index 35b7605..ed6b2da 100644
--- a/src/box/sequence.c
+++ b/src/box/sequence.c
@@ -38,6 +38,7 @@
 #include <small/mempool.h>
 #include <msgpuck/msgpuck.h>
 
+#include "txn.h"
 #include "diag.h"
 #include "error.h"
 #include "errcode.h"
@@ -178,6 +179,30 @@ sequence_update(struct sequence *seq, int64_t value)
 	return 0;
 }
 
+/**
+ * Save generated id in VDBE.
+ *
+ * @param value ID to save in VDBE.
+ * @retval 0 Success.
+ * @retval -1 Error.
+ */
+static inline int
+add_new_id_in_vdbe(int64_t value)
+{
+	struct txn *txn = in_txn();
+	if (txn == NULL || txn->psql_txn == NULL)
+		return 0;
+	assert(txn->psql_txn->vdbe != NULL);
+	struct id_entry *id_entry = malloc(sizeof(*id_entry));
+	if (id_entry == NULL) {
+		diag_set(OutOfMemory, sizeof(*id_entry), "malloc", "id_entry");
+		return -1;
+	}
+	id_entry->id = value;
+	stailq_add_tail_entry(txn->psql_txn->id_list, id_entry, link);
+	return 0;
+}
+
 int
 sequence_next(struct sequence *seq, int64_t *result)
 {
@@ -194,6 +219,8 @@ sequence_next(struct sequence *seq, int64_t *result)
 					  new_data) == light_sequence_end)
 			return -1;
 		*result = def->start;
+		if(add_new_id_in_vdbe(*result) != 0)
+			return -1;
 		return 0;
 	}
 	old_data = light_sequence_get(&sequence_data_index, pos);
@@ -228,6 +255,8 @@ done:
 				   new_data, &old_data) == light_sequence_end)
 		unreachable();
 	*result = value;
+	if(add_new_id_in_vdbe(*result) != 0)
+		return -1;
 	return 0;
 overflow:
 	if (!def->cycle) {
diff --git a/src/box/sql/vdbe.c b/src/box/sql/vdbe.c
index 00ffb0c..cbec9a2 100644
--- a/src/box/sql/vdbe.c
+++ b/src/box/sql/vdbe.c
@@ -3040,7 +3040,7 @@ case OP_TransactionRollback: {
  */
 case OP_TTransaction: {
 	if (!box_txn()) {
-		if (box_txn_begin() != 0) {
+		if (sql_txn_begin(p) != 0) {
 			rc = SQL_TARANTOOL_ERROR;
 			goto abort_due_to_error;
 		}
diff --git a/src/box/sql/vdbe.h b/src/box/sql/vdbe.h
index d5da5d5..9b94183 100644
--- a/src/box/sql/vdbe.h
+++ b/src/box/sql/vdbe.h
@@ -179,10 +179,11 @@ Vdbe *sqlite3VdbeCreate(Parse *);
  * Allocate and initialize SQL-specific struct which completes
  * original Tarantool's txn struct using region allocator.
  *
+ * @param vdbe VDBE that is being prepared or executed.
  * @retval NULL on OOM, new psql_txn struct on success.
  **/
 struct sql_txn *
-sql_alloc_txn();
+sql_alloc_txn(struct Vdbe *vdbe);
 
 /**
  * Prepare given VDBE to execution: initialize structs connected
diff --git a/src/box/sql/vdbeInt.h b/src/box/sql/vdbeInt.h
index ce97f49..ca1794c 100644
--- a/src/box/sql/vdbeInt.h
+++ b/src/box/sql/vdbeInt.h
@@ -368,6 +368,9 @@ struct Vdbe {
 	/** The auto-commit flag. */
 	bool auto_commit;
 
+	/** List of id generated in current VDBE. */
+	struct stailq id_list;
+
 	/* When allocating a new Vdbe object, all of the fields below should be
 	 * initialized to zero or NULL
 	 */
diff --git a/src/box/sql/vdbeaux.c b/src/box/sql/vdbeaux.c
index 54edf1b..98ddef7 100644
--- a/src/box/sql/vdbeaux.c
+++ b/src/box/sql/vdbeaux.c
@@ -57,6 +57,8 @@ sqlite3VdbeCreate(Parse * pParse)
 		return 0;
 	memset(p, 0, sizeof(Vdbe));
 	p->db = db;
+	/* Init list of generated ids. */
+	stailq_create(&p->id_list);
 	if (db->pVdbe) {
 		db->pVdbe->pPrev = p;
 	}
@@ -75,7 +77,7 @@ sqlite3VdbeCreate(Parse * pParse)
 }
 
 struct sql_txn *
-sql_alloc_txn()
+sql_alloc_txn(struct Vdbe *vdbe)
 {
 	struct sql_txn *txn = region_alloc_object(&fiber()->gc,
 						  struct sql_txn);
@@ -86,6 +88,8 @@ sql_alloc_txn()
 	}
 	txn->pSavepoint = NULL;
 	txn->fk_deferred_count = 0;
+	txn->vdbe = vdbe;
+	txn->id_list = &vdbe->id_list;
 	return txn;
 }
 
@@ -103,7 +107,7 @@ sql_vdbe_prepare(struct Vdbe *vdbe)
 		 * check FK violations, at least now.
 		 */
 		if (txn->psql_txn == NULL) {
-			txn->psql_txn = sql_alloc_txn();
+			txn->psql_txn = sql_alloc_txn(vdbe);
 			if (txn->psql_txn == NULL)
 				return -1;
 		}
@@ -2261,7 +2265,7 @@ sql_txn_begin(Vdbe *p)
 	struct txn *ptxn = txn_begin(false);
 	if (ptxn == NULL)
 		return -1;
-	ptxn->psql_txn = sql_alloc_txn();
+	ptxn->psql_txn = sql_alloc_txn(p);
 	if (ptxn->psql_txn == NULL) {
 		box_txn_rollback();
 		return -1;
@@ -2746,6 +2750,11 @@ sqlite3VdbeClearObject(sqlite3 * db, Vdbe * p)
 	vdbeFreeOpArray(db, p->aOp, p->nOp);
 	sqlite3DbFree(db, p->aColName);
 	sqlite3DbFree(db, p->zSql);
+	while (stailq_empty(&p->id_list) == 0) {
+		struct id_entry *id_entry =
+			stailq_shift_entry(&p->id_list, struct id_entry, link);
+		free(id_entry);
+	}
 #ifdef SQLITE_ENABLE_STMT_SCANSTATUS
 	{
 		int i;
diff --git a/src/box/txn.h b/src/box/txn.h
index e74c14d..f8521c4 100644
--- a/src/box/txn.h
+++ b/src/box/txn.h
@@ -117,6 +117,19 @@ struct sql_txn {
 	 * VDBE to the next in the same transaction.
 	 */
 	uint32_t fk_deferred_count;
+	/** Current VDBE. */
+	struct Vdbe *vdbe;
+	/** Pointer to list of generated ids of current VDBE.*/
+	struct stailq *id_list;
+};
+
+/** An element of list of generated ids. */
+struct id_entry
+{
+	/** A link in a generated id list. */
+	struct stailq_entry link;
+	/** Generated id. */
+	int64_t id;
 };
 
 struct txn {
diff --git a/test/sql/iproto.result b/test/sql/iproto.result
index af474bc..163fb6e 100644
--- a/test/sql/iproto.result
+++ b/test/sql/iproto.result
@@ -580,6 +580,77 @@ cn:close()
 box.sql.execute('drop table test')
 ---
 ...
+-- gh-2618 Return generated columns after INSERT in IPROTO.
+-- Return all ids generated in current INSERT statement.
+box.sql.execute('create table test (id integer primary key autoincrement, a integer)')
+---
+...
+cn = remote.connect(box.cfg.listen)
+---
+...
+cn:execute('insert into test values (1, 1)')
+---
+- rowcount: 1
+...
+cn:execute('insert into test values (null, 2)')
+---
+- generated_ids:
+  - 2
+  rowcount: 1
+...
+cn:execute('update test set a = 11 where id == 1')
+---
+- rowcount: 1
+...
+cn:execute('insert into test values (100, 1), (null, 1), (120, 1), (null, 1)')
+---
+- generated_ids:
+  - 101
+  - 121
+  rowcount: 4
+...
+cn:execute('insert into test values (null, 1), (null, 1), (null, 1), (null, 1), (null, 1)')
+---
+- generated_ids:
+  - 122
+  - 123
+  - 124
+  - 125
+  - 126
+  rowcount: 5
+...
+cn:execute('select * from test')
+---
+- metadata:
+  - name: ID
+  - name: A
+  rows:
+  - [1, 11]
+  - [2, 2]
+  - [100, 1]
+  - [101, 1]
+  - [120, 1]
+  - [121, 1]
+  - [122, 1]
+  - [123, 1]
+  - [124, 1]
+  - [125, 1]
+  - [126, 1]
+...
+cn:execute('create table test2 (id integer primary key, a integer)')
+---
+- rowcount: 1
+...
+cn:execute('drop table test2')
+---
+- rowcount: 1
+...
+cn:close()
+---
+...
+box.sql.execute('drop table test')
+---
+...
 box.schema.user.revoke('guest', 'read,write,execute', 'universe')
 ---
 ...
diff --git a/test/sql/iproto.test.lua b/test/sql/iproto.test.lua
index 220331b..6517cde 100644
--- a/test/sql/iproto.test.lua
+++ b/test/sql/iproto.test.lua
@@ -201,6 +201,21 @@ future4:wait_result()
 cn:close()
 box.sql.execute('drop table test')
 
+-- gh-2618 Return generated columns after INSERT in IPROTO.
+-- Return all ids generated in current INSERT statement.
+box.sql.execute('create table test (id integer primary key autoincrement, a integer)')
+cn = remote.connect(box.cfg.listen)
+cn:execute('insert into test values (1, 1)')
+cn:execute('insert into test values (null, 2)')
+cn:execute('update test set a = 11 where id == 1')
+cn:execute('insert into test values (100, 1), (null, 1), (120, 1), (null, 1)')
+cn:execute('insert into test values (null, 1), (null, 1), (null, 1), (null, 1), (null, 1)')
+cn:execute('select * from test')
+cn:execute('create table test2 (id integer primary key, a integer)')
+cn:execute('drop table test2')
+cn:close()
+box.sql.execute('drop table test')
+
 box.schema.user.revoke('guest', 'read,write,execute', 'universe')
 space = nil
 
-- 
2.7.4