added insert_id

petehunt · petehunt · commit fecc4a6cfad5 · 2011-07-29T11:57:52.000-07:00
diff --git a/pymysql/connections.py b/pymysql/connections.py
@@ -508,6 +508,7 @@ def _config(key, default):
         self.encoders = encoders
         self.decoders = conv
 
+        self._result = None
         self._affected_rows = 0
         self.host_info = "Not connected"
 
@@ -685,6 +686,12 @@ def _read_query_result(self):
         self._result = result
         return result.affected_rows
 
+    def insert_id(self):
+        if self._result:
+            return self._result.insert_id
+        else:
+            return 0
+
     def _send_command(self, command, sql):
         #send_data = struct.pack('<i', len(sql) + 1) + command + sql
         # could probably be more efficient, at least it's correct
diff --git a/pymysql/cursors.py b/pymysql/cursors.py
@@ -240,7 +240,6 @@ def _do_get_result(self):
         self.lastrowid = conn._result.insert_id
         self._rows = conn._result.rows
         self._has_next = conn._result.has_next
-        conn._result = None
 
     def __iter__(self):
         return iter(self.fetchone, None)
diff --git a/pymysql/tests/test_issues.py b/pymysql/tests/test_issues.py
@@ -249,7 +249,17 @@ def disabled_test_issue_54(self):
             c.execute("drop table issue54")
 
 class TestGitHubIssues(base.PyMySQLTestCase):
-    pass
+    def test_issue_66(self):
+        conn = self.connections[0]
+        c = conn.cursor()
+        self.assertEquals(0, conn.insert_id())
+        try:
+            c.execute("create table issue66 (id integer primary key auto_increment, x integer)")
+            c.execute("insert into issue66 (x) values (1)")
+            c.execute("insert into issue66 (x) values (1)")
+            self.assertEquals(2, conn.insert_id())
+        finally:
+            c.execute("drop table issue66")
 
 __all__ = ["TestOldIssues", "TestNewIssues", "TestGitHubIssues"]
 
