Coverage for src/lumen/analytics/client.py: 57%

1""" 

2Remote HTTP client SDK for Lumen Analytics. 

3 

4Allows remote Python clients / Jupyter notebooks running on remote machines 

5to query a Lumen FHIR server over HTTP REST and receive pre-flattened 

6Pandas DataFrames automatically. 

7""" 

8 

9from __future__ import annotations 

10 

11from typing import Any 

12 

13import httpx 

14import pandas as pd 

15 

16 

17class LumenClient: 

18 """Remote HTTP Client for querying Lumen FHIR Server DataFrames.""" 

19 

20 def __init__( 

21 self, 

22 base_url: str = "http://localhost:8000", 

23 token: str | None = None, 

24 timeout: float = 30.0, 

25 transport: httpx.BaseTransport | None = None, 

26 verify: bool | str = True, 

27 ) -> None: 

28 self.base_url = base_url.rstrip("/") 

29 headers = {"Accept": "application/fhir+json"} 

30 if token: 

31 headers["Authorization"] = f"Bearer {token}" 

32 self._client = httpx.Client( 

33 base_url=self.base_url, 

34 headers=headers, 

35 timeout=timeout, 

36 transport=transport, 

37 verify=verify, 

38 ) 

39 

40 def search_pandas( 

41 self, 

42 resource_type: str, 

43 params: dict[str, list[str]] | None = None, 

44 **kwargs: Any, 

45 ) -> pd.DataFrame: 

46 """Query the remote Lumen server for a pre-flattened Pandas DataFrame. 

47 

48 The server handles database querying and pandas normalization in-process, 

49 returning a flat JSON stream that is read directly on the client. 

50 """ 

51 query_params: dict[str, Any] = {} 

52 if params: 

53 for k, v in params.items(): 

54 query_params[k] = v 

55 for k, v in kwargs.items(): 

56 norm_k = k.replace("_", "-") if k not in ("_count", "_sort", "_summary") else k 

57 query_params[norm_k] = v 

58 

59 response = self._client.get(f"/{resource_type}/$search-pandas", params=query_params) 

60 response.raise_for_status() 

61 

62 # Parse the pre-flattened JSON list of records directly 

63 import io 

64 

65 return pd.read_json(io.StringIO(response.text), orient="records") 

66 

67 def search_arrow( 

68 self, 

69 resource_type: str, 

70 params: dict[str, Any] | None = None, 

71 **kwargs: Any, 

72 ) -> Any: 

73 """Query the remote Lumen server and receive an Apache Arrow Table.""" 

74 try: 

75 import pyarrow as pa 

76 except ImportError as exc: 

77 raise RuntimeError( 

78 "pyarrow is required for search_arrow(). Install it via: pip install pyarrow" 

79 ) from exc 

80 

81 df = self.search_pandas(resource_type, params=params, **kwargs) 

82 return pa.Table.from_pandas(df) 

83 

84 def query_duckdb( 

85 self, 

86 resource_type: str, 

87 query: str, 

88 params: dict[str, list[str]] | None = None, 

89 **kwargs: Any, 

90 ) -> pd.DataFrame: 

91 """Execute a DuckDB SQL query in-process on the server's DataFrame. 

92 

93 The server builds the DataFrame matching the fhir params, registers it 

94 as the table name (or 'df'), executes the SQL query using DuckDB, 

95 and returns the aggregated result. 

96 """ 

97 fhir_params: dict[str, Any] = {} 

98 if params: 

99 for k, v in params.items(): 

100 fhir_params[k] = v 

101 for k, v in kwargs.items(): 

102 norm_k = k.replace("_", "-") if k not in ("_count", "_sort", "_summary") else k 

103 fhir_params[norm_k] = v 

104 

105 payload = {"query": query, "params": fhir_params} 

106 

107 response = self._client.post(f"/{resource_type}/$search-duckdb", json=payload) 

108 response.raise_for_status() 

109 

110 import io 

111 

112 return pd.read_json(io.StringIO(response.text), orient="records") 

113 

114 def close(self) -> None: 

115 self._client.close() 

116 

117 def __enter__(self) -> LumenClient: 

118 return self 

119 

120 def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: 

121 self.close()