[FIX] Add Partition types for Metastore (#305)

quintoandar · Apr 1, 2021 · 3dcd975 · 3dcd975
1 parent d607297
commit 3dcd975
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/butterfree/migrations/database_migration/metastore_migration.py b/butterfree/migrations/database_migration/metastore_migration.py
@@ -102,8 +102,6 @@ def _get_create_table_query(
             Create table query.
 
         """
-        columns.extend(PARTITION_BY)
-
         parsed_columns = []
         for col in columns:
             parsed_columns.append(f"{col['column_name']} {col['type']}")
@@ -112,9 +110,10 @@ def _get_create_table_query(
         return (
             f"CREATE TABLE IF NOT EXISTS  "
             f"{self.database}.{table_name} ({parsed_columns}) "
-            f"PARTITIONED BY ({PARTITION_BY[0]['column_name']}, "
-            f"{PARTITION_BY[1]['column_name']}, "
-            f"{PARTITION_BY[2]['column_name']});"
+            f"PARTITIONED BY ("
+            f"{PARTITION_BY[0]['column_name']} {PARTITION_BY[0]['type']}, "
+            f"{PARTITION_BY[1]['column_name']} {PARTITION_BY[1]['type']}, "
+            f"{PARTITION_BY[2]['column_name']} {PARTITION_BY[2]['type']});"
         )
 
     def _get_alter_table_drop_query(self, columns: List[Diff], table_name: str) -> str:

diff --git a/tests/unit/butterfree/migrations/database_migration/test_metastore_migration.py b/tests/unit/butterfree/migrations/database_migration/test_metastore_migration.py
@@ -40,9 +40,8 @@ def test_create_table_query(self, fs_schema):
 
         expected_query = [
             "CREATE TABLE IF NOT EXISTS  test.table_name "
-            "(id LongType, timestamp TimestampType, new_feature FloatType, "
-            "feature1__avg_over_1_week_rolling_windows FloatType, year INT, "
-            "month INT, day INT) PARTITIONED BY (year, month, day);"
+            "(id LongType, timestamp TimestampType, new_feature FloatType) "
+            "PARTITIONED BY (year INT, month INT, day INT);"
         ]
 
         query = metastore_migration.create_query(fs_schema, "table_name")