Coverage for arrakis_server/__main__.py: 0.0%

1 import argparse

2 import logging

3 import pathlib

4 import sys

5 from concurrent.futures import ThreadPoolExecutor, wait

6

7 from . import __version__

8 from .backends import BackendType

9 from .constants import DEFAULT_LOCATION

10 from .scope import ScopeMap

11 from .server import ArrakisFlightServer

12

13 logger = logging.getLogger("arrakis")

14

15

16 def 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

26 def 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

149 def _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

156 if __name__ == "__main__":

157 main()