Coverage for dj/api/sql.py: 100%

16 statements  

« prev     ^ index     » next       coverage.py v7.2.3, created at 2023-04-17 20:05 -0700

1""" 

2SQL related APIs. 

3""" 

4 

5import logging 

6from typing import List, Optional 

7 

8from fastapi import APIRouter, Depends, Query 

9from sqlmodel import Session 

10 

11from dj.api.helpers import get_engine, get_query 

12from dj.models.metric import TranslatedSQL 

13from dj.models.query import ColumnMetadata 

14from dj.utils import get_session 

15 

16_logger = logging.getLogger(__name__) 

17router = APIRouter() 

18 

19 

20@router.get("/sql/{node_name}/", response_model=TranslatedSQL) 

21def get_sql( 

22 node_name: str, 

23 dimensions: List[str] = Query([]), 

24 filters: List[str] = Query([]), 

25 *, 

26 session: Session = Depends(get_session), 

27 engine_name: Optional[str] = None, 

28 engine_version: Optional[str] = None, 

29) -> TranslatedSQL: 

30 """ 

31 Return SQL for a node. 

32 """ 

33 engine = ( 

34 get_engine(session, engine_name, engine_version) # type: ignore 

35 if engine_name 

36 else None 

37 ) 

38 query_ast = get_query( 

39 session=session, 

40 node_name=node_name, 

41 dimensions=dimensions, 

42 filters=filters, 

43 engine=engine, 

44 ) 

45 columns = [ 

46 ColumnMetadata(name=col.alias_or_name.name, type=str(col.type)) # type: ignore 

47 for col in query_ast.select.projection 

48 ] 

49 return TranslatedSQL( 

50 sql=str(query_ast), 

51 columns=columns, 

52 dialect=engine.dialect if engine else None, 

53 )