import concurrent.futures
import os
import sys 
from typing import Dict, Tuple

# from datahub.cli import cli_utils
from datahub.cli.cli_utils import post_entity

MAX_WORKERS: int = (os.cpu_count() or 8) * 3 


def reset_aspect(urn: str) -> Tuple[str, Dict]:

    status = post_entity(
        urn=urn,
        aspect_name="upstreamLineage",
        entity_type="dataset",
        aspect_value={"upstreams": []},
    )   
    return (urn, status)


def reset_upstream_aspect(file_name: str) -> None:
    print(f"Using {MAX_WORKERS} workers.")
    with open(file_name, "r") as f:
        with concurrent.futures.ThreadPoolExecutor(
            max_workers=MAX_WORKERS
        ) as async_executor:
            reset_futures = [ 
                async_executor.submit(reset_aspect, urn.strip().strip('"'))
                for urn in f.readlines()
            ]   
            for future in concurrent.futures.as_completed(reset_futures):
                urn, status = future.result()
                print(f"Update succeeded for {urn} with status {status}")


if __name__ == "__main__":
    reset_upstream_aspect(sys.argv[1])

select distinct urn from metadata_aspect_v2 where urn like '%your-pattern-here%'