diff --git a/Makefile b/Makefile index 2156f9b00..1534b80e9 100644 --- a/Makefile +++ b/Makefile @@ -41,7 +41,7 @@ requirements: minimum-requirements dev-requirements requirements-test requiremen .PHONY: ci-install ci-install: @pip install --upgrade pip - @pip install cmake + @pip install "cassandra-driver==3.24.0" @python -m pip install -U -r requirements.test.txt -r requirements.lint.txt -r requirements.dev.txt -r requirements.txt -t ./pip/deps --cache-dir ./pip/cache .PHONY: tests diff --git a/butterfree/clients/__init__.py b/butterfree/clients/__init__.py index 7e8d1a95b..52b171fda 100644 --- a/butterfree/clients/__init__.py +++ b/butterfree/clients/__init__.py @@ -1,7 +1,13 @@ """Holds connection clients.""" from butterfree.clients.abstract_client import AbstractClient -from butterfree.clients.cassandra_client import CassandraClient from butterfree.clients.spark_client import SparkClient -__all__ = ["SparkClient", "CassandraClient", "AbstractClient"] +__all__ = ["SparkClient", "AbstractClient"] + +try: + from butterfree.clients.cassandra_client import CassandraClient # noqa: F401 + + __all__.append("CassandraClient") +except ImportError: + pass diff --git a/butterfree/clients/cassandra_client.py b/butterfree/clients/cassandra_client.py index c4d34521f..693a4a9c9 100644 --- a/butterfree/clients/cassandra_client.py +++ b/butterfree/clients/cassandra_client.py @@ -4,18 +4,24 @@ from ssl import CERT_REQUIRED, PROTOCOL_TLSv1 from typing import Dict, List, Optional, Union -from cassandra.auth import PlainTextAuthProvider -from cassandra.cluster import ( - EXEC_PROFILE_DEFAULT, - Cluster, - ExecutionProfile, - ResponseFuture, - Session, -) -from cassandra.policies import DCAwareRoundRobinPolicy -from cassandra.query import ConsistencyLevel, dict_factory from typing_extensions import TypedDict +try: + from cassandra.auth import PlainTextAuthProvider + from cassandra.cluster import ( + EXEC_PROFILE_DEFAULT, + Cluster, + ExecutionProfile, + ResponseFuture, + Session, + ) + from cassandra.policies import DCAwareRoundRobinPolicy + from cassandra.query import ConsistencyLevel, dict_factory +except ModuleNotFoundError as e: + e.msg = "Cassandra not found. To be able to use this module,you must install butterfree[cassandra] or install cassandra-driver manually." # noqa: E501 + raise + + from butterfree.clients import AbstractClient logger = logging.getLogger(__name__) @@ -168,7 +174,7 @@ def get_schema( if not response: raise RuntimeError( - f"No columns found for table: {table}" f"in key space: {self.keyspace}" + f"No columns found for table: {table}in key space: {self.keyspace}" ) return response @@ -198,7 +204,7 @@ def _get_create_table_query( else: columns_str = joined_parsed_columns - query = f"CREATE TABLE {self.keyspace}.{table} " f"({columns_str}); " + query = f"CREATE TABLE {self.keyspace}.{table} ({columns_str}); " return query diff --git a/butterfree/hooks/schema_compatibility/__init__.py b/butterfree/hooks/schema_compatibility/__init__.py index a00adef8d..855ec15b6 100644 --- a/butterfree/hooks/schema_compatibility/__init__.py +++ b/butterfree/hooks/schema_compatibility/__init__.py @@ -1,10 +1,7 @@ """Holds Schema Compatibility Hooks definitions.""" -from butterfree.hooks.schema_compatibility.cassandra_table_schema_compatibility_hook import ( # noqa - CassandraTableSchemaCompatibilityHook, -) from butterfree.hooks.schema_compatibility.spark_table_schema_compatibility_hook import ( # noqa SparkTableSchemaCompatibilityHook, ) -__all__ = ["SparkTableSchemaCompatibilityHook", "CassandraTableSchemaCompatibilityHook"] +__all__ = ["SparkTableSchemaCompatibilityHook"] diff --git a/butterfree/load/writers/online_feature_store_writer.py b/butterfree/load/writers/online_feature_store_writer.py index bce5a3751..3f3c62e13 100644 --- a/butterfree/load/writers/online_feature_store_writer.py +++ b/butterfree/load/writers/online_feature_store_writer.py @@ -11,7 +11,6 @@ from butterfree.configs.db import AbstractWriteConfig, CassandraConfig from butterfree.constants.columns import TIMESTAMP_COLUMN from butterfree.hooks import Hook -from butterfree.hooks.schema_compatibility import CassandraTableSchemaCompatibilityHook from butterfree.load.writers.writer import Writer from butterfree.transform import FeatureSet @@ -270,9 +269,9 @@ def check_schema( table_name: table name where the dataframe will be saved. database: database name where the dataframe will be saved. """ - if not self.check_schema_hook: - self.check_schema_hook = CassandraTableSchemaCompatibilityHook( - client, table_name - ) + if self.check_schema_hook: + return self.check_schema_hook.run(dataframe) - return self.check_schema_hook.run(dataframe) + raise NotImplementedError( + "Schema check hook not implemented for Online Feature Store Writer" + ) diff --git a/requirements.txt b/requirements.txt index bde1f746e..2b79944d4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,3 @@ -cassandra-driver==3.24.0 mdutils>=1.2.2,<1.7 pandas>=0.24,<2.0 parameters-validation>=1.1.5,<2.0 diff --git a/setup.py b/setup.py index c96888e34..38bb9eb9a 100644 --- a/setup.py +++ b/setup.py @@ -34,7 +34,10 @@ license="Copyright", author="QuintoAndar", install_requires=requirements, - extras_require={"h3": ["h3>=3.7.4,<4"]}, + extras_require={ + "h3": ["h3>=3.7.4,<4"], + "cassandra": ["cassandra-driver==3.24.0"], + }, python_requires=">=3.9, <4", entry_points={"console_scripts": ["butterfree=butterfree._cli.main:app"]}, include_package_data=True, diff --git a/tests/unit/butterfree/hooks/schema_compatibility/test_cassandra_table_schema_compatibility_hook.py b/tests/unit/butterfree/hooks/schema_compatibility/test_cassandra_table_schema_compatibility_hook.py index eccb8d8cc..918abe058 100644 --- a/tests/unit/butterfree/hooks/schema_compatibility/test_cassandra_table_schema_compatibility_hook.py +++ b/tests/unit/butterfree/hooks/schema_compatibility/test_cassandra_table_schema_compatibility_hook.py @@ -3,7 +3,9 @@ import pytest from butterfree.clients import CassandraClient -from butterfree.hooks.schema_compatibility import CassandraTableSchemaCompatibilityHook +from butterfree.hooks.schema_compatibility.cassandra_table_schema_compatibility_hook import ( # noqa: E501 + CassandraTableSchemaCompatibilityHook, +) class TestCassandraTableSchemaCompatibilityHook: