From a994eacd2db59527a1e3c232785eb73c8184c9c3 Mon Sep 17 00:00:00 2001 From: Guillem Borrell Date: Wed, 12 Jun 2024 11:04:28 +0200 Subject: [PATCH] Get available sessions --- src/hellocomputer/main.py | 2 +- src/hellocomputer/routers/analysis.py | 2 +- src/hellocomputer/routers/auth.py | 3 ++- src/hellocomputer/routers/files.py | 2 +- src/hellocomputer/routers/sessions.py | 17 ++++++++++++++++- src/hellocomputer/users.py | 20 ++++++++++++++++++-- test/test_data.py | 2 +- test/test_query.py | 2 +- test/test_user.py | 12 +++++++++--- 9 files changed, 50 insertions(+), 12 deletions(-) diff --git a/src/hellocomputer/main.py b/src/hellocomputer/main.py index a21da9e..5ba6fa9 100644 --- a/src/hellocomputer/main.py +++ b/src/hellocomputer/main.py @@ -2,7 +2,7 @@ import json from pathlib import Path from fastapi import FastAPI -from fastapi.responses import HTMLResponse, RedirectResponse, FileResponse +from fastapi.responses import FileResponse, HTMLResponse, RedirectResponse from fastapi.staticfiles import StaticFiles from starlette.middleware.sessions import SessionMiddleware from starlette.requests import Request diff --git a/src/hellocomputer/routers/analysis.py b/src/hellocomputer/routers/analysis.py index 08c763c..13e9331 100644 --- a/src/hellocomputer/routers/analysis.py +++ b/src/hellocomputer/routers/analysis.py @@ -1,9 +1,9 @@ from fastapi import APIRouter from fastapi.responses import PlainTextResponse -from hellocomputer.sessions import SessionDB from hellocomputer.db import StorageEngines from hellocomputer.extraction import extract_code_block +from hellocomputer.sessions import SessionDB from ..config import settings from ..models import Chat diff --git a/src/hellocomputer/routers/auth.py b/src/hellocomputer/routers/auth.py index aec0b68..aad7883 100644 --- a/src/hellocomputer/routers/auth.py +++ b/src/hellocomputer/routers/auth.py @@ -2,9 +2,10 @@ from authlib.integrations.starlette_client import OAuth, OAuthError from fastapi import APIRouter from fastapi.responses import HTMLResponse, RedirectResponse from starlette.requests import Request + from hellocomputer.config import settings -from hellocomputer.users import UserDB from hellocomputer.db import StorageEngines +from hellocomputer.users import UserDB router = APIRouter() diff --git a/src/hellocomputer/routers/files.py b/src/hellocomputer/routers/files.py index 74542b3..0af6552 100644 --- a/src/hellocomputer/routers/files.py +++ b/src/hellocomputer/routers/files.py @@ -4,9 +4,9 @@ import aiofiles from fastapi import APIRouter, File, UploadFile from fastapi.responses import JSONResponse -from ..sessions import SessionDB from ..config import settings from ..db import StorageEngines +from ..sessions import SessionDB router = APIRouter() diff --git a/src/hellocomputer/routers/sessions.py b/src/hellocomputer/routers/sessions.py index 98ec2c7..adc2532 100644 --- a/src/hellocomputer/routers/sessions.py +++ b/src/hellocomputer/routers/sessions.py @@ -3,8 +3,11 @@ from uuid import uuid4 from fastapi import APIRouter from fastapi.responses import PlainTextResponse from starlette.requests import Request -from hellocomputer.users import OwnershipDB +from typing import List + from hellocomputer.db import StorageEngines +from hellocomputer.users import OwnershipDB + from ..config import settings # Scheme for the Authorization header @@ -32,3 +35,15 @@ async def get_greeting() -> str: "Hi! I'm a helpful assistant. Please upload or select a file " "and I'll try to analyze it following your orders" ) + + +@router.get("/sessions") +async def get_sessions(request: Request) -> List[str]: + user_email = request.session.get("user").get("email") + ownership = OwnershipDB( + StorageEngines.gcs, + gcs_access=settings.gcs_access, + gcs_secret=settings.gcs_secret, + bucket=settings.gcs_bucketname, + ) + return ownership.sessions(user_email) diff --git a/src/hellocomputer/users.py b/src/hellocomputer/users.py index f64d32d..e097aa8 100644 --- a/src/hellocomputer/users.py +++ b/src/hellocomputer/users.py @@ -1,11 +1,12 @@ import json import os +from datetime import datetime from pathlib import Path +from typing import List from uuid import UUID, uuid4 import duckdb import polars as pl -from datetime import datetime from .db import DDB, StorageEngines @@ -79,7 +80,7 @@ class OwnershipDB(DDB): '{sid}' as sid, '{now}' as timestamp ) - TO '{self.path_prefix}/{record_id}.csv' (FORMAT JSON)""" + TO '{self.path_prefix}/{record_id}.csv'""" try: self.db.sql(query) @@ -88,3 +89,18 @@ class OwnershipDB(DDB): self.db.sql(query) return sid + + def sessions(self, user_email: str) -> List[str]: + return ( + self.db.sql(f""" + SELECT + sid + FROM + '{self.path_prefix}/*.csv' + WHERE + email = '{user_email}' + """) + .pl() + .to_series() + .to_list() + ) diff --git a/test/test_data.py b/test/test_data.py index 19ea233..2a8849c 100644 --- a/test/test_data.py +++ b/test/test_data.py @@ -1,8 +1,8 @@ from pathlib import Path import hellocomputer -from hellocomputer.sessions import SessionDB from hellocomputer.db import StorageEngines +from hellocomputer.sessions import SessionDB TEST_STORAGE = StorageEngines.local TEST_XLS_PATH = ( diff --git a/test/test_query.py b/test/test_query.py index e65b9f4..d4d56a2 100644 --- a/test/test_query.py +++ b/test/test_query.py @@ -2,11 +2,11 @@ from pathlib import Path import hellocomputer import pytest -from hellocomputer.sessions import SessionDB from hellocomputer.config import settings from hellocomputer.db import StorageEngines from hellocomputer.extraction import extract_code_block from hellocomputer.models import Chat +from hellocomputer.sessions import SessionDB TEST_XLS_PATH = ( Path(hellocomputer.__file__).parents[2] diff --git a/test/test_user.py b/test/test_user.py index 8acf1f6..df97785 100644 --- a/test/test_user.py +++ b/test/test_user.py @@ -2,7 +2,7 @@ from pathlib import Path import hellocomputer from hellocomputer.db import StorageEngines -from hellocomputer.users import UserDB, OwnershipDB +from hellocomputer.users import OwnershipDB, UserDB TEST_STORAGE = StorageEngines.local TEST_OUTPUT_FOLDER = Path(hellocomputer.__file__).parents[2] / "test" / "output" @@ -28,7 +28,13 @@ def test_user_exists(): def test_assign_owner(): assert ( OwnershipDB(storage_engine=TEST_STORAGE, path=TEST_OUTPUT_FOLDER).set_ownersip( - "something.something@something", "1234", "test" + "something.something@something", "testsession", "test" ) - == "1234" + == "testsession" ) + + +def test_get_sessions(): + assert OwnershipDB(storage_engine=TEST_STORAGE, path=TEST_OUTPUT_FOLDER).sessions( + "something.something@something" + ) == ["testsession"]