Coverage for archdocs/features/sqlalchemy/parser.py: 100%

1import ast 

2import re as py_re 

3import typing 

4 

5from archdocs import prefilter, settings, syntax_tree 

6from archdocs.features.sqlalchemy.const import SQLAlchemyFeatures 

7 

8 

9_ENGINE_FACTORY_NAMES: typing.Final = frozenset(("create_engine", "create_async_engine")) 

10_ASYNC_ENGINE_MODULE: typing.Final = "sqlalchemy.ext.asyncio" 

11_POOL_ARGUMENT_PREFIX: typing.Final = "pool_" 

12_SESSION_ATTRS_NAME: typing.Final = "target_session_attrs" 

13# Inside a string the attribute is a `key=value` pair of a dsn query or of libpq options; outside 

14# of one it is an assignment, a keyword or a dict key, and those the syntax tree spells out itself. 

15_SESSION_ATTRS_IN_STRING_PATTERN: typing.Final = py_re.compile( 

16 r"target_session_attrs\s*=\s*(?P<attrs_value>[\w-]+)", 

17 settings.TYPICAL_RE_FLAGS, 

18) 

19_DSN_HOSTS_PATTERN: typing.Final = py_re.compile(r"://(?:[^/@]*@)?(?P<dsn_hosts>[^/]*)", settings.TYPICAL_RE_FLAGS) 

20_DATABASE_SCHEME_PATTERN: typing.Final = py_re.compile( 

21 r"(?:postgresql|mysql|sqlite|oracle|mssql|mariadb|cockroachdb)(?:\+.*)?", 

22 settings.TYPICAL_RE_FLAGS, 

23) 

24# One union of literals for the whole feature: an engine is found not only by the package 

25# name but by a bare DSN in a string, so the database schemes of `_DATABASE_SCHEME_PATTERN` belong here too. 

26_SQLALCHEMY_LITERALS: typing.Final = ( 

27 "sqlalchemy", 

28 "create_engine", 

29 "create_async_engine", 

30 "target_session_attrs", 

31 "postgresql", 

32 "mysql", 

33 "sqlite", 

34 "oracle", 

35 "mssql", 

36 "mariadb", 

37 "cockroachdb", 

38) 

39_EMPTY_FEATURES: typing.Final = SQLAlchemyFeatures( 

40 async_used=False, 

41 pooling_used=False, 

42 multiple_hosts=False, 

43 target_session_attrs="", 

44 database_type="", 

45) 

46 

47 

48def _read_engine_hosts(engine_call: ast.Call) -> str: 

49 match engine_call.args: 

50 case [ast.Constant(value=str() as engine_dsn), *_]: 

51 hosts_match: typing.Final = _DSN_HOSTS_PATTERN.search(engine_dsn) 

52 return hosts_match.group("dsn_hosts") if hosts_match else "" 

53 case _: 

54 return "" 

55 

56 

57def _read_named_string_values(one_node: ast.AST) -> tuple[tuple[str, str], ...]: 

58 match one_node: 

59 case ast.Assign(targets=every_target, value=ast.Constant(value=str() as assigned_value)): 

60 return tuple( 

61 (one_target.id, assigned_value) for one_target in every_target if isinstance(one_target, ast.Name) 

62 ) 

63 case ast.AnnAssign(target=ast.Name(id=assigned_name), value=ast.Constant(value=str() as assigned_value)): 

64 return ((assigned_name, assigned_value),) 

65 case ast.keyword(arg=str() as keyword_name, value=ast.Constant(value=str() as keyword_value)): 

66 return ((keyword_name, keyword_value),) 

67 case ast.Dict(keys=every_key, values=every_value): 

68 return tuple( 

69 (dict_key.value, dict_value.value) 

70 for dict_key, dict_value in zip(every_key, every_value, strict=True) 

71 if isinstance(dict_key, ast.Constant) 

72 and isinstance(dict_key.value, str) 

73 and isinstance(dict_value, ast.Constant) 

74 and isinstance(dict_value.value, str) 

75 ) 

76 case _: 

77 return () 

78 

79 

80def _read_session_attrs(one_node: ast.AST) -> str: 

81 if isinstance(one_node, ast.Constant) and isinstance(one_node.value, str): 

82 string_match: typing.Final = _SESSION_ATTRS_IN_STRING_PATTERN.search(one_node.value) 

83 return string_match.group("attrs_value") if string_match else "" 

84 return next( 

85 ( 

86 one_value 

87 for one_name, one_value in _read_named_string_values(one_node) 

88 if one_name.lower() == _SESSION_ATTRS_NAME 

89 ), 

90 "", 

91 ) 

92 

93 

94def _read_database_type(one_node: ast.AST) -> str: 

95 match one_node: 

96 case ast.Constant(value=str() as string_value) if _DATABASE_SCHEME_PATTERN.fullmatch(string_value): 

97 return string_value 

98 case _: 

99 return "" 

100 

101 

102def find_sqlalchemy_features(raw_source: str) -> SQLAlchemyFeatures: 

103 if not prefilter.contains_any_literal(raw_source.lower(), _SQLALCHEMY_LITERALS): 

104 return _EMPTY_FEATURES 

105 # A python 2 leftover or a template that only looks like python costs that one file its 

106 # databases: the rest of the project is still drawn. 

107 parsed_tree: typing.Final = syntax_tree.parse_syntax_tree(raw_source) 

108 if parsed_tree is None: 

109 return _EMPTY_FEATURES 

110 every_node: typing.Final = syntax_tree.list_nodes_in_source_order(parsed_tree) 

111 engine_calls: typing.Final = [ 

112 one_node 

113 for one_node in every_node 

114 if isinstance(one_node, ast.Call) and syntax_tree.read_callee_name(one_node) in _ENGINE_FACTORY_NAMES 

115 ] 

116 return SQLAlchemyFeatures( 

117 async_used=any( 

118 one_module.startswith(_ASYNC_ENGINE_MODULE) 

119 for one_node in every_node 

120 for one_module in syntax_tree.read_imported_modules(one_node) 

121 ), 

122 pooling_used=any( 

123 (one_keyword.arg or "").startswith(_POOL_ARGUMENT_PREFIX) 

124 for one_call in engine_calls 

125 for one_keyword in one_call.keywords 

126 ), 

127 multiple_hosts=any("," in _read_engine_hosts(one_call) for one_call in engine_calls), 

128 target_session_attrs=next(filter(None, map(_read_session_attrs, every_node)), ""), 

129 database_type=next(filter(None, map(_read_database_type, every_node)), ""), 

130 )