Braulio F.

@tool
def inspect_db(dataset: str) -> list:
    """
    Allows you to inspect the database.
    Args:
        dataset: The id of the dataset to inspect, without the project name. For example, if the dataset_id is bigquery-public-data.stackoverflow, you should pass stackoverflow.
    Returns:
        A dictionary of table names as keys, and a dictionary of column names and types as values.
    """

    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.getenv("GCP_SERVICE_ACCOUNT_CREDENTIALS")
    client = bigquery.Client()

    # Get dataset reference
    dataset_ref = client.dataset(dataset, project='bigquery-public-data')
    dataset = client.get_dataset(dataset_ref)

    # List tables
    tables = list(client.list_tables(dataset))
    schema_dict = {}
    schema_table = {}
    for table in tables:
        print(f"PROCESSING TABLE:\n{table}")
        # Get schema for each table
        table_ref = dataset_ref.table(table.table_id)
        table = client.get_table(table_ref)
        for schema_field in table.schema:
            schema_table[schema_field.name] = schema_field.field_type
        schema_dict[table.table_id] = schema_table
        schema_table = {}

    schema_str = str(schema_dict)
    return schema_str