""" The Moonstream queries HTTP API """ import logging from typing import Any, Dict, List, Optional, Tuple, Union from uuid import UUID from bugout.data import BugoutResources, BugoutJournalEntryContent, BugoutJournalEntry from bugout.exceptions import BugoutResponseException from fastapi import APIRouter, Body, Request import requests # type: ignore from sqlalchemy import text from .. import data from ..actions import ( get_query_by_name, name_normalization, NameNormalizationException, query_parameter_hash, generate_s3_access_links, ) from ..middleware import MoonstreamHTTPException from ..settings import ( MOONSTREAM_ADMIN_ACCESS_TOKEN, MOONSTREAM_APPLICATION_ID, MOONSTREAM_CRAWLERS_SERVER_URL, MOONSTREAM_CRAWLERS_SERVER_PORT, MOONSTREAM_S3_QUERIES_BUCKET, MOONSTREAM_S3_QUERIES_BUCKET_PREFIX, MOONSTREAM_QUERIES_JOURNAL_ID, ) from ..settings import bugout_client as bc logger = logging.getLogger(__name__) router = APIRouter( prefix="/queries", ) @router.get("/list", tags=["queries"]) async def get_list_of_queries_handler(request: Request) -> List[Dict[str, Any]]: token = request.state.token # Check already existed queries params = { "type": data.BUGOUT_RESOURCE_QUERY_RESOLVER, } try: resources: BugoutResources = bc.list_resources(token=token, params=params) except BugoutResponseException as e: raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) users_queries: List[Dict[str, Any]] = [ resource.resource_data for resource in resources.resources ] return users_queries @router.post("/", tags=["queries"]) async def create_query_handler( request: Request, query_applied: data.PreapprovedQuery = Body(...) ) -> BugoutJournalEntry: """ Create query in bugout journal """ token = request.state.token user = request.state.user # Check already existed queries params = { "type": data.BUGOUT_RESOURCE_QUERY_RESOLVER, } try: resources: BugoutResources = bc.list_resources(token=token, params=params) except BugoutResponseException as e: raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) used_queries: List[str] = [ resource.resource_data["name"] for resource in resources.resources ] try: query_name = name_normalization(query_applied.name) except NameNormalizationException: raise MoonstreamHTTPException( status_code=403, detail=f"Provided query name can't be normalize please select different.", ) if query_name in used_queries: raise MoonstreamHTTPException( status_code=404, detail=f"Provided query name already use. Please remove it or use PUT /{query_name} for update query", ) try: # Put query to journal entry = bc.create_entry( token=MOONSTREAM_ADMIN_ACCESS_TOKEN, journal_id=MOONSTREAM_QUERIES_JOURNAL_ID, title=f"Query:{query_name}", tags=["type:query"], content=query_applied.query, ) except BugoutResponseException as e: raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) try: # create resource query_name_resolver bc.create_resource( token=token, application_id=MOONSTREAM_APPLICATION_ID, resource_data={ "type": data.BUGOUT_RESOURCE_QUERY_RESOLVER, "user_id": str(user.id), "user": str(user.username), "name": query_name, "entry_id": str(entry.id), }, ) except BugoutResponseException as e: logger.error(f"Error creating name resolving resource: {str(e)}") raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) try: bc.update_tags( token=MOONSTREAM_ADMIN_ACCESS_TOKEN, journal_id=MOONSTREAM_QUERIES_JOURNAL_ID, entry_id=entry.id, tags=[f"query_id:{entry.id}", f"preapprove"], ) except BugoutResponseException as e: logger.error(f"Error in applind tags to query entry: {str(e)}") raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) return entry @router.get("/{query_name}/query", tags=["queries"]) async def get_query_handler( request: Request, query_name: str ) -> data.QueryInfoResponse: token = request.state.token try: query_id = get_query_by_name(query_name, token) except NameNormalizationException: raise MoonstreamHTTPException( status_code=403, detail=f"Provided query name can't be normalize please select different.", ) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) try: entry = bc.get_entry( token=MOONSTREAM_ADMIN_ACCESS_TOKEN, journal_id=MOONSTREAM_QUERIES_JOURNAL_ID, entry_id=query_id, ) except BugoutResponseException as e: logger.error(f"Error in get query: {str(e)}") raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) try: if entry.content is None: raise MoonstreamHTTPException( status_code=403, detail=f"Query is empty. Please update it." ) query = text(entry.content) except Exception as e: raise MoonstreamHTTPException( status_code=500, internal_error=e, detail="Error in query parsing" ) query_parameters_names = list(query._bindparams.keys()) tags_dict = { tag.split(":")[0]: (tag.split(":")[1] if ":" in tag else True) for tag in entry.tags } query_parameters: Dict[str, Any] = {} for param in query_parameters_names: if param in tags_dict: query_parameters[param] = tags_dict[param] else: query_parameters[param] = None return data.QueryInfoResponse( query=entry.content, query_id=str(entry.id), preapprove="preapprove" in tags_dict, approved="approved" in tags_dict, parameters=query_parameters, created_at=entry.created_at, updated_at=entry.updated_at, ) @router.put("/{query_name}", tags=["queries"]) async def update_query_handler( request: Request, query_name: str, request_update: data.UpdateQueryRequest = Body(...), ) -> BugoutJournalEntryContent: token = request.state.token try: query_id = get_query_by_name(query_name, token) except NameNormalizationException: raise MoonstreamHTTPException( status_code=403, detail=f"Provided query name can't be normalize please select different.", ) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) try: entry = bc.update_entry_content( token=MOONSTREAM_ADMIN_ACCESS_TOKEN, journal_id=MOONSTREAM_QUERIES_JOURNAL_ID, entry_id=query_id, title=query_name, content=request_update.query, tags=["preapprove"], ) except BugoutResponseException as e: logger.error(f"Error in updating query: {str(e)}") raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) return entry @router.post( "/{query_name}/update_data", tags=["queries"], ) async def update_query_data_handler( request: Request, query_name: str, request_update: data.UpdateDataRequest = Body(...), ) -> Optional[data.QueryPresignUrl]: """ Request update data on S3 bucket """ token = request.state.token try: query_id = get_query_by_name(query_name, token) except NameNormalizationException: raise MoonstreamHTTPException( status_code=403, detail=f"Provided query name can't be normalize please select different.", ) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) try: entries = bc.search( token=MOONSTREAM_ADMIN_ACCESS_TOKEN, journal_id=MOONSTREAM_QUERIES_JOURNAL_ID, query=f"tag:approved tag:query_id:{query_id} !tag:preapprove", limit=1, timeout=5, ) if len(entries.results) == 0: raise MoonstreamHTTPException( status_code=403, detail="Query not approved yet." ) s3_response = None if entries.results[0].content: content = entries.results[0].content tags = entries.results[0].tags file_type = "json" if "ext:csv" in tags: file_type = "csv" responce = requests.post( f"{MOONSTREAM_CRAWLERS_SERVER_URL}:{MOONSTREAM_CRAWLERS_SERVER_PORT}/jobs/{query_id}/query_update", json={ "query": content, "params": request_update.params, "file_type": file_type, }, timeout=5, ) if responce.status_code != 200: raise MoonstreamHTTPException( status_code=responce.status_code, detail=responce.text, ) s3_response = data.QueryPresignUrl(**responce.json()) except BugoutResponseException as e: logger.error(f"Error in updating query: {str(e)}") raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) return s3_response @router.post("/{query_name}", tags=["queries"]) async def get_access_link_handler( request: Request, query_name: str, request_update: data.UpdateDataRequest = Body(...), ) -> Optional[data.QueryPresignUrl]: """ Request S3 presign url """ # get real connect to query_id token = request.state.token try: query_id = get_query_by_name(query_name, token) except NameNormalizationException: raise MoonstreamHTTPException( status_code=403, detail=f"Provided query name can't be normalize please select different.", ) except Exception as e: logger.error(f"Error in get query: {str(e)}") raise MoonstreamHTTPException(status_code=500, internal_error=e) try: entries = bc.search( token=MOONSTREAM_ADMIN_ACCESS_TOKEN, journal_id=MOONSTREAM_QUERIES_JOURNAL_ID, query=f"tag:approved tag:query_id:{query_id} !tag:preapprove", limit=1, timeout=5, ) s3_response = None if entries.results and entries.results[0].content: passed_params = dict(request_update.params) tags = entries.results[0].tags file_type = "json" if "ext:csv" in tags: file_type = "csv" params_hash = query_parameter_hash(passed_params) bucket = MOONSTREAM_S3_QUERIES_BUCKET key = f"{MOONSTREAM_S3_QUERIES_BUCKET_PREFIX}/queries/{query_id}/{params_hash}/data.{file_type}" stats_presigned_url = generate_s3_access_links( method_name="get_object", bucket=bucket, key=key, expiration=300000, http_method="GET", ) s3_response = data.QueryPresignUrl(url=stats_presigned_url) except BugoutResponseException as e: logger.error(f"Error in get access link: {str(e)}") raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: logger.error(f"Error in get access link: {str(e)}") raise MoonstreamHTTPException(status_code=500, internal_error=e) return s3_response @router.delete("/{query_name}", tags=["queries"]) async def remove_query_handler( request: Request, query_name: str, ) -> BugoutJournalEntry: """ Request delete query from journal """ token = request.state.token params = {"type": data.BUGOUT_RESOURCE_QUERY_RESOLVER, "name": query_name} try: resources: BugoutResources = bc.list_resources(token=token, params=params) except BugoutResponseException as e: raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) query_ids: Dict[str, Tuple[UUID, Union[UUID, str]]] = { resource.resource_data["name"]: ( resource.id, resource.resource_data["entry_id"], ) for resource in resources.resources } if len(query_ids) == 0: raise MoonstreamHTTPException(status_code=404, detail="Query does not exists") try: bc.delete_resource(token=token, resource_id=query_ids[query_name][0]) except BugoutResponseException as e: raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) try: entry = bc.delete_entry( token=MOONSTREAM_ADMIN_ACCESS_TOKEN, journal_id=MOONSTREAM_QUERIES_JOURNAL_ID, entry_id=query_ids[query_name][1], ) except BugoutResponseException as e: raise MoonstreamHTTPException(status_code=e.status_code, detail=e.detail) except Exception as e: raise MoonstreamHTTPException(status_code=500, internal_error=e) return entry