credit

app.py

@@ -12,6 +12,7 @@ from fastapi.responses import JSONResponse
 
 from database import init_db
 from routers import auth, blink, general
+from drive_service import DriveService
 
 # Configure logging
 logging.basicConfig(
@@ -20,6 +21,8 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Initialize Drive Service
+drive_service = DriveService()
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -28,9 +31,18 @@ async def lifespan(app: FastAPI):
     Initializes database on startup.
     """
     logger.info("Starting up - initializing database...")
+    
+    # Startup: Download DB from Drive
+    logger.info("Startup: Attempting to download database from Google Drive...")
+    drive_service.download_db()
+    
     await init_db()
     logger.info("Database initialized successfully")
     yield
+    
+    # Shutdown: Upload DB to Drive
+    logger.info("Shutdown: Uploading database to Google Drive...")
+    drive_service.upload_db()
     logger.info("Shutting down...")
 
 

drive_service.py

@@ -0,0 +1,168 @@
+import os
+import logging
+import io
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from google.auth.transport.requests import Request
+from google.oauth2.credentials import Credentials
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload
+from googleapiclient.errors import HttpError
+
+logger = logging.getLogger(__name__)
+
+class DriveService:
+    SCOPES = [
+        'https://www.googleapis.com/auth/gmail.send',
+        'https://www.googleapis.com/auth/drive.file'
+    ]
+    FOLDER_NAME = "apigateway"
+    DB_FILENAME = "blink_data.db"
+
+    def __init__(self):
+        self.creds = None
+        self.service = None
+        self.client_id = os.getenv('GOOGLE_CLIENT_ID')
+        self.client_secret = os.getenv('GOOGLE_CLIENT_SECRET')
+        self.refresh_token = os.getenv('GOOGLE_REFRESH_TOKEN')
+
+    def authenticate(self):
+        """Authenticate using the refresh token."""
+        if not all([self.client_id, self.client_secret, self.refresh_token]):
+            logger.error("Missing Google API credentials for Drive Service")
+            return False
+
+        try:
+            self.creds = Credentials(
+                None,
+                refresh_token=self.refresh_token,
+                token_uri="https://oauth2.googleapis.com/token",
+                client_id=self.client_id,
+                client_secret=self.client_secret,
+                scopes=self.SCOPES
+            )
+            
+            if self.creds and self.creds.expired and self.creds.refresh_token:
+                self.creds.refresh(Request())
+                
+            self.service = build('drive', 'v3', credentials=self.creds)
+            return True
+        except Exception as e:
+            logger.error(f"Failed to authenticate with Drive API: {e}")
+            return False
+
+    def _find_folder(self):
+        """Find the 'apigateway' folder."""
+        try:
+            query = f"mimeType='application/vnd.google-apps.folder' and name='{self.FOLDER_NAME}' and trashed=false"
+            results = self.service.files().list(q=query, spaces='drive', fields='files(id, name)').execute()
+            files = results.get('files', [])
+            if not files:
+                return None
+            return files[0]['id']
+        except HttpError as error:
+            logger.error(f"An error occurred searching for folder: {error}")
+            return None
+
+    def _create_folder(self):
+        """Create the 'apigateway' folder."""
+        try:
+            file_metadata = {
+                'name': self.FOLDER_NAME,
+                'mimeType': 'application/vnd.google-apps.folder'
+            }
+            file = self.service.files().create(body=file_metadata, fields='id').execute()
+            logger.info(f"Created folder with ID: {file.get('id')}")
+            return file.get('id')
+        except HttpError as error:
+            logger.error(f"An error occurred creating folder: {error}")
+            return None
+
+    def _get_folder_id(self):
+        """Get folder ID, creating it if it doesn't exist."""
+        folder_id = self._find_folder()
+        if not folder_id:
+            folder_id = self._create_folder()
+        return folder_id
+
+    def upload_db(self):
+        """Upload the database file to Google Drive."""
+        if not self.service and not self.authenticate():
+            return False
+
+        if not os.path.exists(self.DB_FILENAME):
+            logger.warning(f"Database file {self.DB_FILENAME} not found for upload.")
+            return False
+
+        folder_id = self._get_folder_id()
+        if not folder_id:
+            return False
+
+        try:
+            # Check if file already exists in folder
+            query = f"name='{self.DB_FILENAME}' and '{folder_id}' in parents and trashed=false"
+            results = self.service.files().list(q=query, spaces='drive', fields='files(id)').execute()
+            files = results.get('files', [])
+
+            media = MediaFileUpload(self.DB_FILENAME, mimetype='application/x-sqlite3', resumable=True)
+
+            if files:
+                # Update existing file
+                file_id = files[0]['id']
+                self.service.files().update(fileId=file_id, media_body=media).execute()
+                logger.info(f"Updated database file {self.DB_FILENAME} in Drive (ID: {file_id})")
+            else:
+                # Create new file
+                file_metadata = {
+                    'name': self.DB_FILENAME,
+                    'parents': [folder_id]
+                }
+                self.service.files().create(body=file_metadata, media_body=media, fields='id').execute()
+                logger.info(f"Uploaded new database file {self.DB_FILENAME} to Drive")
+            
+            return True
+        except HttpError as error:
+            logger.error(f"An error occurred uploading DB: {error}")
+            return False
+        except Exception as e:
+            logger.error(f"Unexpected error uploading DB: {e}")
+            return False
+
+    def download_db(self):
+        """Download the database file from Google Drive."""
+        if not self.service and not self.authenticate():
+            return False
+
+        folder_id = self._find_folder() # Don't create if not found, just return
+        if not folder_id:
+            logger.info("No 'apigateway' folder found in Drive. Starting with fresh DB.")
+            return False
+
+        try:
+            query = f"name='{self.DB_FILENAME}' and '{folder_id}' in parents and trashed=false"
+            results = self.service.files().list(q=query, spaces='drive', fields='files(id)').execute()
+            files = results.get('files', [])
+
+            if not files:
+                logger.info(f"No {self.DB_FILENAME} found in Drive folder. Starting with fresh DB.")
+                return False
+
+            file_id = files[0]['id']
+            request = self.service.files().get_media(fileId=file_id)
+            fh = io.FileIO(self.DB_FILENAME, 'wb')
+            downloader = MediaIoBaseDownload(fh, request)
+            done = False
+            while done is False:
+                status, done = downloader.next_chunk()
+                # logger.info(f"Download {int(status.progress() * 100)}%.")
+            
+            logger.info(f"Successfully downloaded {self.DB_FILENAME} from Drive.")
+            return True
+        except HttpError as error:
+            logger.error(f"An error occurred downloading DB: {error}")
+            return False
+        except Exception as e:
+            logger.error(f"Unexpected error downloading DB: {e}")
+            return False

get_gmail_token.py → get_google_token.py

@@ -3,7 +3,11 @@ import json
 from google_auth_oauthlib.flow import InstalledAppFlow
 
 # If modifying these scopes, delete the file token.json.
-SCOPES = ['https://www.googleapis.com/auth/gmail.send']
+# If modifying these scopes, delete the file token.json.
+SCOPES = [
+    'https://www.googleapis.com/auth/gmail.send',
+    'https://www.googleapis.com/auth/drive.file'
+]
 
 def get_refresh_token():
     """Shows basic usage of the Gmail API.

routers/auth.py

@@ -10,8 +10,10 @@ from models import User, AuditLog
 from schemas import CheckRegistrationRequest, RegisterRequest, ValidateRequest, ResetRequest
 from auth_utils import get_password_hash, verify_password, generate_secret_key, send_email
 from dependencies import check_rate_limit
+from drive_service import DriveService
 
 router = APIRouter(prefix="/auth", tags=["auth"])
+drive_service = DriveService()
 
 @router.post("/check-registration")
 async def check_registration(
@@ -109,6 +111,9 @@ Do not share this key with anyone."""
         "Your Secret Key - Credit System",
         email_body
     )
+    
+    # Sync DB to Drive (Async)
+    background_tasks.add_task(drive_service.upload_db)
 
     return {"success": True, "message": "Registration successful. Check your email."}
 
@@ -117,6 +122,7 @@ Do not share this key with anyone."""
 async def validate_key(
     request: ValidateRequest,
     req: Request,
+    background_tasks: BackgroundTasks,
     db: AsyncSession = Depends(get_db)
 ):
     """
@@ -151,6 +157,9 @@ async def validate_key(
         db.add(audit_log)
         await db.commit()
         
+        # Sync DB to Drive (Async) - Optional but good for audit logs
+        background_tasks.add_task(drive_service.upload_db)
+        
         return {
             "valid": True,
             "user_id": valid_user.user_id,
@@ -233,6 +242,10 @@ Current Credits: {user.credits}"""
             "Your New Secret Key",
             email_body
         )
+        
+        # Sync DB to Drive (Async)
+        background_tasks.add_task(drive_service.upload_db)
+        
     else:
         # Log Audit (failed/not found)
         audit_log = AuditLog(