Merge pull request #103 from laughingman7743/add_test_case_for_as_pandas_method

Add test case for as_pandas method

Author: laughingman7743

Pipfile

@@ -7,7 +7,7 @@ name = "pypi"
 allow_prereleases = true
 
 [packages]
-"e1839a8" = {path = ".", extras = ["sqlalchemy"], editable = true}
+"e1839a8" = {path = ".", extras = ["sqlalchemy", "pandas"], editable = true}
 
 [dev-packages]
 awscli = "*"

Pipfile.lock

@@ -3,7 +3,6 @@
 
 import functools
 import threading
-from datetime import datetime
 
 
 def as_pandas(cursor, coerce_float=False):
@@ -39,20 +38,3 @@ def _wrapper(*args, **kwargs):
         return wrapped(*args, **kwargs)
 
     return _wrapper
-
-
-def to_datetime(java_date):
-    if not java_date:
-        return None
-    import jpype
-
-    cal = jpype.java.util.Calendar.getInstance()
-    cal.setTime(java_date)
-    return datetime(
-        cal.get(jpype.java.util.Calendar.YEAR),
-        cal.get(jpype.java.util.Calendar.MONTH) + 1,
-        cal.get(jpype.java.util.Calendar.DAY_OF_MONTH),
-        cal.get(jpype.java.util.Calendar.HOUR_OF_DAY),
-        cal.get(jpype.java.util.Calendar.MINUTE),
-        cal.get(jpype.java.util.Calendar.SECOND),
-    )

pyathenajdbc/util.py

@@ -0,0 +1,3 @@
+true	false
+false
+

scripts/test_data/boolean_na_values.tsv

@@ -5,3 +5,5 @@
 aws s3 rm ${AWS_ATHENA_S3_STAGING_DIR}test_pyathena_jdbc/one_row/one_row.tsv
 aws s3 rm ${AWS_ATHENA_S3_STAGING_DIR}test_pyathena_jdbc/one_row_complex/one_row_complex.tsv
 aws s3 rm ${AWS_ATHENA_S3_STAGING_DIR}test_pyathena_jdbc/many_rows/many_rows.tsv
+aws s3 rm ${AWS_ATHENA_S3_STAGING_DIR}test_pyathena_jdbc/integer_na_values/integer_na_values.tsv
+aws s3 rm ${AWS_ATHENA_S3_STAGING_DIR}test_pyathena_jdbc/boolean_na_values/boolean_na_values.tsv

scripts/test_data/delete_test_data.sh

@@ -0,0 +1,3 @@
+1	2
+1
+

scripts/test_data/integer_na_values.tsv

@@ -5,3 +5,5 @@
 aws s3 cp $(dirname $0)/one_row.tsv ${AWS_ATHENA_S3_STAGING_DIR}test_pyathena_jdbc/one_row/one_row.tsv
 aws s3 cp $(dirname $0)/one_row_complex.tsv ${AWS_ATHENA_S3_STAGING_DIR}test_pyathena_jdbc/one_row_complex/one_row_complex.tsv
 aws s3 cp $(dirname $0)/many_rows.tsv ${AWS_ATHENA_S3_STAGING_DIR}test_pyathena_jdbc/many_rows/many_rows.tsv
+aws s3 cp $(dirname $0)/integer_na_values.tsv ${AWS_ATHENA_S3_STAGING_DIR}test_pyathena_jdbc/integer_na_values/integer_na_values.tsv
+aws s3 cp $(dirname $0)/boolean_na_values.tsv ${AWS_ATHENA_S3_STAGING_DIR}test_pyathena_jdbc/boolean_na_values/boolean_na_values.tsv

scripts/test_data/upload_test_data.sh

@@ -1,4 +1,38 @@
 # -*- coding: utf-8 -*-
+import os
+import random
+import string
+
+from past.builtins.misc import xrange
 from sqlalchemy.dialects import registry
 
 registry.register("awsathena.jdbc", "pyathenajdbc.sqlalchemy_athena", "AthenaDialect")
