adjust fn used by worker to determine timeseries enabled (#975)

Author: adrian-codecov

database/engine.py

@@ -3,13 +3,13 @@
 from decimal import Decimal
 
 from shared.config import get_config
+from shared.timeseries.helpers import is_timeseries_enabled
 from shared.utils.ReportEncoder import ReportEncoder
 from sqlalchemy import create_engine
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 
 import database.events  # noqa: F401
 from database.models.timeseries import TimeseriesBaseModel
-from helpers.timeseries import timeseries_enabled
 
 from .base import Base
 
@@ -44,7 +44,7 @@ def create_session(self):
             json_serializer=json_dumps,
         )
 
-        if timeseries_enabled():
+        if is_timeseries_enabled():
             self.timeseries_engine = create_engine(
                 self.timeseries_database_url,
                 json_serializer=json_dumps,

database/tests/unit/test_engine.py

@@ -8,7 +8,7 @@
 
 class TestDatabaseEngine:
     def test_session_get_bind_timeseries_disabled(self, sqlalchemy_connect_url, mocker):
-        mocker.patch("database.engine.timeseries_enabled", return_value=False)
+        mocker.patch("database.engine.is_timeseries_enabled", return_value=False)
 
         session_factory = SessionFactory(
             database_url=sqlalchemy_connect_url,
@@ -33,7 +33,7 @@ def test_session_get_bind_timeseries_disabled(self, sqlalchemy_connect_url, mock
         assert engine == session_factory.main_engine
 
     def test_session_get_bind_timeseries_enabled(self, sqlalchemy_connect_url, mocker):
-        mocker.patch("database.engine.timeseries_enabled", return_value=True)
+        mocker.patch("database.engine.is_timeseries_enabled", return_value=True)
 
         session_factory = SessionFactory(
             database_url=sqlalchemy_connect_url,

helpers/timeseries.py

@@ -1,9 +1,5 @@
 from shared.config import get_config
 
 
-def timeseries_enabled() -> bool:
-    return get_config("setup", "timeseries", "enabled", default=False)
-
-
 def backfill_max_batch_size() -> int:
     return get_config("setup", "timeseries", "backfill_max_batch_size", default=500)

services/tests/test_timeseries.py

@@ -105,7 +105,7 @@ class TestTimeseriesService(object):
     def test_insert_commit_measurement(
         self, dbsession, sample_report, repository, mocker
     ):
-        mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
+        mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
         mocker.patch(
             "services.report.ReportService.get_existing_report_for_commit",
             return_value=ReadOnlyReport.create_from_report(sample_report),
@@ -140,7 +140,7 @@ def test_insert_commit_measurement(
         assert measurement.value == 60.0
 
     def test_save_commit_measurements_no_report(self, dbsession, repository, mocker):
-        mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
+        mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
         mocker.patch(
             "services.report.ReportService.get_existing_report_for_commit",
             return_value=None,
@@ -167,7 +167,7 @@ def test_save_commit_measurements_no_report(self, dbsession, repository, mocker)
     def test_update_commit_measurement(
         self, dbsession, sample_report, repository, mocker
     ):
-        mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
+        mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
         mocker.patch(
             "services.report.ReportService.get_existing_report_for_commit",
             return_value=ReadOnlyReport.create_from_report(sample_report),
@@ -218,7 +218,7 @@ def test_update_commit_measurement(
     def test_commit_measurement_insert_flags(
         self, dbsession, sample_report, repository, mocker
     ):
-        mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
+        mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
         mocker.patch(
             "services.report.ReportService.get_existing_report_for_commit",
             return_value=ReadOnlyReport.create_from_report(sample_report),
@@ -291,7 +291,7 @@ def test_commit_measurement_insert_flags(
     def test_commit_measurement_update_flags(
         self, dbsession, sample_report, repository, mocker
     ):
-        mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
+        mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
         mocker.patch(
             "services.report.ReportService.get_existing_report_for_commit",
             return_value=ReadOnlyReport.create_from_report(sample_report),
@@ -390,7 +390,7 @@ def test_commit_measurement_update_flags(
     def test_commit_measurement_insert_components(
         self, dbsession, sample_report_for_components, repository, mocker
     ):
-        mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
+        mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
         mocker.patch(
             "services.report.ReportService.get_existing_report_for_commit",
             return_value=ReadOnlyReport.create_from_report(
@@ -575,7 +575,7 @@ def test_commit_measurement_insert_components(
     def test_commit_measurement_update_component(
         self, dbsession, sample_report_for_components, repository, mocker
     ):
-        mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
+        mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
         mocker.patch(
             "services.report.ReportService.get_existing_report_for_commit",
             return_value=ReadOnlyReport.create_from_report(
@@ -645,7 +645,7 @@ def test_commit_measurement_update_component(
         assert measurement.value == 50.0
 
     def test_commit_measurement_no_datasets(self, dbsession, mocker):
-        mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
+        mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
 
         repository = RepositoryFactory.create()
         dbsession.add(repository)
@@ -786,7 +786,7 @@ def test_backfill_batch_size(self, repository, mocker):
         assert batch_size == 100
 
     def test_delete_repository_data(self, dbsession, sample_report, repository, mocker):
-        mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
+        mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
         mocker.patch(
             "services.report.ReportService.get_existing_report_for_commit",
             return_value=ReadOnlyReport.create_from_report(sample_report),
@@ -825,7 +825,7 @@ def test_delete_repository_data(self, dbsession, sample_report, repository, mock
     def test_delete_repository_data_side_effects(
         self, dbsession, sample_report, repository, mocker
     ):
-        mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
+        mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
         mocker.patch(
             "services.report.ReportService.get_existing_report_for_commit",
             return_value=ReadOnlyReport.create_from_report(sample_report),
@@ -901,7 +901,7 @@ def validate_invariants(repository, other_repository):
                 == 16
             )
 
-        mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
+        mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
         mocker.patch(
             "services.report.ReportService.get_existing_report_for_commit",
             return_value=ReadOnlyReport.create_from_report(

services/timeseries.py

@@ -4,13 +4,14 @@
 
 from shared.components import Component
 from shared.reports.readonly import ReadOnlyReport
+from shared.timeseries.helpers import is_timeseries_enabled
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.orm import Session
 
 from database.models import Commit, Dataset, Measurement, MeasurementName
 from database.models.core import Repository
 from database.models.reports import RepositoryFlag
-from helpers.timeseries import backfill_max_batch_size, timeseries_enabled
+from helpers.timeseries import backfill_max_batch_size
 from services.report import ReportService
 from services.yaml import get_repo_yaml
 
@@ -20,7 +21,7 @@
 def save_commit_measurements(
     commit: Commit, dataset_names: Iterable[str] = None
 ) -> None:
-    if not timeseries_enabled():
+    if not is_timeseries_enabled():
         return
 
     if dataset_names is None:

tasks/base.py

@@ -8,6 +8,7 @@
 from django.db import transaction as django_transaction
 from shared.celery_router import route_tasks_based_on_user_plan
 from shared.metrics import Counter, Histogram
+from shared.timeseries.helpers import is_timeseries_enabled
 from sqlalchemy.exc import (
     DataError,
     IntegrityError,
@@ -22,7 +23,6 @@
 from helpers.checkpoint_logger.flows import TestResultsFlow, UploadFlow
 from helpers.log_context import LogContext, set_log_context
 from helpers.telemetry import TimeseriesTimer, log_simple_metric
-from helpers.timeseries import timeseries_enabled
 
 log = logging.getLogger("worker")
 
@@ -179,7 +179,7 @@ def _commit_django(self):
                 extra=dict(e=e),
             )
 
-        if timeseries_enabled():
+        if is_timeseries_enabled():
             try:
                 django_transaction.commit("timeseries")
             except Exception as e:
@@ -199,7 +199,7 @@ def _rollback_django(self):
                 extra=dict(e=e),
             )
 
-        if timeseries_enabled():
+        if is_timeseries_enabled():
             try:
                 django_transaction.rollback("timeseries")
             except Exception as e:

tasks/tests/integration/test_timeseries_backfill.py

@@ -10,8 +10,8 @@
 
 @pytest.mark.integration
 def test_backfill_dataset_run_impl(dbsession, mocker, mock_storage):
-    mocker.patch("services.timeseries.timeseries_enabled", return_value=True)
-    mocker.patch("tasks.timeseries_backfill.timeseries_enabled", return_value=True)
+    mocker.patch("services.timeseries.is_timeseries_enabled", return_value=True)
+    mocker.patch("tasks.timeseries_backfill.is_timeseries_enabled", return_value=True)
     mocked_app = mocker.patch.object(
         TimeseriesBackfillCommitsTask,
         "app",

tasks/tests/unit/test_timeseries_backfill_commits.py

@@ -8,7 +8,7 @@
 
 
 def test_backfill_commits_run_impl(dbsession, mocker):
-    mocker.patch("tasks.timeseries_backfill.timeseries_enabled", return_value=True)
+    mocker.patch("tasks.timeseries_backfill.is_timeseries_enabled", return_value=True)
     mocked_app = mocker.patch.object(
         TimeseriesBackfillCommitsTask,
         "app",
@@ -62,7 +62,7 @@ def test_backfill_commits_run_impl(dbsession, mocker):
 
 
 def test_backfill_commits_run_impl_timeseries_not_enabled(dbsession, mocker):
-    mocker.patch("tasks.timeseries_backfill.timeseries_enabled", return_value=False)
+    mocker.patch("tasks.timeseries_backfill.is_timeseries_enabled", return_value=False)
     mock_group = mocker.patch("tasks.timeseries_backfill.group")
 
     task = TimeseriesBackfillCommitsTask()

tasks/tests/unit/test_timeseries_backfill_dataset.py

@@ -11,7 +11,7 @@
 
 
 def test_backfill_dataset_run_impl(dbsession, mocker):
-    mocker.patch("tasks.timeseries_backfill.timeseries_enabled", return_value=True)
+    mocker.patch("tasks.timeseries_backfill.is_timeseries_enabled", return_value=True)
     mock_group = mocker.patch("tasks.timeseries_backfill.group")
 
     repository = RepositoryFactory.create()
@@ -70,7 +70,7 @@ def test_backfill_dataset_run_impl(dbsession, mocker):
 
 
 def test_backfill_dataset_run_impl_invalid_dataset(dbsession, mocker):
-    mocker.patch("tasks.timeseries_backfill.timeseries_enabled", return_value=True)
+    mocker.patch("tasks.timeseries_backfill.is_timeseries_enabled", return_value=True)
     mock_group = mocker.patch("tasks.timeseries_backfill.group")
 
     repository = RepositoryFactory.create()
@@ -97,7 +97,7 @@ def test_backfill_dataset_run_impl_invalid_dataset(dbsession, mocker):
 
 
 def test_backfill_dataset_run_impl_invalid_repository(dbsession, mocker):
-    mocker.patch("tasks.timeseries_backfill.timeseries_enabled", return_value=True)
+    mocker.patch("tasks.timeseries_backfill.is_timeseries_enabled", return_value=True)
     mock_group = mocker.patch("tasks.timeseries_backfill.group")
 
     repository = RepositoryFactory.create()
@@ -124,7 +124,7 @@ def test_backfill_dataset_run_impl_invalid_repository(dbsession, mocker):
 
 
 def test_backfill_dataset_run_impl_invalid_start_date(dbsession, mocker):
-    mocker.patch("tasks.timeseries_backfill.timeseries_enabled", return_value=True)
+    mocker.patch("tasks.timeseries_backfill.is_timeseries_enabled", return_value=True)
     mock_group = mocker.patch("tasks.timeseries_backfill.group")
 
     repository = RepositoryFactory.create()
@@ -151,7 +151,7 @@ def test_backfill_dataset_run_impl_invalid_start_date(dbsession, mocker):
 
 
 def test_backfill_dataset_run_impl_invalid_end_date(dbsession, mocker):
-    mocker.patch("tasks.timeseries_backfill.timeseries_enabled", return_value=True)
+    mocker.patch("tasks.timeseries_backfill.is_timeseries_enabled", return_value=True)
     mock_group = mocker.patch("tasks.timeseries_backfill.group")
 
     repository = RepositoryFactory.create()
@@ -178,7 +178,7 @@ def test_backfill_dataset_run_impl_invalid_end_date(dbsession, mocker):
 
 
 def test_backfill_dataset_run_impl_timeseries_not_enabled(dbsession, mocker):
-    mocker.patch("tasks.timeseries_backfill.timeseries_enabled", return_value=False)
+    mocker.patch("tasks.timeseries_backfill.is_timeseries_enabled", return_value=False)
     mock_group = mocker.patch("tasks.timeseries_backfill.group")
 
     task = TimeseriesBackfillDatasetTask()

tasks/tests/unit/test_timeseries_delete.py

@@ -4,7 +4,7 @@
 
 
 def test_timeseries_delete_run_impl(dbsession, mocker):
-    mocker.patch("tasks.timeseries_delete.timeseries_enabled", return_value=True)
+    mocker.patch("tasks.timeseries_delete.is_timeseries_enabled", return_value=True)
     delete_repository_data = mocker.patch(
         "tasks.timeseries_delete.delete_repository_data"
     )
@@ -24,7 +24,7 @@ def test_timeseries_delete_run_impl(dbsession, mocker):
 
 
 def test_timeseries_delete_run_impl_invalid_repository(dbsession, mocker):
-    mocker.patch("tasks.timeseries_delete.timeseries_enabled", return_value=True)
+    mocker.patch("tasks.timeseries_delete.is_timeseries_enabled", return_value=True)
 
     task = TimeseriesDeleteTask()
     res = task.run_impl(
@@ -35,7 +35,7 @@ def test_timeseries_delete_run_impl_invalid_repository(dbsession, mocker):
 
 
 def test_timeseries_delete_run_impl_timeseries_not_enabled(dbsession, mocker):
-    mocker.patch("tasks.timeseries_delete.timeseries_enabled", return_value=False)
+    mocker.patch("tasks.timeseries_delete.is_timeseries_enabled", return_value=False)
 
     repository = RepositoryFactory.create()
     dbsession.add(repository)
@@ -50,7 +50,7 @@ def test_timeseries_delete_run_impl_timeseries_not_enabled(dbsession, mocker):
 
 
 def test_timeseries_delete_measurements_only(dbsession, mocker):
-    mocker.patch("tasks.timeseries_delete.timeseries_enabled", return_value=True)
+    mocker.patch("tasks.timeseries_delete.is_timeseries_enabled", return_value=True)
     delete_repository_measurements = mocker.patch(
         "tasks.timeseries_delete.delete_repository_measurements"
     )
@@ -73,7 +73,7 @@ def test_timeseries_delete_measurements_only(dbsession, mocker):
 
 
 def test_timeseries_delete_measurements_only_unsuccessful(dbsession, mocker):
-    mocker.patch("tasks.timeseries_delete.timeseries_enabled", return_value=True)
+    mocker.patch("tasks.timeseries_delete.is_timeseries_enabled", return_value=True)
     delete_repository_measurements = mocker.patch(
         "tasks.timeseries_delete.delete_repository_measurements"
     )