Add PostgresImporter

Author: Gerrrr

hunter/config.py

@@ -8,6 +8,7 @@
 
 from hunter.grafana import GrafanaConfig
 from hunter.graphite import GraphiteConfig
+from hunter.postgres import PostgresConfig
 from hunter.slack import SlackConfig
 from hunter.test_config import TestConfig, create_test_config
 from hunter.util import merge_dict_list
@@ -20,6 +21,7 @@ class Config:
     tests: Dict[str, TestConfig]
     test_groups: Dict[str, List[TestConfig]]
     slack: SlackConfig
+    postgres: PostgresConfig
 
 
 @dataclass
@@ -110,6 +112,27 @@ def load_config_from(config_file: Path) -> Config:
                 bot_token=config["slack"]["token"],
             )
 
+        postgres_config = None
+        if config.get("postgres") is not None:
+            if not config["postgres"]["hostname"]:
+                raise ValueError("postgres.hostname")
+            if not config["postgres"]["port"]:
+                raise ValueError("postgres.port")
+            if not config["postgres"]["username"]:
+                raise ValueError("postgres.username")
+            if not config["postgres"]["password"]:
+                raise ValueError("postgres.password")
+            if not config["postgres"]["database"]:
+                raise ValueError("postgres.database")
+
+            postgres_config = PostgresConfig(
+                hostname=config["postgres"]["hostname"],
+                port=config["postgres"]["port"],
+                username=config["postgres"]["username"],
+                password=config["postgres"]["password"],
+                database=config["postgres"]["database"],
+            )
+
         templates = load_templates(config)
         tests = load_tests(config, templates)
         groups = load_test_groups(config, tests)
@@ -118,6 +141,7 @@ def load_config_from(config_file: Path) -> Config:
             graphite=graphite_config,
             grafana=grafana_config,
             slack=slack_config,
+            postgres=postgres_config,
             tests=tests,
             test_groups=groups,
         )

hunter/importer.py

@@ -9,12 +9,15 @@
 from hunter.config import Config
 from hunter.data_selector import DataSelector
 from hunter.graphite import DataPoint, Graphite, GraphiteError
