[MLOP-647] [MLOP-646] Apply migrations (#300)

* add apply migration method

* add test apply migration

* add migrate actor with tests

* mypy compliant

* fix test interaction with mocked object

* Rebase and some adjusts.

Co-authored-by: Mayara Moromisato <[email protected]>

Author: roelschr

butterfree/_cli/migrate.py

@@ -8,6 +8,7 @@
 import typer
 
 from butterfree._cli import cli_logger
+from butterfree.migrations.database_migration import ALLOWED_DATABASE
 from butterfree.pipelines import FeatureSetPipeline
 
 app = typer.Typer()
@@ -88,6 +89,28 @@ def __fs_objects(path: str) -> Set[FeatureSetPipeline]:
 )
 
 
+class Migrate:
+    """Execute migration operations in a Database based on pipeline Writer.
+
+    Attributes:
+        pipelines: list of Feature Set Pipelines to use to migration.
+    """
+
+    def __init__(self, pipelines: Set[FeatureSetPipeline]) -> None:
+        self.pipelines = pipelines
+
+    def _send_logs_to_s3(self) -> None:
+        """Send all migration logs to S3."""
+        pass
+
+    def run(self) -> None:
+        """Construct and apply the migrations."""
+        for pipeline in self.pipelines:
+            for writer in pipeline.sink.writers:
+                migration = ALLOWED_DATABASE[writer.db_config.database]
+                migration.apply_migration(pipeline.feature_set, writer)
+
+
 @app.callback()
 def migrate(path: str = PATH) -> Set[FeatureSetPipeline]:
     """Scan and run database migrations for feature set pipelines defined under PATH.
@@ -100,5 +123,6 @@ def migrate(path: str = PATH) -> Set[FeatureSetPipeline]:
     All pipelines must be under python modules inside path, so we can dynamically
     import and instantiate them.
     """
-    # TODO call the Migration actor with all feature set pipeline objects
-    return __fs_objects(path)
+    pipe_set = __fs_objects(path)
+    Migrate(pipe_set).run()
+    return pipe_set

butterfree/configs/db/abstract_config.py

@@ -7,6 +7,11 @@
 class AbstractWriteConfig(ABC):
     """Abstract class for database write configurations with spark."""
 
+    @property
+    @abstractmethod
+    def database(self) -> str:
+        """Database name."""
+
     @property
     @abstractmethod
     def mode(self) -> Any:

butterfree/configs/db/cassandra_config.py

@@ -50,6 +50,11 @@ def __init__(
         self.stream_output_mode = stream_output_mode
         self.stream_checkpoint_path = stream_checkpoint_path
 
+    @property
+    def database(self) -> str:
+        """Database name."""
+        return "cassandra"
+
     @property
     def username(self) -> Optional[str]:
         """Username used in connection to Cassandra DB."""

butterfree/configs/db/kafka_config.py

@@ -41,6 +41,11 @@ def __init__(
         self.stream_output_mode = stream_output_mode
         self.stream_checkpoint_path = stream_checkpoint_path
 
+    @property
+    def database(self) -> str:
+        """Database name."""
+        return "kafka"
+
     @property
     def kafka_topic(self) -> Optional[str]:
         """Kafka topic name."""

butterfree/configs/db/metastore_config.py

@@ -35,6 +35,11 @@ def __init__(
         self.format_ = format_
         self.file_system = file_system
 
+    @property
+    def database(self) -> str:
+        """Database name."""
+        return "metastore"
+
     @property
     def path(self) -> Optional[str]:
         """Bucket name."""

butterfree/load/writers/historical_feature_store_writer.py

@@ -1,7 +1,7 @@
 """Holds the Historical Feature Store writer class."""
 
 import os
-from typing import Any, Union
+from typing import Any
 
 from pyspark.sql.dataframe import DataFrame
 from pyspark.sql.functions import dayofmonth, month, year
@@ -106,16 +106,17 @@ class HistoricalFeatureStoreWriter(Writer):
 
     def __init__(
         self,
-        db_config: Union[AbstractWriteConfig, MetastoreConfig] = None,
+        db_config: AbstractWriteConfig = None,
         database: str = None,
         num_partitions: int = None,
         validation_threshold: float = DEFAULT_VALIDATION_THRESHOLD,
         debug_mode: bool = False,
         interval_mode: bool = False,
         check_schema_hook: Hook = None,
     ):
-        super(HistoricalFeatureStoreWriter, self).__init__(debug_mode, interval_mode)
-        self.db_config = db_config or MetastoreConfig()
+        super(HistoricalFeatureStoreWriter, self).__init__(
+            db_config or MetastoreConfig(), debug_mode, interval_mode
+        )
         self.database = database or environment.get_variable(
             "FEATURE_STORE_HISTORICAL_DATABASE"
         )

butterfree/load/writers/online_feature_store_writer.py

@@ -80,15 +80,15 @@ class OnlineFeatureStoreWriter(Writer):
 
     def __init__(
         self,
-        db_config: Union[AbstractWriteConfig, CassandraConfig] = None,
+        db_config: AbstractWriteConfig = None,
         debug_mode: bool = False,
         write_to_entity: bool = False,
         interval_mode: bool = False,
         check_schema_hook: Hook = None,
     ):
-        super(OnlineFeatureStoreWriter, self).__init__(debug_mode, interval_mode)
-        self.db_config = db_config or CassandraConfig()
-        self.write_to_entity = write_to_entity
+        super(OnlineFeatureStoreWriter, self).__init__(
+            db_config or CassandraConfig(), debug_mode, interval_mode, write_to_entity
+        )
         self.check_schema_hook = check_schema_hook
 
     @staticmethod

butterfree/load/writers/writer.py

@@ -7,6 +7,7 @@
 from pyspark.sql.dataframe import DataFrame
 
 from butterfree.clients import SparkClient
+from butterfree.configs.db import AbstractWriteConfig
 from butterfree.hooks import HookableComponent
 from butterfree.transform import FeatureSet
 
@@ -19,11 +20,19 @@ class Writer(ABC, HookableComponent):
 
     """
 
-    def __init__(self, debug_mode: bool = False, interval_mode: bool = False) -> None:
+    def __init__(
+        self,
+        db_config: AbstractWriteConfig,
+        debug_mode: bool = False,
+        interval_mode: bool = False,
+        write_to_entity: bool = False,
+    ) -> None:
         super().__init__()
+        self.db_config = db_config
         self.transformations: List[Dict[str, Any]] = []
         self.debug_mode = debug_mode
         self.interval_mode = interval_mode
+        self.write_to_entity = write_to_entity
 
     def with_(
         self, transformer: Callable[..., DataFrame], *args: Any, **kwargs: Any

butterfree/migrations/__init__.py

@@ -1,4 +1 @@
 """Holds available migrations."""
-from butterfree.migrations.migrate import Migrate
-
-__all__ = ["Migrate"]

butterfree/migrations/database_migration/__init__.py

@@ -9,3 +9,9 @@
 )
 
 __all__ = ["CassandraMigration", "MetastoreMigration", "Diff"]
+
+
+ALLOWED_DATABASE = {
+    "cassandra": CassandraMigration(),
+    "metastore": MetastoreMigration(),
+}