add get_schema to spark client

Author: AlvaroMarquesAndrade

butterfree/clients/spark_client.py

@@ -276,6 +276,16 @@ def add_table_partitions(
 
     @staticmethod
     def _filter_schema(schema: DataFrame) -> List[str]:
+        """Returns filtered schema with the desired information.
+
+        Attributes:
+            schema: desired table.
+
+        Returns:
+            A list of strings in the format
+            ['{"column_name": "example1", type: "Spark_type"}', ...]
+
+        """
         return (
             schema.filter(
                 ~schema.col_name.isin(
@@ -287,12 +297,22 @@ def _filter_schema(schema: DataFrame) -> List[str]:
         )
 
     def _convert_schema(self, schema: DataFrame) -> List[Dict[str, str]]:
+        """Returns schema with the desired information.
+
+        Attributes:
+            schema: desired table.
+
+        Returns:
+            A list of dictionaries in the format
+            [{"column_name": "example1", type: "Spark_type"}, ...]
+
+        """
         schema_list = self._filter_schema(schema)
-        schema = []
+        converted_schema = []
         for row in schema_list:
-            schema.append(json.loads(row))
+            converted_schema.append(json.loads(row))
 
-        return schema
+        return converted_schema
 
     def get_schema(self, table: str, database: str) -> List[Dict[str, str]]:
         """Returns desired table schema.
@@ -301,7 +321,7 @@ def get_schema(self, table: str, database: str) -> List[Dict[str, str]]:
             table: desired table.
 
         Returns:
-            A list dictionaries in the format
+            A list of dictionaries in the format
             [{"column_name": "example1", type: "Spark_type"}, ...]
 
         """

tests/unit/butterfree/clients/test_spark_client.py

@@ -15,6 +15,15 @@ def create_temp_view(dataframe: DataFrame, name: str) -> None:
     dataframe.createOrReplaceTempView(name)
 
 
+def create_db_and_table(spark, database, table, view):
+    spark.sql(f"create database if not exists {database}")
+    spark.sql(f"use {database}")
+    spark.sql(
+        f"create table if not exists {database}.{table} "  # noqa
+        f"as select * from {view}"  # noqa
+    )
+
+
 class TestSparkClient:
     def test_conn(self) -> None:
         # arrange
@@ -293,3 +302,27 @@ def test_add_invalid_partitions(self, mock_spark_sql: Mock, partition):
         # act and assert
         with pytest.raises(ValueError):
             spark_client.add_table_partitions(partition, "table", "db")
+
+    def test_get_schema(
+        self, target_df: DataFrame, spark_session: SparkSession
+    ) -> None:
+        # arrange
+        spark_client = SparkClient()
+        create_temp_view(dataframe=target_df, name="temp_view")
+        create_db_and_table(
+            spark=spark_session,
+            database="test_db",
+            table="test_table",
+            view="temp_view",
+        )
+
+        expected_schema = [
+            {"col_name": "col1", "data_type": "string"},
+            {"col_name": "col2", "data_type": "bigint"},
+        ]
+
+        # act
+        schema = spark_client.get_schema(table="test_table", database="test_db")
+
+        # assert
+        assert schema, expected_schema