Get available sessions

This commit is contained in:
Guillem Borrell 2024-06-12 11:04:28 +02:00
parent 102fc816f8
commit a994eacd2d
9 changed files with 50 additions and 12 deletions

View file

@ -2,7 +2,7 @@ import json
from pathlib import Path from pathlib import Path
from fastapi import FastAPI from fastapi import FastAPI
from fastapi.responses import HTMLResponse, RedirectResponse, FileResponse from fastapi.responses import FileResponse, HTMLResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles from fastapi.staticfiles import StaticFiles
from starlette.middleware.sessions import SessionMiddleware from starlette.middleware.sessions import SessionMiddleware
from starlette.requests import Request from starlette.requests import Request

View file

@ -1,9 +1,9 @@
from fastapi import APIRouter from fastapi import APIRouter
from fastapi.responses import PlainTextResponse from fastapi.responses import PlainTextResponse
from hellocomputer.sessions import SessionDB
from hellocomputer.db import StorageEngines from hellocomputer.db import StorageEngines
from hellocomputer.extraction import extract_code_block from hellocomputer.extraction import extract_code_block
from hellocomputer.sessions import SessionDB
from ..config import settings from ..config import settings
from ..models import Chat from ..models import Chat

View file

@ -2,9 +2,10 @@ from authlib.integrations.starlette_client import OAuth, OAuthError
from fastapi import APIRouter from fastapi import APIRouter
from fastapi.responses import HTMLResponse, RedirectResponse from fastapi.responses import HTMLResponse, RedirectResponse
from starlette.requests import Request from starlette.requests import Request
from hellocomputer.config import settings from hellocomputer.config import settings
from hellocomputer.users import UserDB
from hellocomputer.db import StorageEngines from hellocomputer.db import StorageEngines
from hellocomputer.users import UserDB
router = APIRouter() router = APIRouter()

View file

@ -4,9 +4,9 @@ import aiofiles
from fastapi import APIRouter, File, UploadFile from fastapi import APIRouter, File, UploadFile
from fastapi.responses import JSONResponse from fastapi.responses import JSONResponse
from ..sessions import SessionDB
from ..config import settings from ..config import settings
from ..db import StorageEngines from ..db import StorageEngines
from ..sessions import SessionDB
router = APIRouter() router = APIRouter()

View file

@ -3,8 +3,11 @@ from uuid import uuid4
from fastapi import APIRouter from fastapi import APIRouter
from fastapi.responses import PlainTextResponse from fastapi.responses import PlainTextResponse
from starlette.requests import Request from starlette.requests import Request
from hellocomputer.users import OwnershipDB from typing import List
from hellocomputer.db import StorageEngines from hellocomputer.db import StorageEngines
from hellocomputer.users import OwnershipDB
from ..config import settings from ..config import settings
# Scheme for the Authorization header # 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 " "Hi! I'm a helpful assistant. Please upload or select a file "
"and I'll try to analyze it following your orders" "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)

View file

@ -1,11 +1,12 @@
import json import json
import os import os
from datetime import datetime
from pathlib import Path from pathlib import Path
from typing import List
from uuid import UUID, uuid4 from uuid import UUID, uuid4
import duckdb import duckdb
import polars as pl import polars as pl
from datetime import datetime
from .db import DDB, StorageEngines from .db import DDB, StorageEngines
@ -79,7 +80,7 @@ class OwnershipDB(DDB):
'{sid}' as sid, '{sid}' as sid,
'{now}' as timestamp '{now}' as timestamp
) )
TO '{self.path_prefix}/{record_id}.csv' (FORMAT JSON)""" TO '{self.path_prefix}/{record_id}.csv'"""
try: try:
self.db.sql(query) self.db.sql(query)
@ -88,3 +89,18 @@ class OwnershipDB(DDB):
self.db.sql(query) self.db.sql(query)
return sid 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()
)

View file

@ -1,8 +1,8 @@
from pathlib import Path from pathlib import Path
import hellocomputer import hellocomputer
from hellocomputer.sessions import SessionDB
from hellocomputer.db import StorageEngines from hellocomputer.db import StorageEngines
from hellocomputer.sessions import SessionDB
TEST_STORAGE = StorageEngines.local TEST_STORAGE = StorageEngines.local
TEST_XLS_PATH = ( TEST_XLS_PATH = (

View file

@ -2,11 +2,11 @@ from pathlib import Path
import hellocomputer import hellocomputer
import pytest import pytest
from hellocomputer.sessions import SessionDB
from hellocomputer.config import settings from hellocomputer.config import settings
from hellocomputer.db import StorageEngines from hellocomputer.db import StorageEngines
from hellocomputer.extraction import extract_code_block from hellocomputer.extraction import extract_code_block
from hellocomputer.models import Chat from hellocomputer.models import Chat
from hellocomputer.sessions import SessionDB
TEST_XLS_PATH = ( TEST_XLS_PATH = (
Path(hellocomputer.__file__).parents[2] Path(hellocomputer.__file__).parents[2]

View file

@ -2,7 +2,7 @@ from pathlib import Path
import hellocomputer import hellocomputer
from hellocomputer.db import StorageEngines from hellocomputer.db import StorageEngines
from hellocomputer.users import UserDB, OwnershipDB from hellocomputer.users import OwnershipDB, UserDB
TEST_STORAGE = StorageEngines.local TEST_STORAGE = StorageEngines.local
TEST_OUTPUT_FOLDER = Path(hellocomputer.__file__).parents[2] / "test" / "output" TEST_OUTPUT_FOLDER = Path(hellocomputer.__file__).parents[2] / "test" / "output"
@ -28,7 +28,13 @@ def test_user_exists():
def test_assign_owner(): def test_assign_owner():
assert ( assert (
OwnershipDB(storage_engine=TEST_STORAGE, path=TEST_OUTPUT_FOLDER).set_ownersip( 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"]