[feat] support google sheets

main.py

@@ -1,13 +1,26 @@
-from fastapi import FastAPI, HTTPException, status
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.staticfiles import StaticFiles
-from pydantic import BaseModel
-from typing import List, Dict, Optional
-import hashlib
-import csv
 import os
+import json
+import hashlib
+import logging
+import tempfile
 from datetime import datetime
 
+import pygsheets
+import pandas as pd
+from typing import List
+from pydantic import BaseModel
+
+from fastapi import FastAPI, HTTPException, status, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import JSONResponse
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
 app = FastAPI(title="Data Samples API")
 
 # Add CORS middleware to allow web requests
@@ -32,8 +45,47 @@ class SampleResponse(BaseModel):
     modified_datetime: str
 
 
-# File path for TSV storage
-SAMPLES_FILE = "samples.tsv"
+# Exception handler for Google Sheets issues
+@app.exception_handler(Exception)
+async def general_exception_handler(request: Request, exc: Exception):
+    logger.error(f"Unhandled exception: {exc}")
+    return JSONResponse(
+        status_code=500,
+        content={"detail": f"Internal server error: {str(exc)}"}
+    )
+
+
+# Google Sheets setup - using environment variables for credentials
+def get_worksheet():
+    try:
+        # Check if running in HuggingFace Spaces or local environment
+        if "GOOGLE_SERVICE_JSON_CONTENT" in os.environ:
+            # Get credentials from environment variable
+            service_account_info = json.loads(os.environ["GOOGLE_SERVICE_JSON_CONTENT"])
+            
+            # Create a temporary JSON file with the credentials
+            with tempfile.NamedTemporaryFile(mode='w+', delete=False, suffix='.json') as temp_file:
+                json.dump(service_account_info, temp_file)
+                temp_filepath = temp_file.name
+            
+            try:
+                # Authorize with the temporary file
+                gc = pygsheets.authorize(service_file=temp_filepath)
+                # Open the spreadsheet and return the first worksheet
+                sh = gc.open('drivelology')
+                return sh[0]
+            finally:
+                # Clean up the temporary file
+                if os.path.exists(temp_filepath):
+                    os.unlink(temp_filepath)
+        else:
+            # Local development - use the JSON file directly
+            gc = pygsheets.authorize(service_file='drivelology-1b65510988e8.json')
+            sh = gc.open('drivelology')
+            return sh[0]  # Return the first worksheet
+    except Exception as e:
+        logger.error(f"Error accessing Google Sheets: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to access Google Sheets: {str(e)}")
 
 
 # Helper function to generate ID from text
@@ -47,137 +99,219 @@ def get_current_datetime() -> str:
     return datetime.now().isoformat()
 
 
-# Helper functions for TSV operations
-def load_samples_from_tsv():
-    samples = {}
-    if os.path.exists(SAMPLES_FILE):
-        with open(SAMPLES_FILE, 'r') as file:
-            reader = csv.reader(file, delimiter='\t')
-            for row in reader:
-                if len(row) >= 4:  # id, text, created_datetime, modified_datetime
-                    sample_id = row[0]
-                    samples[sample_id] = SampleResponse(
-                        id=sample_id,
-                        text=row[1],
-                        created_datetime=row[2],
-                        modified_datetime=row[3]
-                    )
-    return samples
-
-
-def save_samples_to_tsv(samples):
-    with open(SAMPLES_FILE, 'w', newline='') as file:
-        writer = csv.writer(file, delimiter='\t')
+# Helper functions for Google Sheets operations
+def load_samples_from_sheets():
+    try:
+        worksheet = get_worksheet()
+        records = worksheet.get_all_records()
+        
+        # Debug log
+        logger.info(f"Retrieved {len(records)} records from Google Sheets")
+        
+        samples = {}
+        for record in records:
+            if record and 'id' in record and 'text' in record:
+                # Handle potential missing fields with defaults
+                sample_id = record['id']
+                created_time = record.get('created_datetime', get_current_datetime())
+                modified_time = record.get('modified_datetime', created_time)
+                
+                samples[sample_id] = SampleResponse(
+                    id=sample_id,
+                    text=record['text'],
+                    created_datetime=created_time,
+                    modified_datetime=modified_time
+                )
+        
+        logger.info(f"Processed {len(samples)} valid samples")
+        return samples
+    except Exception as e:
+        logger.error(f"Error loading samples from sheets: {e}")
+        # Return empty dict on error to avoid crashing
+        return {}
+
+
+def save_samples_to_sheets(samples):
+    try:
+        worksheet = get_worksheet()
+        
+        # Convert samples to DataFrame
+        data = []
         for sample in samples.values():
