-
Notifications
You must be signed in to change notification settings - Fork 36
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
AlvaroMarquesAndrade
committed
Feb 18, 2021
1 parent
a59030d
commit e524c5c
Showing
6 changed files
with
127 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,23 +1,68 @@ | ||
"""Cassandra Migration entity.""" | ||
|
||
import warnings | ||
from typing import Any, Dict, List | ||
|
||
from butterfree.migrations import DatabaseMigration | ||
from butterfree.configs.db import CassandraConfig | ||
from butterfree.migrations.migration import DatabaseMigration | ||
|
||
|
||
class CassandraMigration(DatabaseMigration): | ||
"""Cassandra class for Migrations.""" | ||
|
||
@staticmethod | ||
def _get_alter_table_query(columns: List[Dict[str, Any]], table_name: str) -> str: | ||
parsed_columns = [] | ||
for col in columns: | ||
parsed_columns.append(f"{col['column_name']} {col['type']}") | ||
|
||
parsed_columns = ", ".join(parsed_columns) # type: ignore | ||
|
||
return f"ALTER TABLE {table_name} " f"ADD ({parsed_columns});" | ||
|
||
@staticmethod | ||
def _get_create_table_query(columns: List[Dict[str, Any]], table_name: str,) -> str: | ||
"""Creates CQL statement to create a table.""" | ||
parsed_columns = [] | ||
primary_keys = [] | ||
|
||
for col in columns: | ||
col_str = f"{col['column_name']} {col['type']}" | ||
if col["primary_key"]: | ||
primary_keys.append(col["column_name"]) | ||
parsed_columns.append(col_str) | ||
|
||
joined_parsed_columns = ", ".join(parsed_columns) | ||
|
||
if len(primary_keys) > 0: | ||
joined_primary_keys = ", ".join(primary_keys) | ||
columns_str = ( | ||
f"{joined_parsed_columns}, PRIMARY KEY ({joined_primary_keys})" | ||
) | ||
else: | ||
columns_str = joined_parsed_columns | ||
|
||
keyspace = CassandraConfig().keyspace | ||
|
||
return f"CREATE TABLE {keyspace}.{table_name} " f"({columns_str});" | ||
|
||
def create_query( | ||
self, | ||
fs_schema: List[Dict[str, Any]], | ||
db_schema: List[Dict[str, Any]], | ||
table_name: str, | ||
db_schema: List[Dict[str, Any]] = None, | ||
schema_diff: List[Dict[str, Any]] = None, | ||
) -> Any: | ||
"""Create a query regarding Cassandra. | ||
Returns: | ||
Schema object. | ||
""" | ||
pass | ||
if not schema_diff: | ||
warnings.warn("No migration was performed", UserWarning, stacklevel=1) | ||
return | ||
|
||
if not db_schema: | ||
return self._get_create_table_query(schema_diff, table_name) | ||
|
||
return self._get_alter_table_query(schema_diff, table_name) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,32 @@ | ||
from pytest import fixture | ||
|
||
|
||
@fixture | ||
def dummy_db_schema(): | ||
dummy_db_schema = [ | ||
{"column_name": "id", "type": "int", "primary_key": True}, | ||
{"column_name": "platform", "type": "text", "primary_key": False}, | ||
{"column_name": "ts", "type": "bigint", "primary_key": False}, | ||
] | ||
return dummy_db_schema | ||
|
||
|
||
@fixture | ||
def dummy_schema_diff(): | ||
dummy_schema_diff = [ | ||
{"column_name": "kappa_column", "type": "int", "primary_key": False}, | ||
{"column_name": "pogchamp", "type": "uuid", "primary_key": False}, | ||
] | ||
return dummy_schema_diff | ||
|
||
|
||
@fixture | ||
def dummy_schema(): | ||
dummy_schema = [ | ||
{"column_name": "id", "type": "int", "primary_key": True}, | ||
{"column_name": "platform", "type": "text", "primary_key": False}, | ||
{"column_name": "ts", "type": "bigint", "primary_key": False}, | ||
{"column_name": "kappa_column", "type": "int", "primary_key": False}, | ||
{"column_name": "pogchamp", "type": "uuid", "primary_key": False}, | ||
] | ||
return dummy_schema |
39 changes: 39 additions & 0 deletions
39
tests/unit/butterfree/migrations/test_cassandra_migration.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,39 @@ | ||
from butterfree.migrations import CassandraMigration | ||
|
||
|
||
class TestCassandraMigration: | ||
def test_alter_table_query(self, dummy_db_schema, dummy_schema_diff): | ||
cassandra_migration = CassandraMigration() | ||
|
||
expected_query = "ALTER TABLE test ADD (kappa_column int, pogchamp uuid);" | ||
query = cassandra_migration.create_query( | ||
"test", dummy_db_schema, dummy_schema_diff | ||
) | ||
|
||
assert isinstance(query, str) | ||
assert query, expected_query | ||
|
||
def test_create_table_query(self, dummy_db_schema): | ||
|
||
cassandra_migration = CassandraMigration() | ||
|
||
expected_query = ( | ||
"CREATE TABLE test.test_table " | ||
"(id int, platform text, ts bigint, PRIMARY KEY (id));" | ||
) | ||
query = cassandra_migration.create_query( | ||
table_name="test_table", schema_diff=dummy_db_schema | ||
) | ||
|
||
assert isinstance(query, str) | ||
assert query, expected_query | ||
|
||
def test_no_diff(self, dummy_db_schema): | ||
|
||
cassandra_migration = CassandraMigration() | ||
|
||
query = cassandra_migration.create_query( | ||
table_name="test_table", db_schema=dummy_db_schema, | ||
) | ||
|
||
assert query is None |