API tweaks + query params

This PR adds various little tweaks to make the v2 API compliant with the V1 API.

The one API difference is that in V2 the cursor does not have a `next` method. That is because in Python3 the proper way to do iteration is to create an itterable using `i = iter(cursor)` and then call `i.next()`

* New unit tests
* Using Python inspect module to find differences in API

Author: NiallEgan

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

@@ -1,7 +1,14 @@
+import datetime
+
 from databricks.sql.exc import *
 
+# PEP 249 module globals
+apilevel = '2.0'
+threadsafety = 1  # Threads may share the module, but not connections.
+paramstyle = 'pyformat'  # Python extended format codes, e.g. ...WHERE name=%(name)s
+
 
-class _DBAPITypeObject(object):
+class DBAPITypeObject(object):
     def __init__(self, *values):
         self.values = values
 
@@ -12,18 +19,30 @@ def __repr__(self):
         return "DBAPITypeObject({})".format(self.values)
 
 
-STRING = _DBAPITypeObject('string')
-BINARY = _DBAPITypeObject('binary')
-NUMBER = _DBAPITypeObject('boolean', 'tinyint', 'smallint', 'int', 'bigint', 'float', 'double',
-                          'decimal')
-DATETIME = _DBAPITypeObject('timestamp')
-DATE = _DBAPITypeObject('date')
-ROWID = _DBAPITypeObject()
+STRING = DBAPITypeObject('string')
+BINARY = DBAPITypeObject('binary')
+NUMBER = DBAPITypeObject('boolean', 'tinyint', 'smallint', 'int', 'bigint', 'float', 'double',
+                         'decimal')
+DATETIME = DBAPITypeObject('timestamp')
+DATE = DBAPITypeObject('date')
+ROWID = DBAPITypeObject()
 
 __version__ = "2.0.0rc2"
 USER_AGENT_NAME = "PyDatabricksSqlConnector"
 
+# These two functions are pyhive legacy
+Date = datetime.date
+Timestamp = datetime.datetime
+
+
+def DateFromTicks(ticks):
+    return Date(*time.localtime(ticks)[:3])
+
+
+def TimestampFromTicks(ticks):
+    return Timestamp(*time.localtime(ticks)[:6])
+
 
-def connect(**kwargs):
+def connect(server_hostname, http_path, access_token, **kwargs):
     from .client import Connection
-    return Connection(**kwargs)
+    return Connection(server_hostname, http_path, access_token, **kwargs)

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

@@ -10,7 +10,7 @@
 from databricks.sql import USER_AGENT_NAME, __version__
 from databricks.sql import *
 from databricks.sql.thrift_backend import ThriftBackend
-from databricks.sql.utils import ExecuteResponse
+from databricks.sql.utils import ExecuteResponse, ParamEscaper
 
 logger = logging.getLogger(__name__)
 
@@ -129,6 +129,13 @@ def close(self) -> None:
         for cursor in self._cursors:
             cursor.close()
 