-            writer.writerow([
-                sample.id, 
-                sample.text, 
-                sample.created_datetime, 
-                sample.modified_datetime
-            ])
+            data.append({
+                'id': sample.id,
+                'text': sample.text,
+                'created_datetime': sample.created_datetime,
+                'modified_datetime': sample.modified_datetime
+            })
+        
+        # Create DataFrame and update the worksheet
+        df = pd.DataFrame(data)
+        if not data:
+            # If no samples, create an empty dataframe with the right columns
+            df = pd.DataFrame(columns=['id', 'text', 'created_datetime', 'modified_datetime'])
+        
+        # Clear the worksheet and update with new data
+        worksheet.clear()
+        worksheet.set_dataframe(df, (0, 0))
+        
+        logger.info(f"Successfully saved {len(samples)} samples to Google Sheets")
+        return True
+    except Exception as e:
+        logger.error(f"Error saving samples to sheets: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to save data: {str(e)}")
 
 
-# Get samples from TSV file
+# Get samples from Google Sheets
 def get_samples():
-    return load_samples_from_tsv()
+    return load_samples_from_sheets()
 
 
 @app.post("/samples/", response_model=SampleResponse, status_code=status.HTTP_201_CREATED)
 async def add_sample(sample: Sample):
     """Add a new data sample"""
-    # Load current samples
-    samples = get_samples()
-    
-    # Generate ID based on text content
-    sample_id = generate_id_from_text(sample.text)
-    
-    # Get current datetime
-    current_time = get_current_datetime()
-    
-    # Create a sample response with ID, text, and timestamps
-    sample_response = SampleResponse(
-        id=sample_id, 
-        text=sample.text,
-        created_datetime=current_time,
-        modified_datetime=current_time
-    )
-    
-    # Store the sample
-    samples[sample_id] = sample_response
-    
-    # Save to TSV
-    save_samples_to_tsv(samples)
-    
-    return sample_response
+    try:
+        # Load current samples
+        samples = get_samples()
+        
+        # Generate ID based on text content
+        sample_id = generate_id_from_text(sample.text)
+        
+        # Get current datetime
+        current_time = get_current_datetime()
+        
+        # Create a sample response with ID, text, and timestamps
+        sample_response = SampleResponse(
+            id=sample_id, 
+            text=sample.text,
+            created_datetime=current_time,
+            modified_datetime=current_time
+        )
+        
+        # Store the sample
+        samples[sample_id] = sample_response
+        
+        # Save to Google Sheets
+        save_samples_to_sheets(samples)
+        
+        logger.info(f"Added new sample with ID: {sample_id}")
+        return sample_response
+    except Exception as e:
+        logger.error(f"Error adding sample: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to add sample: {str(e)}")
 
 
 @app.get("/samples/", response_model=List[SampleResponse])
 async def get_all_samples():
     """Get all data samples"""
-    samples = get_samples()
-    return list(samples.values())
+    try:
+        samples = get_samples()
+        result = list(samples.values())
+        logger.info(f"Returning {len(result)} samples")
+        return result
+    except Exception as e:
+        logger.error(f"Error retrieving samples: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to retrieve samples: {str(e)}")
 
 
 @app.get("/samples/{sample_id}", response_model=SampleResponse)
 async def get_sample(sample_id: str):
     """Get a specific data sample by ID"""
-    samples = get_samples()
-    if sample_id not in samples:
-        raise HTTPException(status_code=404, detail="Sample not found")
-    return samples[sample_id]
+    try:
+        samples = get_samples()
+        if sample_id not in samples:
+            logger.warning(f"Sample ID not found: {sample_id}")
+            raise HTTPException(status_code=404, detail="Sample not found")
+        return samples[sample_id]
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error retrieving sample {sample_id}: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to retrieve sample: {str(e)}")
 
 
 @app.put("/samples/{sample_id}", response_model=SampleResponse)
 async def update_sample(sample_id: str, updated_sample: Sample):
     """Update an existing data sample"""
-    samples = get_samples()
-    if sample_id not in samples:
-        raise HTTPException(status_code=404, detail="Sample not found")
-    
-    # Get the existing sample to preserve created_datetime
-    existing_sample = samples[sample_id]
-    
-    # Get current datetime for modified_datetime
-    current_time = get_current_datetime()
-    
-    # Create updated sample with original ID and updated timestamps
-    updated_sample_response = SampleResponse(
-        id=sample_id, 
-        text=updated_sample.text,
-        created_datetime=existing_sample.created_datetime,
-        modified_datetime=current_time
-    )
-    
-    # Store the updated sample
-    samples[sample_id] = updated_sample_response
-    
-    # Save to TSV
-    save_samples_to_tsv(samples)
-    
-    return updated_sample_response
+    try:
+        samples = get_samples()
+        if sample_id not in samples:
+            logger.warning(f"Attempted to update non-existent sample: {sample_id}")
+            raise HTTPException(status_code=404, detail="Sample not found")
+        
+        # Get the existing sample to preserve created_datetime
+        existing_sample = samples[sample_id]
+        
+        # Get current datetime for modified_datetime
+        current_time = get_current_datetime()
+        
+        # Create updated sample with original ID and updated timestamps
+        updated_sample_response = SampleResponse(
+            id=sample_id, 
+            text=updated_sample.text,
+            created_datetime=existing_sample.created_datetime,
+            modified_datetime=current_time
+        )
+        
+        # Store the updated sample
+        samples[sample_id] = updated_sample_response
+        
+        # Save to Google Sheets
+        save_samples_to_sheets(samples)
+        
+        logger.info(f"Updated sample with ID: {sample_id}")
+        return updated_sample_response
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error updating sample {sample_id}: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to update sample: {str(e)}")
 
 
 @app.delete("/samples/{sample_id}", status_code=status.HTTP_204_NO_CONTENT)
 async def delete_sample(sample_id: str):
     """Delete a data sample"""
-    samples = get_samples()
-    if sample_id not in samples:
-        raise HTTPException(status_code=404, detail="Sample not found")
-    
-    # Remove the sample
-    del samples[sample_id]
-    
-    # Save to TSV
-    save_samples_to_tsv(samples)
-    
-    return None
-
-
-# Create empty TSV file if it doesn't exist
-if not os.path.exists(SAMPLES_FILE):
-    save_samples_to_tsv({})
+    try:
+        samples = get_samples()
+        if sample_id not in samples:
+            logger.warning(f"Attempted to delete non-existent sample: {sample_id}")
+            raise HTTPException(status_code=404, detail="Sample not found")
+        
+        # Remove the sample
+        del samples[sample_id]
+        
+        # Save to Google Sheets
+        save_samples_to_sheets(samples)
+        
+        logger.info(f"Deleted sample with ID: {sample_id}")
+        return None
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error deleting sample {sample_id}: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to delete sample: {str(e)}")
+
+
+# Initialize the Google Sheet with headers if needed
+def initialize_sheet():
+    try:
+        worksheet = get_worksheet()
+        # Check if the worksheet has headers
+        records = worksheet.get_all_records()
+        # If empty, add headers
+        if not records:
+            logger.info("Initializing empty Google Sheet with headers")
+            df = pd.DataFrame(columns=['id', 'text', 'created_datetime', 'modified_datetime'])
+            worksheet.set_dataframe(df, (0, 0))
+            logger.info("Google Sheet initialized successfully")
+    except Exception as e:
+        logger.error(f"Error initializing sheet: {e}")
+        # Don't raise exception here to allow the app to start even if sheet init fails
+        pass
+
+
+# Initialize the sheet when the application starts
+initialize_sheet()
 
 # Mount the static files directory at the end after defining all API routes
-# This serves the index.html file at root (/) while preserving API routes
-# app.mount("/", StaticFiles(directory=".", html=True), name="static")
 app.mount("/", StaticFiles(directory="static", html=True), name="static")

requirements.txt

@@ -1,3 +1,5 @@
 fastapi==0.104.0
 pydantic==2.11.3
-uvicorn==0.23.2
+uvicorn==0.23.2
+pandas==2.2.3
+pygsheets==2.0.6