def datahub_api_should_use_rest_sink_as_default(
cls, v: Optional[DatahubClientConfig], values: Dict[str, Any],
**kwargs: Any) -> Optional[DatahubClientConfig]:
if v is None:
if values["sink"].type is not None:
sink_type = values["sink"].type
if sink_type == "datahub-rest":
sink_config = values["sink"].config
v = DatahubClientConfig.parse_obj(sink_config)
return v
logging.basicConfig(level=logging.INFO)
# Inputs -> owner, ownership_type, dataset
owner_to_add = make_user_urn("jdoe")
ownership_type = OwnershipTypeClass.DATAOWNER
dataset_urn = make_dataset_urn(platform="hive",
name="realestate_db.sales",
env="PROD")
# Some objects to help with conditional pathways later
owner_class_to_add = OwnerClass(owner=owner_to_add, type=ownership_type)
ownership_to_add = OwnershipClass(owners=[owner_class_to_add])
# First we get the current owners
gms_endpoint = "http://localhost:8080"
graph = DataHubGraph(DatahubClientConfig(server=gms_endpoint))
current_owners: Optional[OwnershipClass] = graph.get_aspect_v2(
entity_urn=dataset_urn,
aspect="ownership",
aspect_type=OwnershipClass,
)
need_write = False
if current_owners:
if (owner_to_add, ownership_type) not in [(x.owner, x.type)
for x in current_owners.owners]:
# owners exist, but this owner is not present in the current owners
current_owners.owners.append(owner_class_to_add)
need_write = True
else:
class DatahubIngestionReportingProviderConfig(IngestionReportingProviderConfig
):
datahub_api: Optional[DatahubClientConfig] = DatahubClientConfig()
# 1. Datahub Checkpointing State Provider Config test params
checkpointing_provider_config_test_params: Dict[
str, Tuple[Type[DatahubIngestionStateProviderConfig], Dict[str, Any],
Optional[DatahubIngestionStateProviderConfig], bool, ], ] = {
# Full custom-config
"checkpointing_valid_full_config": (
DatahubIngestionStateProviderConfig,
{
"datahub_api": datahub_client_configs["full"],
},
DatahubIngestionStateProviderConfig(
datahub_api=DatahubClientConfig(
server="http://localhost:8080",
token="dummy_test_tok",
timeout_sec=10,
extra_headers={},
max_threads=1,
), ),
False,
),
# Simple config
"checkpointing_valid_simple_config": (
DatahubIngestionStateProviderConfig,
{
"datahub_api": datahub_client_configs["simple"],
},
DatahubIngestionStateProviderConfig(
datahub_api=DatahubClientConfig(
server="http://localhost:8080", ), ),
False,
class DatahubIngestionStateProviderConfig(IngestionCheckpointingProviderConfig
):
datahub_api: Optional[DatahubClientConfig] = DatahubClientConfig()
class DatahubIngestionStateProviderConfig(ConfigModel):
datahub_api: Optional[DatahubClientConfig] = DatahubClientConfig()