Merge pull request #10 from laughingman7743/retrieve_query_id

Add property to retrieve query_id (fix #9)

pyathenajdbc/cursor.py

@@ -75,6 +75,12 @@ def description(self):
         ]
         return self._description
 
+    @property
+    def query_id(self):
+        if not self.has_result_set:
+            raise ProgrammingError('No result set.')
+        return self._result_set.getClient().getQueryExecutionId()
+
     def close(self):
         self._meta_data = None
         if self._result_set and not self._result_set.isClosed():

tests/test_pyathenajdbc.py

@@ -4,6 +4,7 @@
 import contextlib
 import os
 import random
+import re
 import string
 import unittest
 from datetime import datetime, date
@@ -196,6 +197,14 @@ def test_description(self, cursor):
         cursor.execute('SELECT 1 AS foobar FROM one_row')
         self.assertEqual(cursor.description, [('foobar', 4, 11, None, 10, 0, 2)])
 
+    @with_cursor
+    def test_query_id(self, cursor):
+        cursor.execute('SELECT * from one_row')
+        # query_id is UUID v4
+        expected_pattern = \
+            r'^[0-9a-f]{8}-[0-9a-f]{4}-[4][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$'
+        self.assertTrue(re.match(expected_pattern, cursor.query_id))
+
     @with_cursor
     def test_complex(self, cursor):
         # TODO DECIMAL type failed to fetch with java.lang.IndexOutOfBoundsException
@@ -311,6 +320,7 @@ def test_cursor_is_closed(self):
             self.assertRaises(ProgrammingError, lambda: cursor.fetchmany())
             self.assertRaises(ProgrammingError, lambda: cursor.fetchall())
             self.assertRaises(ProgrammingError, lambda: cursor.cancel())
+            self.assertRaises(ProgrammingError, lambda: cursor.query_id)
 
     def test_no_ops(self):
         conn = self.connect()