+from hunter.postgres import Postgres
 from hunter.series import Metric, Series
 from hunter.test_config import (
     CsvMetric,
     CsvTestConfig,
     GraphiteTestConfig,
     HistoStatTestConfig,
+    PostgresMetric,
+    PostgresTestConfig,
     TestConfig,
 )
 from hunter.util import (
@@ -431,17 +434,122 @@ def fetch_all_metric_names(self, test: HistoStatTestConfig) -> List[str]:
         return metric_names
 
 
+class PostgresImporter:
+    __postgres: Postgres
+
+    def __init__(self, postgres: Postgres):
+        self.__postgres = postgres
+
+    @staticmethod
+    def __selected_metrics(
+        defined_metrics: Dict[str, PostgresMetric], selected_metrics: Optional[List[str]]
+    ) -> Dict[str, PostgresMetric]:
+
+        if selected_metrics is not None:
+            return {name: defined_metrics[name] for name in selected_metrics}
+        else:
+            return defined_metrics
+
+    def fetch_data(self, test_conf: TestConfig, selector: DataSelector = DataSelector()) -> Series:
+        if not isinstance(test_conf, PostgresTestConfig):
+            raise ValueError("Expected PostgresTestConfig")
+
+        if selector.branch:
+            raise ValueError("Postgres tests don't support branching yet")
+
+        since_time = selector.since_time
+        until_time = selector.until_time
+        if since_time.timestamp() > until_time.timestamp():
+            raise DataImportError(
+                f"Invalid time range: ["
+                f"{format_timestamp(int(since_time.timestamp()))}, "
+                f"{format_timestamp(int(until_time.timestamp()))}]"
+            )
+        metrics = self.__selected_metrics(test_conf.metrics, selector.metrics)
+
+        columns, rows = self.__postgres.fetch_data(test_conf.query)
+
+        # Decide which columns to fetch into which components of the result:
+        try:
+            time_index: int = columns.index(test_conf.time_column)
+            attr_indexes: List[int] = [columns.index(c) for c in test_conf.attributes]
+            metric_names = [m.name for m in metrics.values()]
+            metric_columns = [m.column for m in metrics.values()]
+            metric_indexes: List[int] = [columns.index(c) for c in metric_columns]
+        except ValueError as err:
+            raise DataImportError(f"Column not found {err.args[0]}")
+
+        time: List[int] = []
+        data: Dict[str, List[float]] = {}
+        for n in metric_names:
+            data[n] = []
+        attributes: Dict[str, List[str]] = {}
+        for i in attr_indexes:
+            attributes[columns[i]] = []
+
+        for row in rows:
+            ts: datetime = row[time_index]
+            if since_time is not None and ts < since_time:
+                continue
+            if until_time is not None and ts >= until_time:
+                continue
+            time.append(int(ts.timestamp()))
+
+            # Read metric values. Note we can still fail on conversion to float,
+            # because the user is free to override the column selection and thus
+            # they may select a column that contains non-numeric data:
+            for (name, i) in zip(metric_names, metric_indexes):
+                try:
+                    data[name].append(float(row[i]))
+                except ValueError as err:
+                    raise DataImportError(
+                        "Could not convert value in column " + headers[i] + ": " + err.args[0]
+                    )
+
+            # Attributes are just copied as-is, with no conversion:
+            for i in attr_indexes:
+                attributes[columns[i]].append(row[i])
+
+        # Convert metrics to series.Metrics
+        metrics = {m.name: Metric(m.direction, m.scale) for m in metrics.values()}
+
+        # Leave last n points:
+        time = time[-selector.last_n_points :]
+        tmp = data
+        data = {}
+        for k, v in tmp.items():
+            data[k] = v[-selector.last_n_points :]
+        tmp = attributes
+        attributes = {}
+        for k, v in tmp.items():
+            attributes[k] = v[-selector.last_n_points :]
+
+        return Series(
+            test_conf.name,
+            branch=None,
+            time=time,
+            metrics=metrics,
+            data=data,
+            attributes=attributes,
+        )
+
+    def fetch_all_metric_names(self, test_conf: PostgresTestConfig) -> List[str]:
+        return [m for m in test_conf.metrics.keys()]
+
+
 class Importers:
     __config: Config
     __csv_importer: Optional[CsvImporter]
     __graphite_importer: Optional[GraphiteImporter]
     __histostat_importer: Optional[HistoStatImporter]
+    __postgres_importer: Optional[PostgresImporter]
 
     def __init__(self, config: Config):
         self.__config = config
         self.__csv_importer = None
         self.__graphite_importer = None
         self.__histostat_importer = None
+        self.__postgres_importer = None
 
     def csv_importer(self) -> CsvImporter:
         if self.__csv_importer is None:
@@ -458,12 +566,19 @@ def histostat_importer(self) -> HistoStatImporter:
             self.__histostat_importer = HistoStatImporter()
         return self.__histostat_importer
 
+    def postgres_importer(self) -> PostgresImporter:
+        if self.__postgres_importer is None:
+            self.__postgres_importer = PostgresImporter(Postgres(self.__config.postgres))
+        return self.__postgres_importer
+
     def get(self, test: TestConfig) -> Importer:
         if isinstance(test, CsvTestConfig):
             return self.csv_importer()
         elif isinstance(test, GraphiteTestConfig):
             return self.graphite_importer()
         elif isinstance(test, HistoStatTestConfig):
             return self.histostat_importer()
+        elif isinstance(test, PostgresTestConfig):
+            return self.postgres_importer()
         else:
             raise ValueError(f"Unsupported test type {type(test)}")

hunter/postgres.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+
+import psycopg2
+
+
+@dataclass
+class PostgresConfig:
+    hostname: str
+    port: int
+    username: str
+    password: str
+    database: str
+
+
+class Postgres:
+    __conn = None
+    __config = None
+
+    def __init__(self, config: PostgresConfig):
+        self.__config = config
+
+    def __get_conn(self) -> psycopg2.extensions.connection:
+        if self.__conn is None:
+            self.__conn = psycopg2.connect(
+                host=self.__config.hostname,
+                port=self.__config.port,
+                user=self.__config.username,
+                password=self.__config.password,
+                database=self.__config.database,
+            )
+        return self.__conn
+
+    def fetch_data(self, query: str):
+        cursor = self.__get_conn().cursor()
+        cursor.execute(query)
+        columns = [c.name for c in cursor.description]
+        return (columns, cursor.fetchall())

hunter/test_config.py

@@ -121,12 +121,45 @@ def fully_qualified_metric_names(self):
         return HistoStatImporter().fetch_all_metric_names(self)
 
 
+@dataclass
+class PostgresMetric:
+    name: str
+    direction: int
+    scale: float
+    column: str
+
+
+@dataclass
+class PostgresTestConfig(TestConfig):
+    query: str
+    time_column: str
+    attributes: List[str]
+    metrics: Dict[str, PostgresMetric]
+
+    def __init__(
+        self,
+        name: str,
+        query: str,
+        time_column: str = "time",
+        metrics: List[PostgresMetric] = None,
+        attributes: List[str] = None,
+    ):
+        self.name = name
+        self.query = query
+        self.time_column = time_column
+        self.metrics = {m.name: m for m in metrics} if metrics else {}
+        self.attributes = attributes
+
+    def fully_qualified_metric_names(self) -> List[str]:
+        return list(self.metrics.keys())
+
+
 def create_test_config(name: str, config: Dict) -> TestConfig:
     """
     Loads properties of a test from a dictionary read from hunter's config file
     This dictionary must have the `type` property to determine the type of the test.
     Other properties depend on the type.
-    Currently supported test types are `fallout`, `graphite` and `csv`.
+    Currently supported test types are `fallout`, `graphite`, `csv`, and `psql`.
     """
     test_type = config.get("type")
     if test_type == "csv":
@@ -135,6 +168,8 @@ def create_test_config(name: str, config: Dict) -> TestConfig:
         return create_graphite_test_config(name, config)
     elif test_type == "histostat":
         return create_histostat_test_config(name, config)
+    elif test_type == "postgres":
+        return create_postgres_test_config(name, config)
     elif test_type is None:
         raise TestConfigError(f"Test type not set for test {name}")
     else:
@@ -226,3 +261,32 @@ def create_histostat_test_config(name: str, test_info: Dict) -> HistoStatTestCon
             f"Configuration referenced histostat file which does not exist: {file}"
         )
     return HistoStatTestConfig(name, file)
+
+
+def create_postgres_test_config(test_name: str, test_info: Dict) -> PostgresTestConfig:
+    try:
+        time_column = test_info.get("time_column", "time")
+        attributes = test_info.get("attributes", [])
+        metrics_info = test_info.get("metrics")
+        query = test_info["query"]
+
+        metrics = []
+        if isinstance(metrics_info, List):
+            for name in metrics_info:
+                metrics.append(CsvMetric(name, 1, 1.0))
+        elif isinstance(metrics_info, Dict):
+            for (metric_name, metric_conf) in metrics_info.items():
+                metrics.append(
+                    PostgresMetric(
+                        name=metric_name,
+                        column=metric_conf.get("column", metric_name),
+                        direction=int(metric_conf.get("direction", "1")),
+                        scale=float(metric_conf.get("scale", "1")),
+                    )
+                )
+        else:
+            raise TestConfigError(f"Metrics of the test {test_name} must be a list or dictionary")
+
+        return PostgresTestConfig(test_name, query, time_column, metrics, attributes)
+    except KeyError as e:
+        raise TestConfigError(f"Configuration key not found in test {test_name}: {e.args[0]}")