Coverage for arrakis_server/__main__.py: 0.0%

1import argparse 

2import logging 

3import pathlib 

4import sys 

5from concurrent.futures import ThreadPoolExecutor, wait 

6 

7from . import __version__ 

8from .backends import BackendType 

9from .constants import DEFAULT_LOCATION 

10from .scope import ScopeMap 

11from .server import ArrakisFlightServer 

12 

13logger = logging.getLogger("arrakis") 

14 

15 

16def get_log_level(args: argparse.Namespace) -> int: 

17 """Determine the log level from logging options.""" 

18 if args.quiet: 

19 return logging.WARNING 

20 elif args.verbose: 

21 return logging.DEBUG 

22 else: 

23 return logging.INFO 

24 

25 

26def main() -> None: 

27 parser = argparse.ArgumentParser( 

28 prog="arrakis-server", 

29 description="Arrakis Arrow Flight server.", 

30 ) 

31 parser.add_argument( 

32 "--version", 

33 action="version", 

34 version=__version__, 

35 ) 

36 group = parser.add_mutually_exclusive_group() 

37 group.add_argument( 

38 "-q", 

39 "--quiet", 

40 action="store_true", 

41 help="terse logging, show only warnings and errors", 

42 ) 

43 group.add_argument( 

44 "-v", 

45 "--verbose", 

46 action="store_true", 

47 help="verbose logging", 

48 ) 

49 parser.add_argument( 

50 "-u", 

51 "--url", 

52 type=str, 

53 default=DEFAULT_LOCATION, 

54 help=( 

55 f"serve requests at this URL, default: {DEFAULT_LOCATION}." 

56 " specify '0' to pick a random port on localhost" 

57 ), 

58 ) 

59 parser.add_argument( 

60 "-s", 

61 "--scope-map-file", 

62 "--scope-map", 

63 type=pathlib.Path, 

64 metavar="SCOPE.yaml", 

65 help="scope map YAML file", 

66 ) 

67 

68 # backend subparsers 

69 subparsers = parser.add_subparsers( 

70 title="Backends", 

71 description="""Backends determine how to access the data to be 

72 served to clients. Each backend defines its own arguments, 

73 see \"<BACKEND NAME> -h\" for more info. If a backend is not 

74 specified then the server will act as an information server 

75 only.""", 

76 metavar="Available backends:", 

77 dest="backend", 

78 ) 

79 for _backend in BackendType: 

80 if _backend.value is None: 

81 continue 

82 bparser = subparsers.add_parser( 

83 _backend.name, 

84 aliases=[_backend.name.lower()], 

85 formatter_class=argparse.RawDescriptionHelpFormatter, 

86 help=_backend.value.__doc__.splitlines()[0], 

87 description=_backend.value.__doc__, 

88 ) 

89 bparser.set_defaults(backend=_backend) 

90 _backend.value.add_arguments(bparser) 

91 

92 # parse command line args 

93 args = parser.parse_args() 

94 

95 # set up logger 

96 logger = logging.getLogger("arrakis") 

97 log_level = get_log_level(args) 

98 logger.setLevel(log_level) 

99 handler = logging.StreamHandler(sys.stderr) 

100 handler.setLevel(log_level) 

101 formatter = logging.Formatter("%(asctime)s | arrakis : %(levelname)s : %(message)s") 

102 handler.setFormatter(formatter) 

103 logger.addHandler(handler) 

104 

105 ############################## 

106 

107 logger.info("Arrakis server %s", __version__) 

108 

109 # load scope map 

110 scope_map = None 

111 if args.scope_map_file: 

112 logger.info("loading global scope map %s", args.scope_map_file) 

113 scope_map = ScopeMap.load(args.scope_map_file) 

114 for loc, info in scope_map.servers.items(): 

115 logger.info(" %s: %s", loc, info.domains) 

116 else: 

117 logger.info("no scope map specified.") 

118 

119 # initialize the backend 

120 if args.backend: 

121 logger.info("initializing %s backend...", args.backend.name) 

122 backend = args.backend.value.from_args(args) 

123 else: 

124 logger.info("no backend specified.") 

125 if not args.scope_map_file: 

126 parser.exit( 

127 1, "error: Nothing to serve, must specify scope map and/or backend.\n" 

128 ) 

129 backend = None 

130 

131 # serve requests 

132 logger.info("initializing Flight server...") 

133 server = ArrakisFlightServer( 

134 url=args.url, 

135 backend=backend, 

136 scope_map=scope_map, 

137 ) 

138 logger.info("Flight server initialized") 

139 

140 # serve requests 

141 with ThreadPoolExecutor(max_workers=1) as executor: 

142 try: 

143 future = executor.submit(_run_until_shutdown, server) 

144 wait([future]) 

145 except KeyboardInterrupt: 

146 server.shutdown() 

147 

148 

149def _run_until_shutdown(server: ArrakisFlightServer) -> None: 

150 """Run an Arrakis server instance until a shutdown request is received.""" 

151 logger.info("serving...") 

152 with server: 

153 server.wait_until_shutdown() 

154 

155 

156if __name__ == "__main__": 

157 main()