def test_make_handler_metadata(
db_keys: Dict[str, str],
handler_type: str,
expected_type: Any,
kwargs: Dict[str, str],
):
base_config = {"username": "******", "password": "******", "database": "baz"}
config = {**base_config, **db_keys}
with pytest.raises(KeyError):
ds.make_handler_metadata(table="foo", schema="bar", config=config, **kwargs)
config["handler_type"] = handler_type
# Validate for all required keys
for k, v in {**base_config, **db_keys}.items():
config.pop(k)
with pytest.raises(ValidationError):
ds.make_handler_metadata(table="foo", schema="bar", config=config, **kwargs)
config[k] = v
handler = ds.make_handler_metadata(
table="foo", schema="bar", config=config, **kwargs
)
assert isinstance(handler, ds.Handler)
assert isinstance(handler.metadata, expected_type)
for k, v in {**base_config, **db_keys, **kwargs}.items():
assert handler.metadata.dict()[k] == v
def data_sources_enroll_iterator(
client: "ImmutaClient",
schema_table_mapping: Dict[str, List[Dict[str, str]]],
schema_obj: Dict[str, str],
config: Dict[str, Any],
) -> Iterator[Tuple[DataSource, Handler]]:
LOGGER.info("Processing schema_prefix: %s", schema_obj["schema_prefix"])
matches_prefix = partial(fnmatch.fnmatch, pat=schema_obj["schema_prefix"])
for schema, tables in keyfilter(matches_prefix, schema_table_mapping).items():
LOGGER.info("Processing schema: %s", schema)
for table in tables:
if not fnmatch.fnmatch(table["tableName"], schema_obj["table_prefix"]):
continue
LOGGER.info("Processing table: %s.%s", schema, table["tableName"])
handler = make_handler_metadata(
config=config, table=table["tableName"], schema=schema
)
columns = client.get_column_types(
config=config, data_source_type=config["handler_type"], handler=handler
)
data_source, handler = to_immuta_objects(
schema=schema, table=table["tableName"], columns=columns, config=config
)
yield (data_source, handler)