Add Glue.get_table_dtypes and Glue.get_table_python_types.

Author: igorborgest

awswrangler/athena.py

@@ -9,26 +9,24 @@
 class Athena:
     def __init__(self, session):
         self._session = session
+        self._client_athena = session.boto3_session.client(
+            service_name="athena", config=session.botocore_config)
 
     def run_query(self, query, database, s3_output):
-        client = self._session.boto3_session.client(
-            service_name="athena", config=self._session.botocore_config)
-        response = client.start_query_execution(
+        response = self._client_athena.start_query_execution(
             QueryString=query,
             QueryExecutionContext={"Database": database},
             ResultConfiguration={"OutputLocation": s3_output},
         )
         return response["QueryExecutionId"]
 
     def wait_query(self, query_execution_id):
-        client = self._session.boto3_session.client(
-            service_name="athena", config=self._session.botocore_config)
         final_states = ["FAILED", "SUCCEEDED", "CANCELLED"]
-        response = client.get_query_execution(
+        response = self._client_athena.get_query_execution(
             QueryExecutionId=query_execution_id)
         while (response.get("QueryExecution").get("Status").get("State") not in
                final_states):
             sleep(QUERY_WAIT_POLLING_DELAY)
-            response = client.get_query_execution(
+            response = self._client_athena.get_query_execution(
                 QueryExecutionId=query_execution_id)
         return response

awswrangler/glue.py

@@ -10,6 +10,53 @@
 class Glue:
     def __init__(self, session):
         self._session = session
+        self._client_glue = session.boto3_session.client(
+            service_name="glue", config=session.botocore_config)
+
+    def get_table_dtypes(self, database, table):
+        """
+        Get all columns names and the related data types
+        :param database: Glue database's name
+        :param table: Glue table's name
+        :return: A dictionary as {"col name": "col dtype"}
+        """
+        response = self._client_glue.get_table(DatabaseName=database,
+                                               Name=table)
+        logger.debug(f"get_table response:\n{response}")
+        dtypes = {}
+        for col in response["Table"]["StorageDescriptor"]["Columns"]:
+            dtypes[col["Name"]] = col["Type"]
+        for par in response["Table"]["PartitionKeys"]:
+            dtypes[par["Name"]] = par["Type"]
+        return dtypes
+
+    def get_table_python_types(self, database, table):
+        """
+        Get all columns names and the related python types
+        :param database: Glue database's name
+        :param table: Glue table's name
+        :return: A dictionary as {"col name": "col python type"}
+        """
+        dtypes = self.get_table_dtypes(database=database, table=table)
+        return {k: Glue._type_athena2python(v) for k, v in dtypes.items()}
+
+    @staticmethod
+    def _type_athena2python(dtype):
+        dtype = dtype.lower()
+        if dtype == "int":
+            return int
+        elif dtype == "bigint":
+            return int
+        elif dtype == "float":
+            return float
+        elif dtype == "double":
+            return float
+        elif dtype == "boolean":
+            return bool
+        elif dtype == "string":
+            return str
+        else:
+            raise UnsupportedType(f"Unsupported Athena type: {dtype}")
 
     def metadata_to_glue(
             self,
@@ -53,20 +100,16 @@ def metadata_to_glue(
             )
 
     def delete_table_if_exists(self, database, table):
-        client = self._session.boto3_session.client(
-            service_name="glue", config=self._session.botocore_config)
         try:
-            client.delete_table(DatabaseName=database, Name=table)
-        except client.exceptions.EntityNotFoundException:
+            self._client_glue.delete_table(DatabaseName=database, Name=table)
+        except self._client_glue.exceptions.EntityNotFoundException:
             pass
 
     def does_table_exists(self, database, table):
-        client = self._session.boto3_session.client(
-            service_name="glue", config=self._session.botocore_config)
         try:
-            client.get_table(DatabaseName=database, Name=table)
+            self._client_glue.get_table(DatabaseName=database, Name=table)
             return True
-        except client.exceptions.EntityNotFoundException:
+        except self._client_glue.exceptions.EntityNotFoundException:
             return False
 
     def create_table(self,
@@ -76,8 +119,6 @@ def create_table(self,
                      path,
                      file_format,
                      partition_cols=None):
-        client = self._session.boto3_session.client(
-            service_name="glue", config=self._session.botocore_config)
         if file_format == "parquet":
             table_input = Glue.parquet_table_definition(
                 table, partition_cols, schema, path)
@@ -86,11 +127,10 @@ def create_table(self,
                                                     schema, path)
         else:
             raise UnsupportedFileFormat(file_format)
-        client.create_table(DatabaseName=database, TableInput=table_input)
+        self._client_glue.create_table(DatabaseName=database,
+                                       TableInput=table_input)
 
     def add_partitions(self, database, table, partition_paths, file_format):
-        client = self._session.boto3_session.client(
-            service_name="glue", config=self._session.botocore_config)
         if not partition_paths:
             return None
         partitions = list()
@@ -106,15 +146,13 @@ def add_partitions(self, database, table, partition_paths, file_format):
         for _ in range(pages_num):
             page = partitions[:100]
             del partitions[:100]
-            client.batch_create_partition(DatabaseName=database,
-                                          TableName=table,
-                                          PartitionInputList=page)
+            self._client_glue.batch_create_partition(DatabaseName=database,
+                                                     TableName=table,
+                                                     PartitionInputList=page)
 
     def get_connection_details(self, name):
-        client = self._session.boto3_session.client(
-            service_name="glue", config=self._session.botocore_config)
-        return client.get_connection(Name=name,
-                                     HidePassword=False)["Connection"]
+        return self._client_glue.get_connection(
+            Name=name, HidePassword=False)["Connection"]
 
     @staticmethod
     def _build_schema(dataframe, partition_cols, preserve_index):
@@ -155,7 +193,7 @@ def _type_pandas2athena(dtype):
         elif dtype[:10] == "datetime64":
             return "timestamp"
         else:
-            raise UnsupportedType("Unsupported Pandas type: " + dtype)
+            raise UnsupportedType(f"Unsupported Pandas type: {dtype}")
 
     @staticmethod
     def _parse_table_name(path):

testing/test_awswrangler/test_glue.py

@@ -0,0 +1,86 @@
+import logging
+
+import pytest
+import boto3
+import pandas
+
+from awswrangler import Session
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="[%(asctime)s][%(levelname)s][%(name)s][%(funcName)s] %(message)s")
+logging.getLogger("awswrangler").setLevel(logging.DEBUG)
+
+
+@pytest.fixture(scope="module")
+def cloudformation_outputs():
+    response = boto3.client("cloudformation").describe_stacks(
+        StackName="aws-data-wrangler-test-arena")
+    outputs = {}
+    for output in response.get("Stacks")[0].get("Outputs"):
+        outputs[output.get("OutputKey")] = output.get("OutputValue")
+    yield outputs
+
+
+@pytest.fixture(scope="module")
+def session():
+    yield Session()
+
+
+@pytest.fixture(scope="module")
+def bucket(session, cloudformation_outputs):
+    if "BucketName" in cloudformation_outputs:
+        bucket = cloudformation_outputs.get("BucketName")
+        session.s3.delete_objects(path=f"s3://{bucket}/")
+    else:
+        raise Exception("You must deploy the test infrastructure using SAM!")
+    yield bucket
+    session.s3.delete_objects(path=f"s3://{bucket}/")
+
+
+@pytest.fixture(scope="module")
+def database(cloudformation_outputs):
+    if "GlueDatabaseName" in cloudformation_outputs:
+        database = cloudformation_outputs.get("GlueDatabaseName")
+    else:
+        raise Exception("You must deploy the test infrastructure using SAM!")
+    yield database
+
+
+@pytest.fixture(scope="module")
+def table(
+        session,
+        bucket,
+        database,
+):
+    dataframe = pandas.read_csv("data_samples/micro.csv")
+    path = f"s3://{bucket}/test/"
+    table = "test"
+    session.pandas.to_parquet(dataframe=dataframe,
+                              database=database,
+                              table=table,
+                              path=path,
+                              preserve_index=False,
+                              mode="overwrite",
+                              procs_cpu_bound=1,
+                              partition_cols=["name", "date"])
+    yield table
+    session.glue.delete_table_if_exists(database=database, table=table)
+    session.s3.delete_objects(path=path)
+
+
+def test_get_table_dtypes(session, database, table):
+    dtypes = session.glue.get_table_dtypes(database=database, table=table)
+    assert dtypes["id"] == "bigint"
+    assert dtypes["value"] == "double"
+    assert dtypes["name"] == "string"
+    assert dtypes["date"] == "string"
+
+
+def test_get_table_python_types(session, database, table):
+    ptypes = session.glue.get_table_python_types(database=database,
+                                                 table=table)
+    assert ptypes["id"] == int
+    assert ptypes["value"] == float
+    assert ptypes["name"] == str
+    assert ptypes["date"] == str