openrvdas.logger.writers.google_sheets_writer

  1#!/usr/bin/env python3
  2"""
  3GoogleSheetsWriter - A Python class for writing dictionaries to Google Sheets
  4
  5This module provides a simple interface for writing Python dictionaries as rows
  6to Google Sheets, with automatic column creation and management.
  7
  8Author: Assistant
  9License: MIT
 10"""
 11
 12import os
 13
 14# Don't barf if they don't have google api packages installed. Only complain
 15# if they actually try to use it, below.
 16try:
 17    from googleapiclient.discovery import build
 18    from google.oauth2.service_account import Credentials
 19    from google.auth.transport.requests import Request
 20    from google.oauth2.credentials import Credentials as UserCredentials
 21    GOOGLE_SHEETS_ENABLED = True
 22except ModuleNotFoundError:
 23    GOOGLE_SHEETS_ENABLED = False
 24
 25
 26from logger.utils.das_record import DASRecord  # noqa: E402
 27from logger.writers.writer import Writer  # noqa: E402
 28
 29
 30class GoogleSheetsWriter(Writer):
 31    """
 32    A class for writing Python dictionaries as rows to Google Sheets
 33    with automatic column management.
 34
 35    This class handles authentication, column creation, and data writing to
 36    Google Sheets. Each dictionary key becomes a column header, and each
 37    dictionary becomes a row. New columns are automatically created when new
 38    keys are encountered. Numeric values are preserved as numbers in the
 39    spreadsheet.
 40
 41    SETUP INSTRUCTIONS:
 42
 43    1. Install Required Packages:
 44       pip install google-api-python-client google-auth-httplib2 \
 45           google-auth-oauthlib
 46
 47    2. Set Up Google Cloud Project:
 48       - Go to https://console.cloud.google.com/
 49       - Create a new project or select existing one
 50       - Note your project ID
 51
 52    3. Enable Google Sheets API:
 53       - In Cloud Console, go to "APIs & Services" > "Library"
 54       - Search for "Google Sheets API"
 55       - Click on it and press "Enable"
 56
 57    4. Create Service Account:
 58       - Go to "APIs & Services" > "Credentials"
 59       - Click "Create Credentials" > "Service Account"
 60       - Enter service account name (e.g., "sheets-writer-bot")
 61       - Click "Create and Continue"
 62       - Skip role assignment (click "Continue")
 63       - Skip user access (click "Done")
 64
 65    5. Generate JSON Key File:
 66       - Find your service account in the credentials list
 67       - Click on it to open details
 68       - Go to "Keys" tab
 69       - Click "Add Key" > "Create New Key"
 70       - Select "JSON" format and click "Create"
 71       - Save the downloaded JSON file securely (never commit to version
 72         control!)
 73
 74    6. Share Your Spreadsheet:
 75       - Open your Google Sheets document
 76       - Click "Share" button
 77       - Enter the service account email (client_email from JSON file)
 78       - Give "Editor" permissions
 79       - Uncheck "Notify people"
 80       - Click "Share"
 81
 82    EXAMPLE USAGE:
 83
 84        # Initialize writer with force_create enabled
 85        writer = GoogleSheetsWriter(
 86            sheet_name_or_id="1ABC123def456...",  # Your spreadsheet ID
 87            auth_key_path="path/to/service-account-key.json",
 88            use_service_account=True,
 89            worksheet_name="MyData",
 90            force_create=True  # Create "MyData" worksheet if needed
 91        )
 92
 93        # Write single dictionary (timestamp automatically goes first)
 94        writer.write({'timestamp': 1234567890, 'name': 'John', 'age': 30})
 95
 96        # Write single DASRecord (uses DASRecord.timestamp automatically)
 97        from das_record import DASRecord
 98        das_record = DASRecord(timestamp=1234567891,
 99                               fields={'name': 'Jane', 'age': 25})
100        writer.write(das_record)
101
102        # Write multiple records (mixed types)
103        records = [
104            {'timestamp': 1234567892, 'name': 'Bob', 'city': 'NYC',
105             'salary': 75000.50},
106            DASRecord(timestamp=1234567893,
107                      fields={'name': 'Alice', 'department': 'Engineering',
108                              'rating': 4.8})
109        ]
110        writer.write(records)
111
112    SECURITY NOTES:
113    - Never commit the JSON key file to version control
114    - Store the key file securely and use environment variables in production
115    - Limit service account permissions to only what's needed
116    - Consider key rotation for long-term production use
117
118    Attributes:
119        sheet_id (str): The Google Sheets spreadsheet ID
120        headers (list): Current column headers in the sheet
121        worksheet_name (str): Name of the worksheet/tab being used
122        force_create (bool): Whether to create the worksheet if it doesn't exist
123        service: Google Sheets API service object
124    """
125
126    def __init__(self, sheet_name_or_id, auth_key_path=None,
127                 use_service_account=True, worksheet_name="Sheet1",
128                 force_create=False, **kwargs):
129        """
130        Initialize GoogleSheetsWriter with spreadsheet name/ID and auth.
131
132        Args:
133            sheet_name_or_id (str): Google Sheets spreadsheet ID, full URL,
134                                    or name.
135            auth_key_path (str): Path to service account JSON key file or
136                                 OAuth credentials file. If None and
137                                 use_service_account=False, looks for
138                                 'token.json'.
139            use_service_account (bool): True for service account auth, False
140                                        for OAuth user auth.
141            worksheet_name (str): Name of the specific worksheet/tab to write
142                                  to (default: "Sheet1").
143            force_create (bool): If True, create the worksheet if it doesn't
144                                 exist (default: False).
145
146        Raises:
147            ValueError: If authentication parameters are invalid
148            Exception: If unable to authenticate or access the spreadsheet
149        """
150        if not GOOGLE_SHEETS_ENABLED:
151            raise ModuleNotFoundError(
152                'GoogleSheetsWriter requires google-api-python-client, '
153                'google-auth-httplib2 and google-auth-oauthlib.\n'
154                'Please run: pip install google-api-python-client '
155                'google-auth-httplib2 google-auth-oauthlib'
156            )
157        super().__init__(**kwargs)  # processes 'quiet', encoding and hints
158
159        self.sheet_name_or_id = sheet_name_or_id
160        self.auth_key_path = auth_key_path
161        self.use_service_account = use_service_account
162        self.worksheet_name = worksheet_name
163        self.force_create = force_create
164        self.service = None
165        self.sheet_id = None
166        self.headers = []
167
168        # Initialize the service
169        self._authenticate()
170        self._get_sheet_id()
171        self._ensure_worksheet_exists()
172        self._load_existing_headers()
173
174    def _authenticate(self):
175        """
176        Authenticate and build the Google Sheets service.
177
178        Sets up authentication using either service account credentials or
179        OAuth user credentials. Creates the Google Sheets API service object.
180
181        Raises:
182            ValueError: If authentication credentials are missing or invalid
183        """
184        scopes = ['https://www.googleapis.com/auth/spreadsheets']
185
186        if self.use_service_account:
187            # Service account authentication
188            if not self.auth_key_path:
189                raise ValueError("Service account key file path is required")
190
191            credentials = Credentials.from_service_account_file(
192                self.auth_key_path, scopes=scopes
193            )
194        else:
195            # OAuth user credentials
196            creds = None
197            token_path = self.auth_key_path or 'token.json'
198
199            if os.path.exists(token_path):
200                creds = UserCredentials.from_authorized_user_file(
201                    token_path, scopes)
202
203            if not creds or not creds.valid:
204                if creds and creds.expired and creds.refresh_token:
205                    creds.refresh(Request())
206                else:
207                    raise ValueError("Valid OAuth credentials not found. "
208                                     "Run OAuth flow first.")
209
210            credentials = creds
211
212        self.service = build('sheets', 'v4', credentials=credentials)
213
214    def _get_sheet_id(self):
215        """
216        Extract or find the spreadsheet ID from various input formats.
217
218        Handles different input formats:
219        - Direct spreadsheet ID: "1ABC123def456..."
220        - Full URL: "https://docs.google.com/spreadsheets/d/1ABC.../edit"
221        - Assumes other inputs are spreadsheet IDs
222
223        Sets self.sheet_id to the extracted spreadsheet ID.
224        """
225        # If it looks like a spreadsheet ID (long alphanumeric string)
226        if len(self.sheet_name_or_id) > 20 and '/' not in self.sheet_name_or_id:
227            self.sheet_id = self.sheet_name_or_id
228        else:
229            # If it's a full URL, extract the ID
230            if 'docs.google.com/spreadsheets/d/' in self.sheet_name_or_id:
231                part = self.sheet_name_or_id.split('/d/')[1]
232                self.sheet_id = part.split('/')[0]
233            else:
234                # Assume it's a spreadsheet ID
235                self.sheet_id = self.sheet_name_or_id
236
237    def _ensure_worksheet_exists(self):
238        """
239        Ensure the specified worksheet exists, creating it if necessary.
240
241        Checks if the worksheet exists in the spreadsheet. If force_create
242        is True and the worksheet doesn't exist, creates it.
243
244        Raises:
245            Exception: If worksheet doesn't exist and force_create is False,
246                       or if creation fails.
247        """
248        try:
249            # Get spreadsheet metadata to check existing worksheets
250            spreadsheet = self.service.spreadsheets().get(
251                spreadsheetId=self.sheet_id
252            ).execute()
253
254            # Check if the worksheet already exists
255            existing_sheets = [sheet['properties']['title']
256                               for sheet in spreadsheet.get('sheets', [])]
257
258            if self.worksheet_name in existing_sheets:
259                # Worksheet exists, nothing to do
260                return
261
262            if not self.force_create:
263                # Worksheet doesn't exist and we're not allowed to create it
264                raise Exception(f"Worksheet '{self.worksheet_name}' does not "
265                                "exist in spreadsheet. Set force_create=True "
266                                "to create it automatically.")
267
268            # Create the worksheet
269            requests = [{
270                'addSheet': {
271                    'properties': {
272                        'title': self.worksheet_name
273                    }
274                }
275            }]
276
277            body = {'requests': requests}
278
279            self.service.spreadsheets().batchUpdate(
280                spreadsheetId=self.sheet_id,
281                body=body
282            ).execute()
283
284            print(f"Created worksheet '{self.worksheet_name}' in spreadsheet")
285
286        except Exception as e:
287            if "force_create=True" in str(e):
288                # Re-raise our custom error message
289                raise e
290            else:
291                raise Exception(f"Failed to ensure worksheet exists: {str(e)}")
292
293    def _load_existing_headers(self):
294        """
295        Load existing column headers from the first row of the worksheet.
296
297        Reads the first row of the specified worksheet to get current column
298        headers. If the sheet is empty or doesn't exist, initializes with
299        empty headers list.
300
301        Sets self.headers to the list of existing column headers.
302        """
303        try:
304            # Get the first row to see existing headers
305            range_name = f"{self.worksheet_name}!1:1"
306            result = self.service.spreadsheets().values().get(
307                spreadsheetId=self.sheet_id,
308                range=range_name
309            ).execute()
310
311            values = result.get('values', [])
312            if values:
313                self.headers = values[0]
314            else:
315                self.headers = []
316
317        except Exception as e:
318            print(f"Warning: Could not load existing headers: {str(e)}")
319            self.headers = []
320
321    def _update_headers(self, new_keys):
322        """
323        Add new column headers for any keys not already present.
324
325        Args:
326            new_keys (iterable): Keys from a dictionary that should be headers
327
328        Side Effects:
329            - Updates self.headers with any new column names
330            - Updates the first row of the spreadsheet with new headers
331        """
332        # Find keys that aren't already headers
333        new_headers = [key for key in new_keys if key not in self.headers]
334
335        if new_headers:
336            # Add new headers to our list
337            self.headers.extend(new_headers)
338
339            # Update the header row in the sheet
340            range_name = f"{self.worksheet_name}!1:1"
341            body = {
342                'values': [self.headers]
343            }
344
345            self.service.spreadsheets().values().update(
346                spreadsheetId=self.sheet_id,
347                range=range_name,
348                valueInputOption='RAW',
349                body=body
350            ).execute()
351
352    def _get_next_row(self):
353        """
354        Find the next empty row number to write data to.
355
356        Scans column A to find the last row with data and returns the next
357        row number.
358
359        Returns:
360            int: The row number (1-indexed) where the next data should be
361                 written. Returns 2 if unable to determine (assumes headers).
362        """
363        try:
364            # Get all data to find the last row with content
365            range_name = f"{self.worksheet_name}!A:A"
366            result = self.service.spreadsheets().values().get(
367                spreadsheetId=self.sheet_id,
368                range=range_name
369            ).execute()
370
371            values = result.get('values', [])
372            return len(values) + 1  # Next row after the last one with data
373
374        except Exception:
375            return 2  # Start at row 2 (after headers) if we can't determine
376
377    def _normalize_record(self, record):
378        """
379        Convert a record (dict or DASRecord) to a standardized dictionary.
380
381        Args:
382            record: Either a dictionary or a DASRecord object
383
384        Returns:
385            dict: Standardized dictionary with timestamp and other fields
386        """
387        if isinstance(record, DASRecord):
388            result_dict = record.fields.copy()
389            result_dict['timestamp'] = record.timestamp
390            return result_dict
391        elif isinstance(record, dict):
392            # It's already a dictionary
393            return record.copy()
394        else:
395            raise ValueError(f"Unsupported record type: {type(record)}")
396
397    def _ensure_timestamp_first(self, keys):
398        """
399        Ensure 'timestamp' is the first column, followed by other keys.
400
401        Args:
402            keys: Iterable of column names
403
404        Returns:
405            list: Ordered list with 'timestamp' first, then other keys
406        """
407        keys_list = list(keys)
408        ordered_keys = []
409
410        # Always put timestamp first if it exists
411        if 'timestamp' in keys_list:
412            ordered_keys.append('timestamp')
413            keys_list.remove('timestamp')
414
415        # Add remaining keys in order they appear
416        ordered_keys.extend(keys_list)
417
418        return ordered_keys
419
420    def _format_value_for_sheets(self, value):
421        """
422        Format a value for Google Sheets while preserving numeric types.
423
424        Args:
425            value: The value to format
426
427        Returns:
428            The value formatted for Google Sheets:
429            - Numbers (int, float) are returned as-is
430            - None values are converted to empty string
431            - Booleans are converted to strings
432            - Everything else is converted to string
433        """
434        if value is None:
435            return ''
436        elif isinstance(value, bool):
437            # Convert booleans to strings to avoid confusion with 1/0
438            return str(value)
439        elif isinstance(value, (int, float)):
440            # Preserve numeric types - don't convert to string
441            return value
442        else:
443            # Convert everything else to string
444            return str(value)
445
446    def write_dict(self, record_dict):
447        """
448        Legacy method: Write a single dictionary as a new row.
449        Use write() instead.
450        """
451        return self.write(record_dict)
452
453    def write_dicts(self, record_dicts):
454        """
455        Legacy method: Write multiple dictionaries as rows.
456        Use write() instead.
457        """
458        return self.write(record_dicts)
459
460    def write(self, records):
461        """
462        Write record(s) to the Google Sheet with automatic column management.
463
464        Accepts either a single record or a list of records. Each record can
465        be either a dictionary or a DASRecord object. The 'timestamp' field
466        is always placed in the first column. Numeric values are preserved.
467
468        Args:
469            records: Single record (dict or DASRecord) or list of records.
470
471        Returns:
472            dict: Response from the Google Sheets API update operation,
473                  or None if records is empty.
474
475        Raises:
476            Exception: If unable to write to the spreadsheet
477            ValueError: If record type is not supported
478        """
479        if not records:
480            return None
481
482        # Normalize input to list of records
483        if not isinstance(records, list):
484            records = [records]
485
486        try:
487            # Convert all records to dictionaries
488            normalized_records = []
489            all_keys = set()
490
491            for record in records:
492                normalized = self._normalize_record(record)
493                normalized_records.append(normalized)
494                all_keys.update(normalized.keys())
495
496            # Ensure timestamp is first in the column order
497            ordered_keys = self._ensure_timestamp_first(all_keys)
498
499            # Update headers with any new keys (maintaining order)
500            new_headers = []
501            for key in ordered_keys:
502                if key not in self.headers:
503                    new_headers.append(key)
504
505            if new_headers:
506                # Insert new headers in the correct position
507                if 'timestamp' in new_headers and 'timestamp' not in self.headers:
508                    # If timestamp is new, it goes first
509                    self.headers.insert(0, 'timestamp')
510                    new_headers.remove('timestamp')
511
512                # Add remaining new headers to the end
513                self.headers.extend(new_headers)
514
515                # Update the header row in the sheet
516                range_name = f"{self.worksheet_name}!1:1"
517                body = {'values': [self.headers]}
518
519                self.service.spreadsheets().values().update(
520                    spreadsheetId=self.sheet_id,
521                    range=range_name,
522                    valueInputOption='RAW',
523                    body=body
524                ).execute()
525
526            # Create rows data with proper value formatting
527            rows_data = []
528            for record_dict in normalized_records:
529                row_data = []
530                for header in self.headers:
531                    value = record_dict.get(header)
532                    formatted_value = self._format_value_for_sheets(value)
533                    row_data.append(formatted_value)
534                rows_data.append(row_data)
535
536            # Find next available row
537            next_row = self._get_next_row()
538            last_col_char = chr(65 + len(self.headers) - 1)
539
540            if len(rows_data) == 1:
541                # Single row
542                range_name = (f"{self.worksheet_name}!A{next_row}:"
543                              f"{last_col_char}{next_row}")
544            else:
545                # Multiple rows
546                end_row = next_row + len(rows_data) - 1
547                range_name = (f"{self.worksheet_name}!A{next_row}:"
548                              f"{last_col_char}{end_row}")
549
550            # Write the data using USER_ENTERED to preserve numeric values
551            body = {'values': rows_data}
552
553            result = self.service.spreadsheets().values().update(
554                spreadsheetId=self.sheet_id,
555                range=range_name,
556                valueInputOption='USER_ENTERED',
557                body=body
558            ).execute()
559
560            return result
561
562        except Exception as e:
563            raise Exception(f"Failed to write to spreadsheet: {str(e)}")
564
565    def get_headers(self):
566        """
567        Return the current column headers in the sheet.
568
569        Returns:
570            list: A copy of the current column headers list.
571        """
572        return self.headers.copy()
573
574    def clear_sheet(self):
575        """
576        Clear all data from the worksheet.
577
578        Removes all content including headers and data. Resets internal
579        headers list.
580
581        Raises:
582            Exception: If unable to clear the sheet
583        """
584        try:
585            range_name = f"{self.worksheet_name}!A:Z"
586            self.service.spreadsheets().values().clear(
587                spreadsheetId=self.sheet_id,
588                range=range_name
589            ).execute()
590            self.headers = []
591        except Exception as e:
592            raise Exception(f"Failed to clear sheet: {str(e)}")
593
594
595# Example usage:
596if __name__ == "__main__":
597    # Initialize the writer
598    writer = GoogleSheetsWriter(
599        sheet_name_or_id="sheet_id_here",
600        auth_key_path="path_to_auth_key_here.json",
601        use_service_account=True,
602        worksheet_name="Sheet1"
603    )
604
605    # Write single dictionary with numeric values
606    record1 = {
607        'timestamp': 1234567890,
608        'name': 'John Doe',
609        'age': 30,
610        'city': 'New York',
611        'salary': 75000.50,
612        'rating': 4.8
613    }
614    writer.write(record1)
615
616    # Write another dict with new columns including numeric ones
617    record2 = {
618        'timestamp': 1234567891,
619        'name': 'Jane Smith',
620        'age': 25,
621        'city': 'Los Angeles',
622        'occupation': 'Engineer',
623        'salary': 85000,
624        'bonus': 10000.25
625    }
626    writer.write(record2)
627
628    # Write multiple dictionaries at once with mixed numeric types
629    records = [
630        {'timestamp': 1234567892, 'name': 'Bob', 'age': 35,
631         'city': 'Chicago', 'salary': 75000, 'score': 89.5},
632        {'timestamp': 1234567893, 'name': 'Alice', 'age': 28,
633         'occupation': 'Designer', 'salary': 65000, 'score': 92.1}
634    ]
635    writer.write(records)
636
637    # Write a DASRecord with numeric fields
638    record = DASRecord(
639        timestamp=1234567894,
640        fields={'name': 'Hal', 'favorite_color': 'green',
641                'years_experience': 5, 'rating': 4.2}
642    )
643    writer.write(record)
644
645    # Check current headers
646    print("Current headers:", writer.get_headers())