mirror of
https://github.com/open-metadata/OpenMetadata.git
synced 2025-12-03 19:16:10 +00:00
e7c9d6aa7f
* Implement Initial MultiThread suggestion * Update all the ingestion sources to use the new ContextManager * Fix missing wraps on decorator * Fix Unittests * Fix linters * Fix linters * Fix BigQuery UnitTests * Add UnitTests to the newly created code * Fix unittest * change the threads from table to schemas * Update README.md * Small change suggested by Sonar * Slight change to test a different way to multithread over tables * Debug changes * More multithread tests * Remove uneeded wait time * Testing * refactor code based on removal of time.sleep * Fix wrong paste * Improve ExecutionTimeContextManager * Fix missing .get() and unit tests * Fix conflicting changes * Update Multithread logic with the incremental extraction * Fix linters * Fix unittest * Remove commented code * Fix Unittests * Fix checkstyle * Change default to threads = 1
293 lines
11 KiB
Python
293 lines
11 KiB
Python
# Copyright 2021 Collate
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""
|
|
Test MongoDB using the topology
|
|
"""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from unittest import TestCase
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
from metadata.generated.schema.api.data.createTable import CreateTableRequest
|
|
from metadata.generated.schema.entity.data.database import Database
|
|
from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
|
|
from metadata.generated.schema.entity.data.table import (
|
|
Column,
|
|
ConstraintType,
|
|
DataType,
|
|
TableConstraint,
|
|
TableType,
|
|
)
|
|
from metadata.generated.schema.entity.services.databaseService import (
|
|
DatabaseConnection,
|
|
DatabaseService,
|
|
DatabaseServiceType,
|
|
)
|
|
from metadata.generated.schema.metadataIngestion.workflow import (
|
|
OpenMetadataWorkflowConfig,
|
|
)
|
|
from metadata.generated.schema.type.basic import SourceUrl
|
|
from metadata.generated.schema.type.entityReference import EntityReference
|
|
from metadata.ingestion.ometa.ometa_api import OpenMetadata
|
|
from metadata.ingestion.source.database.bigtable.metadata import BigtableSource
|
|
|
|
mock_file_path = (
|
|
Path(__file__).parent.parent.parent / "resources/datasets/glue_db_dataset.json"
|
|
)
|
|
with open(mock_file_path) as file:
|
|
mock_data: dict = json.load(file)
|
|
|
|
mock_bigtable_config = {
|
|
"source": {
|
|
"type": "bigtable",
|
|
"serviceName": "local_bigtable",
|
|
"serviceConnection": {
|
|
"config": {
|
|
"type": "BigTable",
|
|
"credentials": {
|
|
"gcpConfig": {
|
|
"type": "service_account",
|
|
"projectId": "my-gcp-project",
|
|
"privateKeyId": "private_key_id",
|
|
# this is a valid key that was generated on a local machine and is not used for any real project
|
|
"privateKey": "-----BEGIN RSA PRIVATE KEY-----\nMIIEpQIBAAKCAQEAw3vHG9fDIkcYB0xi2Mv4fS2gUzKR9ZRrcVNeKkqGFTT71AVB\nOzgIqYVe8b2aWODuNye6sipcrqTqOt05Esj+sxhk5McM9bE2RlxXC5QH/Bp9zxMP\n/Yksv9Ov7fdDt/loUk7sTXvI+7LDJfmRYU6MtVjyyLs7KpQIB2xBWEToU1xZY+v0\ndRC1NA+YWc+FjXbAiFAf9d4gXkYO8VmU5meixVh4C8nsjokEXk0T/HEItpZCxadk\ndZ7LKUE/HDmWCO2oNG6sCf4ET2crjSdYIfXuREopX1aQwnk7KbI4/YIdlRz1I369\nAz3+Hxlf9lLJVH3+itN4GXrR9yWWKWKDnwDPbQIDAQABAoIBAQC3X5QuTR7SN8iV\niBUtc2D84+ECSmza5shG/UJW/6N5n0Mf53ICgBS4GNEwiYCRISa0/ILIgK6CcVb7\nsuvH8F3kWNzEMui4TO0x4YsR5GH9HkioCCS224frxkLBQnL20HIIy9ok8Rpe6Zjg\nNZUnp4yczPyqSeA9l7FUbTt69uDM2Cx61m8REOpFukpnYLyZGbmNPYmikEO+rq9r\nwNID5dkSeVuQYo4MQdRavOGFUWvUYXzkEQ0A6vPyraVBfolESX8WaLNVjic7nIa3\nujdSNojnJqGJ3gslntcmN1d4JOfydc4bja4/NdNlcOHpWDGLzY1QnaDe0Koxn8sx\nLT9MVD2NAoGBAPy7r726bKVGWcwqTzUuq1OWh5c9CAc4N2zWBBldSJyUdllUq52L\nWTyva6GRoRzCcYa/dKLLSM/k4eLf9tpxeIIfTOMsvzGtbAdm257ndMXNvfYpxCfU\nK/gUFfAUGHZ3MucTHRY6DTkJg763Sf6PubA2fqv3HhVZDK/1HGDtHlTPAoGBAMYC\npdV7O7lAyXS/d9X4PQZ4BM+P8MbXEdGBbPPlzJ2YIb53TEmYfSj3z41u9+BNnhGP\n4uzUyAR/E4sxrA2+Ll1lPSCn+KY14WWiVGfWmC5j1ftdpkbrXstLN8NpNYzrKZwx\njdR0ZkwvZ8B5+kJ1hK96giwWS+SJxJR3TohcQ18DAoGAJSfmv2r//BBqtURnHrd8\nwq43wvlbC8ytAVg5hA0d1r9Q4vM6w8+vz+cuWLOTTyobDKdrG1/tlXrd5r/sh9L0\n15SIdkGm3kPTxQbPNP5sQYRs8BrV1tEvoao6S3B45DnEBwrdVN42AXOvpcNGoqE4\nuHpahyeuiY7s+ZV8lZdmxSsCgYEAolr5bpmk1rjwdfGoaKEqKGuwRiBX5DHkQkxE\n8Zayt2VOBcX7nzyRI05NuEIMrLX3rZ61CktN1aH8fF02He6aRaoE/Qm9L0tujM8V\nNi8WiLMDeR/Ifs3u4/HAv1E8v1byv0dCa7klR8J257McJ/ID4X4pzcxaXgE4ViOd\nGOHNu9ECgYEApq1zkZthEQymTUxs+lSFcubQpaXyf5ZC61cJewpWkqGDtSC+8DxE\nF/jydybWuoNHXymnvY6QywxuIooivbuib6AlgpEJeybmnWlDOZklFOD0abNZ+aNO\ndUk7XVGffCakXQ0jp1kmZA4lGsYK1h5dEU5DgXqu4UYJ88Vttax2W+Y=\n-----END RSA PRIVATE KEY-----\n",
|
|
"clientEmail": "gcpuser@project_id.iam.gserviceaccount.com",
|
|
"clientId": "client_id",
|
|
"authUri": "https://accounts.google.com/o/oauth2/auth",
|
|
"tokenUri": "https://oauth2.googleapis.com/token",
|
|
"authProviderX509CertUrl": "https://www.googleapis.com/oauth2/v1/certs",
|
|
"clientX509CertUrl": "https://www.googleapis.com/oauth2/v1/certs",
|
|
}
|
|
},
|
|
},
|
|
},
|
|
"sourceConfig": {
|
|
"config": {
|
|
"type": "DatabaseMetadata",
|
|
"schemaFilterPattern": {"includes": ["my_instance"]},
|
|
"tableFilterPattern": {"includes": ["random_table"]},
|
|
}
|
|
},
|
|
},
|
|
"sink": {"type": "metadata-rest", "config": {}},
|
|
"workflowConfig": {
|
|
"openMetadataServerConfig": {
|
|
"hostPort": "http://localhost:8585/api",
|
|
"authProvider": "openmetadata",
|
|
"securityConfig": {
|
|
"jwtToken": "eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg"
|
|
},
|
|
}
|
|
},
|
|
}
|
|
|
|
MOCK_DATABASE_SERVICE = DatabaseService(
|
|
id="85811038-099a-11ed-861d-0242ac120002",
|
|
name="local_bigtable",
|
|
connection=DatabaseConnection(),
|
|
serviceType=DatabaseServiceType.Glue,
|
|
)
|
|
|
|
MOCK_DATABASE = Database(
|
|
id="2aaa012e-099a-11ed-861d-0242ac120002",
|
|
name="my-gcp-project",
|
|
fullyQualifiedName="local_bigtable.my-gcp-project",
|
|
displayName="my-gcp-project",
|
|
description="",
|
|
service=EntityReference(
|
|
id="85811038-099a-11ed-861d-0242ac120002",
|
|
type="databaseService",
|
|
),
|
|
)
|
|
|
|
MOCK_DATABASE_SCHEMA = DatabaseSchema(
|
|
id="2aaa012e-099a-11ed-861d-0242ac120056",
|
|
name="my_instance",
|
|
fullyQualifiedName="local_bigtable.my-gcp-project.my_instance",
|
|
displayName="default",
|
|
description="",
|
|
database=EntityReference(
|
|
id="2aaa012e-099a-11ed-861d-0242ac120002",
|
|
type="database",
|
|
),
|
|
service=EntityReference(
|
|
id="85811038-099a-11ed-861d-0242ac120002",
|
|
type="databaseService",
|
|
),
|
|
)
|
|
|
|
|
|
MOCK_CREATE_TABLE = CreateTableRequest(
|
|
name="random_table",
|
|
tableType=TableType.Regular,
|
|
columns=[
|
|
Column(
|
|
name="row_key",
|
|
displayName="row_key",
|
|
dataType=DataType.BYTES,
|
|
dataTypeDisplay=DataType.BYTES.value,
|
|
),
|
|
Column(
|
|
name="cf1.col1",
|
|
displayName="cf1.col1",
|
|
dataType=DataType.BYTES,
|
|
dataTypeDisplay=DataType.BYTES.value,
|
|
),
|
|
Column(
|
|
name="cf2.col2",
|
|
displayName="cf2.col2",
|
|
dataType=DataType.BYTES,
|
|
dataTypeDisplay=DataType.BYTES.value,
|
|
),
|
|
],
|
|
tableConstraints=[
|
|
TableConstraint(constraintType=ConstraintType.PRIMARY_KEY, columns=["row_key"])
|
|
],
|
|
databaseSchema="local_bigtable.my-gcp-project.my_instance",
|
|
sourceUrl=SourceUrl(
|
|
__root__="https://console.cloud.google.com/bigtable/instances/my_instance/tables/random_table/overview?project=my-gcp-project"
|
|
),
|
|
)
|
|
|
|
|
|
EXPECTED_DATABASE_NAMES = ["my-gcp-project"]
|
|
|
|
EXPECTED_DATABASE_SCHEMA_NAMES = [
|
|
"my_instance",
|
|
]
|
|
|
|
MOCK_DATABASE_SCHEMA_NAMES = [
|
|
"my_instance",
|
|
"random1_schema",
|
|
]
|
|
|
|
EXPECTED_TABLE_NAMES = [
|
|
("random_table", TableType.Regular),
|
|
]
|
|
|
|
|
|
def custom_column_compare(self, other):
|
|
return (
|
|
self.name == other.name
|
|
and self.description == other.description
|
|
and self.children == other.children
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_bigtable_row():
|
|
mock = Mock()
|
|
cell = Mock()
|
|
cell.value = b"cell_value"
|
|
cell.timestamp = 1234567890
|
|
mock.cells = {"cf1": {b"col1": [cell]}, "cf2": {b"col2": [cell]}}
|
|
mock.row_key = b"row_key"
|
|
yield mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_bigtable_table(mock_bigtable_row):
|
|
mock = Mock()
|
|
mock.table_id = "random_table"
|
|
mock.list_column_families.return_value = {"cf1": None, "cf2": None}
|
|
mock.read_rows.return_value = [mock_bigtable_row]
|
|
yield mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_bigtable_instance(mock_bigtable_table):
|
|
mock = Mock()
|
|
mock.instance_id = "my_instance"
|
|
mock.project_id = "my-gcp-project"
|
|
mock.list_tables.return_value = [mock_bigtable_table]
|
|
yield mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_google_cloud_client(mock_bigtable_instance):
|
|
with patch("google.cloud.bigtable.Client") as mock_client:
|
|
mock_client.list_instances.return_value = [[], []]
|
|
mock_client().list_instances.return_value = [[mock_bigtable_instance], []]
|
|
yield mock_client
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_test_connection():
|
|
with patch.object(BigtableSource, "test_connection") as mock_test_connection:
|
|
mock_test_connection.return_value = True
|
|
yield mock_test_connection
|
|
|
|
|
|
class BigTableUnitTest(TestCase):
|
|
@pytest.fixture(autouse=True)
|
|
def setup(
|
|
self,
|
|
monkeypatch,
|
|
mock_google_cloud_client,
|
|
mock_test_connection,
|
|
mock_bigtable_instance,
|
|
mock_bigtable_table,
|
|
):
|
|
self.config = OpenMetadataWorkflowConfig.parse_obj(mock_bigtable_config)
|
|
self.bigtable_source = BigtableSource.create(
|
|
mock_bigtable_config["source"],
|
|
OpenMetadata(self.config.workflowConfig.openMetadataServerConfig),
|
|
)
|
|
self.bigtable_source.context.get().__dict__[
|
|
"database_service"
|
|
] = MOCK_DATABASE_SERVICE.name.__root__
|
|
self.bigtable_source.context.get().__dict__[
|
|
"database"
|
|
] = MOCK_DATABASE.name.__root__
|
|
self.bigtable_source.context.get().__dict__[
|
|
"database_schema"
|
|
] = MOCK_DATABASE_SCHEMA.name.__root__
|
|
self.bigtable_source.instances = {
|
|
"my-gcp-project": {
|
|
mock_bigtable_instance.instance_id: mock_bigtable_instance
|
|
}
|
|
}
|
|
self.bigtable_source.tables = {
|
|
"my-gcp-project": {
|
|
mock_bigtable_instance.instance_id: {
|
|
mock_bigtable_table.table_id: mock_bigtable_table
|
|
}
|
|
}
|
|
}
|
|
|
|
def test_database_names(self):
|
|
assert (
|
|
list(self.bigtable_source.get_database_names()) == EXPECTED_DATABASE_NAMES
|
|
)
|
|
|
|
def test_database_schema_names(self):
|
|
assert (
|
|
list(self.bigtable_source.get_database_schema_names())
|
|
== EXPECTED_DATABASE_SCHEMA_NAMES
|
|
)
|
|
|
|
def test_table_names(self):
|
|
assert (
|
|
list(self.bigtable_source.get_tables_name_and_type())
|
|
== EXPECTED_TABLE_NAMES
|
|
)
|
|
|
|
def test_yield_tables(self):
|
|
Column.__eq__ = custom_column_compare
|
|
result = next(self.bigtable_source.yield_table(EXPECTED_TABLE_NAMES[0]))
|
|
assert result.left is None
|
|
assert result.right.name.__root__ == "random_table"
|
|
assert result.right == MOCK_CREATE_TABLE
|