workercommon 0.4.1__py3-none-any.whl

workercommon/database/pg.py

@@ -0,0 +1,278 @@
+#!/usr/bin/env python3
+import sys
+
+if sys.version_info < (3, 8):
+    raise RuntimeError("At least Python 3.8 is required")
+
+from . import base
+from typing import (
+    Optional,
+    Any,
+    List,
+    Dict,
+    Generator,
+    Iterable,
+    Tuple,
+    Union,
+    Iterator,
+    Sequence,
+)
+from threading import Lock
+import logging, weakref
+import psycopg2, psycopg2.extensions
+
+psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
+psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY)
+
+
+class BaseConnection(base.Connection):
+    def get(self) -> psycopg2.extensions.connection:
+        raise NotImplementedError
+
+
+class Cursor(base.Cursor):
+    def __init__(self, connection: BaseConnection):
+        self.connection = connection
+        self.cursor = None
+
+    def get(self) -> psycopg2.extensions.cursor:
+        if self.cursor is None:
+            raise RuntimeError("Cursor is not available")
+        return self.cursor
+
+    def execute_direct(self, query, *args) -> psycopg2.extensions.cursor:
+        cursor = self.get()
+        logging.debug(f"execute_direct({query}, {args})")
+        cursor.execute(query, args)
+        return cursor
+
+    @classmethod
+    def iter_inner(cls, cursor) -> Generator[Dict[str, Any], None, None]:
+        columns = [desc[0] for desc in cursor.description]
+        while (row := cursor.fetchone()) is not None:
+            yield dict(zip(columns, row))
+
+    def execute(self, query, *args) -> Iterator[Dict[str, Any]]:
+        cursor = self.get()
+        logging.debug(f"execute({query}, {args})")
+        cursor.execute(query, args)
+        return iter(self.iter_inner(cursor))
+
+    def __iter__(self) -> Iterator[Dict[str, Any]]:
+        return iter(self.iter_inner(self.get()))
+
+    @property
+    def rowcount(self) -> int:
+        return self.get().rowcount
+
+    @property
+    def columns(self) -> List[str]:
+        return [desc[0] for desc in self.get().description]
+
+    def init_cursor(self) -> None:
+        self.cursor = self.connection.get().cursor()
+
+    def __enter__(self) -> Any:
+        self.init_cursor()
+        return self
+
+    def __exit__(self, type, value, traceback) -> None:
+        if self.cursor is not None:
+            self.cursor = None
+            if value is None:
+                self.connection.commit()
+            else:
+                self.connection.rollback()
+
+
+class Connection(BaseConnection):
+    def __init__(
+        self,
+        host: Optional[str],
+        port: Optional[int],
+        username: str,
+        password: Optional[str],
+        database: str,
+    ):
+        self.lock = Lock()
+        port_number = "5432"
+        if port is not None and port > 0:
+            port_number = str(port)
+        self.connection: Optional[psycopg2.extensions.connection] = psycopg2.connect(
+            host=host,
+            port=port_number,
+            user=username,
+            password=password,
+            database=database,
+        )
+
+    def get(self) -> psycopg2.extensions.connection:
+        with self.lock:
+            if self.connection is None:
+                raise RuntimeError("Connection has been closed")
+            return self.connection
+
+    def close(self) -> None:
+        with self.lock:
+            if self.connection is not None:
+                self.connection.close()
+                self.connection = None
+
+    def commit(self) -> None:
+        self.get().commit()
+
+    def rollback(self) -> None:
+        self.get().rollback()
+
+    def cursor(self) -> Cursor:
+        return Cursor(self)
+
+
+class BaseTable(base.BaseTable):
+    ORDERINGS = frozenset(["ASC", "DESC"])
+
+    def __init__(
+        self,
+        name: str,
+        primary_key: str,
+        columns: Sequence[str],
+        ordering: Optional[str] = "ASC",
+    ):
+        super().__init__(name, primary_key, columns, ordering)
+
+    @staticmethod
+    def convert_cursor(cursor: base.Cursor) -> Cursor:
+        if not isinstance(cursor, Cursor):
+            raise RuntimeError(f"Not a PgCursor: {cursor}")
+        return cursor
+
+    def insert(self, base_cursor: base.Cursor, **input_values: Any) -> Any:
+        cursor = self.convert_cursor(base_cursor)
+        self.validate(input_values.keys())
+        columns, values = self.unpack(input_values)
+
+        try:
+            column_string = ", ".join(columns)
+            value_string = ", ".join(("%s" for i in columns))
+            row = next(
+                cursor.execute(
+                    f"INSERT INTO {self.name}({column_string}) VALUES({value_string}) RETURNING {self.primary_key}",
+                    *values,
+                )
+            )
+            return row[self.primary_key]
+
+        except StopIteration:
+            raise RuntimeError(f"Failed to insert values into {self.name}")
+
+    def upsert_on_column(
+        self,
+        base_cursor: base.Cursor,
+        column_name: str,
+        column_value: Any,
+        **input_values: Any,
+    ) -> Any:
+        cursor = self.convert_cursor(base_cursor)
+        self.validate(input_values.keys())
+
+        if self.primary_key in input_values:
+            input_values[self.primary_key] = column_value
+
+        try:
+            if column_value is not None:
+                row = next(
+                    cursor.execute(
+                        f"SELECT {self.primary_key} FROM {self.name} WHERE {column_name} = %s",
+                        column_value,
+                    )
+                )
+            else:
+                row = next(
+                    cursor.execute(
+                        f"SELECT {self.primary_key} FROM {self.name} WHERE {column_name} IS NULL"
+                    )
+                )
+
+            if cursor.rowcount > 1:
+                raise RuntimeError(f"Multiple results were found in {self.name}")
+
+            pkey_value = row[self.primary_key]
+
+            if not input_values or (
+                len(input_values) == 1 and input_values == {column_name: column_value}
+            ):
+                # No changes are necessary
+                return pkey_value
+
+            # Update existing row
+            columns, values = self.unpack(input_values)
+            set_string = ", ".join((f"{column} = %s" for column in columns))
+            cursor.execute(
+                f"UPDATE {self.name} SET {set_string} WHERE {self.primary_key} = %s",
+                *(values + (pkey_value,)),
+            )
+
+            try:
+                return input_values[self.primary_key]
+            except KeyError:
+                return pkey_value
+
+        except StopIteration:
+            if column_name not in input_values:
+                # Specifying the column multiple times seems silly.
+                input_values[column_name] = column_value
+
+            return self.insert(cursor, **input_values)
+
+    def select_by_columns(
+        self, base_cursor: base.Cursor, limit: Optional[int] = None, **input_values: Any
+    ) -> Iterator[Dict[str, Any]]:
+        cursor = self.convert_cursor(base_cursor)
+        self.validate(input_values.keys())
+        fixed_input_values = list(input_values.items())
+        columns, values = self.unpack(fixed_input_values)
+
+        column_string = ", ".join(self.columns)
+        where_string = " AND ".join(
+            (
+                (f"{column} = %s" if value is not None else f"{column} IS NULL")
+                for column, value in fixed_input_values
+            )
+        )
+
+        cleaned_values = [v for v in values if v is not None]
+
+        limit_str = f"LIMIT {limit}" if limit and limit >= 0 else ""
+        return cursor.execute(
+            f"SELECT {column_string} FROM {self.name} WHERE {where_string} ORDER BY {self.primary_key} {self.ordering} {limit_str}",
+            *cleaned_values,
+        )
+
+    def select_custom_where(
+        self, base_cursor: base.Cursor, where: str, *values
+    ) -> Iterator[Dict[str, Any]]:
+        cursor = self.convert_cursor(base_cursor)
+        column_string = ", ".join(self.columns)
+        return cursor.execute(
+            f"SELECT {column_string} FROM {self.name} WHERE {where}", *values
+        )
+
+    def select_by_id(self, base_cursor: base.Cursor, pkey_value: Any) -> Dict[str, Any]:
+        cursor = self.convert_cursor(base_cursor)
+        column_string = ", ".join(self.columns)
+        try:
+            return next(
+                cursor.execute(
+                    f"SELECT {column_string} FROM {self.name} WHERE {self.primary_key} = %s",
+                    pkey_value,
+                )
+            )
+        except StopIteration:
+            raise KeyError(pkey_value)
+
+    def delete(self, base_cursor: base.Cursor, pkey_value: Any) -> bool:
+        cursor = self.convert_cursor(base_cursor)
+        cursor.execute(
+            f"DELETE FROM {self.name} WHERE {self.primary_key} = %s", pkey_value
+        )
+        return cursor.rowcount > 0

