Merge branch 'develop'

mogui · Oct 10, 2015 · f48eb05 · f48eb05
2 parents 83267b1 + eb40386
commit f48eb05
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 2 deletions.
diff --git a/pyorient/orient.py b/pyorient/orient.py
@@ -92,7 +92,10 @@ def close(self):
         self.connected = False
 
     def write(self, buff):
-        return self._socket.send(buff)
+        count = 0
+        while count < len(buff):
+            count += self._socket.send(buff[count:])
+        return count
 
     # The man page for recv says: The receive calls normally return
     #   any data available, up to the requested amount, rather than waiting

diff --git a/pyorient/types.py b/pyorient/types.py
@@ -23,6 +23,14 @@ def __str__(self):
             rep = rep + ",'rid':'" + str(self.__rid) + "'"
         return '{' + rep + '}'
 
+    @staticmethod
+    def addslashes(string):
+        l = [ "\\", '"', "'", "\0", ]
+        for i in l:
+            if i in string:
+                string = string.replace( i, '\\' + i )
+        return string
+
     def __init__(self, content=None):
 
         self.__rid = None
@@ -45,7 +53,7 @@ def __init__(self, content=None):
                 # { '@my_class': { 'accommodation': 'hotel' } }
                 self.__o_class = key[1:]
                 for _key, _value in content[key].items():
-                    self.__o_storage[_key] = _value
+                    self.__o_storage[_key] = self.addslashes( _value )
             elif key == '__o_storage':
                 self.__o_storage = content[key]
             else:

diff --git a/pyorient/utils.py b/pyorient/utils.py
@@ -75,6 +75,8 @@ def parse_cluster_id(cluster_id):
             cluster_id = cluster_id.decode("utf-8")
         elif isinstance( cluster_id, OrientRecordLink ):
             cluster_id = cluster_id.get()
+        elif sys.version_info[0] < 3 and isinstance( cluster_id, unicode ):
+            cluster_id = cluster_id.encode('utf-8')
 
         _cluster_id, _position = cluster_id.split( ':' )
         if _cluster_id[0] is '#':

diff --git a/tests/test_serializations.py b/tests/test_serializations.py
@@ -46,3 +46,67 @@ def test_csv_encoding(self):
         assert 'specie:"rodent"' in raw
 
         # TODO: add several more complex tests to have more coverage
+
+    def test_csv_escape(self):
+        import pyorient
+        DB = pyorient.OrientDB("localhost", 2424)
+        DB.connect("root", "root")
+
+        db_name = "test_escape"
+        try:
+            DB.db_drop(db_name)
+        except pyorient.PyOrientCommandException as e:
+            print(e)
+        finally:
+            db = DB.db_create(db_name, pyorient.DB_TYPE_GRAPH,
+                              pyorient.STORAGE_TYPE_MEMORY)
+            pass
+
+        cluster_info = DB.db_open(
+            db_name, "admin", "admin", pyorient.DB_TYPE_GRAPH, ""
+        )
+
+        cluster_id = DB.command( "CREATE CLASS MyModel EXTENDS V" )[0]
+
+        data0 = {'key': '"""'}
+        DB.record_create( cluster_id, {'@MyModel': data0} )
+
+        data1 = {'key': "'''"}
+        DB.record_create( cluster_id, {'@MyModel': data1} )
+
+        data2 = {'key': '\\'}
+        DB.record_create( cluster_id, {'@MyModel': data2} )
+
+        data3 = {'key': '\0'}
+        DB.record_create( cluster_id, {'@MyModel': data3} )
+
+        data4 = {'key': '""'}
+        DB.record_create( cluster_id, {'@MyModel': data4} )
+
+        data5 = {'key': '\'\'""\0 \\ execution'}
+        DB.record_create( cluster_id, {'@MyModel': data5} )
+
+        rec0 = DB.record_load( "#" + cluster_id.decode() + ":0" )
+        assert rec0._class == "MyModel"
+        assert rec0.oRecordData == data0
+
+        rec1 = DB.record_load( "#" + cluster_id.decode('utf-8') + ":1" )
+        assert rec1._class == "MyModel"
+        assert rec1.oRecordData == data1
+
+        rec2 = DB.record_load( "#" + cluster_id.decode('utf-8') + ":2")
+        assert rec2._class == "MyModel"
+        assert rec2.oRecordData == data2
+
+        rec3 = DB.record_load( "#" + cluster_id.decode('utf-8') + ":3" )
+        assert rec3._class == "MyModel"
+        assert rec3.oRecordData == data3
+
+        rec4 = DB.record_load( "#" + cluster_id.decode('utf-8') + ":4" )
+        assert rec4._class == "MyModel"
+        assert rec4.oRecordData == data4
+
+        rec5 = DB.record_load( "#" + cluster_id.decode('utf-8') + ":5" )
+        assert rec5._class == "MyModel"
+        assert rec5.oRecordData == data5
+