Testing changes

Author: Niall Egan <niall.egan@databricks.com>

Author: Niall Egan

cmdexec/clients/python/src/databricks/sql/__init__.py

@@ -1,14 +1,21 @@
+from databricks.sql.errors import *
+
+
 class _DBAPITypeObject(object):
     def __init__(self, *values):
         self.values = values
 
     def __eq__(self, other):
         return other in self.values
 
+    def __repr__(self):
+        return "DBAPITypeObject(%s)" % self.values
+
 
 STRING = _DBAPITypeObject('string')
 BINARY = _DBAPITypeObject('binary')
-NUMBER = _DBAPITypeObject('boolean', 'byte', 'short', 'integer', 'long', 'double', 'decimal')
+NUMBER = _DBAPITypeObject('boolean', 'tinyint', 'smallint', 'int', 'bigint', 'float', 'double',
+                          'decimal')
 DATETIME = _DBAPITypeObject('timestamp')
 DATE = _DBAPITypeObject('date')
 ROWID = _DBAPITypeObject()

cmdexec/clients/python/src/databricks/sql/client.py

@@ -1,4 +1,6 @@
 import base64
+import datetime
+from decimal import Decimal
 import logging
 import re
 import time
@@ -7,14 +9,19 @@
 import grpc
 import pyarrow
 
-from databricks.sql.errors import OperationalError, InterfaceError, DatabaseError, Error
+from databricks.sql.errors import OperationalError, InterfaceError, DatabaseError, Error, DataError
 from databricks.sql.api import messages_pb2
 from databricks.sql.api.sql_cmd_service_pb2_grpc import SqlCommandServiceStub
 from databricks.sql import USER_AGENT_NAME, __version__
 
 logger = logging.getLogger(__name__)
 
 DEFAULT_RESULT_BUFFER_SIZE_BYTES = 10485760
+DEFAULT_ARRAY_SIZE = 100000
+
+_TIMESTAMP_PATTERN = re.compile(r'(\d+-\d+-\d+ \d+:\d+:\d+(\.\d{,6})?)')
+
+TYPES_CONVERTER = {"decimal": Decimal}
 
 
 class Connection:
@@ -101,11 +108,13 @@ def _http_path_to_routing_header(self, http_path):
         else:
             raise ValueError("Please provide a valid http_path")
 
-    def cursor(self, buffer_size_bytes=DEFAULT_RESULT_BUFFER_SIZE_BYTES):
+    def cursor(self,
+               arraysize=DEFAULT_ARRAY_SIZE,
+               buffer_size_bytes=DEFAULT_RESULT_BUFFER_SIZE_BYTES):
         if not self.open:
             raise Error("Cannot create cursor from closed connection")
 
-        cursor = Cursor(self, buffer_size_bytes)
+        cursor = Cursor(self, arraysize=arraysize, result_buffer_size_bytes=buffer_size_bytes)
         self._cursors.append(cursor)
         return cursor
 
@@ -122,12 +131,12 @@ class Cursor:
     def __init__(self,
                  connection,
                  result_buffer_size_bytes=DEFAULT_RESULT_BUFFER_SIZE_BYTES,
-                 arraysize=10000):
+                 arraysize=DEFAULT_ARRAY_SIZE):
         self.connection = connection
         self.rowcount = -1
-        self.arraysize = arraysize
         self.buffer_size_bytes = result_buffer_size_bytes
         self.active_result_set = None
+        self.arraysize = arraysize
         # Note that Cursor closed => active result set closed, but not vice versa
         self.open = True
         self.executing_command_id = None
@@ -351,9 +360,19 @@ def _fill_results_buffer(self):
             self.description = description
 
     @staticmethod
-    def _convert_arrow_table(table):
+    def parse_type(type_, value):
+        converter = TYPES_CONVERTER.get(type_)
+        if converter:
+            return converter(value)
+        else:
+            return value
+
+    def _convert_arrow_table(self, table):
         n_rows, _ = table.shape
-        list_repr = [[col[i].as_py() for col in table.itercolumns()] for i in range(n_rows)]
+        list_repr = [[
+            self.parse_type(self.description[col_index][1], col[row_index].as_py())
+            for col_index, col in enumerate(table.itercolumns())
+        ] for row_index in range(n_rows)]
         return list_repr
 
     def fetchmany_arrow(self, n_rows):
@@ -397,7 +416,6 @@ def fetchone(self):
         Fetch the next row of a query result set, returning a single sequence,
         or None when no more data is available.
         """
-        self._row_index += 1
         res = self._convert_arrow_table(self.fetchmany_arrow(1))
         if len(res) > 0:
             return res[0]

cmdexec/clients/python/tests/test_fetches.py

@@ -40,7 +40,7 @@ def make_arrow_queue(batch):
     def make_dummy_result_set_from_initial_results(initial_results):
         # If the initial results have been set, then we should never try and fetch more
         arrow_ipc_stream = FetchTests.make_arrow_ipc_stream(initial_results)
-        return client.ResultSet(
+        rs = client.ResultSet(
             connection=None,
             command_id=None,
             status=None,
@@ -49,6 +49,9 @@ def make_dummy_result_set_from_initial_results(initial_results):
             arrow_ipc_stream=arrow_ipc_stream,
             num_valid_rows=len(initial_results),
             schema_message=MagicMock())
+        num_cols = len(initial_results[0]) if initial_results else 0
+        rs.description = [('', 'integer', None, None, None, None, None) * num_cols]
+        return rs
 
     @staticmethod
     def make_dummy_result_set_from_batch_list(batch_list):
@@ -63,7 +66,10 @@ def _fetch_and_deserialize_results(self):
                 return results, batch_index < len(batch_list), \
                        [('id', 'integer', None, None, None, None, None)]
 
-        return SemiFakeResultSet(None, None, None, False, False)
+        rs = SemiFakeResultSet(None, None, None, False, False)
+        num_cols = len(batch_list[0][0]) if batch_list and batch_list[0] else 0
+        rs.description = [('', 'integer', None, None, None, None, None) * num_cols]
+        return rs
 
     def test_fetchmany_with_initial_results(self):
         # Fetch all in one go