mypy compliant

Author: roelschr

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

@@ -28,8 +28,6 @@ class CassandraConfig(AbstractWriteConfig):
 
     """
 
-    database = "cassandra"
-
     def __init__(
         self,
         username: str = None,
@@ -52,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

@@ -23,8 +23,6 @@ class MetastoreConfig(AbstractWriteConfig):
 
     """
 
-    database = "metastore"
-
     def __init__(
         self,
         path: str = None,
@@ -37,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/database_migration/database_migration.py

@@ -3,12 +3,9 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import Enum, auto
-from typing import Any, Dict, List, Set, Union
+from typing import Any, Dict, List, Set
 
-from butterfree.load.writers import (
-    HistoricalFeatureStoreWriter,
-    OnlineFeatureStoreWriter,
-)
+from butterfree.load.writers.writer import Writer
 from butterfree.transform import FeatureSet
 
 
@@ -129,11 +126,7 @@ def _get_schema(self, table_name: str) -> List[Dict[str, Any]]:
             db_schema = []
         return db_schema
 
-    def apply_migration(
-        self,
-        feature_set: FeatureSet,
-        writer: Union[HistoricalFeatureStoreWriter, OnlineFeatureStoreWriter],
-    ) -> None:
+    def apply_migration(self, feature_set: FeatureSet, writer: Writer,) -> None:
         """Apply the migration in the respective database.
 
         Args: