Browse Source

[indexer] Add yelp example

Localize unlocalized timestamps
Abraham Elmahrek 11 năm trước cách đây
mục cha
commit
50aa91949e

+ 3 - 0
desktop/libs/indexer/src/indexer/controller.py

@@ -50,6 +50,9 @@ class CollectionManagerController(object):
         field[flags[1]] = field['flags'][index] == FLAGS[index][0]
     return fields
 
+  def collection_exists(self, collection):
+    return collection in self.get_collections()
+
   def get_collections(self):
     try:
       solr_collections = SolrApi(SOLR_URL.get(), self.user, SECURITY_ENABLED.get()).collections()

+ 16 - 6
desktop/libs/indexer/src/indexer/management/commands/indexer_install_examples.py

@@ -38,10 +38,10 @@ class Command(NoArgsCommand):
   def handle_noargs(self, **options):
     self.user = install_sample_user()
     self.fs = cluster.get_hdfs()
-
-    twitter_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../../../../../apps/search/examples/collections/solr_configs_twitter_demo/index_data.csv'))
+    self.searcher = controller.CollectionManagerController(self.user)
 
     LOG.info(_("Installing twitter collection"))
+    twitter_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../../../../../apps/search/examples/collections/solr_configs_twitter_demo/index_data.csv'))
     self._setup_collection_from_csv({
       'name': 'twitter_example',
       'fields': self._parse_fields(twitter_path),
@@ -49,11 +49,21 @@ class Command(NoArgsCommand):
     }, twitter_path)
     LOG.info(_("Twitter collection successfully installed"))
 
+    LOG.info(_("Installing yelp collection"))
+    twitter_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../../../../../apps/search/examples/collections/solr_configs_yelp_demo/index_data.csv'))
+    self._setup_collection_from_csv({
+      'name': 'yelp_example',
+      'fields': self._parse_fields(twitter_path),
+      'uniqueKeyField': 'id'
+    }, twitter_path)
+    LOG.info(_("Yelp collection successfully installed"))
+
   def _setup_collection_from_csv(self, collection, path, separator=',', quote_character='"'):
-    searcher = controller.CollectionManagerController(self.user)
+    if self.searcher.collection_exists(collection['name']):
+      self.searcher.delete_collection(collection['name'])
 
     # Create instance directory, collection, and add fields
-    searcher.create_collection(collection['name'], collection['fields'], collection['uniqueKeyField'])
+    self.searcher.create_collection(collection['name'], collection['fields'], collection['uniqueKeyField'])
 
     try:
       hdfs_path = '/tmp/%s' % uuid.uuid4()
@@ -67,7 +77,7 @@ class Command(NoArgsCommand):
         self.fs.do_as_user(self.fs.DEFAULT_USER, self.fs.create, hdfs_path, data=fh.read(), overwrite=overwrite)
 
       # Index data
-      searcher.update_data_from_hdfs(self.fs,
+      self.searcher.update_data_from_hdfs(self.fs,
                                      collection['name'],
                                      collection['fields'],
                                      hdfs_path,
@@ -76,7 +86,7 @@ class Command(NoArgsCommand):
                                      quote_character=quote_character)
 
     except:
-      searcher.delete_collection(collection['name'])
+      self.searcher.delete_collection(collection['name'])
       raise
     finally:
       # Remove HDFS file

+ 4 - 0
desktop/libs/indexer/src/indexer/utils.py

@@ -26,6 +26,7 @@ import StringIO
 import tempfile
 from dateutil.parser import parse
 
+from django.conf import settings
 from django.utils.translation import ugettext as _
 
 from desktop.lib.i18n import force_unicode, smart_str
@@ -200,9 +201,12 @@ def field_values_from_separated_file(fh, delimiter, quote_character, fields=None
 
       # Parse dates
       if timestamp_fields:
+        tzinfo = pytz.timezone(settings.TIME_ZONE)
         for key in timestamp_fields:
           if key in row:
             dt = parse(row[key])
+            if not dt.tzinfo:
+              dt = tzinfo.localize(dt)
             row[key] = dt.astimezone(pytz.utc).strftime('%Y-%m-%dT%H:%M:%SZ')
 
       # Parse decimal