+    def commit(self):
+        """No-op because Databricks does not support transactions"""
+        pass
+
+    def rollback(self):
+        raise NotSupportedError("Transactions are not supported on Databricks")
+
 
 class Cursor:
     def __init__(self,
@@ -144,7 +151,7 @@ def __init__(self,
         visible by other cursors or connections.
         """
         self.connection = connection
-        self.rowcount = -1
+        self.rowcount = -1  # Return -1 as this is not supported
         self.buffer_size_bytes = result_buffer_size_bytes
         self.active_result_set = None
         self.arraysize = arraysize
@@ -153,6 +160,8 @@ def __init__(self,
         self.executing_command_id = None
         self.thrift_backend = thrift_backend
         self.active_op_handle = None
+        self.escaper = ParamEscaper()
+        self.lastrowid = None
 
     def __enter__(self):
         return self
@@ -178,26 +187,23 @@ def _check_not_closed(self):
         if not self.open:
             raise Error("Attempting operation on closed cursor")
 
-    def execute(self,
-                operation: str,
-                query_params: Optional[Dict[str, str]] = None,
-                metadata: Optional[Dict[str, str]] = None) -> "Cursor":
+    def execute(self, operation: str, parameters: Optional[Dict[str, str]] = None) -> "Cursor":
         """
         Execute a query and wait for execution to complete.
-
+        Parameters should be given in extended param format style: %(...)<s|d|f>.
+        For example:
+            operation = "SELECT * FROM %(table_name)s"
+            parameters = {"table_name": "my_table_name"}
+            Will result in the query "SELECT * FROM 'my_table_name' being sent to the server
         :returns self
         """
-        if query_params is None:
-            sql = operation
-        else:
-            # TODO(https://databricks.atlassian.net/browse/SC-88829) before public release
-            logger.error("query param substitution currently un-implemented")
-            sql = operation
+        if parameters is not None:
+            operation = operation % self.escaper.escape_args(parameters)
 
         self._check_not_closed()
         self._close_and_clear_active_result_set()
         execute_response = self.thrift_backend.execute_command(
-            operation=sql,
+            operation=operation,
             session_handle=self.connection._session_handle,
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
@@ -206,6 +212,19 @@ def execute(self,
                                            self.buffer_size_bytes, self.arraysize)
         return self
 
+    def executemany(self, operation, seq_of_parameters):
+        """
+        Prepare a database operation (query or command) and then execute it against all parameter
+        sequences or mappings found in the sequence ``seq_of_parameters``.
+
+        Only the final result set is retained.
+
+        :returns self
+        """
+        for parameters in seq_of_parameters:
+            self.execute(operation, parameters)
+        return self
+
     def catalogs(self) -> "Cursor":
         """
         Get all available catalogs.
@@ -327,7 +346,7 @@ def fetchone(self) -> Tuple:
         else:
             raise Error("There is no active result set")
 
-    def fetchmany(self, n_rows: int) -> List[Tuple]:
+    def fetchmany(self, size: int) -> List[Tuple]:
         """
         Fetch the next set of rows of a query result, returning a sequence of sequences (e.g. a
         list of tuples).
@@ -345,7 +364,7 @@ def fetchmany(self, n_rows: int) -> List[Tuple]:
         """
         self._check_not_closed()
         if self.active_result_set:
-            return self.active_result_set.fetchmany(n_rows)
+            return self.active_result_set.fetchmany(size)
         else:
             raise Error("There is no active result set")
 
@@ -356,10 +375,10 @@ def fetchall_arrow(self):
         else:
             raise Error("There is no active result set")
 
-    def fetchmany_arrow(self, n_rows):
+    def fetchmany_arrow(self, size):
         self._check_not_closed()
         if self.active_result_set:
-            return self.active_result_set.fetchmany_arrow(n_rows)
+            return self.active_result_set.fetchmany_arrow(size)
         else:
             raise Error("There is no active result set")
 
@@ -407,6 +426,24 @@ def description(self) -> Optional[List[Tuple]]:
         else:
             return None
 
+    @property
+    def rownumber(self):
+        """This read-only attribute should provide the current 0-based index of the cursor in the
+        result set.
+
+        The index can be seen as index of the cursor in a sequence (the result set). The next fetch
+        operation will fetch the row indexed by ``rownumber`` in that sequence.
+        """
+        return self.active_result_set.rownumber if self.active_result_set else 0
+
+    def setinputsizes(self, sizes):
+        """Does nothing by default"""
+        pass
+
+    def setoutputsize(self, size, column=None):
+        """Does nothing by default"""
+        pass
+
 
 class ResultSet:
     def __init__(self,
@@ -468,16 +505,20 @@ def _convert_arrow_table(self, table):
                      for row_index in range(n_rows)]
         return list_repr
 
-    def fetchmany_arrow(self, n_rows: int) -> pyarrow.Table:
+    @property
+    def rownumber(self):
+        return self._next_row_index
+
+    def fetchmany_arrow(self, size: int) -> pyarrow.Table:
         """
         Fetch the next set of rows of a query result, returning a PyArrow table.
 
         An empty sequence is returned when no more rows are available.
         """
-        if n_rows < 0:
-            raise ValueError("n_rows argument for fetchmany is %s but must be >= 0", n_rows)
-        results = self.results.next_n_rows(n_rows)
-        n_remaining_rows = n_rows - results.num_rows
+        if size < 0:
+            raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
+        results = self.results.next_n_rows(size)
+        n_remaining_rows = size - results.num_rows
         self._next_row_index += results.num_rows
 
         while n_remaining_rows > 0 and not self.has_been_closed_server_side and self.has_more_rows:
@@ -519,13 +560,13 @@ def fetchall(self) -> List[Tuple]:
         """
         return self._convert_arrow_table(self.fetchall_arrow())
 
-    def fetchmany(self, n_rows: int) -> List[Tuple]:
+    def fetchmany(self, size: int) -> List[Tuple]:
         """
         Fetch the next set of rows of a query result, returning a list of lists.
 
         An empty sequence is returned when no more rows are available.
         """
-        return self._convert_arrow_table(self.fetchmany_arrow(n_rows))
+        return self._convert_arrow_table(self.fetchmany_arrow(size))
 
     def close(self) -> None:
         """

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

@@ -1,4 +1,5 @@
-from collections import namedtuple, OrderedDict
+from collections import namedtuple, OrderedDict, Iterable
+import datetime
 from enum import Enum
 
 import pyarrow
@@ -110,3 +111,58 @@ def user_friendly_error_message(self, no_retry_reason, attempt, elapsed):
                 user_friendly_error_message, elapsed)
 
         return user_friendly_error_message
+
+
+# Taken from PyHive
+class ParamEscaper:
+    _DATE_FORMAT = "%Y-%m-%d"
+    _TIME_FORMAT = "%H:%M:%S.%f"
+    _DATETIME_FORMAT = "{} {}".format(_DATE_FORMAT, _TIME_FORMAT)
+
+    def escape_args(self, parameters):
+        if isinstance(parameters, dict):
+            return {k: self.escape_item(v) for k, v in parameters.items()}
+        elif isinstance(parameters, (list, tuple)):
+            return tuple(self.escape_item(x) for x in parameters)
+        else:
+            raise exc.ProgrammingError("Unsupported param format: {}".format(parameters))
+
+    def escape_number(self, item):
+        return item
+
+    def escape_string(self, item):
+        # Need to decode UTF-8 because of old sqlalchemy.
+        # Newer SQLAlchemy checks dialect.supports_unicode_binds before encoding Unicode strings
+        # as byte strings. The old version always encodes Unicode as byte strings, which breaks
+        # string formatting here.
+        if isinstance(item, bytes):
+            item = item.decode('utf-8')
+        # This is good enough when backslashes are literal, newlines are just followed, and the way
+        # to escape a single quote is to put two single quotes.
+        # (i.e. only special character is single quote)
+        return "'{}'".format(item.replace("'", "''"))
+
+    def escape_sequence(self, item):
+        l = map(str, map(self.escape_item, item))
+        return '(' + ','.join(l) + ')'
+
+    def escape_datetime(self, item, format, cutoff=0):
+        dt_str = item.strftime(format)
+        formatted = dt_str[:-cutoff] if cutoff and format.endswith(".%f") else dt_str
+        return "'{}'".format(formatted)
+
+    def escape_item(self, item):
+        if item is None:
+            return 'NULL'
+        elif isinstance(item, (int, float)):
+            return self.escape_number(item)
+        elif isinstance(item, str):
+            return self.escape_string(item)
+        elif isinstance(item, Iterable):
+            return self.escape_sequence(item)
+        elif isinstance(item, datetime.datetime):
+            return self.escape_datetime(item, self._DATETIME_FORMAT)
+        elif isinstance(item, datetime.date):
+            return self.escape_datetime(item, self._DATE_FORMAT)
+        else:
+            raise exc.ProgrammingError("Unsupported object {}".format(item))