{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-tenant Authentication with API Keys - SDK Usage\n",
    "\n",
    "This notebook demonstrates how to use TimeDB's authentication system using the **Python SDK** directly.\n",
    "\n",
    "## What we'll cover:\n",
    "1. Setting up the database schema with users table using SDK\n",
    "2. Creating tenants and users using Python functions\n",
    "3. Using API keys to authenticate requests\n",
    "4. Demonstrating tenant isolation (each user only sees their own data)\n",
    "5. Testing authentication enforcement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import uuid\n",
    "import requests\n",
    "import pandas as pd\n",
    "import time\n",
    "import psycopg\n",
    "from datetime import datetime, timezone, timedelta\n",
    "import timedb as td\n",
    "from timedb.db import users as db_users\n",
    "\n",
    "# Use TEST_DATABASE_URL for testing (falls back to DATABASE_URL if not set)\n",
    "DATABASE_URL = os.environ.get(\"TEST_DATABASE_URL\") or os.environ.get(\"DATABASE_URL\")\n",
    "if not DATABASE_URL:\n",
    "    raise ValueError(\"Please set TEST_DATABASE_URL or DATABASE_URL environment variable\")\n",
    "\n",
    "# IMPORTANT: Override both TIMEDB_DSN and DATABASE_URL so the API server uses the same database\n",
    "os.environ[\"TIMEDB_DSN\"] = DATABASE_URL\n",
    "os.environ[\"DATABASE_URL\"] = DATABASE_URL\n",
    "\n",
    "# API base URL\n",
    "API_BASE_URL = \"http://127.0.0.1:8000\"\n",
    "\n",
    "print(f\"Database URL configured: {DATABASE_URL[:50]}...\")\n",
    "print(f\"Environment variables set for API server\")\n",
    "print(f\"API URL: {API_BASE_URL}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Setup Database Schema with Users Table\n",
    "\n",
    "First, we'll create the database schema including the users table for authentication using the SDK."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete existing schema (for a clean start)\n",
    "td.delete(dsn=DATABASE_URL)\n",
    "print(\"Deleted existing tables\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create schema with users table using SDK\n",
    "td.create(dsn=DATABASE_URL, create_users_table=True)\n",
    "print(\"Created schema with users table\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Create Tenants and Users\n",
    "\n",
    "Now we'll create two separate tenants using the `timedb.db.users` module. Each user gets a unique API key that is tied to their tenant_id."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate two tenant IDs\n",
    "tenant_a_id = uuid.uuid4()\n",
    "tenant_b_id = uuid.uuid4()\n",
    "\n",
    "print(f\"Tenant A ID: {tenant_a_id}\")\n",
    "print(f\"Tenant B ID: {tenant_b_id}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create User A for Tenant A using SDK\n",
    "with psycopg.connect(DATABASE_URL) as conn:\n",
    "    user_a = db_users.create_user(\n",
    "        conn,\n",
    "        tenant_id=tenant_a_id,\n",
    "        email=\"alice@tenant-a.com\"\n",
    "    )\n",
    "    conn.commit()\n",
    "\n",
    "api_key_a = user_a['api_key']\n",
    "print(f\"Created User A:\")\n",
    "print(f\"  User ID: {user_a['user_id']}\")\n",
    "print(f\"  Tenant ID: {user_a['tenant_id']}\")\n",
    "print(f\"  Email: {user_a['email']}\")\n",
    "print(f\"  API Key: {api_key_a}\")\n",
    "print(f\"  Is Active: {user_a['is_active']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create User B for Tenant B using SDK\n",
    "with psycopg.connect(DATABASE_URL) as conn:\n",
    "    user_b = db_users.create_user(\n",
    "        conn,\n",
    "        tenant_id=tenant_b_id,\n",
    "        email=\"bob@tenant-b.com\"\n",
    "    )\n",
    "    conn.commit()\n",
    "\n",
    "api_key_b = user_b['api_key']\n",
    "print(f\"Created User B:\")\n",
    "print(f\"  User ID: {user_b['user_id']}\")\n",
    "print(f\"  Tenant ID: {user_b['tenant_id']}\")\n",
    "print(f\"  Email: {user_b['email']}\")\n",
    "print(f\"  API Key: {api_key_b}\")\n",
    "print(f\"  Is Active: {user_b['is_active']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List all users to verify\n",
    "with psycopg.connect(DATABASE_URL) as conn:\n",
    "    all_users = db_users.list_users(conn, include_inactive=True)\n",
    "\n",
    "print(f\"\\nTotal users: {len(all_users)}\\n\")\n",
    "for user in all_users:\n",
    "    status = \"Active\" if user['is_active'] else \"Inactive\"\n",
    "    print(f\"Email: {user['email']:30} | Tenant ID: {user['tenant_id']} | Status: {status}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Start the API Server\n",
    "\n",
    "Now we start the API server using the SDK. With the users_table present, authentication will be **required** for all endpoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Kill any existing API server first\n",
    "import subprocess\n",
    "result = subprocess.run([\"pkill\", \"-f\", \"uvicorn.*timedb\"], capture_output=True)\n",
    "if result.returncode == 0:\n",
    "    print(\"Killed existing API server\")\n",
    "    time.sleep(2)  # Wait for port to be released\n",
    "else:\n",
    "    print(\"No existing API server found\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start API server in background using SDK\n",
    "td.start_api_background(host=\"127.0.0.1\", port=8000)\n",
    "print(\"API server started in background\")\n",
    "\n",
    "# Give the server time to start\n",
    "time.sleep(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify API is running\n",
    "try:\n",
    "    response = requests.get(f\"{API_BASE_URL}/\")\n",
    "    print(f\"API is running: {response.json()['name']}\")\n",
    "    print(f\"Version: {response.json().get('version', 'unknown')}\")\n",
    "except Exception as e:\n",
    "    print(f\"API check failed: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 4: Test Authentication Enforcement\n",
    "\n",
    "Let's verify that authentication is required when the users_table exists."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test request WITHOUT API key - should return 401\n",
    "response = requests.get(f\"{API_BASE_URL}/values\")\n",
    "print(f\"Request without API key:\")\n",
    "print(f\"  Status: {response.status_code}\")\n",
    "print(f\"  Response: {response.json()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test request with INVALID API key - should return 401\n",
    "response = requests.get(\n",
    "    f\"{API_BASE_URL}/values\",\n",
    "    headers={\"X-API-Key\": \"invalid-key-12345\"}\n",
    ")\n",
    "print(f\"Request with invalid API key:\")\n",
    "print(f\"  Status: {response.status_code}\")\n",
    "print(f\"  Response: {response.json()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test request with VALID API key - should return 200\n",
    "response = requests.get(\n",
    "    f\"{API_BASE_URL}/values\",\n",
    "    headers={\"X-API-Key\": api_key_a}\n",
    ")\n",
    "print(f\"Request with valid API key (User A):\")\n",
    "print(f\"  Status: {response.status_code}\")\n",
    "print(f\"  Response: {response.json()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 5: Demonstrate Tenant Isolation\n",
    "\n",
    "Now let's upload data from both users and verify that each user can only see their own tenant's data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# User A uploads temperature data\n",
    "base_time = datetime(2025, 1, 1, 0, 0, tzinfo=timezone.utc)\n",
    "\n",
    "value_rows_a = [\n",
    "    {\"valid_time\": (base_time + timedelta(hours=i)).isoformat(), \n",
    "     \"value_key\": \"temperature_a\", \n",
    "     \"value\": 20.0 + i}\n",
    "    for i in range(6)\n",
    "]\n",
    "\n",
    "response = requests.post(\n",
    "    f\"{API_BASE_URL}/upload\",\n",
    "    json={\n",
    "        \"run_start_time\": datetime.now(timezone.utc).isoformat(),\n",
    "        \"value_rows\": value_rows_a\n",
    "    },\n",
    "    headers={\"X-API-Key\": api_key_a, \"Content-Type\": \"application/json\"}\n",
    ")\n",
    "print(f\"User A upload status: {response.status_code}\")\n",
    "print(f\"Response: {response.json()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# User B uploads humidity data\n",
    "value_rows_b = [\n",
    "    {\"valid_time\": (base_time + timedelta(hours=i)).isoformat(), \n",
    "     \"value_key\": \"humidity_b\", \n",
    "     \"value\": 50.0 + i}\n",
    "    for i in range(6)\n",
    "]\n",
    "\n",
    "response = requests.post(\n",
    "    f\"{API_BASE_URL}/upload\",\n",
    "    json={\n",
    "        \"run_start_time\": datetime.now(timezone.utc).isoformat(),\n",
    "        \"value_rows\": value_rows_b\n",
    "    },\n",
    "    headers={\"X-API-Key\": api_key_b, \"Content-Type\": \"application/json\"}\n",
    ")\n",
    "print(f\"User B upload status: {response.status_code}\")\n",
    "print(f\"Response: {response.json()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# User A reads data - should ONLY see their own data (temperature_a)\n",
    "response = requests.get(\n",
    "    f\"{API_BASE_URL}/values\",\n",
    "    params={\n",
    "        \"start_valid\": base_time.isoformat(),\n",
    "        \"end_valid\": (base_time + timedelta(hours=6)).isoformat(),\n",
    "        \"mode\": \"flat\"\n",
    "    },\n",
    "    headers={\"X-API-Key\": api_key_a}\n",
    ")\n",
    "data_a = response.json()\n",
    "print(f\"User A sees {data_a['count']} records:\")\n",
    "if data_a['count'] > 0:\n",
    "    df_a = pd.DataFrame(data_a['data'])\n",
    "    print(df_a[['valid_time', 'series_key', 'value']])\n",
    "    print(f\"\\nSeries keys visible to User A: {df_a['series_key'].unique().tolist()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# User B reads data - should ONLY see their own data (humidity_b)\n",
    "response = requests.get(\n",
    "    f\"{API_BASE_URL}/values\",\n",
    "    params={\n",
    "        \"start_valid\": base_time.isoformat(),\n",
    "        \"end_valid\": (base_time + timedelta(hours=6)).isoformat(),\n",
    "        \"mode\": \"flat\"\n",
    "    },\n",
    "    headers={\"X-API-Key\": api_key_b}\n",
    ")\n",
    "data_b = response.json()\n",
    "print(f\"User B sees {data_b['count']} records:\")\n",
    "if data_b['count'] > 0:\n",
    "    df_b = pd.DataFrame(data_b['data'])\n",
    "    print(df_b[['valid_time', 'series_key', 'value']])\n",
    "    print(f\"\\nSeries keys visible to User B: {df_b['series_key'].unique().tolist()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify tenant isolation\n",
    "print(\"Tenant Isolation Verification:\")\n",
    "print(f\"  User A (Tenant A) sees: {data_a['count']} records\")\n",
    "print(f\"  User B (Tenant B) sees: {data_b['count']} records\")\n",
    "print(\"\")\n",
    "\n",
    "# Check that they see different data\n",
    "if data_a['count'] > 0 and data_b['count'] > 0:\n",
    "    series_a = set(pd.DataFrame(data_a['data'])['series_key'].unique())\n",
    "    series_b = set(pd.DataFrame(data_b['data'])['series_key'].unique())\n",
    "    overlap = series_a & series_b\n",
    "    \n",
    "    if len(overlap) == 0:\n",
    "        print(\"SUCCESS: No overlapping series between tenants!\")\n",
    "        print(f\"  Tenant A series: {series_a}\")\n",
    "        print(f\"  Tenant B series: {series_b}\")\n",
    "    else:\n",
    "        print(f\"WARNING: Found overlapping series: {overlap}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 6: User Deactivation\n",
    "\n",
    "Let's demonstrate how deactivating a user revokes their API access using the SDK."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deactivate User B using SDK\n",
    "with psycopg.connect(DATABASE_URL) as conn:\n",
    "    success = db_users.deactivate_user(conn, email=\"bob@tenant-b.com\")\n",
    "    conn.commit()\n",
    "\n",
    "if success:\n",
    "    print(\"Successfully deactivated user: bob@tenant-b.com\")\n",
    "else:\n",
    "    print(\"User not found\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try to access API with deactivated user's key - should return 401\n",
    "response = requests.get(\n",
    "    f\"{API_BASE_URL}/values\",\n",
    "    headers={\"X-API-Key\": api_key_b}\n",
    ")\n",
    "print(f\"Request with deactivated user's API key:\")\n",
    "print(f\"  Status: {response.status_code}\")\n",
    "print(f\"  Response: {response.json()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reactivate User B using SDK\n",
    "with psycopg.connect(DATABASE_URL) as conn:\n",
    "    success = db_users.activate_user(conn, email=\"bob@tenant-b.com\")\n",
    "    conn.commit()\n",
    "\n",
    "if success:\n",
    "    print(\"Successfully reactivated user: bob@tenant-b.com\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify access is restored\n",
    "response = requests.get(\n",
    "    f\"{API_BASE_URL}/values\",\n",
    "    headers={\"X-API-Key\": api_key_b}\n",
    ")\n",
    "print(f\"Request after reactivation:\")\n",
    "print(f\"  Status: {response.status_code}\")\n",
    "if response.status_code == 200:\n",
    "    print(f\"  SUCCESS: Access restored for User B\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 7: Additional SDK Functions\n",
    "\n",
    "Let's explore a few more user management functions available in the SDK."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get user by email\n",
    "with psycopg.connect(DATABASE_URL) as conn:\n",
    "    user = db_users.get_user_by_email(conn, email=\"alice@tenant-a.com\")\n",
    "\n",
    "if user:\n",
    "    print(\"User found:\")\n",
    "    print(f\"  Email: {user['email']}\")\n",
    "    print(f\"  Tenant ID: {user['tenant_id']}\")\n",
    "    print(f\"  User ID: {user['user_id']}\")\n",
    "    print(f\"  Is Active: {user['is_active']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Regenerate API key for User B\n",
    "with psycopg.connect(DATABASE_URL) as conn:\n",
    "    new_api_key = db_users.regenerate_api_key(conn, email=\"bob@tenant-b.com\")\n",
    "    conn.commit()\n",
    "\n",
    "if new_api_key:\n",
    "    print(f\"New API key generated for bob@tenant-b.com:\")\n",
    "    print(f\"  {new_api_key}\")\n",
    "    print(f\"\\nOld API key: {api_key_b[:20]}...\")\n",
    "    print(f\"New API key: {new_api_key[:20]}...\")\n",
    "    \n",
    "    # Update the variable\n",
    "    api_key_b_old = api_key_b\n",
    "    api_key_b = new_api_key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify old key no longer works\n",
    "response = requests.get(\n",
    "    f\"{API_BASE_URL}/values\",\n",
    "    headers={\"X-API-Key\": api_key_b_old}\n",
    ")\n",
    "print(f\"Request with old API key:\")\n",
    "print(f\"  Status: {response.status_code}\")\n",
    "\n",
    "# Verify new key works\n",
    "response = requests.get(\n",
    "    f\"{API_BASE_URL}/values\",\n",
    "    headers={\"X-API-Key\": api_key_b}\n",
    ")\n",
    "print(f\"\\nRequest with new API key:\")\n",
    "print(f\"  Status: {response.status_code}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List users for a specific tenant\n",
    "with psycopg.connect(DATABASE_URL) as conn:\n",
    "    tenant_a_users = db_users.list_users(conn, tenant_id=tenant_a_id)\n",
    "\n",
    "print(f\"Users in Tenant A ({tenant_a_id}):\")\n",
    "for user in tenant_a_users:\n",
    "    print(f\"  {user['email']} - Active: {user['is_active']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrated TimeDB's multi-tenant authentication system using the **Python SDK**.\n",
    "\n",
    "### Key SDK Functions Used:\n",
    "\n",
    "1. **Schema Management** (`timedb` module):\n",
    "   - `td.create(dsn=url, create_users_table=True)` - Create schema with users table\n",
    "   - `td.delete(dsn=url)` - Delete schema\n",
    "   - `td.start_api_background(host=..., port=...)` - Start API server\n",
    "\n",
    "2. **User Management** (`timedb.db.users` module):\n",
    "   - `db_users.create_user(conn, tenant_id=..., email=...)` - Create user with API key\n",
    "   - `db_users.list_users(conn, tenant_id=None, include_inactive=False)` - List users\n",
    "   - `db_users.get_user_by_email(conn, email=...)` - Get user by email\n",
    "   - `db_users.deactivate_user(conn, email=...)` - Revoke API access\n",
    "   - `db_users.activate_user(conn, email=...)` - Restore API access\n",
    "   - `db_users.regenerate_api_key(conn, email=...)` - Generate new API key\n",
    "   - `db_users.get_user_by_api_key(conn, api_key=...)` - Validate API key\n",
    "\n",
    "### Key Differences from CLI:\n",
    "\n",
    "- No need for subprocess calls or parsing command output\n",
    "- Direct Python function calls with return values\n",
    "- Full access to user dictionaries with all fields\n",
    "- More programmatic control over database operations\n",
    "- Requires managing database connections with `psycopg.connect()`\n",
    "\n",
    "### Authentication Concepts:\n",
    "\n",
    "- API keys are generated securely using `secrets.token_urlsafe(32)`\n",
    "- Each user's API key is tied to a tenant_id for data isolation\n",
    "- When users_table exists, authentication is **required** for all endpoints\n",
    "- Use `X-API-Key` header with all API requests\n",
    "- Invalid or missing API keys return 401 Unauthorized"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}