adding method to add columns to a table (#18)

examples/add_columns_to_table.py

@@ -0,0 +1,22 @@
+"""
+    Use the builders for creating Column thrift object.
+    Some arguments are optional when creating a thrift object.
+    Check Builder constructor for more information.
+"""
+
+from hive_metastore_client.builders.column_builder import ColumnBuilder
+from hive_metastore_client.hive_mestastore_client import HiveMetastoreClient
+
+HIVE_HOST = "<ADD_HIVE_HOST_HERE>"
+HIVE_PORT = 9083
+
+# You must create a list with the columns
+columns = [
+    ColumnBuilder(name="quantity", type="int", comment="item's quantity").build()
+]
+
+with HiveMetastoreClient(HIVE_HOST, HIVE_PORT) as hive_client:
+    # Adding columns to the table
+    hive_client.add_columns_to_table(
+        db_name="store", table_name="order", columns=columns
+    )

examples/create_database.py

@@ -1,12 +1,18 @@
+"""
+    Use the builders for creating Database thrift object.
+    Some arguments are optional when creating a thrift object.
+    Check Builder constructor for more information.
+"""
+
 from hive_metastore_client.builders.database_builder import DatabaseBuilder
 from hive_metastore_client.hive_mestastore_client import HiveMetastoreClient
 
 HIVE_HOST = "<ADD_HIVE_HOST_HERE>"
 HIVE_PORT = 9083
 
-with HiveMetastoreClient(HIVE_HOST, HIVE_PORT) as hive_client:
-
-    database_builder = DatabaseBuilder("database_name")
-    database = database_builder.build()
+# Creating database object using builder
+database = DatabaseBuilder("database_name").build()
 
+with HiveMetastoreClient(HIVE_HOST, HIVE_PORT) as hive_client:
+    # Creating new database from thrift table object
     hive_client.create_database(database)

hive_metastore_client/hive_mestastore_client.py

@@ -1,10 +1,11 @@
 """Hive Metastore Client main class."""
 from thrift.protocol import TBinaryProtocol
 from thrift.transport import TSocket, TTransport
-
+from typing import List
 from thrift_files.libraries.thrift_hive_metastore_client.ThriftHiveMetastore import (  # type: ignore # noqa: E501
     Client as ThriftClient,
 )
+from thrift_files.libraries.thrift_hive_metastore_client.ttypes import FieldSchema  # type: ignore # noqa: E501
 
 
 class HiveMetastoreClient(ThriftClient):
@@ -58,3 +59,21 @@ def __enter__(self) -> "HiveMetastoreClient":
     def __exit__(self, exc_type: str, exc_val: str, exc_tb: str) -> None:
         """Handles the conn closing after the code inside 'with' block is ended."""
         self.close()
+
+    def add_columns_to_table(
+        self, db_name: str, table_name: str, columns: List[FieldSchema]
+    ) -> None:
+        """
+        Adds columns to a table.
+
+        :param db_name: database name of the table
+        :param table_name: table name
+        :param columns: columns to be added to the table
+        """
+        table = self.get_table(dbname=db_name, tbl_name=table_name)
+
+        # add more columns to the list of columns
+        table.sd.cols.extend(columns)
+
+        # call alter table to add columns
+        self.alter_table(dbname=db_name, tbl_name=table_name, new_tbl=table)

tests/unit/hive_metastore_client/test_hive_mestastore_client.py

@@ -67,3 +67,27 @@ def test__exit__(self, mocked_close, hive_metastore_client):
 
         # assert
         mocked_close.assert_called_once_with()
+
+    @mock.patch.object(HiveMetastoreClient, "get_table")
+    @mock.patch.object(HiveMetastoreClient, "alter_table")
+    def test_add_columns_to_table(
+        self, mocked_alter_table, mocked_get_table, hive_metastore_client
+    ):
+        # arrange
+        db_name = "db_name"
+        table_name = "table_name"
+        cols = [Mock(), Mock()]
+
+        mocked_return_get_table = Mock()
+        mocked_get_table.return_value = mocked_return_get_table
+
+        # act
+        hive_metastore_client.add_columns_to_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=mocked_return_get_table
+        )