Setting the file. One moment. scripts/with_server.py
Python·106 lines·4 KB
"""
16
17import subprocess
18import socket
19import time
20import sys
21import argparse
22
23def is_server_ready(port, timeout=30):
24 """Wait for server to be ready by polling the port."""
25 start_time = time.time()
26 while time.time() - start_time < timeout:
27 try:
28 with socket.create_connection(('localhost', port), timeout=1):
29 return True
30 except (socket.error, ConnectionRefusedError):
31 time.sleep(0.5)
32 return False
33
34
35def main():
36 parser = argparse.ArgumentParser(description='Run command with one or more servers')
37 parser.add_argument('--server', action='append', dest='servers', required=True, help='Server command (can be repeated)')
38 parser.add_argument('--port', action='append', dest='ports', type=int, required=True, help='Port for each server (must match --server count)')
39 parser.add_argument('--timeout', type=int, default=30, help='Timeout in seconds per server (default: 30)')
40 parser.add_argument('command', nargs=argparse.REMAINDER, help='Command to run after server(s) ready')
41
42 args = parser.parse_args()
43
44 # Remove the '--' separator if present
45 if args.command and args.command[0] == '--':
46 args.command = args.command[1:]
47
48 if not args.command:
49 print("Error: No command specified to run")
50 sys.exit(1)
51
52 # Parse server configurations
53 if len(args.servers) != len(args.ports):
54 print("Error: Number of --server and --port arguments must match")
55 sys.exit(1)
56
57 servers = []
58 for cmd, port in zip(args.servers, args.ports):
59 servers.append({'cmd': cmd, 'port': port})
60
61 server_processes = []
62
63 try:
64 # Start all servers
65 for i, server in enumerate(servers):
66 print(f"Starting server {i+1}/{len(servers)}: {server['cmd']}")
67
68 # Use shell=True to support commands with cd and &&
69 process = subprocess.Popen(
70 server['cmd'],
71 shell=True,
72 stdout=subprocess.PIPE,
73 stderr=subprocess.PIPE
74 )
75 server_processes.append(process)
76
77 # Wait for this server to be ready
78 print(f"Waiting for server on port {server['port']}...")
79 if not is_server_ready(server['port'], timeout=args.timeout):
80 raise RuntimeError(f"Server failed to start on port {server['port']} within {args.timeout}s")
81
82 print(f"Server ready on port {server['port']}")
83
84 print(f"\nAll {len(servers)} server(s) ready")
85
86 # Run the command
87 print(f"Running: {' '.join(args.command)}\n")
88 result = subprocess.run(args.command)
89 sys.exit(result.returncode)
90
91 finally:
92 # Clean up all servers
93 print(f"\nStopping {len(server_processes)} server(s)...")
94 for i, process in enumerate(server_processes):
95 try:
96 process.terminate()
97 process.wait(timeout=5)
98 except subprocess.TimeoutExpired:
99 process.kill()
100 process.wait()
101 print(f"Server {i+1} stopped")
102 print("All servers stopped")
103
104
105if __name__ == '__main__':
106 main()