from config import Config
import io
from airflow.hooks.base import BaseHook
from airflow.providers.postgres.hooks.postgres import PostgresHook
from typing import Dict, Any, Optional, List
[docs]
class GDriveUtil:
"""
Google Drive utilities for create folder, upload file, download file.
The user must share folder or file with the service account email address
:return: _description_
:rtype: _type_
"""
[docs]
@staticmethod
def get_mimetype(service, file_id: str) -> str:
"""
Get the mimetype of a Google Drive file.
:param service: Google Drive API service object
:type service: object
:param file_id: the ID of the Google Drive file
:type file_id: str
:return: the mimetype of the file
:rtype: str
"""
file_metadata = (
service.files()
.get(
fileId=file_id,
fields="mimeType",
supportsAllDrives=True,
supportsTeamDrives=True,
)
.execute()
)
return file_metadata["mimeType"]
[docs]
@staticmethod
def is_drive_folder(service, google_drive_file_id: str) -> bool:
"""
Check if a Google Drive file is a folder.
:param service: Google Drive API service object
:type service: object
:param google_drive_file_id: the ID of the Google Drive file
:type google_drive_file_id: str
:return: whether the file is a folder
:rtype: bool
"""
file_metadata = (
service.files()
.get(
fileId=google_drive_file_id,
fields="mimeType",
supportsAllDrives=True,
supportsTeamDrives=True,
)
.execute()
)
return file_metadata["mimeType"] == "application/vnd.google-apps.folder"
[docs]
@staticmethod
def get_all_gdrive_files(
service, folder_id: str, recursive: bool = True
) -> List[str]:
"""
Get all file IDs from a Google Drive folder.
:param service: Google Drive API service object
:type service: object
:param folder_id: the ID of the Google Drive folder
:type folder_id: str
:param recursive: whether to recursively traverse subfolders (default: True)
:type recursive: bool
:return: the list of file IDs
:rtype: List[str]
"""
file_objects = []
try:
# Get all items (files and folders) in the current folder
query = f"'{folder_id}' in parents and trashed = false"
page_token = Config.INITIAL_GDRIVE_PAGE_TOKEN
while page_token:
print("Query:", query)
print("Page Token:", page_token)
results = (
service.files()
.list(
q=query,
fields="nextPageToken, files(id, name, mimeType)",
pageSize=1000,
pageToken=(
page_token
if page_token != Config.INITIAL_GDRIVE_PAGE_TOKEN
else None
),
supportsAllDrives=True,
corpora="allDrives",
includeItemsFromAllDrives=True,
)
.execute()
)
for item in results.get("files", []):
print(f'Item Name: {item["name"]}')
if item["mimeType"] == "application/vnd.google-apps.folder":
if recursive:
file_objects.extend(
GDriveUtil.get_all_gdrive_files(
service, item["id"], recursive=True
)
)
else:
file_objects.append(item)
page_token = results.get("nextPageToken")
return file_objects
except Exception as e:
print(f"Error accessing folder {folder_id}: {str(e)}")
return []
[docs]
@staticmethod
def get_gcp_credential(
service_account_conn_id: str = Config.GDRIVE_SERVICE_ACCOUNT,
scopes=[
"https://www.googleapis.com/auth/drive",
"https://www.googleapis.com/auth/spreadsheets",
],
) -> dict:
"""
Get the Google Cloud Platform credentials from the connection id.
The service account email address must be added to the folder or file to be accessed.
default service account is stored in AWS Secrets Manager.
default service account email address is:googledrive@inner-period-265200.iam.gserviceaccount.com
:param service_account_conn_id: the connection id for the service account (stored in AWS Secrets Manager)
:type service_account_conn_id: str
:return: the Google Cloud Platform credentials
:rtype: dict
"""
from google.oauth2 import service_account
import json
gdrive_hook = BaseHook.get_connection(service_account_conn_id)
gcp_json_credentials_dict = json.loads(gdrive_hook.extra)
credentials = service_account.Credentials.from_service_account_info(
gcp_json_credentials_dict, scopes=scopes
)
return credentials
[docs]
@staticmethod
def get_gspread_service(
service_account_conn_id: str = Config.GDRIVE_SERVICE_ACCOUNT,
) -> object:
"""
Construct a Google Spread Sheet service object using the sevice account.
The service account email address must be added to the folder or file to be accessed.
default service account is stored in aws secrets manager.
default service account email address is:googledrive@inner-period-265200.iam.gserviceaccount.com
:param service_account_conn_id: the connection id for the service account (stored in aws secrets manager)
:type service_account_conn_id: str
:return: a Google Sheet service object
:rtype: object
"""
import gspread
credentials = GDriveUtil.get_gcp_credential(service_account_conn_id)
gsp = gspread.authorize(credentials)
return gsp
[docs]
@staticmethod
def construct_gdrive_service(
service_account_conn_id: str = Config.GDRIVE_SERVICE_ACCOUNT,
) -> object:
"""Construct a Google Drive service object using the sevice account.
The service account email address must be added to the folder or file to be accessed.
default service account is stored in aws secrets manager.
default service account email address is:googledrive@inner-period-265200.iam.gserviceaccount.com
:param service_account_conn_id: the connection id for the service account (stored in aws secrets manager)
:type service_account_conn_id: str
:return: a Google Drive service object
:rtype: object
"""
from googleapiclient.discovery import build
credentials = GDriveUtil.get_gcp_credential(service_account_conn_id)
drive_service = build(
"drive",
"v3",
credentials=credentials,
)
return drive_service
[docs]
@staticmethod
def gdrive_create_folder(
folder_name: str,
drive_service: object,
parent_folder_id: str = None,
pageSize: int = 1000,
) -> str:
"""
Create a folder in Google Drive and return its ID. If already created, return the ID of the existing folder.
:param folder_name: the name of the folder to create under the parent folder
:type folder_name: str
:param drive_service: the Google Drive service object
:type drive_service: object
:param parent_folder_id: the parent folder ID. This can be None if the folder is to be created at the root level
:type parent_folder_id: str
:param pageSize: the page size for the list request
:type pageSize: int
:return: the ID of the created folder
:rtype: str
"""
folder_metadata = {
"name": folder_name,
"mimeType": "application/vnd.google-apps.folder",
"parents": [parent_folder_id] if parent_folder_id else [],
}
# check if folder exists
items = GDriveUtil.gdrive_list_folder(
drive_service,
parent_folder_id,
pageSize,
)
print(f"WARNING: currently only checking the {pageSize} items")
if not items:
print("No folders or files found in Google Drive parent folder.")
created_folder = (
drive_service.files()
.create(
body=folder_metadata,
fields="id",
supportsAllDrives=True,
)
.execute()
)
else:
print("Folders and files in Google Drive existing in the parent folder")
names_existing = [item["name"] for item in items]
created_folder = (
items[names_existing.index(folder_name)]
if folder_name in names_existing
else (
drive_service.files()
.create(
body=folder_metadata,
fields="id",
supportsAllDrives=True,
)
.execute()
)
)
print(f'Created Folder ID: {created_folder["id"]}')
return created_folder["id"]
[docs]
@staticmethod
def gdrive_delete_files(
drive_service: object,
file_or_folder_id: str,
) -> None:
"""
Delete a file or folder in Google Drive by ID.
:param drive_service: the Google Drive service object
:type drive_service: object
:param file_or_folder_id: the ID of the file or folder to delete
:type file_or_folder_id: str
:return: None
:rtype: None
"""
try:
drive_service.files().delete(
fileId=file_or_folder_id, supportsAllDrives=True
).execute()
print(f"Successfully deleted file/folder with ID: {file_or_folder_id}")
except Exception as e:
print(f"Error deleting file/folder with ID: {file_or_folder_id}")
print(f"Error details: {str(e)}")
[docs]
@staticmethod
def gdrive_list_folder(
drive_service: object,
parent_folder_id: str = None,
pageSize: int = 1000,
) -> list:
"""
List sub folders and files in a Google Drive parent folder.
:param drive_service: the Google Drive service object
:type drive_service: object
:param parent_folder_id: the ID of the parent folder. This can be None if the folder is to be created at the root level
:type parent_folder_id: str
:param pageSize: the page size for the list request. Default is 1000
:type pageSize: int
:return: the list of folders and files
:rtype: list
"""
results = (
drive_service.files()
.list(
q=(
f"'{parent_folder_id}' in parents and trashed=false"
if parent_folder_id
else None
),
pageSize=pageSize, # default is 100
fields="nextPageToken, files(id, name, mimeType)",
supportsAllDrives=True,
corpora="allDrives",
includeItemsFromAllDrives=True,
)
.execute()
)
items = results.get("files", [])
if not items:
print("No files or subfolders found in this Google Drive parent folder.")
else:
print("Folders and files in this Google Drive parent folder:")
for item in items:
print(
f"Name: {item['name']}, ID: {item['id']}, Type: {item['mimeType']}"
)
return items
[docs]
@staticmethod
def gdrive_download_file(
drive_service: object,
file_id: str,
destination_path: str,
) -> None:
"""
Download a file from Google Drive by its ID.
WARINING: needs to add the service account to the drive as a user with access to the file first. Risk of out of storage error.
:param drive_service: the Google Drive service object
:type drive_service: object
:param file_id: the ID of the file to download
:type file_id: str
:param destination_path: the destination path to save the file
:type destination_path: str
:return: None
:rtype: None
"""
from googleapiclient.http import MediaIoBaseDownload
request = drive_service.files().get_media(fileId=file_id)
fh = io.FileIO(destination_path, mode="wb")
downloader = MediaIoBaseDownload(fh, request)
done = False
while not done:
status, done = downloader.next_chunk()
print(f"Download {int(status.progress() * 100)}%.")
[docs]
@staticmethod
def gdrive_upload_file(
drive_service: object,
file_path: str,
file_name: str,
parent_folder_id: str = None,
overwrite: bool = False,
) -> str:
"""
Upload a file to Google Drive.
:param drive_service: the Google Drive service object
:type drive_service: object
:param file_path: the file path of the file to upload
:type file_path: str
:param file_name: the name of the file to upload
:type file_name: str
:param parent_folder_id: the ID of the parent folder. This can be None if the file is to be added to the root level.
:type parent_folder_id: str
:param overwrite: whether to overwrite the file if it already exists. Default is False, meaning a new file with exact same name will be created in same folder!
:type overwrite: bool
:return: the ID of the created file
:rtype: str
"""
from googleapiclient.http import MediaFileUpload
print(
f"""
- Now uploading file: {file_name} to Google Drive.
- Parent folder ID: {parent_folder_id}
- Overwrite: {overwrite}
"""
)
media = MediaFileUpload(file_path, resumable=True)
file_metadata = {
"name": file_name,
"parents": [parent_folder_id] if parent_folder_id else [],
}
if overwrite:
# STEP 1. check whether file_name already exists in the parent_folder
items = GDriveUtil.gdrive_list_folder(
drive_service,
parent_folder_id,
)
names_existing = [item["name"] for item in items]
if file_name in names_existing:
print(f"File {file_name} already exists in the parent folder")
file_id = items[names_existing.index(file_name)]["id"]
# STEP 2. Update the existing file
update_file = (
drive_service.files()
.update(
fileId=file_id,
body={"name": file_name},
media_body=media,
fields="id",
supportsAllDrives=True,
)
.execute()
)
print(
"File with same name exists. Updated file for the ID: ",
update_file["id"],
)
print(f'https://drive.google.com/file/d/{update_file["id"]}/view')
return update_file["id"]
else:
print("File with same name does not exist. Creating a new file.")
created_file = (
drive_service.files()
.create(
body=file_metadata,
media_body=media,
fields="id",
supportsAllDrives=True,
)
.execute()
)
print(f'Created a new File ID: {created_file["id"]}')
print(f'https://drive.google.com/file/d/{created_file["id"]}/view')
return created_file["id"]
else:
print(
"Developer does not want to overwrite the file. Creating a new file. "
)
created_file = (
drive_service.files()
.create(
body=file_metadata,
media_body=media,
fields="id",
supportsAllDrives=True,
)
.execute()
)
print(f'Created a new File ID: {created_file["id"]}')
return created_file["id"]
[docs]
@staticmethod
def gdrive_share_with_emails(
drive_service: object,
file_or_folder_id: str,
emails: list,
role: str = "reader",
send_notification_email: bool = False,
) -> None:
"""
Grant a list of emails access to a Google Drive file or folder.
Best-effort: a failure to share with one email (e.g. an invalid address
or a domain-restricted Shared Drive) is logged and skipped rather than
failing the whole task.
:param drive_service: the Google Drive service object
:type drive_service: object
:param file_or_folder_id: the ID of the file/folder to share
:type file_or_folder_id: str
:param emails: the email addresses to grant access to
:type emails: list
:param role: the Drive permission role (reader/commenter/writer). Default reader
:type role: str
:param send_notification_email: whether Google should send its own share
notification. Default False (the DAG sends its own email).
:type send_notification_email: bool
:return: None
:rtype: None
"""
for email in emails:
if not email:
continue
try:
drive_service.permissions().create(
fileId=file_or_folder_id,
body={"type": "user", "role": role, "emailAddress": email},
sendNotificationEmail=send_notification_email,
supportsAllDrives=True,
fields="id",
).execute()
print(f"Shared {file_or_folder_id} with {email} as {role}")
except Exception as e:
print(f"Error sharing {file_or_folder_id} with {email}: {str(e)}")
[docs]
@staticmethod
def gsheet_to_dataframe(
url,
sheet_name,
headerrow=1,
) -> object:
"""
Read the Google Sheet by URL and return a pandas DataFrame. The Sheet must be shared with the service account email address.
:param url: the URL of the Google Sheet
:type url: str
:param sheet_name: the name of the sheet to read
:type sheet_name: str
:param headerrow: the row number of the header. Default is 1
:type headerrow: int
:return df_spread: the pandas DataFrame of the Google Sheet
:rtype: object
:return worksheet: the Google Sheet worksheet object
:rtype: object
"""
import pandas as pd
gsp = GDriveUtil.get_gspread_service()
SHEET_ID = url.split("/")[5]
spreadsheet = gsp.open_by_key(SHEET_ID)
worksheet = spreadsheet.worksheet(sheet_name)
columns = worksheet.row_values(headerrow)
print(columns)
rows = worksheet.get_all_values()[headerrow:]
print(rows[:1])
df_spread = pd.DataFrame(rows, columns=columns)
return df_spread, worksheet
[docs]
@staticmethod
def create_new_gsheet(
url,
sheet_name,
rows,
cols,
) -> object:
"""
Create a new sheet in the Google Sheet by URL and return the worksheet object.
:param url: the URL of the Google Sheet
:type url: str
:param sheet_name: the name of the new sheet to create
:type sheet_name: str
:param rows: the number of rows in the new sheet
:type rows: int
:param cols: the number of columns in the new sheet
:type cols: int
:return worksheet: the Google Sheet worksheet object
:rtype: object
:return: the worksheet object
:rtype: object
"""
gsp = GDriveUtil.get_gspread_service()
SHEET_ID = url.split("/")[5]
spreadsheet = gsp.open_by_key(SHEET_ID)
worksheet = spreadsheet.add_worksheet(
title=sheet_name,
rows=rows,
cols=cols,
)
return worksheet
[docs]
@staticmethod
def export_query_to_gdrive(
query,
export_file_name,
gdrive_url="",
sub_gdrive_id="",
rename_dict={},
database_type="athena",
database=Config.DATA_LAKE_DB,
returning_df=False,
overwrite_dest_file=True,
) -> object:
"""
Export a query result to Google Drive as a CSV file.
:param query: the query to execute
:type query: str
:param export_file_name: the name of the file to export
:type export_file_name: str
:param gdrive_url: the Google Drive parent folder URL
:type gdrive_url: str
:param sub_gdrive_id: the Google Drive parent folder ID
:type sub_gdrive_id: str
:param rename_dict: the dictionary to rename the columns
:type rename_dict: dict
:param database_type: the database type (athena or postgres)
:type database_type: str
:param database: the database name
:type database: str
:param returning_df: whether to return the DataFrame
:type returning_df: bool
:param overwrite_dest_file: whether to overwrite the destination file if it already exists
:type overwrite_dest_file: bool
:return: the DataFrame if returning_df is True
:rtype: object
"""
from tempfile import NamedTemporaryFile
print(
f"""LOGGING ************************************************
Step 1: Querying {database_type}.{database} into dataframe with query:
{query}
"""
)
ATHENA_DB_TYPE, POSTGRES_DB_TYPE = "athena", "postgres"
if database_type == ATHENA_DB_TYPE:
import awswrangler as wr
df = wr.athena.read_sql_query(query, database=database)
elif database_type == POSTGRES_DB_TYPE:
import pandas as pd
pgeng = PostgresHook(postgres_conn_id=database).get_sqlalchemy_engine()
df = pd.read_sql(query, pgeng)
else:
raise ValueError(
f"database_type must be either '{POSTGRES_DB_TYPE}' or '{ATHENA_DB_TYPE}'; instead, input was: ",
database_type,
)
if rename_dict:
df.rename(columns=rename_dict, inplace=True)
# GDRIVE PATH FOR THIS PARENT ID: "Ready <> SBO File Transfer/general_query"
GDRIVE_DIR = (
"https://drive.google.com/drive/folders/1IQYsA1cF09ryLomA7lwwVFs9_DtlsZn5"
)
gdrive_url = gdrive_url if gdrive_url else GDRIVE_DIR
print(
f"""LOGGING ************************************************
Step 2: Upload dataframe to...
GDrive parent folder: {gdrive_url=}
File: {export_file_name=}
"""
)
drive_service = GDriveUtil.construct_gdrive_service()
if not sub_gdrive_id:
sub_gdrive_id = GDriveUtil.gdrive_create_folder(
drive_service=drive_service,
folder_name=(
Config.pendulum.now("US/Eastern").format("YYYY_MM_DD") + database
),
parent_folder_id=gdrive_url.split("/")[-1],
)
with NamedTemporaryFile() as temp:
df.to_csv(temp.name, index=False)
GDriveUtil.gdrive_upload_file(
drive_service=drive_service,
file_name=export_file_name,
file_path=temp.name,
parent_folder_id=sub_gdrive_id,
overwrite=overwrite_dest_file,
)
print(
"""
LOGGING ************************************************
FINISHED EXPORTING QUERY TO GDRIVE
"""
)
if returning_df:
return df
[docs]
@staticmethod
def write_df_to_gsheet(
df,
file_name: str,
sheet_name: str = "Sheet1",
parent_folder_id: str = Config.GDRIVE_FOLDER_GSHEET,
overwrite: bool = True,
) -> str:
import gspread
MAX_CELL_CHARS = 50000
def split_large_cells(df):
new_df = df.copy()
new_columns = []
for col in df.columns:
max_length = df[col].astype(str).map(len).max()
if max_length > MAX_CELL_CHARS:
parts = (max_length // MAX_CELL_CHARS) + 1
for i in range(parts):
if i == 0:
new_df[col] = df[col].astype(str).str[:MAX_CELL_CHARS]
else:
new_col = f"{col}_extra{i}"
new_df[new_col] = (
df[col]
.astype(str)
.str[i * MAX_CELL_CHARS : (i + 1) * MAX_CELL_CHARS]
)
new_columns.append(new_col)
return new_df, new_columns
if df.empty:
print("Error: Empty DataFrame provided.")
return None
df, new_columns = split_large_cells(df)
print(f"Writing DataFrame {df.shape} as '{file_name}' to Google Drive.")
print(f"Parent folder ID: {parent_folder_id}, Overwrite: {overwrite}")
print(f"Added {len(new_columns)} new columns to handle large cells.")
gsp = GDriveUtil.get_gspread_service()
drive_service = GDriveUtil.construct_gdrive_service()
try:
query = f"name='{file_name}' and '{parent_folder_id}' in parents and mimeType='application/vnd.google-apps.spreadsheet' and trashed=false"
files = (
drive_service.files()
.list(
q=query,
fields="files(id, name)",
supportsAllDrives=True,
includeItemsFromAllDrives=True,
)
.execute()
.get("files", [])
)
if files:
print(f"File {file_name} exists - open it")
spreadsheet = gsp.open_by_key(files[0]["id"])
try:
worksheet = spreadsheet.worksheet(sheet_name)
if overwrite:
worksheet.clear()
else:
sheet_name = f"{sheet_name}_{Config.pendulum.now('US/Eastern').format('YYYY-MM-DD HH:mm:ss zz')}"
worksheet = spreadsheet.add_worksheet(
title=sheet_name, rows=1, cols=1
)
except gspread.exceptions.WorksheetNotFound:
worksheet = spreadsheet.add_worksheet(
title=sheet_name, rows=1, cols=1
)
else:
print(f"File {file_name} doesn't exist - create new one")
spreadsheet = gsp.create(file_name)
drive_service.files().update(
fileId=spreadsheet.id,
addParents=parent_folder_id,
removeParents="root",
fields="id, parents",
supportsAllDrives=True,
).execute()
worksheet = spreadsheet.sheet1
worksheet.update_title(sheet_name)
# Update the worksheet with the new data
worksheet.resize(rows=df.shape[0] + 1, cols=df.shape[1])
# Convert DataFrame to list of lists
data = [df.columns.tolist()] + df.values.tolist()
# Update in batches to avoid API limits
batch_size = 1000
for i in range(0, len(data), batch_size):
worksheet.update(
f"A{i+1}", data[i : i + batch_size], value_input_option="RAW"
)
if sheet_name != "Sheet1":
try:
sheet1 = spreadsheet.worksheet("Sheet1")
if not sheet1.get_all_values():
spreadsheet.del_worksheet(sheet1)
except gspread.exceptions.WorksheetNotFound:
pass
print(
f"Spreadsheet URL: https://docs.google.com/spreadsheets/d/{spreadsheet.id}/edit"
)
return spreadsheet.id
except Exception as e:
print(f"Error: {str(e)}")
return None
[docs]
@staticmethod
def style_google_sheet(worksheet) -> None:
"""
Apply professional styling to a Google Sheets worksheet.
Args:
worksheet: gspread worksheet object
"""
# Get the number of columns
header_row = worksheet.row_values(1)
num_columns = len(header_row)
total_rows = worksheet.row_count
# Format header row
header_format = {
"backgroundColor": {"red": 0.27, "green": 0.45, "blue": 0.77}, # Blue color
"textFormat": {
"foregroundColor": {
"red": 1.0,
"green": 1.0,
"blue": 1.0,
}, # White text
"fontSize": 11,
"bold": True,
},
"horizontalAlignment": "CENTER",
"verticalAlignment": "MIDDLE",
"wrapStrategy": "WRAP",
}
# Apply header formatting
worksheet.format("1:1", header_format)
# Format data rows
data_format = {
"textFormat": {"fontSize": 10},
"horizontalAlignment": "LEFT",
"verticalAlignment": "MIDDLE",
}
# Apply data formatting
if total_rows > 1:
worksheet.format(f"2:{total_rows}", data_format)
# Remove any existing banding first (separate batch update)
delete_requests = []
try:
# Get all banded ranges
sheet_metadata = worksheet.spreadsheet.fetch_sheet_metadata()
for sheet in sheet_metadata["sheets"]:
if sheet["properties"]["sheetId"] == worksheet.id:
if "bandedRanges" in sheet:
for banded_range in sheet["bandedRanges"]:
delete_requests.append(
{
"deleteBanding": {
"bandedRangeId": banded_range["bandedRangeId"]
}
}
)
# Execute delete requests if any exist
if delete_requests:
worksheet.spreadsheet.batch_update({"requests": delete_requests})
print(f"Removed {len(delete_requests)} existing banded range(s)")
except Exception as e:
print(f"Note: Could not check/remove existing banding: {e}")
# Build batch update requests
requests = []
# Set header row height
requests.append(
{
"updateDimensionProperties": {
"range": {
"sheetId": worksheet.id,
"dimension": "ROWS",
"startIndex": 0,
"endIndex": 1,
},
"properties": {"pixelSize": 40},
"fields": "pixelSize",
}
}
)
# Add alternating row colors (banding) - FIXED VERSION
requests.append(
{
"addBanding": {
"bandedRange": {
"range": {
"sheetId": worksheet.id,
"startRowIndex": 1, # Start after header (row 2)
"endRowIndex": total_rows,
"startColumnIndex": 0,
"endColumnIndex": num_columns,
},
"rowProperties": {
# Remove headerColor - we don't want it in the data rows
"firstBandColor": {
"red": 1.0,
"green": 1.0,
"blue": 1.0,
}, # White
"secondBandColor": {
"red": 0.95,
"green": 0.95,
"blue": 0.95,
}, # Light gray
},
}
}
}
)
# Set column widths
column_widths = {
"version_date": 150,
"environment": 150,
"status": 120,
"updated_at": 200,
"created_at": 200,
"dp-ro relationship": 180,
"sa-relationship": 180,
"source_folder": 350,
}
for idx, header in enumerate(header_row):
width = column_widths.get(header, 150)
requests.append(
{
"updateDimensionProperties": {
"range": {
"sheetId": worksheet.id,
"dimension": "COLUMNS",
"startIndex": idx,
"endIndex": idx + 1,
},
"properties": {"pixelSize": width},
"fields": "pixelSize",
}
}
)
# Freeze header row
requests.append(
{
"updateSheetProperties": {
"properties": {
"sheetId": worksheet.id,
"gridProperties": {"frozenRowCount": 1},
},
"fields": "gridProperties.frozenRowCount",
}
}
)
# Add borders to all cells
requests.append(
{
"updateBorders": {
"range": {
"sheetId": worksheet.id,
"startRowIndex": 0,
"endRowIndex": total_rows,
"startColumnIndex": 0,
"endColumnIndex": num_columns,
},
"top": {
"style": "SOLID",
"width": 1,
"color": {"red": 0.83, "green": 0.83, "blue": 0.83},
},
"bottom": {
"style": "SOLID",
"width": 1,
"color": {"red": 0.83, "green": 0.83, "blue": 0.83},
},
"left": {
"style": "SOLID",
"width": 1,
"color": {"red": 0.83, "green": 0.83, "blue": 0.83},
},
"right": {
"style": "SOLID",
"width": 1,
"color": {"red": 0.83, "green": 0.83, "blue": 0.83},
},
"innerHorizontal": {
"style": "SOLID",
"width": 1,
"color": {"red": 0.83, "green": 0.83, "blue": 0.83},
},
"innerVertical": {
"style": "SOLID",
"width": 1,
"color": {"red": 0.83, "green": 0.83, "blue": 0.83},
},
}
}
)
# Execute all requests in a single batch
worksheet.spreadsheet.batch_update({"requests": requests})
[docs]
@staticmethod
def append_row_to_google_sheets(
spreadsheet_id_or_url: str,
row_data: Dict[str, Any],
expected_column_order: Optional[List[str]] = None,
sheet_names: Optional[List[str]] = None,
style_sheet: bool = True,
) -> None:
"""
Append a single row to one or more sheets in a Google Sheets file.
Args:
spreadsheet_id_or_url: Google Sheets ID or full URL
row_data: Dictionary with column names as keys and values to append
sheet_names: List of sheet names to append to. If None, appends to all sheets.
credentials_file: Path to Google service account credentials JSON file.
If None, uses default credentials.
Example:
row_data = {
'version_date': '2024-01-08',
'environment': 'production',
'status': 'active',
'updated_at': '2024-01-08 10:30:00',
'created_at': '2024-01-08 10:30:00',
'dp-ro relationship': '1DP-1RP',
'sa-relationship': 'linked',
'source_folder': 's3://bucket/folder'
}
append_row_to_google_sheets(
'1ABC123xyz...',
row_data,
sheet_names=['AR', 'MN'],
credentials_file='credentials.json'
)
"""
from datetime import datetime
# Authorize the client
client = GDriveUtil.get_gspread_service()
# Open the spreadsheet
if "docs.google.com" in spreadsheet_id_or_url:
# Extract ID from URL
spreadsheet = client.open_by_url(spreadsheet_id_or_url)
else:
spreadsheet = client.open_by_key(spreadsheet_id_or_url)
# Get all worksheets or specific ones
all_worksheets = spreadsheet.worksheets()
if sheet_names:
worksheets_to_update = [
ws for ws in all_worksheets if ws.title in sheet_names
]
else:
worksheets_to_update = all_worksheets
missing_sheets = set(sheet_names or []) - {
ws.title for ws in worksheets_to_update
}
if missing_sheets:
print(
f"Warning: The following sheets were not found: {', '.join(missing_sheets)}"
)
# create missing sheets
for sheet_name in missing_sheets:
worksheets_to_update.append(
spreadsheet.add_worksheet(title=sheet_name, rows="100", cols="20")
)
for worksheet in worksheets_to_update:
# Get header row
header_row = worksheet.row_values(1)
# if header is empty, set it
if not header_row:
worksheet.append_row(
expected_column_order, value_input_option="USER_ENTERED"
)
header_row = expected_column_order
# Build row values matching header order
row_values = []
for col in header_row:
value = row_data.get(col)
# Convert datetime objects to strings
if isinstance(value, datetime):
value = value.strftime("%Y-%m-%d %H:%M:%S")
row_values.append(value if value is not None else "")
# Append the row
worksheet.append_row(row_values, value_input_option="USER_ENTERED")
print(f"Appended row to sheet '{worksheet.title}'")
if style_sheet:
GDriveUtil.style_google_sheet(worksheet)
print(f"Successfully updated {len(worksheets_to_update)} sheet(s)")