def generate_sql_and_query()

in backend_apis/app/utils_codey.py [0:0]

• 35 lines of code
• 2 McCabe index (conditional complexity)

def generate_sql_and_query(
        llm,
        datacatalog_client,
        prompt_template: str,
        query_metadata: str,
        question: str,
        project_id: str,
        dataset_id: str,
        tag_template_name: str,
        bqclient: bigquery.Client):
    """Generates a GoogleSQL query and executes it against a BigQuery dataset.

    Args:
        query: 
            The initial query text.
        project_id: 
            The ID of the BigQuery project.
        dataset_id: 
            The ID of the BigQuery dataset.
        tag_template_name: 
            The name of the tag template to use for the query.
        bqclient: 
            A BigQuery client object.

    Returns:
        A DataFrame containing the results of the query.

    Raises:
        NotFoundError: If the dataset or table is not found.
        BadRequestError: If the query is invalid.
    """

    metadata = get_metadata_from_dataset(
        bqclient=bqclient,
        datacatalog_client=datacatalog_client,
        query=query_metadata, 
        project_id=project_id, 
        dataset_id=dataset_id,
        tag_template_name=tag_template_name)

    prompt = generate_prompt(
        question, 
        metadata,
        prompt_template,
        project_id)

    gen_code = llm.generate_content(
        prompt,
        generation_config={ "temperature": 0.3,"max_output_tokens": 1024}
    ).text.replace("```","")
    
    gen_code = gen_code[gen_code.find("SELECT"):]
    result = []
    try:
        result_job = bqclient.query(gen_code)
        for row in result_job:
            result.append(dict(row.items()))
    except:
        raise Exception(f"We couldn't complete your request due to an error in the generated query. Query: {gen_code} ")
    
    return result, gen_code, prompt