+
+BASE_PATH = os.path.dirname(os.path.abspath(__file__))
+S3_PREFIX = "test_pyathena_jdbc"
+WORK_GROUP = "test-pyathena-jdbc"
+SCHEMA = "test_pyathena_jdbc_" + "".join(
+    [random.choice(string.ascii_lowercase + string.digits) for i in xrange(10)]
+)
+
+
+class Env(object):
+    def __init__(self):
+        self.region_name = os.getenv("AWS_DEFAULT_REGION", None)
+        assert (
+            self.region_name
+        ), "Required environment variable `AWS_DEFAULT_REGION` not found."
+        self.s3_staging_dir = os.getenv("AWS_ATHENA_S3_STAGING_DIR", None)
+        assert (
+            self.s3_staging_dir
+        ), "Required environment variable `AWS_ATHENA_S3_STAGING_DIR` not found."
+
+
+ENV = Env()
+
+
+class WithConnect(object):
+    def connect(self, work_group=None):
+        from pyathenajdbc import connect
+
+        return connect(schema_name=SCHEMA, work_group=work_group)

tests/__init__.py

@@ -3,22 +3,12 @@
 
 import contextlib
 import os
-import random
-import string
 
 import pytest
-from past.builtins.misc import xrange
 
 from pyathenajdbc import connect
-from tests.util import Env, read_query
-
-ENV = Env()
-BASE_PATH = os.path.dirname(os.path.abspath(__file__))
-S3_PREFIX = "test_pyathena_jdbc"
-WORK_GROUP = "test-pyathena-jdbc"
-SCHEMA = "test_pyathena_jdbc_" + "".join(
-    [random.choice(string.ascii_lowercase + string.digits) for i in xrange(10)]
-)
+from tests import BASE_PATH, ENV, S3_PREFIX, SCHEMA
+from tests.util import read_query
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -54,12 +44,24 @@ def _create_table(cursor):
     location_one_row_complex = "{0}{1}/{2}/".format(
         ENV.s3_staging_dir, S3_PREFIX, "one_row_complex"
     )
+    location_partition_table = "{0}{1}/{2}/".format(
+        ENV.s3_staging_dir, S3_PREFIX, "partition_table"
+    )
+    location_integer_na_values = "{0}{1}/{2}/".format(
+        ENV.s3_staging_dir, S3_PREFIX, "integer_na_values"
+    )
+    location_boolean_na_values = "{0}{1}/{2}/".format(
+        ENV.s3_staging_dir, S3_PREFIX, "boolean_na_values"
+    )
     for q in read_query(os.path.join(BASE_PATH, "sql", "create_table.sql")):
         cursor.execute(
             q.format(
                 schema=SCHEMA,
                 location_one_row=location_one_row,
                 location_many_rows=location_many_rows,
                 location_one_row_complex=location_one_row_complex,
+                location_partition_table=location_partition_table,
+                location_integer_na_values=location_integer_na_values,
+                location_boolean_na_values=location_boolean_na_values,
             )
         )

tests/conftest.py

@@ -30,3 +30,27 @@ CREATE EXTERNAL TABLE IF NOT EXISTS {schema}.one_row_complex (
 )
 ROW FORMAT DELIMITED FIELDS TERMINATED BY '\t' LINES TERMINATED BY '\n' STORED AS TEXTFILE
 LOCATION '{location_one_row_complex}';
+
+DROP TABLE IF EXISTS {schema}.partition_table;
+CREATE EXTERNAL TABLE IF NOT EXISTS {schema}.partition_table (
+    a STRING
+)
+PARTITIONED BY (b INT)
+ROW FORMAT DELIMITED FIELDS TERMINATED BY '\t' LINES TERMINATED BY '\n' STORED AS TEXTFILE
+LOCATION '{location_partition_table}';
+
+DROP TABLE IF EXISTS {schema}.integer_na_values;
+CREATE EXTERNAL TABLE IF NOT EXISTS {schema}.integer_na_values (
+    a INT,
+    b INT
+)
+ROW FORMAT DELIMITED FIELDS TERMINATED BY '\t' LINES TERMINATED BY '\n' STORED AS TEXTFILE
+LOCATION '{location_integer_na_values}';
+
+DROP TABLE IF EXISTS {schema}.boolean_na_values;
+CREATE EXTERNAL TABLE IF NOT EXISTS {schema}.boolean_na_values (
+    a BOOLEAN,
+    b BOOLEAN
+)
+ROW FORMAT DELIMITED FIELDS TERMINATED BY '\t' LINES TERMINATED BY '\n' STORED AS TEXTFILE
+LOCATION '{location_boolean_na_values}';