Fix tests to pass against Presto 0.76 and fix test setup issues

- Modify test script to make more databases
- Make tests pass against Presto 0.76
- Default to current user instead of hard-coded "hadoop" user
- Squash together Hive setup commands and use local JT to make tests faster
- Convert binary type from base64 to byte strings

Author: jingw

README.rst

@@ -48,9 +48,10 @@ Testing
 Run the following in an environment with Hive/Presto::
 
     ./scripts/make_test_tables.sh
-    virtualenv env
+    virtualenv --no-site-packages env
     source env/bin/activate
     pip install -r dev_requirements.txt
     py.test
 
-WARNING: This drops/creates tables named ``one_row``, ``one_row_complex``, and ``many_rows``.
+WARNING: This drops/creates tables named ``one_row``, ``one_row_complex``, and ``many_rows``, plus a
+database called ``pyhive_test_database``.

dev_requirements.txt

@@ -1,6 +1,7 @@
 mock>=1.0.0
-pytest
-pytest-cov
+# SQLAlchemy's test suite breaks with pytest 2.6
+pytest==2.5.2
+pytest-cov==1.7.0
 requests>=1.0.0
 sasl>=0.1.3
 sqlalchemy>=0.9.4

pyhive/presto.py

@@ -11,6 +11,7 @@
 from pyhive.common import DBAPITypeObject
 # Make all exceptions visible in this module per DB-API
 from pyhive.exc import *
+import base64
 import getpass
 import logging
 import requests
@@ -181,6 +182,14 @@ def _fetch_more(self):
         """Fetch the next URI and update state"""
         self._process_response(requests.get(self._nextUri))
 
