fixes key loader

unicef · Sep 19, 2019 · 7f64bac · 7f64bac
1 parent 1ef7c95
commit 7f64bac
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 13 deletions.
diff --git a/src/etools_datamart/apps/data/loader.py b/src/etools_datamart/apps/data/loader.py
@@ -17,7 +17,10 @@
 
 
 class EToolsLoaderOptions(BaseLoaderOptions):
-    pass
+    DEFAULT_KEY = lambda loader, record: dict(schema_name=loader.context['country'].schema_name,
+                                              source_id=record.pk)
+
+
 #     def __init__(self, base=None):
 #         super().__init__(base)
 #         self.key = lambda loader, record: dict(schema_name=loader.context['country'].schema_name,
@@ -175,6 +178,10 @@ def load(self, *, verbosity=0, stdout=None, ignore_dependencies=False, max_recor
         return self.results
 
 
+class CommonSchemaLoaderOptions(BaseLoaderOptions):
+    DEFAULT_KEY = lambda loader, record: dict(source_id=record.pk)
+
+
 class CommonSchemaLoader(EtoolsLoader):
     def get_mart_values(self, record=None):
         ret = {'seen': self.context['today']}

diff --git a/src/etools_datamart/apps/data/models/base.py b/src/etools_datamart/apps/data/models/base.py
@@ -5,7 +5,8 @@
 
 from celery.local import class_property
 
-from etools_datamart.apps.data.loader import EtoolsLoader, EToolsLoaderOptions
+from etools_datamart.apps.data.loader import (CommonSchemaLoader, CommonSchemaLoaderOptions,
+                                              EtoolsLoader, EToolsLoaderOptions,)
 from etools_datamart.apps.etl.base import DataMartModelBase
 
 
@@ -40,12 +41,17 @@ def truncate(self, reset=True):
                                                                       restart))
 
 
-class EToolsDataMartModelBase(DataMartModelBase):
+class CommonDataMartModelModelBase(DataMartModelBase):
+    loader_option_class = CommonSchemaLoaderOptions
+    loader_class = CommonSchemaLoader
+
+
+class EToolsDataMartModelBase(CommonDataMartModelModelBase):
     loader_option_class = EToolsLoaderOptions
     loader_class = EtoolsLoader
 
 
-class CommonDataMartModel(models.Model, metaclass=DataMartModelBase):
+class CommonDataMartModel(models.Model, metaclass=CommonDataMartModelModelBase):
     source_id = models.IntegerField(blank=True, null=True, db_index=True)
     last_modify_date = models.DateTimeField(blank=True, auto_now=True)
     seen = models.DateTimeField(blank=True, null=True)

diff --git a/src/etools_datamart/apps/etl/loader.py b/src/etools_datamart/apps/etl/loader.py
@@ -107,10 +107,11 @@ class MaxRecordsException(Exception):
     pass
 
 
-DEFAULT_KEY = lambda loader, record: dict(source_id=record.pk)
+undefined = object()
 
 
 class BaseLoaderOptions:
+    DEFAULT_KEY = lambda loader, record: dict(source_id=record.pk)
     __attrs__ = ['mapping', 'celery', 'source', 'last_modify_field',
                  'queryset', 'key', 'locks', 'filters', 'sync_deleted_records', 'truncate',
                  'depends', 'timeout', 'lock_key', 'always_update', 'fields_to_compare']
@@ -123,7 +124,7 @@ def __init__(self, base=None):
         self.always_update = False
         self.source = None
         self.lock_key = None
-        self.key = DEFAULT_KEY
+        self.key = undefined
         self.timeout = None
         self.depends = ()
         self.filters = None
@@ -141,6 +142,8 @@ def __init__(self, base=None):
                         setattr(self, attr, n)
                     else:
                         setattr(self, attr, getattr(base, attr, getattr(self, attr)))
+        if self.key == undefined:
+            self.key = type(self).DEFAULT_KEY
 
         if self.truncate:
             self.sync_deleted_records = lambda loader: False

diff --git a/src/etools_datamart/apps/rapidpro/loader.py b/src/etools_datamart/apps/rapidpro/loader.py
@@ -10,16 +10,11 @@
 
 logger = get_task_logger(__name__)
 
-DEFAULT_KEY = lambda loader, record: dict(uuid=record.uuid,
-                                          organization=loader.context['organization'])
-
 
 class TembaLoaderOptions(BaseLoaderOptions):
     __attrs__ = BaseLoaderOptions.__attrs__ + ['host', 'temba_object']
-
-    def __init__(self, base=None):
-        super().__init__(base)
-        self.key = DEFAULT_KEY
+    DEFAULT_KEY = lambda loader, record: dict(uuid=record.uuid,
+                                              organization=loader.context['organization'])
 
 
 class TembaLoader(BaseLoader):