-}
-"""
-
-usdts = map(lambda pid: USDT(pid=pid), args.pids)
-for usdt in usdts:
- usdt.enable_probe("query__start", "probe_start")
- usdt.enable_probe("query__done", "probe_end")
+};
+""".replace("DEFINE_USDT", "#define USDT" if mode == "USDT" else "") \
+ .replace("DEFINE_MYSQL56", "#define MYSQL56" if mode == "MYSQL56" else "") \
+ .replace("DEFINE_MYSQL57", "#define MYSQL57" if mode == "MYSQL57" else "") \
+ .replace("DEFINE_THRESHOLD", ("#define THRESHOLD " + str(threshold_ns)) if threshold_ns > 0 else "")
+
+if mode.startswith("MYSQL"):
+ # Uprobes mode
+ bpf = BPF(text=program)
+ bpf.attach_uprobe(name=args.path, sym=mysql_func_name, fn_name="query_start")
+ bpf.attach_uretprobe(name=args.path, sym=mysql_func_name, fn_name="query_end")
+else:
+ # USDT mode
+ if not args.pids or len(args.pids) == 0:
+ if args.db == "mysql":
+ args.pids = map(int, subprocess.check_output(
+ "pidof mysqld".split()).split())
+ elif args.db == "postgres":
+ args.pids = map(int, subprocess.check_output(
+ "pidof postgres".split()).split())
+
+ usdts = map(lambda pid: USDT(pid=pid), args.pids)
+ for usdt in usdts:
+ usdt.enable_probe("query__start", "query_start")
+ usdt.enable_probe("query__done", "query_end")
+ if args.verbose:
+ print('\n'.join(map(lambda u: u.get_text(), usdts)))
+
+ bpf = BPF(text=program, usdt_contexts=usdts)