[MLOP-640] Create CLI with migrate command (#298)

Author: roelschr

butterfree/_cli/__init__.py

@@ -0,0 +1,10 @@
+import logging
+
+
+def __logger(name: str) -> logging.Logger:
+    format_ = "%(name)s:%(asctime)-15s:%(levelname)s:< %(message)s >"
+    logging.basicConfig(format=format_, level=logging.INFO)
+    return logging.getLogger(name)
+
+
+cli_logger = __logger("butterfree")

butterfree/_cli/main.py

@@ -3,7 +3,7 @@
 from butterfree._cli import migrate
 
 app = typer.Typer()
-app.add_typer(migrate.app, name="migrate")
+app.add_typer(migrate.app)
 
 if __name__ == "__main__":
     app()

butterfree/_cli/migrate.py

@@ -1,24 +1,16 @@
-import datetime
 import importlib
 import inspect
-import os
 import pkgutil
 import sys
 from typing import Set
 
-import boto3
 import setuptools
 import typer
-from botocore.exceptions import ClientError
 
-from butterfree.configs import environment
-from butterfree.configs.logger import __logger
-from butterfree.migrations.database_migration import ALLOWED_DATABASE
+from butterfree._cli import cli_logger
 from butterfree.pipelines import FeatureSetPipeline
 
-app = typer.Typer(help="Apply the automatic migrations in a database.")
-
-logger = __logger("migrate", True)
+app = typer.Typer()
 
 
 def __find_modules(path: str) -> Set[str]:
@@ -40,19 +32,18 @@ def __find_modules(path: str) -> Set[str]:
 
 
 def __fs_objects(path: str) -> Set[FeatureSetPipeline]:
-    logger.info(f"Looking for python modules under {path}...")
+    cli_logger.info(f"Looking for python modules under {path}...")
     modules = __find_modules(path)
     if not modules:
-        logger.error(f"Path: {path} not found!")
         return set()
 
-    logger.info(f"Importing modules...")
+    cli_logger.info(f"Importing modules...")
     package = ".".join(path.strip("/").split("/"))
     imported = set(
         importlib.import_module(f".{name}", package=package) for name in modules
     )
 
-    logger.info(f"Scanning modules...")
+    cli_logger.info(f"Scanning modules...")
     content = {
         module: set(
             filter(
@@ -88,91 +79,17 @@ def __fs_objects(path: str) -> Set[FeatureSetPipeline]:
 
             instances.add(value)
 
-    logger.info("Creating instances...")
+    cli_logger.info("Creating instances...")
     return set(value() for value in instances)
 
 
 PATH = typer.Argument(
     ..., help="Full or relative path to where feature set pipelines are being defined.",
 )
 
-GENERATE_LOGS = typer.Option(
-    False, help="To generate the logs in local file 'logging.json'."
-)
-
-DEBUG_MODE = typer.Option(
-    False,
-    help="To view the queries resulting from the migration, DON'T apply the migration.",
-)
-
 
-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, file_local: bool, debug_mode: bool) -> None:
-        """Send all migration logs to S3."""
-        file_name = "../logging.json"
-
-        if not file_local and os.path.exists(file_name):
-            s3_client = boto3.client("s3")
-
-            timestamp = datetime.datetime.now()
-
-            if debug_mode:
-                object_name = (
-                    f"logs/migrate-debug-mode/"
-                    f"{timestamp.strftime('%Y-%m-%d')}"
-                    f"/logging-{timestamp.strftime('%H:%M:%S')}.json"
-                )
-            else:
-                object_name = (
-                    f"logs/migrate/"
-                    f"{timestamp.strftime('%Y-%m-%d')}"
-                    f"/logging-{timestamp.strftime('%H:%M:%S')}.json"
-                )
-            bucket = environment.get_variable("FEATURE_STORE_S3_BUCKET")
-
-            try:
-                s3_client.upload_file(
-                    file_name,
-                    bucket,
-                    object_name,
-                    ExtraArgs={"ACL": "bucket-owner-full-control"},
-                )
-            except ClientError:
-                raise
-
-            os.remove(file_name)
-        elif os.path.exists(file_name):
-            print("Logs written to ../logging.json")
-        else:
-            print("No logs were generated.")
-
-    def run(self, generate_logs: bool = False, debug_mode: bool = False) -> None:
-        """Construct and apply the migrations."""
-        for pipeline in self.pipelines:
-            for writer in pipeline.sink.writers:
-                db = writer.db_config.database
-                if db == "cassandra":
-                    migration = ALLOWED_DATABASE[db]
-                    migration.apply_migration(pipeline.feature_set, writer, debug_mode)
-                else:
-                    logger.warning(f"Butterfree not supporting {db} Migrations yet.")
-
-        self._send_logs_to_s3(generate_logs, debug_mode)
-
-
-@app.command("apply")
-def migrate(
-    path: str = PATH, generate_logs: bool = GENERATE_LOGS, debug_mode: bool = DEBUG_MODE
-) -> Set[FeatureSetPipeline]:
+@app.callback()
+def migrate(path: str = PATH) -> Set[FeatureSetPipeline]:
     """Scan and run database migrations for feature set pipelines defined under PATH.
 
     Butterfree will scan a given path for classes that inherit from its
@@ -183,6 +100,5 @@ def migrate(
     All pipelines must be under python modules inside path, so we can dynamically
     import and instantiate them.
     """
-    pipe_set = __fs_objects(path)
-    Migrate(pipe_set).run(generate_logs, debug_mode)
-    return pipe_set
+    # TODO call the Migration actor with all feature set pipeline objects
+    return __fs_objects(path)

requirements.txt

@@ -4,6 +4,4 @@ pandas>=0.24,<1.1
 parameters-validation>=1.1.5,<2.0
 pyspark==3.*
 typer>=0.3,<0.4
-setuptools>=41,<42
-typing-extensions==3.7.4.3
-boto3==1.17.*
+setuptools>=41,<42

setup.py

@@ -36,4 +36,6 @@
     install_requires=requirements,
     extras_require={"h3": ["cmake==3.16.3", "h3==3.4.2"]},
     python_requires=">=3.7, <4",
+    entry_points={"console_scripts": ["butterfree=butterfree._cli.main:app"]},
+    include_package_data=True,
 )

tests/unit/butterfree/_cli/test_migrate.py

@@ -1,43 +1,8 @@
-from unittest.mock import call
-
-from typer.testing import CliRunner
-
 from butterfree._cli import migrate
-from butterfree._cli.main import app
-from butterfree.migrations.database_migration import CassandraMigration
 from butterfree.pipelines import FeatureSetPipeline
 
-runner = CliRunner()
-
-
-class TestMigrate:
-    def test_migrate_success(self, mocker):
-        mocker.patch.object(migrate.Migrate, "run")
-        all_fs = migrate.migrate("tests/mocks/entities/")
-        assert all(isinstance(fs, FeatureSetPipeline) for fs in all_fs)
-        assert sorted([fs.feature_set.name for fs in all_fs]) == ["first", "second"]
-
-    def test_migrate_run_methods(self, mocker):
-        mocker.patch.object(CassandraMigration, "apply_migration")
-        mocker.patch.object(migrate.Migrate, "_send_logs_to_s3")
-
-        all_fs = migrate.migrate("tests/mocks/entities/", False, False)
-
-        assert CassandraMigration.apply_migration.call_count == 2
-
-        cassandra_pairs = [
-            call(pipe.feature_set, pipe.sink.writers[1], False) for pipe in all_fs
-        ]
-        CassandraMigration.apply_migration.assert_has_calls(
-            cassandra_pairs, any_order=True
-        )
-        migrate.Migrate._send_logs_to_s3.assert_called_once()
-
-    def test_app_cli(self):
-        result = runner.invoke(app, "migrate")
-        assert result.exit_code == 0
 
-    def test_app_migrate(self, mocker):
-        mocker.patch.object(migrate.Migrate, "run")
-        result = runner.invoke(app, ["migrate", "apply", "tests/mocks/entities/"])
-        assert result.exit_code == 0
+def test_migrate_success():
+    all_fs = migrate.migrate("tests/mocks/entities/")
+    assert all(isinstance(fs, FeatureSetPipeline) for fs in all_fs)
+    assert sorted([fs.feature_set.name for fs in all_fs]) == ["first", "second"]