From 0f3ba1f69ff36625709b3da1af90e244198b6818 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Krier?= <ced@b2ck.com>
Date: Sun, 21 Aug 2011 10:20:23 +0200
Subject: [PATCH] Use RETURNING in INSERT if available

review21003
---
 CHANGELOG                              |  1 +
 trytond/backend/database.py            |  8 +++++++
 trytond/backend/postgresql/database.py | 30 ++++++++++++++++++++++----
 trytond/model/modelsql.py              | 27 +++++++++++++++--------
 4 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG b/CHANGELOG
index 72c3b15b42..dbb65d52b3 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,3 +1,4 @@
+* Add has_returning on Cursor
 * Remove use of python code in workflow XML
 * Use SQL sequence on PostgreSQL for ir.sequence
 * Allow to specify more than one interface
diff --git a/trytond/backend/database.py b/trytond/backend/database.py
index b61c88baa2..363b656150 100644
--- a/trytond/backend/database.py
+++ b/trytond/backend/database.py
@@ -225,3 +225,11 @@ class CursorInterface(object):
         :param value: the auto_increment value
         '''
         pass
+
+    def has_returning(self):
+        '''
+        Return True if database implements RETURNING clause in INSERT or UPDATE
+        statements.
+
+        :return: a boolean
+        '''
diff --git a/trytond/backend/postgresql/database.py b/trytond/backend/postgresql/database.py
index be7f01d7ec..2fa694ddbd 100644
--- a/trytond/backend/postgresql/database.py
+++ b/trytond/backend/postgresql/database.py
@@ -26,6 +26,7 @@ from trytond.protocols.datatype import Float
 
 RE_FROM = re.compile('.* from "?([a-zA-Z_0-9]+)"?.*$')
 RE_INTO = re.compile('.* into "?([a-zA-Z_0-9]+)"?.*$')
+RE_VERSION = re.compile(r'\S+ (\d+)\.(\d+)')
 
 
 class Database(DatabaseInterface):
@@ -34,6 +35,7 @@ class Database(DatabaseInterface):
     _connpool = None
     _list_cache = None
     _list_cache_timestamp = None
+    _version_cache = {}
 
     def __new__(cls, database_name='template1'):
         if database_name in cls._databases:
@@ -69,7 +71,7 @@ class Database(DatabaseInterface):
             conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
         else:
             conn.set_isolation_level(ISOLATION_LEVEL_SERIALIZABLE)
-        return Cursor(self._connpool, conn, self.database_name)
+        return Cursor(self._connpool, conn, self)
 
     def close(self):
         if self._connpool is None:
@@ -86,6 +88,14 @@ class Database(DatabaseInterface):
         cursor.execute('DROP DATABASE "' + database_name + '"')
         Database._list_cache = None
 
+    def get_version(self, cursor):
+        if self.database_name not in self._version_cache:
+            cursor.execute('SELECT version()')
+            version, = cursor.fetchone()
+            self._version_cache[self.database_name] = map(int,
+                RE_VERSION.search(version).groups())
+        return self._version_cache[self.database_name]
+
     @staticmethod
     def dump(database_name):
         from trytond.tools import exec_pg_command_pipe
@@ -270,12 +280,11 @@ class _Cursor(PsycopgCursor):
 
 class Cursor(CursorInterface):
 
-    def __init__(self, connpool, conn, database_name):
+    def __init__(self, connpool, conn, database):
         super(Cursor, self).__init__()
         self._connpool = connpool
         self._conn = conn
-        self.database_name = database_name
-        self.dbname = self.database_name #XXX to remove
+        self._database = database
         self.cursor = conn.cursor(cursor_factory=_Cursor)
         self.commit()
         self.sql_from_log = {}
@@ -285,6 +294,15 @@ class Cursor(CursorInterface):
             'into': 0,
         }
 
+    @property
+    def database_name(self):
+        return self._database.database_name
+
+    # TODO to remove
+    @property
+    def dbname(self):
+        return self.database_name
+
     def __getattr__(self, name):
         return getattr(self.cursor, name)
 
@@ -401,6 +419,10 @@ class Cursor(CursorInterface):
             select += ' OFFSET %d' % offset
         return select
 
+    def has_returning(self):
+        # RETURNING clause is available since PostgreSQL 8.2
+        return self._database.get_version(self) >= (8, 2)
+
 register_type(UNICODE)
 if PYDATE:
     register_type(PYDATE)
diff --git a/trytond/model/modelsql.py b/trytond/model/modelsql.py
index 89e7cba646..97a317cc97 100644
--- a/trytond/model/modelsql.py
+++ b/trytond/model/modelsql.py
@@ -333,16 +333,25 @@ class ModelSQL(ModelStorage):
         upd2.append(Transaction().user)
         upd2.append(datetime.datetime.now())
         try:
-            id_new = cursor.nextid(self._table)
-            if id_new:
-                cursor.execute('INSERT INTO "' + self._table + '" ' \
-                        '(id' + upd0 + ') ' \
-                        'VALUES (' + str(id_new) + upd1 + ')', tuple(upd2))
+            if cursor.has_returning():
+                    cursor.execute('INSERT INTO "' + self._table + '" '
+                        '(' + upd0[1:] + ') '
+                        'VALUES (' + upd1[1:] + ') RETURNING id',
+                        tuple(upd2))
+                    id_new, = cursor.fetchone()
             else:
-                cursor.execute('INSERT INTO "' + self._table + '" ' \
-                        '(' + upd0[1:] + ') ' \
-                        'VALUES (' + upd1[1:] + ')', tuple(upd2))
-                id_new = cursor.lastid()
+                id_new = cursor.nextid(self._table)
+                if id_new:
+                    cursor.execute('INSERT INTO "' + self._table + '" '
+                        '(id' + upd0 + ') '
+                        'VALUES (' + str(id_new) + upd1 + ')',
+                        tuple(upd2))
+                else:
+                    cursor.execute('INSERT INTO "' + self._table + '" '
+                        '(' + upd0[1:] + ') '
+                        'VALUES (' + upd1[1:] + ')',
+                        tuple(upd2))
+                    id_new = cursor.lastid()
         except DatabaseIntegrityError, exception:
             with contextlib.nested(Transaction().new_cursor(),
                     Transaction().set_user(0)):