workercommon/locking.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+import sys
+
+if sys.version_info < (3, 8):
+    raise RuntimeError("At least Python 3.8 is required")
+
+from sys import stdout, stderr
+from os import fchmod
+import fcntl, logging
+from time import sleep
+from typing import Union
+from pathlib import Path
+
+
+class LockFile(object):
+    def __init__(self, location: Union[str, Path], exclusive: bool = True):
+        self.location = str(location)
+        self.fd = None
+        self.exclusive = exclusive
+
+    def __enter__(self):
+        if self.fd is not None:
+            raise RuntimeError("Cannot lock twice")
+        self.fd = open(self.location, "w+b")
+        if self.fd is None:
+            raise RuntimeError(f"Failed to open {self.location}")
+        try:
+            fchmod(self.fd.fileno(), 0o666)
+        except OSError:
+            pass
+
+        for i in range(10):
+            try:
+                mode = fcntl.LOCK_NB
+                if self.exclusive:
+                    mode |= fcntl.LOCK_EX
+                else:
+                    mode |= fcntl.LOCK_SH
+
+                logging.debug(f"Attempting to lock file {self.location}")
+                fcntl.lockf(self.fd, mode)
+                return
+            except IOError:
+                logging.warning("Waiting for existing lock...")
+                sleep(1)
+                continue
+        raise RuntimeError("The existing bot did not exit in time; dying")
+
+    def __exit__(self, type, value, traceback):
+        fcntl.lockf(self.fd, fcntl.LOCK_UN)
+        self.fd.close()
+        self.fd = None