def generate_pk_query()

in nl2sql/datasets/custom.py [0:0]


def generate_pk_query(dataset_id: str, tablename: str, primary_key_column: str) -> str:
    """
    Generate DDL queries to add associated primary key columns to respective
    tables.

    Args:
        dataset_id (str): Bigquery dataset id.
        tablename (str): Bigquery table name.
        primary_key_column (str): Name of the primary key column in table.

    Returns:
        query (str): DDL query to add primary key to table.
    """
    query = (
        f"ALTER TABLE `{dataset_id}.{tablename}` "
        f"ADD PRIMARY KEY({primary_key_column}) NOT ENFORCED;"
    )
    return query