#!/usr/bin/env python
# Copyright (c) TripleBlind Holdings, Inc. Confidential and Proprietary. All rights reserved.

import tripleblind as tb


tb.initialize(api_token=tb.config.example_user2["token"], example=True)

# A second organization (the current user's org) accesses the report which was
# defined by the first organization.
report_asset = tb.Asset.find("Company Payroll Diversity Report by Ethnicity")
if not report_asset:
    print("ERROR: Report not found.")
    raise SystemExit("You must run 1_owner_create_query first")

# NOTE:
# The report could also be shared without listing the asset if the first
# organization just sends the asset ID to the second organization.  Then they
# could just use an asset id sent by the data owner:
#
# report_asset = tb.Asset("a1b2c3d4-e5f6-a1b2-c3d4-e5f6a1b2c3d4")


# Create a job which runs the report
job = tb.create_job(
    job_name="Execute report",
    operation=tb.Operation.BLIND_QUERY,
    dataset=report_asset,
    preprocessor=tb.TabularPreprocessor.builder().all_columns(True),
)

if job.submit():
    job.wait_for_completion()

    if job.success:
        print("\n\nQuery results:")
        # tb.TableAsset.cast(job.result.asset).print_content()
        table = tb.TableAsset.cast(job.result.asset)
        table.pretty_print(col_names=["Ethnicity", "Count", "Avg Salary"])

        # Also save for this example
        filename = table.retrieve(save_as="out.csv", overwrite=True)
        print("\nResults also saved in file:", filename)
    else:
        raise SystemExit("Blind Query failed.")
