Coverage for gws-app/gws/base/database/connection.py: 96%

1"""Database connection wrapper.""" 

2 

3import gws 

4import gws.lib.sa as sa 

5 

6 

7class Object(gws.DatabaseConnection): 

8 """Database connection.""" 

9 

10 db: gws.DatabaseProvider 

11 """Provider this connection belongs to.""" 

12 saConn: sa.Connection 

13 """The underlying SQLAlchemy connection.""" 

14 

15 def __init__(self, db: gws.DatabaseProvider, conn: sa.Connection): 

16 """Create a connection wrapper. 

17 

18 Args: 

19 db: Provider this connection belongs to. 

20 conn: The SQLAlchemy connection. 

21 """ 

22 self.db = db 

23 self.saConn = conn 

24 

25 def execute(self, stmt, params=None, execution_options=None): 

26 if isinstance(stmt, str): 

27 stmt = sa.text(stmt) 

28 return self.saConn.execute(stmt, params, execution_options=execution_options) 

29 

30 def fetch_all(self, stmt, **params): 

31 return [r._asdict() for r in self.execute(stmt, params)] 

32 

33 def fetch_first(self, stmt, **params): 

34 res = self.execute(stmt, params) 

35 r = res.first() 

36 return r._asdict() if r else None 

37 

38 def fetch_scalars(self, stmt, **params): 

39 res = self.execute(stmt, params) 

40 return list(res.scalars().all()) 

41 

42 def fetch_strings(self, stmt, **params): 

43 res = self.execute(stmt, params) 

44 return [_to_str(s) for s in res.scalars().all()] 

45 

46 def fetch_ints(self, stmt, **params): 

47 res = self.execute(stmt, params) 

48 return [_to_int(s) for s in res.scalars().all()] 

49 

50 def fetch_scalar(self, stmt, **params): 

51 res = self.execute(stmt, params) 

52 return res.scalar() 

53 

54 def fetch_string(self, stmt, **params): 

55 res = self.execute(stmt, params) 

56 s = res.scalar() 

57 return _to_str(s) if s is not None else None 

58 

59 def fetch_int(self, stmt, **params): 

60 res = self.execute(stmt, params) 

61 s = res.scalar() 

62 return _to_int(s) if s is not None else None 

63 

64 

65## 

66 

67 

68def _to_int(s) -> int: 

69 """Return the value if it is an int, raise ``ValueError`` otherwise.""" 

70 if isinstance(s, int): 

71 return s 

72 raise ValueError(f'db: expected int, got {s=}') 

73 

74 

75def _to_str(s) -> str: 

76 """Convert a value to a string, ``None`` to an empty string.""" 

77 if s is None: 

78 return '' 

79 return str(s)