notion_handler.py

import os
from typing import List, Dict, Any
from notion_client import Client
from datetime import datetime
from dotenv import load_dotenv

load_dotenv()  # Load environment variables from .env file

# Check for required environment variables
if not os.getenv("NOTION_TOKEN"):
    raise ValueError("NOTION_TOKEN environment variable is not set")
if not os.getenv("NOTION_PARENT_PAGE_ID"):
    raise ValueError("NOTION_PARENT_PAGE_ID environment variable is not set")

notion = Client(auth=os.getenv("NOTION_TOKEN"))
parent_page_id = os.getenv("NOTION_PARENT_PAGE_ID")

def create_tasks_database() -> str:
    """
    Create a new Notion database for tasks.
    
    Returns:
        str: The ID of the created database
    
    Raises:
        Exception: If database creation fails
    """
    try:
        response = notion.databases.create(
            parent={"type": "page_id", "page_id": parent_page_id},
            title=[{"type": "text", "text": {"content": f"NotionTaskSense Tasks - {datetime.now().strftime('%Y-%m-%d')}"}}],
            properties={
                "Task": {"title": {}},
                "Category": {
                    "select": {
                        "options": [
                            {"name": "Career", "color": "blue"},
                            {"name": "Learning", "color": "green"},
                            {"name": "Personal", "color": "orange"},
                            {"name": "Outreach", "color": "purple"},
                            {"name": "Health", "color": "red"},
                            {"name": "Finance", "color": "yellow"}
                        ]
                    }
                },
                "Priority": {
                    "select": {
                        "options": [
                            {"name": "High", "color": "red"},
                            {"name": "Medium", "color": "yellow"},
                            {"name": "Low", "color": "blue"}
                        ]
                    }
                },
                "Due Date": {"date": {}},
                "Status": {
                    "select": {
                        "options": [
                            {"name": "To Do", "color": "gray"},
                            {"name": "In Progress", "color": "blue"},
                            {"name": "Done", "color": "green"}
                        ]
                    }
                },
                "Notes": {"rich_text": {}}
            }
        )
        return response["id"]
    except Exception as e:
        raise Exception(f"Failed to create Notion database: {str(e)}")

def push_tasks_to_notion(tasks: List[Dict[str, Any]], db_id: str) -> None:
    """
    Add or update tasks in the Notion database.
    
    Args:
        tasks (List[Dict[str, Any]]): List of task dictionaries
        db_id (str): Notion database ID
    
    Raises:
        Exception: If pushing tasks fails
    """
    try:
        for task in tasks:
            properties = {
                "Task": {"title": [{"text": {"content": task.get("task", "Untitled Task")}}]},
                "Category": {"select": {"name": task.get("category", "Personal")}},
                "Priority": {"select": {"name": task.get("priority", "Medium")}},
                "Status": {"select": {"name": "To Do"}},
            }

            if "dueDate" in task and task["dueDate"]:
                properties["Due Date"] = {"date": {"start": task["dueDate"]}}

            if "notes" in task and task["notes"]:
                properties["Notes"] = {"rich_text": [{"text": {"content": task["notes"]}}]}

            notion.pages.create(
                parent={"database_id": db_id},
                properties=properties
            )
    except Exception as e:
        raise Exception(f"Failed to push tasks to Notion: {str(e)}")