Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/sqlalchemy_utils/functions/mock.py: 18%

1import contextlib 

2import datetime 

3import inspect 

4import io 

5import re 

6 

7import sqlalchemy as sa 

8 

9 

10def create_mock_engine(bind, stream=None): 

11 """Create a mock SQLAlchemy engine from the passed engine or bind URL. 

12 

13 :param bind: A SQLAlchemy engine or bind URL to mock. 

14 :param stream: Render all DDL operations to the stream. 

15 """ 

16 

17 if not isinstance(bind, str): 

18 bind_url = str(bind.url) 

19 

20 else: 

21 bind_url = bind 

22 

23 if stream is not None: 

24 

25 def dump(sql, *args, **kwargs): 

26 class Compiler(type(sql._compiler(engine.dialect))): 

27 def visit_bindparam(self, bindparam, *args, **kwargs): 

28 return self.render_literal_value(bindparam.value, bindparam.type) 

29 

30 def render_literal_value(self, value, type_): 

31 if isinstance(value, int): 

32 return str(value) 

33 

34 elif isinstance(value, (datetime.date, datetime.datetime)): 

35 return "'%s'" % value 

36 

37 return super().render_literal_value(value, type_) 

38 

39 text = str(Compiler(engine.dialect, sql).process(sql)) 

40 text = re.sub(r'\n+', '\n', text) 

41 text = text.strip('\n').strip() 

42 

43 stream.write('\n%s;' % text) 

44 

45 else: 

46 

47 def dump(*args, **kw): 

48 return None 

49 

50 engine = sa.create_mock_engine(bind_url, executor=dump) 

51 return engine 

52 

53 

54@contextlib.contextmanager 

55def mock_engine(engine, stream=None): 

56 """Mocks out the engine specified in the passed bind expression. 

57 

58 Note this function is meant for convenience and protected usage. Do NOT 

59 blindly pass user input to this function as it uses exec. 

60 

61 :param engine: A python expression that represents the engine to mock. 

62 :param stream: Render all DDL operations to the stream. 

63 """ 

64 

65 # Create a stream if not present. 

66 

67 if stream is None: 

68 stream = io.StringIO() 

69 

70 # Navigate the stack and find the calling frame that allows the 

71 # expression to execute. 

72 

73 for frame in inspect.stack()[1:]: 

74 try: 

75 frame = frame[0] 

76 expression = '__target = %s' % engine 

77 exec(expression, frame.f_globals, frame.f_locals) 

78 target = frame.f_locals['__target'] 

79 break 

80 

81 except Exception: 

82 pass 

83 

84 else: 

85 raise ValueError('Not a valid python expression', engine) 

86 

87 # Evaluate the expression and get the target engine. 

88 

89 frame.f_locals['__mock'] = create_mock_engine(target, stream) 

90 

91 # Replace the target with our mock. 

92 

93 exec('%s = __mock' % engine, frame.f_globals, frame.f_locals) 

94 

95 # Give control back. 

96 

97 yield stream 

98 

99 # Put the target engine back. 

100 

101 frame.f_locals['__target'] = target 

102 exec('%s = __target' % engine, frame.f_globals, frame.f_locals) 

103 exec('del __target', frame.f_globals, frame.f_locals) 

104 exec('del __mock', frame.f_globals, frame.f_locals)