[DAE-94] Add get_table_partition_keys and keys_names methods (#43)

Author: Felipe Miquelim

docs/source/getstarted.md

@@ -28,6 +28,10 @@ Click on the following links to open the [examples](https://github.com/quintoand
 
 **[#5 Add partitions to a table](https://github.com/quintoandar/hive-metastore-client/blob/main/examples/add_partitions.py)**
 
+**[#6 Get partition keys objects from a table](https://github.com/quintoandar/hive-metastore-client/blob/main/examples/get_partition_keys_objects.py)**
+
+**[#7 Get partition keys names from a table](https://github.com/quintoandar/hive-metastore-client/blob/main/examples/get_partition_keys_names.py)**
+
 ## Available methods
 
 You can see all the Hive Metastore server available methods by looking at the 
@@ -41,4 +45,6 @@ the [`HiveMetastoreClient`](https://github.com/quintoandar/hive-metastore-client
 - [`drop_columns_from_table`](https://hive-metastore-client.readthedocs.io/en/latest/hive_metastore_client.html#hive_metastore_client.hive_metastore_client.HiveMetastoreClient.drop_columns_from_table)
 - [`add_partitions_if_not_exists`](https://hive-metastore-client.readthedocs.io/en/latest/hive_metastore_client.html#hive_metastore_client.hive_metastore_client.HiveMetastoreClient.add_partitions_if_not_exists)
 - [`create_database_if_not_exists`](https://hive-metastore-client.readthedocs.io/en/latest/hive_metastore_client.html#hive_metastore_client.hive_metastore_client.HiveMetastoreClient.create_database_if_not_exists)
-- [`create_external_table`](https://hive-metastore-client.readthedocs.io/en/latest/hive_metastore_client.html#hive_metastore_client.hive_metastore_client.HiveMetastoreClient.create_external_table)
+- [`create_external_table`](https://hive-metastore-client.readthedocs.io/en/latest/hive_metastore_client.html#hive_metastore_client.hive_metastore_client.HiveMetastoreClient.create_external_table)
+- [`get_partition_keys_objects`](https://hive-metastore-client.readthedocs.io/en/latest/hive_metastore_client.html#hive_metastore_client.hive_metastore_client.HiveMetastoreClient.get_partition_keys_objects)
+- [`get_partition_keys_names`](https://hive-metastore-client.readthedocs.io/en/latest/hive_metastore_client.html#hive_metastore_client.hive_metastore_client.HiveMetastoreClient.get_partition_keys_names)

examples/get_partition_keys_names.py

@@ -0,0 +1,11 @@
+from hive_metastore_client import HiveMetastoreClient
+
+HIVE_HOST = "<ADD_HIVE_HOST_HERE>"
+HIVE_PORT = 9083
+
+DATABASE_NAME = "database_name"
+TABLE_NAME = "table_name"
+
+with HiveMetastoreClient(HIVE_HOST, HIVE_PORT) as hive_client:
+    # Retrieving the partition keys names via table schema
+    hive_client.get_partition_keys_names(DATABASE_NAME, TABLE_NAME)

examples/get_partition_keys_objects.py

@@ -0,0 +1,11 @@
+from hive_metastore_client import HiveMetastoreClient
+
+HIVE_HOST = "<ADD_HIVE_HOST_HERE>"
+HIVE_PORT = 9083
+
+DATABASE_NAME = "database_name"
+TABLE_NAME = "table_name"
+
+with HiveMetastoreClient(HIVE_HOST, HIVE_PORT) as hive_client:
+    # Retrieving the partition keys via table schema
+    hive_client.get_partition_keys_objects(DATABASE_NAME, TABLE_NAME)

hive_metastore_client/hive_metastore_client.py

@@ -229,3 +229,33 @@ def _validate_lists_length(list_a: List[Any], list_b: List[Any]) -> None:
                 "m=_validate_lists_length, msg=The length of the two provided "
                 "lists does not match"
             )
+
+    def get_partition_keys_objects(
+        self, db_name: str, table_name: str
+    ) -> List[FieldSchema]:
+        """
+        Gets the partition keys objects, containing the metadata, from a table.
+
+        An empty list will be returned when no table is found or
+        when the table has no partitions
+
+        :param db_name: database name where the table is at
+        :param table_name: table name which the partition keys belong to
+        """
+        table = self.get_table(dbname=db_name, tbl_name=table_name)
+        return list(table.partitionKeys) if table else []
+
+    def get_partition_keys_names(self, db_name: str, table_name: str) -> List[str]:
+        """
+        Gets the partition keys names from a table.
+
+        An empty list will be returned when no table is found or
+        when the table has no partitions
+
+        :param db_name: database name where the table is at
+        :param table_name: table name which the partition keys belong to
+        """
+        partition_keys = self.get_partition_keys_objects(
+            db_name=db_name, table_name=table_name
+        )
+        return [partition.name for partition in partition_keys]

tests/unit/hive_metastore_client/test_hive_metastore_client.py

@@ -317,3 +317,120 @@ def test_create_external_table(self, mocked_create_table, hive_metastore_client)
 
         # assert
         mocked_create_table.assert_called_once_with(updated_table)
+
+    @mock.patch.object(HiveMetastoreClient, "get_table", return_value=None)
+    def test_get_partition_keys_objects_with_invalid_table(
+        self, mocked_get_table, hive_metastore_client,
+    ):
+        # arrange
+        table_name = "table_name"
+        database_name = "database_name"
+
+        # act
+        returned_value = hive_metastore_client.get_partition_keys_objects(
+            database_name, table_name
+        )
+
+        # assert
+        assert returned_value == []
+        mocked_get_table.assert_called_once_with(
+            dbname=database_name, tbl_name=table_name
+        )
+
+    @mock.patch.object(HiveMetastoreClient, "get_table")
+    def test_get_partition_keys_objects_with_not_partitioned_table(
+        self, mocked_get_table, hive_metastore_client,
+    ):
+        # arrange
+        table_name = "table_name"
+        database_name = "database_name"
+        mocked_table = Mock()
+        # the default return from hive metastore for not partitioned tables is an empty list
+        mocked_table.partitionKeys = []
+        mocked_get_table.return_value = mocked_table
+
+        # act
+        returned_value = hive_metastore_client.get_partition_keys_objects(
+            database_name, table_name
+        )
+
+        # assert
+        assert returned_value == []
+        mocked_get_table.assert_called_once_with(
+            dbname=database_name, tbl_name=table_name
+        )
+
+    @mock.patch.object(HiveMetastoreClient, "get_table")
+    def test_get_partition_keys_objects_with_partitioned_table(
+        self, mocked_get_table, hive_metastore_client,
+    ):
+        # arrange
+        table_name = "table_name"
+        database_name = "database_name"
+        mocked_table = Mock()
+        mocked_partition_a = Mock()
+        mocked_partition_b = Mock()
+        mocked_table.partitionKeys = [mocked_partition_a, mocked_partition_b]
+        mocked_get_table.return_value = mocked_table
+
+        # act
+        returned_value = hive_metastore_client.get_partition_keys_objects(
+            database_name, table_name
+        )
+
+        # assert
+        assert returned_value == [mocked_partition_a, mocked_partition_b]
+        mocked_get_table.assert_called_once_with(
+            dbname=database_name, tbl_name=table_name
+        )
+
+    @mock.patch.object(
+        HiveMetastoreClient, "get_partition_keys_objects", return_value=[]
+    )
+    def test_get_partition_keys_names_with_invalid_or_not_partitioned_table(
+        self, mocked_get_partition_keys_objects, hive_metastore_client,
+    ):
+        # arrange
+        table_name = "table_name"
+        database_name = "database_name"
+
+        # act
+        returned_value = hive_metastore_client.get_partition_keys_names(
+            database_name, table_name
+        )
+
+        # assert
+        assert returned_value == []
+        mocked_get_partition_keys_objects.assert_called_once_with(
+            db_name=database_name, table_name=table_name
+        )
+
+    @mock.patch.object(
+        HiveMetastoreClient, "get_partition_keys_objects", return_value=[]
+    )
+    def test_get_partition_keys_names_with_partitioned_table(
+        self, mocked_get_partition_keys_objects, hive_metastore_client,
+    ):
+        # arrange
+        table_name = "table_name"
+        database_name = "database_name"
+        mocked_partition_a = Mock()
+        mocked_partition_a.name = "mocked_partition_a"
+        mocked_partition_b = Mock()
+        mocked_partition_b.name = "mocked_partition_b"
+        mocked_get_partition_keys_objects.return_value = [
+            mocked_partition_a,
+            mocked_partition_b,
+        ]
+        expected_return = ["mocked_partition_a", "mocked_partition_b"]
+
+        # act
+        returned_value = hive_metastore_client.get_partition_keys_names(
+            database_name, table_name
+        )
+
+        # assert
+        assert returned_value == expected_return
+        mocked_get_partition_keys_objects.assert_called_once_with(
+            db_name=database_name, table_name=table_name
+        )