Adding method to drop columns from a table (#30)

Author: jufreire

examples/drop_columns_from_table.py

@@ -0,0 +1,13 @@
+from hive_metastore_client import HiveMetastoreClient
+
+HIVE_HOST = "<ADD_HIVE_HOST_HERE>"
+HIVE_PORT = 9083
+
+# You must create a list with the columns' names to drop
+columns = ["quantity"]
+
+with HiveMetastoreClient(HIVE_HOST, HIVE_PORT) as hive_client:
+    # Dropping columns from table
+    hive_client.drop_columns_from_table(
+        db_name="store", table_name="order", columns=columns
+    )

hive_metastore_client/hive_mestastore_client.py

@@ -84,6 +84,32 @@ def add_columns_to_table(
         # call alter table to add columns
         self.alter_table(dbname=db_name, tbl_name=table_name, new_tbl=table)
 
+    def drop_columns_from_table(
+        self, db_name: str, table_name: str, columns: List[str]
+    ) -> None:
+        """
+        Drops columns from a table.
+
+        It encapsulates the logic of calling alter table with removed columns from
+        the list of columns, since hive does not have a drop command.
+
+        :param db_name: database name of the table
+        :param table_name: table name
+        :param columns: names of the columns to be dropped from the table
+        """
+        if columns:
+            table = self.get_table(dbname=db_name, tbl_name=table_name)
+
+            # remove columns from the list of columns in table object
+            cols = []
+            for col in table.sd.cols:
+                if col.name not in columns:
+                    cols.append(col)
+            table.sd.cols = cols
+
+            # call alter table to drop columns removed from list of table columns
+            self.alter_table(dbname=db_name, tbl_name=table_name, new_tbl=table)
+
     def add_partitions_to_table(
         self, db_name: str, table_name: str, partition_list: List[Partition]
     ) -> None:

tests/unit/hive_metastore_client/test_hive_mestastore_client.py

@@ -3,6 +3,7 @@
 
 import pytest
 from hive_metastore_client import HiveMetastoreClient
+from thrift_files.libraries.thrift_hive_metastore_client.ttypes import FieldSchema
 
 
 class TestHiveMetastoreClient:
@@ -93,6 +94,38 @@ def test_add_columns_to_table(
             dbname=db_name, tbl_name=table_name, new_tbl=mocked_return_get_table
         )
 
+    @mock.patch.object(HiveMetastoreClient, "get_table")
+    @mock.patch.object(HiveMetastoreClient, "alter_table")
+    def test_drop_columns_from_table(
+        self, mocked_alter_table, mocked_get_table, hive_metastore_client
+    ):
+        # arrange
+        db_name = "db_name"
+        table_name = "table_name"
+        cols = ["col1", "col2"]
+
+        mocked_return_get_table = Mock()
+        mocked_return_get_table.sd.cols = [
+            FieldSchema(name="col1"),
+            FieldSchema(name="col2"),
+            FieldSchema(name="col3"),
+        ]
+        mocked_get_table.return_value = mocked_return_get_table
+        expected_table_column = [FieldSchema(name="col3")]
+        expected_mocked_table = mocked_return_get_table
+        expected_mocked_table.sd.cols = expected_table_column
+
+        # act
+        hive_metastore_client.drop_columns_from_table(
+            db_name=db_name, table_name=table_name, columns=cols
+        )
+
+        # assert
+        mocked_get_table.assert_called_once_with(dbname=db_name, tbl_name=table_name)
+        mocked_alter_table.assert_called_once_with(
+            dbname=db_name, tbl_name=table_name, new_tbl=expected_mocked_table
+        )
+
     def test__validate_lists_length_with_diff_lens(self, hive_metastore_client):
         # arrange
         list_a = [1, 2, 3]