dbworkload 0.8.0__tar.gz → 0.8.3__tar.gz

{dbworkload-0.8.0 → dbworkload-0.8.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: dbworkload
-Version: 0.8.0
+Version: 0.8.3
 Summary: Workload framework
 License: GPLv3+
 Author: Fabio Ghirardello

{dbworkload-0.8.0 → dbworkload-0.8.3}/dbworkload/cli/dep.py

@@ -1,8 +1,9 @@
 #!/usr/bin/python
 
-from .. import __version__
 import typer
 
+from .. import __version__
+
 EPILOG = "Docs: <https://dbworkload.github.io/dbworkload/>"
 
 

{dbworkload-0.8.0 → dbworkload-0.8.3}/dbworkload/cli/main.py

@@ -1,23 +1,26 @@
 #!/usr/bin/python
 
-from .. import __version__
-from dbworkload.cli.dep import Param, EPILOG, ConnInfo
+import json
+import logging
+import os
+import platform
+import sys
 from enum import Enum
 from pathlib import Path
 from typing import Optional
 from urllib.parse import urlparse
+
+import pandas as pd
+import typer
+import yaml
+
 import dbworkload.cli.util
 import dbworkload.models.run
 import dbworkload.models.util
 import dbworkload.utils.common
-import json
-import logging
-import os
-import platform
-import sys
-import typer
-import yaml
-import pandas as pd
+from dbworkload.cli.dep import EPILOG, ConnInfo, Param
+
+from .. import __version__
 
 logger = logging.getLogger("dbworkload")
 
@@ -145,6 +148,11 @@ def run(
         "--schedule",
         help="schedule JSON string or filepath to the schedule file.",
     ),
+    histogram_bins: str = typer.Option(
+        "5,10,25,50,75,100,125,250,500,750,1000",
+        "--bins",
+        help="comma separated list of ints defining the histogram bins.",
+    ),
     log_level: LogLevel = Param.LogLevel,
 ):
     logger.setLevel(log_level.upper())
