save_data.py

import os
from google.oauth2 import service_account
from googleapiclient.discovery import build
import json


import os

def load_envs(local=False):
    """Load the environment variables."""
    if local:
        from dotenv import load_dotenv
        load_dotenv(override=True)
    service_account_info = json.loads(os.environ['GOOGLE_APPLICATION_CREDENTIALS_JSON'])
    SHEET_ID1 = os.environ['SHEET_ID1'] # human-ai-sequential
    SHEET_ID2 = os.environ['SHEET_ID2'] # human-ai-parallel
    SHEET_ID3 = os.environ['SHEET_ID3'] # ai-human-sequential
    SHEET_IDs = [SHEET_ID1, SHEET_ID2, SHEET_ID3]
    return service_account_info, SHEET_IDs


def get_sheet_service(local=False):
    """Get the google sheet service object."""
    service_account_info, SHEET_ID = load_envs(local=local)
    # verify the service_account_info
    credentials = service_account.Credentials.from_service_account_info(
        service_account_info,
        scopes=['https://www.googleapis.com/auth/spreadsheets']
    )

    # build the service object
    service = build('sheets', 'v4', credentials=credentials)
    return service, SHEET_ID

def col_letter(col_num):
    """Convert a column number to a column letter (1-indexed)."""
    letter = ''
    while col_num > 0:
        col_num, remainder = divmod(col_num - 1, 26)
        letter = chr(65 + remainder) + letter
    return letter

def col_letter(col_num):
    """Convert a column number to its corresponding Excel-style letter."""
    letter = ''
    while col_num > 0:
        col_num, remainder = divmod(col_num - 1, 26)
        letter = chr(65 + remainder) + letter
    return letter

def add_new_data(new_row, service, SPREADSHEET_ID, num_of_columns = 5):
    """Add new data to the spreadsheet.
    new_row: list of data to be added. """
    # read the existing data
    range_to_read = f'Sheet1!A:{col_letter(num_of_columns)}'

    result = service.spreadsheets().values().get(
        spreadsheetId=SPREADSHEET_ID,
        range=range_to_read
    ).execute()
 

    # search for first columns match
    match_found = False
    update_idx = None
    for idx, row in enumerate(result.get('values', [])):
        if row[0] == new_row[0]:
            match_found = True
            range_to_write = f'Sheet1!A{idx + 1}:{col_letter(num_of_columns)}{idx + 1}'
            update_idx = idx + 1
            break

    values = result.get('values', [])
    number_of_rows = len(values)
    new_row = [new_row]
    if not match_found:
        range_to_write = f'Sheet1!A{number_of_rows + 1}'

    request_body = {
        'values': new_row
    }
    if match_found:
        service.spreadsheets().values().clear(
            spreadsheetId=SPREADSHEET_ID,
            range=range_to_write,
            body={},
        ).execute()
        response = service.spreadsheets().values().update(
            spreadsheetId=SPREADSHEET_ID,
            range=range_to_write,
            valueInputOption='RAW',
            body=request_body
        ).execute()
        print(f"Updated row at position {update_idx}")
    else:
        response = service.spreadsheets().values().append(
            spreadsheetId=SPREADSHEET_ID,
            range=range_to_write,
            valueInputOption='RAW',
            insertDataOption='INSERT_ROWS',
            body=request_body
        ).execute()
        print(f"Added new row at position {number_of_rows + 1}")



if __name__ == "__main__":
    service, SHEET_ID = get_sheet_service(local=True)
    new_row = ["test1", "test2", "test3", "test4", "test5"]
    add_new_data(new_row, service, SHEET_ID)