Skip to content
Snippets Groups Projects
Commit bb40503b authored by Tim Repke's avatar Tim Repke
Browse files

add count nql

parent 07aa9730
No related branches found
No related tags found
1 merge request!77Master
Pipeline #2556 canceled
from typing import TYPE_CHECKING
import httpx import httpx
from nacsos_data.db.schemas import Project, ItemType
from pydantic import BaseModel from pydantic import BaseModel
from fastapi import APIRouter, Depends from fastapi import APIRouter, Depends
import sqlalchemy.sql.functions as func import sqlalchemy.sql.functions as func
from sqlalchemy import select from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from nacsos_data.db.engine import ensure_session
from nacsos_data.db.schemas import Project, ItemType
from nacsos_data.util.nql import NQLQuery, NQLFilter
from nacsos_data.util.academic.openalex import query_async, SearchResult from nacsos_data.util.academic.openalex import query_async, SearchResult
from nacsos_data.models.items import AcademicItemModel, FullLexisNexisItemModel, GenericItemModel from nacsos_data.models.items import AcademicItemModel, FullLexisNexisItemModel, GenericItemModel
from nacsos_data.models.openalex.solr import SearchField, DefType, OpType from nacsos_data.models.openalex.solr import SearchField, DefType, OpType
from nacsos_data.util.nql import NQLQuery, NQLFilter
from server.util.security import UserPermissionChecker, UserPermissions from server.util.security import UserPermissionChecker, UserPermissions
from server.util.logging import get_logger from server.util.logging import get_logger
from server.util.config import settings from server.util.config import settings
from server.data import db_engine from server.data import db_engine
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession # noqa F401
router = APIRouter() router = APIRouter()
logger = get_logger('nacsos.api.route.search') logger = get_logger('nacsos.api.route.search')
...@@ -88,22 +85,34 @@ class QueryResult(BaseModel): ...@@ -88,22 +85,34 @@ class QueryResult(BaseModel):
docs: list[AcademicItemModel] | list[FullLexisNexisItemModel] | list[GenericItemModel] docs: list[AcademicItemModel] | list[FullLexisNexisItemModel] | list[GenericItemModel]
@ensure_session
async def _get_query(session: AsyncSession, query: NQLFilter, project_id: str) -> NQLQuery:
project_type: ItemType | None = (
await session.scalar(select(Project.type).where(Project.project_id == project_id)))
if project_type is None:
raise KeyError(f'Found no matching project for {project_id}. This should NEVER happen!')
return NQLQuery(query, project_id=str(project_id), project_type=project_type)
@router.post('/nql/query', response_model=QueryResult) @router.post('/nql/query', response_model=QueryResult)
async def nql_query(query: NQLFilter, async def nql_query(query: NQLFilter,
page: int = 1, page: int = 1,
limit: int = 20, limit: int = 20,
permissions: UserPermissions = Depends(UserPermissionChecker('dataset_read'))) -> QueryResult: permissions: UserPermissions = Depends(UserPermissionChecker('dataset_read'))) -> QueryResult:
async with db_engine.session() as session: # type: AsyncSession async with db_engine.session() as session: # type: AsyncSession
project_id = permissions.permissions.project_id nql = await _get_query(session=session, query=query, project_id=permissions.permissions.project_id)
project_type: ItemType | None = (
await session.scalar(select(Project.type).where(Project.project_id == project_id)))
if project_type is None:
raise KeyError(f'Found no matching project for {project_id}. This should NEVER happen!')
nql = NQLQuery(query, project_id=str(project_id), project_type=project_type)
n_docs = (await session.execute(func.count(nql.stmt.subquery().c.item_id))).scalar() n_docs = (await session.execute(func.count(nql.stmt.subquery().c.item_id))).scalar()
docs = await nql.results_async(session=session, limit=limit, offset=(page - 1) * limit) docs = await nql.results_async(session=session, limit=limit, offset=(page - 1) * limit)
return QueryResult(n_docs=n_docs, docs=docs) # type: ignore[arg-type] return QueryResult(n_docs=n_docs, docs=docs) # type: ignore[arg-type]
@router.post('/nql/count', response_model=int)
async def nql_query_count(query: NQLFilter,
permissions: UserPermissions = Depends(UserPermissionChecker('dataset_read'))) -> QueryResult:
async with db_engine.session() as session: # type: AsyncSession
nql = await _get_query(session=session, query=query, project_id=permissions.permissions.project_id)
return (await session.execute(func.count(nql.stmt.subquery().c.item_id))).scalar()
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment