Merge pull request #699 from tisnik/lcore-741-refactored-code-to-connect-to-databases

LCORE-741: refactored code used to connect to databases with quota

Author: tisnik

src/quota/connect_pg.py

@@ -0,0 +1,31 @@
+"""PostgreSQL connection handler."""
+
+from typing import Any
+import psycopg2
+
+from log import get_logger
+from models.config import PostgreSQLDatabaseConfiguration
+
+logger = get_logger(__name__)
+
+
+def connect_pg(config: PostgreSQLDatabaseConfiguration) -> Any:
+    """Initialize connection to PostgreSQL database."""
+    logger.info("Connecting to PostgreSQL storage")
+    try:
+        connection = psycopg2.connect(
+            host=config.host,
+            port=config.port,
+            user=config.user,
+            password=config.password.get_secret_value(),
+            dbname=config.db,
+            sslmode=config.ssl_mode,
+            # sslrootcert=config.ca_cert_path,
+            gssencmode=config.gss_encmode,
+        )
+        if connection is not None:
+            connection.autocommit = True
+        return connection
+    except psycopg2.Error as e:
+        logger.exception("Error connecting to PostgreSQL database:\n%s", e)
+        raise

src/quota/connect_sqlite.py

@@ -0,0 +1,25 @@
+"""SQLite connection handler."""
+
+import sqlite3
+from typing import Any
+
+from log import get_logger
+from models.config import SQLiteDatabaseConfiguration
+
+logger = get_logger(__name__)
+
+
+def connect_sqlite(config: SQLiteDatabaseConfiguration) -> Any:
+    """Initialize connection to database."""
+    logger.info("Connecting to SQLite storage")
+    # make sure the connection will have known state
+    # even if SQLite is not alive
+    connection = None
+    try:
+        connection = sqlite3.connect(database=config.db_path)
+        if connection is not None:
+            connection.autocommit = True
+        return connection
+    except sqlite3.Error as e:
+        logger.exception("Error initializing SQLite cache:\n%s", e)
+        raise

src/runners/quota_scheduler.py

@@ -4,19 +4,18 @@
 from threading import Thread
 from time import sleep
 
-import sqlite3
-import psycopg2
 
 import constants
 from log import get_logger
 from models.config import (
     Configuration,
     QuotaHandlersConfiguration,
     QuotaLimiterConfiguration,
-    PostgreSQLDatabaseConfiguration,
-    SQLiteDatabaseConfiguration,
 )
 
+from quota.connect_pg import connect_pg
+from quota.connect_sqlite import connect_sqlite
+
 from quota.sql import (
     CREATE_QUOTA_TABLE,
     INCREASE_QUOTA_STATEMENT_PG,
@@ -211,41 +210,6 @@ def connect(config: QuotaHandlersConfiguration) -> Any:
     return None
 
 
-def connect_pg(config: PostgreSQLDatabaseConfiguration) -> Any:
-    """Initialize connection to PostgreSQL database."""
-    logger.info("Connecting to PostgreSQL storage")
-    connection = psycopg2.connect(
-        host=config.host,
-        port=config.port,
-        user=config.user,
-        password=config.password.get_secret_value(),
-        dbname=config.db,
-        sslmode=config.ssl_mode,
-        # sslrootcert=config.ca_cert_path,
-        gssencmode=config.gss_encmode,
-    )
-    if connection is not None:
-        connection.autocommit = True
-    return connection
-
-
-def connect_sqlite(config: SQLiteDatabaseConfiguration) -> Any:
-    """Initialize connection to database."""
-    logger.info("Connecting to SQLite storage")
-    # make sure the connection will have known state
-    # even if SQLite is not alive
-    connection = None
-    try:
-        connection = sqlite3.connect(database=config.db_path)
-    except sqlite3.Error as e:
-        if connection is not None:
-            connection.close()
-        logger.exception("Error initializing SQLite cache:\n%s", e)
-        raise
-    connection.autocommit = True
-    return connection
-
-
 def init_tables(connection: Any) -> None:
     """Initialize tables used by quota limiter."""
     logger.info("Initializing tables for quota limiter")