@@ -231,6 +239,7 @@ def run(
 
     args = load_args(args)
 
+    histogram_bins = histogram_bins.split(",")
     schedule = load_schedule(schedule)
 
     dbworkload.models.run.run(
@@ -249,6 +258,7 @@ def run(
         quiet,
         save,
         schedule,
+        histogram_bins,
         log_level.upper(),
     )
 

{dbworkload-0.8.0 → dbworkload-0.8.3}/dbworkload/cli/util.py

@@ -1,13 +1,15 @@
 #!/usr/bin/python
 
-from pathlib import Path
 from enum import Enum
+from pathlib import Path
 from typing import Optional
+
+import typer
+
 import dbworkload.models.run
 import dbworkload.models.util
 import dbworkload.utils.common
-from dbworkload.cli.dep import Param, EPILOG
-import typer
+from dbworkload.cli.dep import EPILOG, Param
 
 
 class Compression(str, Enum):

{dbworkload-0.8.0 → dbworkload-0.8.3}/dbworkload/models/run.py

@@ -1,24 +1,23 @@
 #!/usr/bin/python
 
-from contextlib import contextmanager
-from dbworkload.cli.dep import ConnInfo
-import dbworkload.utils.common
+import errno
 import logging
-import logging.handlers
 import multiprocessing as mp
-import numpy as np
+import os
 import queue
 import random
 import signal
 import sys
-import sys
-import tabulate
-from threading import Thread
 import time
 import traceback
-import os
-import errno
+from contextlib import contextmanager
+from threading import Thread
+
+import numpy as np
+import tabulate
 
+import dbworkload.utils.common
+from dbworkload.cli.dep import ConnInfo
 
 # from cassandra.cluster import Cluster, ExecutionProfile, EXEC_PROFILE_DEFAULT, Session
 # from cassandra.policies import (
@@ -37,6 +36,7 @@ FIFO = "dbworkload.pipe"
 
 logger = logging.getLogger("dbworkload")
 
+force_exit = False
 
 HEADERS: list = [
     "elapsed",
@@ -96,7 +96,13 @@ def signal_handler(sig, frame):
         frame (_type_):
     """
     logger.info("KeyboardInterrupt signal detected. Stopping processes...")
+    global force_exit
+    if force_exit:
+        logger.warning("Forcibly quitting. You're rude!")
+        sys.exit(1)
 
+    force_exit = True
+    
     # send the poison pill to each proc.
     # if dbworkload cannot graceful shutdown due
     # to processes being still in the init phase
@@ -105,13 +111,15 @@ def signal_handler(sig, frame):
     # and raise the queue.Full exception, forcing to quit.
     for q in queues.values():
         try:
-            q.put("proc_end", timeout=0.1)
+            q.put("proc_end")
         except queue.Full:
             logger.error("Timed out")
             sys.exit(1)
 
     logger.debug("Sent poison pill to all procs")
-    os.remove(FIFO)
+
+    if os.path.exists(FIFO):
+        os.remove(FIFO)
 
 
 def cycle(iterable, backwards=False):
@@ -178,21 +186,20 @@ def run(
     quiet: bool,
     save: bool,
     schedule: list,
+    histogram_bins: list,
     log_level: str,
 ):
     def gracefully_shutdown(by_keyinterrupt: bool = False):
-        """
-        wait for final stat reports to come in,
-        then print final stats and quit
-        """
+
+        logger.info("Gracefully shutting down...")
 
         end_time = int(time.time())
-        _s = stats_received
+        # _s = stats_received
 
         if not by_keyinterrupt:
             for q in queues.values():
                 try:
-                    q.put("proc_end", timeout=0.1)
+                    q.put("proc_end")
                 except queue.Full:
                     logger.error("Timed out")
                     sys.exit(1)
@@ -201,20 +208,23 @@ def run(
                 if x.is_alive():
                     x.join()
 
-        while True:
-            try:
-                msg = to_main_q.get(block=True, timeout=2.0)
-                if isinstance(msg, list):
-                    _s += 1
-                    stats.add_tds(msg)
-                    if _s >= active_connections:
-                        break
-                else:
-                    logger.error("Timed out, quitting")
-                    sys.exit(1)
-
-            except queue.Empty:
-                break
+        # Commenting the below as in theory there shouldn't be any stats that 
+        # comes in *after* the PROC returns. That is, all threads send stats
+        # when all threads are returned, then the supervisor returns.
+        # while True:
+        #     try:
+        #         msg = to_main_q.get(block=True, timeout=2.0)
+        #         if isinstance(msg, list):
+        #             _s += 1
+        #             stats.add_tds(msg)
+        #             if _s >= active_connections:
+        #                 break
+        #         else:
+        #             logger.error("Timed out, quitting")
+        #             sys.exit(1)
+
+        #     except queue.Empty:
+        #         break
 
         # now that we have all stat reports, calculate the stats one last time.
         report = stats.calculate_stats(active_connections, end_time)
@@ -330,7 +340,7 @@ def run(
 
     stats = dbworkload.utils.common.Stats(start_time)
 
-    prom = dbworkload.utils.common.Prom(prom_port)
+    prom = dbworkload.utils.common.Prom(prom_port, stats, histogram_bins)
 
     to_main_q = mp.Queue()
 
@@ -478,6 +488,7 @@ def run(
                     active_connections -= 1
                 elif msg == "proc_returned":
                     returned_procs += 1
+                    logger.info(f"Stopped processes: {returned_procs}/{procs} ")
                 elif msg == "task_done":
                     returned_threads += 1
             except queue.Empty:
@@ -580,7 +591,7 @@ def run(
                 report_time += FREQUENCY
 
             # pause briefly to prevent the loop from overheating the CPU
-            time.sleep(0.1)
+            time.sleep(.001)
 
     gracefully_shutdown()
 

{dbworkload-0.8.0 → dbworkload-0.8.3}/dbworkload/models/util.py

@@ -1,28 +1,29 @@
 #!/usr/bin/python
 
-from io import TextIOWrapper
-from jinja2 import Environment, PackageLoader
-from pathlib import PosixPath
-from plotly.subplots import make_subplots
-from pytdigest import TDigest
 import datetime as dt
-import dbworkload
-import dbworkload.utils.common
-import dbworkload.utils.simplefaker
 import gzip
 import itertools
 import logging
-import numpy as np
 import os
+import shutil
+import sys
+from io import TextIOWrapper
+from pathlib import PosixPath
+
+import numpy as np
 import pandas as pd
 import plotext as plt
 import plotly.graph_objects as go
 import plotly.io as pio
-import shutil
 import sqlparse
-import sys
 import yaml
+from jinja2 import Environment, PackageLoader
+from plotly.subplots import make_subplots
+from pytdigest import TDigest
 
+import dbworkload
+import dbworkload.utils.common
+import dbworkload.utils.simplefaker
 
 logger = logging.getLogger("dbworkload")
 logger.setLevel(logging.INFO)

{dbworkload-0.8.0 → dbworkload-0.8.3}/dbworkload/utils/common.py

@@ -2,14 +2,17 @@
 
 import importlib
 import logging
-import numpy as np
 import os
 import random
 import sys
 import time
 import urllib.parse
-import yaml
+
+import numpy as np
 import prometheus_client as prom
+import yaml
+from prometheus_client.core import REGISTRY, HistogramMetricFamily
+from prometheus_client.registry import Collector
 from pytdigest import TDigest
 
 RESERVED_WORDS = [
@@ -45,57 +48,8 @@ NOT_NULL_MAX = 40
 
 logger = logging.getLogger("dbworkload")
 
-
-class Prom:
-    def __init__(self, prom_port: int = 26260):
-        self.prom_latency: dict[str, list[prom.Gauge]] = {}
-
-        # don't stop just because prom server can't start
-        try:
-            prom.start_http_server(prom_port)
-        except OSError as e:
-            logger.warning(f"Cannot start prometheus server: {e}")
-
-        self.threads = prom.Gauge(
-            "threads", "count of connection threads to the database."
-        )
-
-    def publish(self, report: list):
-        for row in report:
-            id = row[1]
-
-            if id not in self.prom_latency:
-                self.prom_latency[id] = []
-                self.prom_latency[id].append(
-                    prom.Gauge(f"{id}__tot_ops", "total count of ops")
-                )
-                self.prom_latency[id].append(
-                    prom.Gauge(
-                        f"{id}__tot_ops_s", "derived value from tot_ops / elapsed"
-                    )
-                )
-                self.prom_latency[id].append(
-                    prom.Gauge(f"{id}__period_ops", "ops count for the recent window")
-                )
-                self.prom_latency[id].append(
-                    prom.Gauge(
-                        f"{id}__period_ops_s",
-                        "derived value from period_ops / window duration",
-                    )
-                )
-                self.prom_latency[id].append(prom.Gauge(f"{id}__mean_ms", "mean_ms"))
-                self.prom_latency[id].append(prom.Gauge(f"{id}__p50_ms", "p50_ms"))
-                self.prom_latency[id].append(prom.Gauge(f"{id}__p90_ms", "p90_ms"))
-                self.prom_latency[id].append(prom.Gauge(f"{id}__p95_ms", "p95_ms"))
-                self.prom_latency[id].append(prom.Gauge(f"{id}__p99_ms", "p99_ms"))
-                self.prom_latency[id].append(prom.Gauge(f"{id}__max_ms", "max_ms"))
-
-            for idx, v in enumerate(row[3:]):
-                self.prom_latency[id][idx].set(v)
-
-        # threads value is the same for all rows
-        if report:
-            self.threads.set(report[0][2])
+from prometheus_client.core import REGISTRY, HistogramMetricFamily
+from prometheus_client.registry import Collector
 
 
 class Stats:
@@ -216,6 +170,84 @@ class WorkerStats:
         ]
 
 
+class CustomHistogram(Collector):
+    def __init__(self, name: str, stats: Stats, bins: list):
+        self.name = name
+        self.stats = stats
+        self.bins = bins
+
+    def get_buckets(self, name):
+        td = self.stats.cumulative_counts.get(name)
+        if td is None:
+            return [["+Inf", 0]]
+
+        # create buckets from 10 ... 180
+        td_hist = [[x, int(td.cdf((int(x) + 1) / 1000) * td.weight)] for x in self.bins]
+        td_hist.append(["+Inf", td.weight])
+
+        return td.mean * 1000 * td.weight, td_hist
+
+    def collect(self):
+        sum_value, buckets = self.get_buckets(self.name)
+        yield HistogramMetricFamily(
+            f"{self.name}_latency_ms",
+            f"Latency in ms for {self.name}",
+            buckets,
+            sum_value,
+        )
+
+
+class Prom:
+    def __init__(self, prom_port: int = 26260, stats: Stats = None, bins: list = []):
+        self.prom_latency: dict[str, list[prom.Gauge]] = {}
+        self.stats = stats
+        self.bins = bins
+
+        # don't stop just because prom server can't start
+        try:
+            prom.start_http_server(prom_port)
+        except OSError as e:
+            logger.warning(f"Cannot start prometheus server: {e}")
+
+        self.threads = prom.Gauge(
+            "threads", "count of connection threads to the database."
+        )
+
+    def publish(self, report: list, td: dict = {}):
+        for row in report:
+            id = row[1]
+
+            if id not in self.prom_latency:
+                self.prom_latency[id] = []
+
+                REGISTRY.register(CustomHistogram(id, self.stats, self.bins))
+
+                self.prom_latency[id].append(
+                    prom.Gauge(f"{id}__tot_ops", "total count of ops")
+                )
+                self.prom_latency[id].append(
+                    prom.Gauge(
+                        f"{id}__tot_ops_s", "derived value from tot_ops / elapsed"
+                    )
+                )
+                self.prom_latency[id].append(
+                    prom.Gauge(f"{id}__period_ops", "ops count for the recent window")
+                )
+                self.prom_latency[id].append(
+                    prom.Gauge(
+                        f"{id}__period_ops_s",
+                        "derived value from period_ops / window duration",
+                    )
+                )
+
+            for idx, v in enumerate(row[3:6]):
+                self.prom_latency[id][idx].set(v)
+
+        # threads value is the same for all rows
+        if report:
+            self.threads.set(report[0][2])
+
+
 def get_driver_from_scheme(scheme: str):
     return {
         "postgres": "postgres",
@@ -654,23 +686,6 @@ def ddl_to_yaml(ddl: str):
             elif within_brackets > 0 and i == ",":
                 col_def += ":"
 
-        # process the content within parenthesis in the
-        # CREATE TABLE stmt char by char to distinguish
-        # the comma for separating columns vs the comma
-        # included in single quote strings such as those in DEFAULT
-        # eg: mycol STRING NULL DEFAULT 'corporate, inc'
-        within_quote = False
-        col_def_str = col_def
-        col_def = ""
-        for i in col_def_str:
-            if i == "'":
-                within_quote = not within_quote
-                continue
-            if within_quote:
-                continue
-            else:
-                col_def += i
-
         col_def = [x.strip().lower() for x in col_def.split(",")]
 
         ll = []

{dbworkload-0.8.0 → dbworkload-0.8.3}/dbworkload/utils/simplefaker.py

@@ -1,12 +1,14 @@
+import builtins
 import csv
 import datetime as dt
 import logging
 import multiprocessing as mp
 import os
-import pandas as pd
-import uuid
 import random
-import builtins
+import uuid
+
+import pandas as pd
+
 from .common import import_class_at_runtime
 
 logger = logging.getLogger("dbworkload")

{dbworkload-0.8.0 → dbworkload-0.8.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "dbworkload"
-version = "0.8.0"
+version = "0.8.3"
 description = "Workload framework"
 authors = ["Fabio Ghirardello"]
 license = "GPLv3+"