Add sql execute to Connection

Closes #159

Co-authored-by Denis Ignatenko <[email protected]>

Author: artembo

tarantool/connection.py

@@ -34,7 +34,8 @@
     RequestSubscribe,
     RequestUpdate,
     RequestUpsert,
-    RequestAuthenticate
+    RequestAuthenticate,
+    RequestExecute
 )
 from tarantool.space import Space
 from tarantool.const import (
@@ -257,17 +258,18 @@ def _read_response(self):
 
     def _send_request_wo_reconnect(self, request):
         '''
-        :rtype: `Response` instance
+        :rtype: `Response` instance or subclass
 
         :raise: NetworkError
         '''
-        assert isinstance(request, Request)
+        if not isinstance(request, Request):
+            raise TypeError("Expected Request instance")
 
         response = None
         while True:
             try:
                 self._socket.sendall(bytes(request))
-                response = Response(self, self._read_response())
+                response = request.response_class(self, self._read_response())
                 break
             except SchemaReloadException as e:
                 self.update_schema(e.schema_version)
@@ -792,3 +794,22 @@ def generate_sync(self):
         Need override for async io connection
         '''
         return 0
+
+    def execute(self, query, params=None):
+        '''
+        Execute SQL request.
+
+        :param query: SQL syntax query
+        :type query: str
+
+        :param params: Bind values to use in query
+        :type params: list, dict
+
+        :return: query result data
+        :rtype: `Response` instance
+        '''
+        if not params:
+            params = []
+        request = RequestExecute(self, query, params)
+        response = self._send_request(request)
+        return response

tarantool/const.py

@@ -29,6 +29,13 @@
 #
 IPROTO_DATA = 0x30
 IPROTO_ERROR = 0x31
+#
+IPROTO_METADATA = 0x32
+IPROTO_SQL_TEXT = 0x40
+IPROTO_SQL_BIND = 0x41
+IPROTO_SQL_INFO = 0x42
+IPROTO_SQL_INFO_ROW_COUNT = 0x00
+IPROTO_SQL_INFO_AUTOINCREMENT_IDS = 0x01
 
 IPROTO_GREETING_SIZE = 128
 IPROTO_BODY_MAX_LEN = 2147483648
@@ -44,6 +51,7 @@
 REQUEST_TYPE_EVAL = 8
 REQUEST_TYPE_UPSERT = 9
 REQUEST_TYPE_CALL = 10
+REQUEST_TYPE_EXECUTE = 11
 REQUEST_TYPE_PING = 64
 REQUEST_TYPE_JOIN = 65
 REQUEST_TYPE_SUBSCRIBE = 66

tarantool/request.py

@@ -4,10 +4,17 @@
 Request types definitions
 '''
 
+import collections
 import msgpack
 import hashlib
 
+try:
+    collectionsAbc = collections.abc
+except AttributeError:
+    collectionsAbc = collections
 
+
+from tarantool.error import DatabaseError
 from tarantool.const import (
     IPROTO_CODE,
     IPROTO_SYNC,
@@ -27,6 +34,8 @@
     IPROTO_OPS,
     # IPROTO_INDEX_BASE,
     IPROTO_SCHEMA_ID,
+    IPROTO_SQL_TEXT,
+    IPROTO_SQL_BIND,
     REQUEST_TYPE_OK,
     REQUEST_TYPE_PING,
     REQUEST_TYPE_SELECT,
@@ -37,11 +46,13 @@
     REQUEST_TYPE_UPSERT,
     REQUEST_TYPE_CALL16,
     REQUEST_TYPE_CALL,
+    REQUEST_TYPE_EXECUTE,
     REQUEST_TYPE_EVAL,
     REQUEST_TYPE_AUTHENTICATE,
     REQUEST_TYPE_JOIN,
     REQUEST_TYPE_SUBSCRIBE
 )
+from tarantool.response import Response, ResponseExecute
 from tarantool.utils import (
     strxor,
     binary_types
@@ -64,6 +75,7 @@ def __init__(self, conn):
         self.conn = conn
         self._sync = None
         self._body = ''
+        self.response_class = Response
 
         packer_kwargs = dict()
 
@@ -360,3 +372,25 @@ def __init__(self, conn, sync):
         request_body = self._dumps({IPROTO_CODE: self.request_type,
                                     IPROTO_SYNC: sync})
         self._body = request_body
+
+
+class RequestExecute(Request):
+    '''
+    Represents EXECUTE SQL request
+    '''
+    request_type = REQUEST_TYPE_EXECUTE
+
+    # pylint: disable=W0231
+    def __init__(self, conn, sql, args):
+        super(RequestExecute, self).__init__(conn)
+        if isinstance(args, collectionsAbc.Mapping):
+            args = [{":%s" % name: value} for name, value in args.items()]
+        elif not isinstance(args, collectionsAbc.Sequence):
+            raise TypeError("Parameter type '%s' is not supported. "
+                            "Must be a mapping or sequence" % type(args))
+
+        request_body = self._dumps({IPROTO_SQL_TEXT: sql,
+                                    IPROTO_SQL_BIND: args})
+
+        self._body = request_body
+        self.response_class = ResponseExecute

tarantool/response.py

@@ -17,7 +17,10 @@
     IPROTO_ERROR,
     IPROTO_SYNC,
     IPROTO_SCHEMA_ID,
-    REQUEST_TYPE_ERROR
+    REQUEST_TYPE_ERROR,
+    IPROTO_SQL_INFO,
+    IPROTO_SQL_INFO_ROW_COUNT,
+    IPROTO_SQL_INFO_AUTOINCREMENT_IDS
 )
 from tarantool.error import (
     DatabaseError,
@@ -268,3 +271,51 @@ def __str__(self):
         return ''.join(output)
 
     __repr__ = __str__
+
+
+class ResponseExecute(Response):
+    @property
+    def autoincrement_ids(self):
+        """
+        Returns the list of tuples containing
+        values of autoincrement fields for a DML request
+        and None for a DQL request (NOT result set size)
+
+        :rtype: list or None
+        """
+        if self._body is None:
+            raise InterfaceError("Trying to access data, when there's no data")
+        info = self._body.get(IPROTO_SQL_INFO)
+
+        if info is None:
+            return None
+
+        autoincrement_ids = info.get(IPROTO_SQL_INFO_AUTOINCREMENT_IDS)
+
+        return autoincrement_ids
+
+    @property
+    def affected_row_count(self):
+        """
+        Returns the number of affected rows for responses
+        to DML requests and None for other requests
+        :rtype: int
+        """
+        if self._body is None:
+            raise InterfaceError("Trying to access data, when there's no data")
+
+        info = self._body.get(IPROTO_SQL_INFO)
+
+        if info is None:
+            return None
+
+        return info.get(IPROTO_SQL_INFO_ROW_COUNT)
+
+    @property
+    def rows(self):
+        """
+        :rtype: list or None
+        """
+        if self._return_code != 0:
+            return None
+        return self._body.get(IPROTO_DATA)

unit/suites/__init__.py

@@ -10,11 +10,12 @@
 from .test_protocol import TestSuite_Protocol
 from .test_reconnect import TestSuite_Reconnect
 from .test_mesh import TestSuite_Mesh
+from .test_execute import TestSuite_Execute
 
 test_cases = (TestSuite_Schema_UnicodeConnection,
               TestSuite_Schema_BinaryConnection,
               TestSuite_Request, TestSuite_Protocol, TestSuite_Reconnect,
-              TestSuite_Mesh)
+              TestSuite_Mesh, TestSuite_Execute)
 
 def load_tests(loader, tests, pattern):
     suite = unittest.TestSuite()

unit/suites/test_execute.py

@@ -0,0 +1,88 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import print_function
+
+import sys
+import unittest
+
+import tarantool
+from .lib.tarantool_server import TarantoolServer
+
+
+class TestSuite_Execute(unittest.TestCase):
+    ddl = 'create table %s (id INTEGER PRIMARY KEY AUTOINCREMENT, ' \
+           'name varchar(20))'
+
+    ddl_params = [
+        {'id': None, 'name': 'Michael'},
+        {'id': None, 'name': 'Mary'},
+        {'id': None, 'name': 'John'},
+        {'id': None, 'name': 'Ruth'},
+        {'id': None, 'name': 'Rachel'}
+    ]
+
+    @classmethod
+    def setUpClass(self):
+        print(' EXECUTE '.center(70, '='), file=sys.stderr)
+        print('-' * 70, file=sys.stderr)
+        self.srv = TarantoolServer()
+        self.srv.script = 'unit/suites/box.lua'
+        self.srv.start()
+        self.con = tarantool.Connection(self.srv.host, self.srv.args['primary'])
+
+    def setUp(self):
+        # prevent a remote tarantool from clean our session
+        if self.srv.is_started():
+            self.srv.touch_lock()
+        self.con.flush_schema()
+
+        # grant full access to guest
+        self.srv.admin("box.schema.user.grant('guest', 'create,read,write,"
+                       "execute', 'universe')")
+
+    @classmethod
+    def tearDownClass(self):
+        self.con.close()
+        self.srv.stop()
+        self.srv.clean()
+
+    def _populate_data(self, table_name):
+        query = "insert into %s values (:id, :name)" % table_name
+        for param in self.ddl_params:
+            self.con.execute(query, param)
+
+    def _create_table(self, table_name):
+        return self.con.execute(self.ddl % table_name)
+
+    def test_dml_response(self):
+        table_name = 'foo'
+        response = self._create_table(table_name)
+        self.assertEqual(response.autoincrement_ids, None)
+        self.assertEqual(response.affected_row_count, 1)
+        self.assertEqual(response.data, None)
+
+        query = "insert into %s values (:id, :name)" % table_name
+
+        for num, param in enumerate(self.ddl_params, start=1):
+            response = self.con.execute(query, param)
+            self.assertEqual(response.autoincrement_ids[0], num)
+            self.assertEqual(response.affected_row_count, 1)
+            self.assertEqual(response.data, None)
+
+        query = "delete from %s where id in (4, 5)" % table_name
+        response = self.con.execute(query)
+        self.assertEqual(response.autoincrement_ids, None)
+        self.assertEqual(response.affected_row_count, 2)
+        self.assertEqual(response.data, None)
+
+    def test_dql_response(self):
+        table_name = 'bar'
+        self._create_table(table_name)
+        self._populate_data(table_name)
+
+        select_query = "select name from %s where id in (1, 3, 5)" % table_name
+        response = self.con.execute(select_query)
+        self.assertEqual(response.autoincrement_ids, None)
+        self.assertEqual(response.affected_row_count, None)
+        expected_data = [['Michael'], ['John'], ['Rachel']]
+        self.assertListEqual(response.data, expected_data)