def test_multiple_indexes(self):
index_metadata_without_type = IndexMetadata(
name="test_index",
columns=["test_column_1", "test_column_2"],
)
index_metadata_with_type = IndexMetadata(
name="test_index",
columns=["test_column_1", "test_column_2"],
index_type="primary",
architecture="clustered",
constraint="unique",
)
table_indexes_metadata = TableIndexesMetadata(
database="test_database",
cluster="test_cluster",
schema="test_schema",
table="test_table",
indexes=[index_metadata_without_type, index_metadata_with_type],
)
expected_list = [
"* [] `test_index` [`test_column_1`, `test_column_2`]",
"* [primary, unique, clustered] `test_index` [`test_column_1`, `test_column_2`]",
]
self.assertEqual(table_indexes_metadata.format_for_markdown(),
"\n".join(expected_list))
def test_index_without_any_type(self):
index_metadata = IndexMetadata(
name="test_index",
columns=["test_column_1", "test_column_2"],
)
expected = "* [] `test_index` [`test_column_1`, `test_column_2`]"
self.assertEqual(index_metadata.format_for_markdown(), expected)
def test_index_with_type(self):
index_metadata = IndexMetadata(
name="test_index",
columns=["test_column_1", "test_column_2"],
index_type="primary",
architecture="clustered",
constraint="unique",
)
expected = "* [primary, unique, clustered] `test_index` [`test_column_1`, `test_column_2`]"
self.assertEqual(index_metadata.format_for_markdown(), expected)
def test_load_index_metadata(patched_config):
index_metadata = IndexMetadata(
name="mock_index",
columns=["mock_column_1", "mock_column_2"],
)
record = TableIndexesMetadata(
database="mock_database",
cluster="mock_catalog",
schema="mock_schema",
table="mock_table",
indexes=[index_metadata],
)
loader = whale_loader.WhaleLoader()
loader.init(patched_config)
loader.load(record)
loader.close()
file_path = os.path.join(
patched_config.get("base_directory"),
"mock_database/mock_catalog.mock_schema.mock_table.md",
)
with open(file_path, "r") as f:
written_record = f.read()
assert "mock_schema" in written_record
assert "mock_table" in written_record
assert "mock_catalog" in written_record
assert "mock_database" in written_record
assert "mock_index" in written_record
def _get_extract_iter(self):
for table_key, table_group in groupby(self._get_raw_extract_iter(),
self._get_table_key):
indexes = []
for index_key, index_group in groupby(table_group,
self._get_index_key):
columns = []
for row in index_group:
last_row = row
columns.append(row["column_name"])
indexes.append(
IndexMetadata(
name=last_row["index_name"],
index_type="primary"
if last_row["is_primary"] else None,
architecture="clustered"
if last_row["is_clustered"] else None,
constraint="unique" if last_row["is_unique"] else None,
columns=columns,
))
yield TableIndexesMetadata(
database=self._database,
cluster=last_row["cluster"],
schema=last_row["schema"],
table=last_row["table"],
indexes=indexes,
)
def test_single_index_without_any_type(self):
index_metadata = IndexMetadata(
name="test_index",
columns=["test_column_1", "test_column_2"],
)
table_indexes_metadata = TableIndexesMetadata(
database="test_database",
cluster="test_cluster",
schema="test_schema",
table="test_table",
indexes=[index_metadata],
)
expected = "* [] `test_index` [`test_column_1`, `test_column_2`]"
self.assertEqual(table_indexes_metadata.format_for_markdown(),
expected)
def test_single_index_with_type(self):
index_metadata = IndexMetadata(
name="test_index",
columns=["test_column_1", "test_column_2"],
index_type="primary",
architecture="clustered",
constraint="unique",
)
table_indexes_metadata = TableIndexesMetadata(
database="test_database",
cluster="test_cluster",
schema="test_schema",
table="test_table",
indexes=[index_metadata],
)
expected = "* [primary, unique, clustered] `test_index` [`test_column_1`, `test_column_2`]"
self.assertEqual(table_indexes_metadata.format_for_markdown(),
expected)
def test_extraction_with_multiple_result(self) -> None:
with patch.object(SQLAlchemyExtractor,
"_get_connection") as mock_connection:
connection = MagicMock()
mock_connection.return_value = connection
sql_execute = MagicMock()
connection.execute = sql_execute
table_1 = {
"cluster": self.conf[PostgresIndexExtractor.CLUSTER_KEY],
"schema": "test_schema_1",
"table": "test_table_1",
}
table_2 = {
"cluster": self.conf[PostgresIndexExtractor.CLUSTER_KEY],
"schema": "test_schema_2",
"table": "test_table_2",
}
sql_execute.return_value = [
self._union(
{
"index_name": "idx_1",
"is_primary": False,
"is_clustered": False,
"is_unique": True,
"column_name": "idx_1_column_1",
},
table_1,
),
self._union(
{
"index_name": "idx_1",
"is_primary": False,
"is_clustered": False,
"is_unique": True,
"column_name": "idx_1_column_2",
},
table_1,
),
self._union(
{
"index_name": "idx_2",
"is_primary": False,
"is_clustered": False,
"is_unique": True,
"column_name": "idx_2_column_1",
},
table_1,
),
self._union(
{
"index_name": "idx_3",
"is_primary": False,
"is_clustered": False,
"is_unique": True,
"column_name": "idx_3_column_1",
},
table_2,
),
self._union(
{
"index_name": "idx_3",
"is_primary": False,
"is_clustered": False,
"is_unique": True,
"column_name": "idx_3_column_2",
},
table_2,
),
]
extractor = PostgresIndexExtractor()
extractor.init(self.conf)
expected_index_1 = IndexMetadata(
name="idx_1",
description=None,
index_type=None,
architecture=None,
constraint="unique",
columns=["idx_1_column_1", "idx_1_column_2"],
tags=None,
)
expected_index_2 = IndexMetadata(
name="idx_2",
description=None,
index_type=None,
architecture=None,
constraint="unique",
columns=["idx_2_column_1"],
tags=None,
)
expected = TableIndexesMetadata(
database="postgres",
cluster="MY_CLUSTER",
schema="test_schema_1",
table="test_table_1",
indexes=[expected_index_1, expected_index_2],
)
self.assertEqual(expected.__repr__(),
extractor.extract().__repr__())
expected_index = IndexMetadata(
name="idx_3",
description=None,
index_type=None,
architecture=None,
constraint="unique",
columns=["idx_3_column_1", "idx_3_column_2"],
tags=None,
)
expected = TableIndexesMetadata(
database="postgres",
cluster="MY_CLUSTER",
schema="test_schema_2",
table="test_table_2",
indexes=[expected_index],
)
self.assertEqual(expected.__repr__(),
extractor.extract().__repr__())
self.assertIsNone(extractor.extract())