feat: unify database class (#1068)

Author: davidgamez

api/src/feeds/impl/feeds_api_impl.py

@@ -45,7 +45,7 @@
 from feeds.impl.error_handling import raise_http_error, raise_http_validation_error, convert_exception
 from middleware.request_context import is_user_email_restricted
 from utils.date_utils import valid_iso_date
-from utils.logger import Logger
+from shared.common.logging_utils import Logger
 
 T = TypeVar("T", bound="BasicFeed")
 
@@ -66,7 +66,7 @@ def __init__(self) -> None:
     def get_feed(self, id: str, db_session: Session) -> BasicFeed:
         """Get the specified feed from the Mobility Database."""
         is_email_restricted = is_user_email_restricted()
-        self.logger.info(f"User email is restricted: {is_email_restricted}")
+        self.logger.debug(f"User email is restricted: {is_email_restricted}")
 
         feed = (
             FeedFilter(stable_id=id, provider__ilike=None, producer_url__ilike=None, status=None)
@@ -98,7 +98,7 @@ def get_feeds(
     ) -> List[BasicFeed]:
         """Get some (or all) feeds from the Mobility Database."""
         is_email_restricted = is_user_email_restricted()
-        self.logger.info(f"User email is restricted: {is_email_restricted}")
+        self.logger.debug(f"User email is restricted: {is_email_restricted}")
         feed_filter = FeedFilter(
             status=status, provider__ilike=provider, producer_url__ilike=producer_url, stable_id=None
         )
@@ -137,7 +137,6 @@ def _get_gtfs_feed(
         query = get_gtfs_feeds_query(
             db_session=db_session, stable_id=stable_id, include_options_for_joinedload=include_options_for_joinedload
         )
-        self.logger.debug("Query: %s", str(query.statement.compile(compile_kwargs={"literal_binds": True})))
         results = query.all()
         if len(results) == 0:
             return None

api/src/feeds/impl/models/validation_report_impl.py

@@ -1,6 +1,6 @@
 from shared.database_gen.sqlacodegen_models import Validationreport
 from feeds_gen.models.validation_report import ValidationReport
-from utils.logger import Logger
+from shared.common.logging_utils import Logger
 
 DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%S"
 

api/src/scripts/load_dataset_on_create.py

@@ -10,7 +10,7 @@
 from google.cloud.pubsub_v1.futures import Future
 
 from shared.database_gen.sqlacodegen_models import Feed
-from utils.logger import Logger
+from shared.common.logging_utils import Logger
 
 # Lazy create so we won't try to connect to google cloud when the file is imported.
 pubsub_client = None

api/src/scripts/populate_db.py

@@ -9,7 +9,7 @@
 
 from shared.database.database import Database
 from shared.database_gen.sqlacodegen_models import Feed, Gtfsrealtimefeed, Gtfsfeed, Gbfsfeed
-from utils.logger import Logger
+from shared.common.logging_utils import Logger
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session

api/src/scripts/populate_db_test_data.py

@@ -17,7 +17,7 @@
     Officialstatushistory,
 )
 from scripts.populate_db import set_up_configs, DatabasePopulateHelper
-from utils.logger import Logger
+from shared.common.logging_utils import Logger
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:

api/src/shared/common/logging_utils.py

@@ -0,0 +1,35 @@
+import logging
+import os
+
+
+def get_env_logging_level():
+    """
+    Get the logging level from the environment via OS variable LOGGING_LEVEL. Returns INFO if not set.
+    """
+    return os.getenv("LOGGING_LEVEL", "INFO")
+
+
+class Logger:
+    """
+    Util class for logging information, errors or warnings
+    """
+
+    def __init__(self, name):
+        """
+        Initialize the logger
+        """
+        formatter = logging.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s")
+
+        console_handler = logging.StreamHandler()
+        console_handler.setFormatter(formatter)
+
+        self.logger = logging.getLogger(name)
+        self.logger.addHandler(console_handler)
+        self.logger.setLevel(get_env_logging_level())
+
+    def get_logger(self):
+        """
+        Get the logger instance
+        :return: the logger instance
+        """
+        return self.logger

api/src/shared/database/database.py

@@ -5,12 +5,14 @@
 import uuid
 from typing import Type, Callable
 from dotenv import load_dotenv
-from sqlalchemy import create_engine
-from sqlalchemy.orm import load_only, Query, class_mapper, Session
+from sqlalchemy import create_engine, text, event
+from sqlalchemy.orm import load_only, Query, class_mapper, Session, mapper
 from shared.database_gen.sqlacodegen_models import Base, Feed, Gtfsfeed, Gtfsrealtimefeed, Gbfsfeed
 from sqlalchemy.orm import sessionmaker
 import logging
 
+from shared.common.logging_utils import get_env_logging_level
+
 
 def generate_unique_id() -> str:
     """
@@ -42,7 +44,48 @@ def configure_polymorphic_mappers():
     gbfsfeed_mapper.polymorphic_identity = Gbfsfeed.__tablename__.lower()
 
 
-def with_db_session(func):
+def set_cascade(mapper, class_):
+    """
+    Set cascade for relationships in Gtfsfeed.
+    This allows to delete/add the relationships when their respective relation array changes.
+    """
+    if class_.__name__ == "Gtfsfeed":
+        for rel in class_.__mapper__.relationships:
+            if rel.key in [
+                "redirectingids",
+                "redirectingids_",
+                "externalids",
+                "externalids_",
+            ]:
+                rel.cascade = "all, delete-orphan"
+
+
+def mapper_configure_listener(mapper, class_):
+    """
+    Mapper configure listener
+    """
+    set_cascade(mapper, class_)
+    configure_polymorphic_mappers()
+
+
+# Add the mapper_configure_listener to the mapper_configured event
+event.listen(mapper, "mapper_configured", mapper_configure_listener)
+
+
+def refresh_materialized_view(session: "Session", view_name: str) -> bool:
+    """
+    Refresh Materialized view by name.
+    @return: True if the view was refreshed successfully, False otherwise
+    """
+    try:
+        session.execute(text(f"REFRESH MATERIALIZED VIEW CONCURRENTLY {view_name}"))
+        return True
+    except Exception as error:
+        logging.error(f"Error raised while refreshing view: {error}")
+        return False
+
+
+def with_db_session(func=None, db_url: str | None = None):
     """
     Decorator to handle the session management for the decorated function.
 
@@ -58,12 +101,15 @@ def with_db_session(func):
           exception occurs, and closed in either case.
         - The session is then passed to the decorated function as the 'db_session' keyword argument.
         - If 'db_session' is already provided, it simply calls the decorated function with the existing session.
+        - The echoed SQL queries will be logged if the environment variable LOGGING_LEVEL is set to DEBUG.
     """
+    if func is None:
+        return lambda f: with_db_session(f, db_url=db_url)
 
     def wrapper(*args, **kwargs):
         db_session = kwargs.get("db_session")
         if db_session is None:
-            db = Database()
+            db = Database(echo_sql=get_env_logging_level() == "DEBUG", feeds_database_url=db_url)
             with db.start_db_session() as session:
                 kwargs["db_session"] = session
                 return func(*args, **kwargs)
@@ -89,12 +135,16 @@ def __new__(cls, *args, **kwargs):
                     cls.instance = object.__new__(cls)
         return cls.instance
 
-    def __init__(self, echo_sql=False):
+    def __init__(self, echo_sql=False, feeds_database_url: str | None = None):
         """
         Initializes the database instance
-        :param echo_sql: whether to echo the SQL queries or not
-        echo_sql set to False reduces the amount of information and noise going to the logs.
-        In case of errors, the exceptions will still contain relevant information about the failing queries.
+
+        :param echo_sql: whether to echo the SQL queries or not echo_sql.
+            False reduces the amount of information and noise going to the logs.
+            In case of errors, the exceptions will still contain relevant information about the failing queries.
+
+        :param feeds_database_url: The URL of the target database.
+            If it's None the URL will be assigned from the environment variable FEEDS_DATABASE_URL.
         """
 
         # This init function is called each time we call Database(), but in the case of a singleton, we only want to
@@ -107,7 +157,7 @@ def __init__(self, echo_sql=False):
             load_dotenv()
             self.logger = logging.getLogger(__name__)
             self.connection_attempts = 0
-            database_url = os.getenv("FEEDS_DATABASE_URL")
+            database_url = feeds_database_url if feeds_database_url else os.getenv("FEEDS_DATABASE_URL")
             if database_url is None:
                 raise Exception("Database URL not provided.")
             self.pool_size = int(os.getenv("DB_POOL_SIZE", 10))

api/src/utils/logger.py

@@ -125,29 +125,3 @@ async def async_emit(self, record):
             jsonPayload=json_payload,
         )
         self.logger.info(json.dumps(log_record.__dict__))
-
-
-class Logger:
-    """
-    Util class for logging information, errors or warnings
-    """
-
-    def __init__(self, name):
-        """
-        Initialize the logger
-        """
-        formatter = logging.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s")
-
-        console_handler = logging.StreamHandler()
-        console_handler.setFormatter(formatter)
-
-        self.logger = logging.getLogger(name)
-        self.logger.addHandler(console_handler)
-        self.logger.setLevel(logging.DEBUG)
-
-    def get_logger(self):
-        """
-        Get the logger instance
-        :return: the logger instance
-        """
-        return self.logger

api/tests/utils/test_logger.py

@@ -1,7 +1,8 @@
 import unittest
 from unittest.mock import patch
 
-from utils.logger import HttpRequest, LogRecord, AsyncStreamHandler, GCPLogHandler, Logger
+from utils.logger import HttpRequest, LogRecord, AsyncStreamHandler, GCPLogHandler
+from shared.common.logging_utils import Logger
 
 
 class TestLogger(unittest.TestCase):

functions-python/backfill_dataset_service_date_range/.coveragerc

@@ -1,9 +1,8 @@
 [run]
 omit =
     */test*/*
-    */database_gen/*
-    */dataset_service/*
     */helpers/*
+    */shared/*
 
 [report]
 exclude_lines =