BUG23550743: Add close method to XSession and NodeSession

amitab · amitab · commit def54743e9c6 · 2016-07-15T14:33:18.000+05:30
Added test cases for regression
diff --git a/lib/mysqlx/connection.py b/lib/mysqlx/connection.py
@@ -26,7 +26,7 @@
 import socket
 
 from .authentication import MySQL41AuthPlugin
-from .errors import InterfaceError
+from .errors import InterfaceError, OperationalError
 from .crud import Schema
 from .protocol import Protocol, MessageReaderWriter
 from .result import Result, RowResult, DocResult
@@ -45,6 +45,9 @@ def connect(self, host, port):
         self._socket.connect((host, port,))
 
     def read(self, count):
+        if self._socket == None:
+            raise OperationalError("MySQLx Connection not available")
+
         buf = ""
         while count > 0:
             data = self._socket.recv(count)
@@ -55,8 +58,15 @@ def read(self, count):
         return buf
 
     def sendall(self, data):
+        if self._socket == None:
+            raise OperationalError("MySQLx Connection not available")
+
         self._socket.sendall(data)
 
+    def close(self):
+        self._socket.close()
+        self._socket = None
+
 
 class Connection(object):
     def __init__(self, settings):
@@ -128,6 +138,13 @@ def get_row_result(self, cmd, *args):
         self.protocol.send_execute_statement("xplugin", cmd, args)
         return RowResult(self)
 
+    def close(self):
+        if self._active_result is not None:
+            self._active_result.fetch_all()
+        self.protocol.send_close()
+        self.protocol.read_ok()
+        self.stream.close()
+
 
 class BaseSession(object):
     """Base functionality for Session classes through the X Protocol.
@@ -148,6 +165,9 @@ def __init__(self, settings):
         self._connection = Connection(self._settings)
         self._connection.connect()
 
+    def is_open(self):
+        return self._connection.stream._socket is not None
+
     def get_schema(self, name):
         """Retrieves a Schema object from the current session by it's name.
 
@@ -196,6 +216,9 @@ def rollback(self):
         """
         self._connection.execute_nonquery("sql", "ROLLBACK", True)
 
+    def close(self):
+        self._connection.close()
+
 
 class XSession(BaseSession):
     """Enables interaction with a X Protocol enabled MySQL Product.
diff --git a/lib/mysqlx/protocol.py b/lib/mysqlx/protocol.py
@@ -55,6 +55,7 @@
     (MySQLx.ServerMessages.RESULTSET_FETCH_DONE, MySQLxResultset.FetchDone),
     (MySQLx.ServerMessages.RESULTSET_FETCH_DONE_MORE_RESULTSETS,
      MySQLxResultset.FetchDoneMoreResultsets),
+    (MySQLx.ServerMessages.OK, MySQLx.Ok),
 ]
 
 
@@ -333,3 +334,15 @@ def arg_object_to_expr(self, value, allow_relational):
 
     def arg_object_to_scalar(self, value, allow_relational):
         return self.arg_object_to_expr(value, allow_relational).literal
+
+    def read_ok(self):
+        msg = self._reader.read_message()
+        if isinstance(msg, MySQLx.Error):
+            raise InterfaceError(msg.msg)
+
+        if not isinstance(msg, MySQLx.Ok):
+            raise InterfaceError("Unexpected message encountered")
+
+    def send_close(self):
+        msg = MySQLxSession.Close()
+        self._writer.write_message(MySQLx.ClientMessages.SESS_CLOSE, msg)
diff --git a/tests/test_mysqlx_connection.py b/tests/test_mysqlx_connection.py
@@ -68,6 +68,13 @@ def test_create_schema(self):
         schema = self.session.create_schema(self.schema_name)
         self.assertTrue(schema.exists_in_database())
 
+    def test_close(self):
+        session = mysqlx.get_session(self.connect_kwargs)
+        schema = session.get_schema(self.schema_name)
+
+        session.close()
+        self.assertRaises(mysqlx.OperationalError, schema.exists_in_database)
+
 
 @unittest.skipIf(tests.MYSQL_VERSION < (5, 7, 12), "XPlugin not compatible")
 class MySQLxNodeSessionTests(tests.MySQLxTests):