+    def _decode_binary(self, rows):
+        # As of Presto 0.69, binary data is returned as the varbinary type in base64 format
+        # This function decodes base64 data in place
+        for i, col in enumerate(self.description):
+            if col[1] == 'varbinary':
+                for row in rows:
+                    row[i] = base64.b64decode(row[i])
+
     def _process_response(self, response):
         """Given the JSON response from Presto's REST API, update the internal state with the next
         URI and any data from the response
@@ -194,7 +203,11 @@ def _process_response(self, response):
         assert self._state == self._STATE_RUNNING, "Should be running if processing response"
         self._nextUri = response_json.get('nextUri')
         self._columns = response_json.get('columns')
-        self._data += response_json.get('data', [])
+        if 'data' in response_json:
+            assert self._columns
+            new_data = response_json['data']
+            self._decode_binary(new_data)
+            self._data += new_data
         if 'nextUri' not in response_json:
             self._state = self._STATE_FINISHED
         if 'error' in response_json:

pyhive/sqlalchemy_presto.py

@@ -179,16 +179,15 @@ def get_columns(self, connection, table_name, schema=None, **kw):
         rows = self._get_table_columns(connection, table_name, None)
         result = []
         for row in rows:
-            name, coltype, nullable, _is_partition_key = row
             try:
-                coltype = _type_map[coltype]
+                coltype = _type_map[row.Type]
             except KeyError:
-                util.warn("Did not recognize type '%s' of column '%s'" % (coltype, name))
+                util.warn("Did not recognize type '%s' of column '%s'" % (row.Type, row.Column))
                 coltype = types.NullType
             result.append({
-                'name': name,
+                'name': row.Column,
                 'type': coltype,
-                'nullable': nullable,
+                'nullable': row.Null,
                 'default': None,
             })
         return result

pyhive/tests/test_hive.py

@@ -21,7 +21,7 @@ class TestHive(unittest.TestCase, DBAPITestCase):
     __test__ = True
 
     def connect(self):
-        return hive.connect(host=_HOST, username='hadoop')
+        return hive.connect(host=_HOST, configuration={'mapred.job.tracker': 'local'})
 
     @with_cursor
     def test_description(self, cursor):

pyhive/tests/test_presto.py

@@ -40,8 +40,8 @@ def test_complex(self, cursor):
             ('float', 'double', None, None, None, None, True),
             ('double', 'double', None, None, None, None, True),
             ('string', 'varchar', None, None, None, None, True),
-            ('timestamp', 'bigint', None, None, None, None, True),
-            ('binary', 'varchar', None, None, None, None, True),
+            ('timestamp', 'timestamp', None, None, None, None, True),
+            ('binary', 'varbinary', None, None, None, None, True),
             ('array', 'varchar', None, None, None, None, True),
             ('map', 'varchar', None, None, None, None, True),
             ('struct', 'varchar', None, None, None, None, True),
@@ -57,7 +57,7 @@ def test_complex(self, cursor):
             0.5,
             0.25,
             'a string',
-            0,
+            '1970-01-01 00:00:00.000',
             '123',
             '[1,2]',
             '{"1":2,"3":4}',  # Presto converts all keys to strings so that they're valid JSON

pyhive/tests/test_sqlalchemy_hive.py

@@ -33,7 +33,7 @@
 
 class TestSqlAlchemyHive(unittest.TestCase, SqlAlchemyTestCase):
     def create_engine(self):
-        return create_engine('hive://hadoop@localhost:10000/default')
+        return create_engine('hive://localhost:10000/default')
 
     @with_engine_connection
     def test_reflect_select(self, engine, connection):
@@ -64,7 +64,7 @@ def test_reserved_words(self, engine, connection):
         self.assertNotIn('"map"', query)
 
     def test_switch_database(self):
-        engine = create_engine('hive://hadoop@localhost:10000/pyhive_test_database')
+        engine = create_engine('hive://localhost:10000/pyhive_test_database')
         try:
             with contextlib.closing(engine.connect()) as connection:
                 self.assertEqual(

pyhive/tests/test_sqlalchemy_presto.py

@@ -33,7 +33,7 @@ def test_reflect_select(self, engine, connection):
             0.5,
             0.25,
             'a string',
-            0,
+            '1970-01-01 00:00:00.000',
             '123',
             '[1,2]',
             '{"1":2,"3":4}',  # Presto converts all keys to strings so that they're valid JSON

scripts/make_many_rows.sh

@@ -1,14 +1,15 @@
 #!/bin/bash
 
-hive -e 'DROP TABLE IF EXISTS many_rows'
+temp_file=/tmp/pyhive_test_data_many_rows.tsv
+seq 0 9999 > $temp_file
+
 hive -e "
+DROP TABLE IF EXISTS many_rows;
 CREATE TABLE many_rows (
     a INT
 ) PARTITIONED BY (
     b STRING
-) ROW FORMAT DELIMITED FIELDS TERMINATED BY '\t' STORED AS TEXTFILE"
-
-temp_file=/tmp/pyhive_test_data_many_rows.tsv
-seq 0 9999 > $temp_file
-hive -e "LOAD DATA LOCAL INPATH '$temp_file' INTO TABLE many_rows PARTITION (b='blah')"
+) ROW FORMAT DELIMITED FIELDS TERMINATED BY '\t' STORED AS TEXTFILE;
+LOAD DATA LOCAL INPATH '$temp_file' INTO TABLE many_rows PARTITION (b='blah');
+"
 rm -f $temp_file

scripts/make_one_row.sh

@@ -1,4 +1,7 @@
 #!/bin/bash -eux
-hive -e 'DROP TABLE IF EXISTS one_row'
-hive -e 'CREATE TABLE one_row (number_of_rows INT)'
-hive -e 'INSERT OVERWRITE TABLE one_row SELECT COUNT(*) + 1 FROM one_row'
+hive -e '
+set mapred.job.tracker=local;
+DROP TABLE IF EXISTS one_row;
+CREATE TABLE one_row (number_of_rows INT);
+INSERT OVERWRITE TABLE one_row SELECT COUNT(*) + 1 FROM one_row;
+'