Fix migration query. (#318)

quintoandar · Apr 26, 2021 · c46f171 · c46f171
1 parent 3b18b5a
commit c46f171
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 16 deletions.
diff --git a/butterfree/migrations/database_migration/cassandra_migration.py b/butterfree/migrations/database_migration/cassandra_migration.py
@@ -75,7 +75,7 @@ def _get_alter_table_add_query(self, columns: List[Diff], table_name: str) -> st
 
         return f"ALTER TABLE {table_name} ADD ({parsed_columns});"
 
-    def _get_alter_column_type_query(self, columns: List[Diff], table_name: str) -> str:
+    def _get_alter_column_type_query(self, column: Diff, table_name: str) -> str:
         """Creates CQL statement to alter columns' types.
 
         Args:
@@ -86,9 +86,9 @@ def _get_alter_column_type_query(self, columns: List[Diff], table_name: str) ->
             Alter column type query.
 
         """
-        parsed_columns = self._get_parsed_columns(columns)
+        parsed_columns = self._get_parsed_columns([column])
 
-        return f"ALTER TABLE {table_name} ALTER ({parsed_columns});"
+        return f"ALTER TABLE {table_name} ALTER {parsed_columns};"
 
     @staticmethod
     def _get_create_table_query(columns: List[Dict[str, Any]], table_name: str) -> str:

diff --git a/butterfree/migrations/database_migration/database_migration.py b/butterfree/migrations/database_migration/database_migration.py
@@ -93,7 +93,7 @@ def _get_alter_table_drop_query(self, columns: List[Diff], table_name: str) -> s
         pass
 
     @abstractmethod
-    def _get_alter_column_type_query(self, columns: List[Diff], table_name: str) -> str:
+    def _get_alter_column_type_query(self, column: Diff, table_name: str) -> str:
         """Creates desired statement to alter columns' types.
 
         Args:
@@ -152,10 +152,11 @@ def _get_queries(
                 )
                 queries.append(drop_columns_query)
         if alter_type_items:
-            alter_column_types_query = self._get_alter_column_type_query(
-                alter_type_items, table_name
-            )
-            queries.append(alter_column_types_query)
+            for item in alter_type_items:
+                alter_column_types_query = self._get_alter_column_type_query(
+                    item, table_name
+                )
+                queries.append(alter_column_types_query)
         if alter_key_items:
             logger.info("This operation is not supported by Spark.")
 

diff --git a/butterfree/migrations/database_migration/metastore_migration.py b/butterfree/migrations/database_migration/metastore_migration.py
@@ -74,7 +74,7 @@ def _get_alter_table_add_query(self, columns: List[Diff], table_name: str) -> st
             f"ADD IF NOT EXISTS columns ({parsed_columns});"
         )
 
-    def _get_alter_column_type_query(self, columns: List[Diff], table_name: str) -> str:
+    def _get_alter_column_type_query(self, column: Diff, table_name: str) -> str:
         """Creates SQL statement to alter columns' types.
 
         Args:
@@ -85,9 +85,9 @@ def _get_alter_column_type_query(self, columns: List[Diff], table_name: str) ->
             Alter column type query.
 
         """
-        parsed_columns = self._get_parsed_columns(columns)
+        parsed_columns = self._get_parsed_columns([column])
 
-        return f"ALTER TABLE {table_name} ALTER COLUMN ({parsed_columns});"
+        return f"ALTER TABLE {table_name} ALTER COLUMN {parsed_columns};"
 
     def _get_create_table_query(
         self, columns: List[Dict[str, Any]], table_name: str

diff --git a/setup.py b/setup.py
@@ -1,7 +1,7 @@
 from setuptools import find_packages, setup
 
 __package_name__ = "butterfree"
-__version__ = "1.2.0.dev13"
+__version__ = "1.2.0.dev14"
 __repository_url__ = "https://github.com/quintoandar/butterfree"
 
 with open("requirements.txt") as f:

diff --git a/tests/unit/butterfree/migrations/database_migration/test_cassandra_migration.py b/tests/unit/butterfree/migrations/database_migration/test_cassandra_migration.py
@@ -8,7 +8,7 @@ def test_queries(self, fs_schema, db_schema):
             "ALTER TABLE table_name ADD (new_feature FloatType);",
             "ALTER TABLE table_name DROP (feature1__avg_over_2_days_rolling_windows);",
             "ALTER TABLE table_name ALTER "
-            "(feature1__avg_over_1_week_rolling_windows FloatType);",
+            "feature1__avg_over_1_week_rolling_windows FloatType;",
         ]
         query = cassandra_migration.create_query(fs_schema, "table_name", db_schema)
 
@@ -19,7 +19,7 @@ def test_queries_on_entity(self, fs_schema, db_schema):
         expected_query = [
             "ALTER TABLE table_name ADD (new_feature FloatType);",
             "ALTER TABLE table_name ALTER "
-            "(feature1__avg_over_1_week_rolling_windows FloatType);",
+            "feature1__avg_over_1_week_rolling_windows FloatType;",
         ]
         query = cassandra_migration.create_query(
             fs_schema, "table_name", db_schema, True

diff --git a/tests/unit/butterfree/migrations/database_migration/test_metastore_migration.py b/tests/unit/butterfree/migrations/database_migration/test_metastore_migration.py
@@ -11,7 +11,7 @@ def test_queries(self, fs_schema, db_schema):
             "ALTER TABLE table_name DROP IF EXISTS "
             "(feature1__avg_over_2_days_rolling_windows None);",
             "ALTER TABLE table_name ALTER COLUMN "
-            "(feature1__avg_over_1_week_rolling_windows FloatType);",
+            "feature1__avg_over_1_week_rolling_windows FloatType;",
         ]
 
         query = metastore_migration.create_query(fs_schema, "table_name", db_schema)
@@ -25,7 +25,7 @@ def test_queries_on_entity(self, fs_schema, db_schema):
             "ALTER TABLE test.table_name ADD IF NOT EXISTS "
             "columns (new_feature FloatType);",
             "ALTER TABLE table_name ALTER COLUMN "
-            "(feature1__avg_over_1_week_rolling_windows FloatType);",
+            "feature1__avg_over_1_week_rolling_windows FloatType;",
         ]
 
         query = metastore_migration.create_query(