Python context_managers.switch_db函数代码示例

 def test_undoes_db_changes_when_error_2(self):
     """
     Tests that load_from_config.run() removes only the documents it saved to the database and leaves
     alone any documents that were already there when it comes across an unexpected error
     """
     expected_experi_model.switch_db(TEST_DB_ALIAS)
     expected_experi_model.save()
     expected_ds_model.switch_db(TEST_DB_ALIAS)
     expected_ds_model.save()
     json_config = open(path_string_json_full, 'w')
     json_config.write(breaking_json)
     json_config.close()
     with self.assertRaises(KeyError):
         load_from_config.run()
     with switch_db(Experiment, TEST_DB_ALIAS) as TestEx:
         query = TestEx.objects.all()
         self.assertEqual(len(query), 1)
         self.document_compare(query[0], expected_experi_model)
     with switch_db(DataSource, TEST_DB_ALIAS) as TestDs:
         query = TestDs.objects.all()
         self.assertEqual(len(query), 1)
         self.document_compare(query[0], expected_ds_model)
     with switch_db(Genotype, TEST_DB_ALIAS) as TestGen:
         query = TestGen.objects.all()
         self.assertEqual(len(query), 0)

def create_update_footnotes(data, db_alias):
    total_created = 0
    total_updated = 0

    logging.info('Processing %d footnotes' % len(data))

    from usda_mongo.models import Footnote, Nutrient, Food

    with switch_db(Footnote, db_alias) as Footnote:
        with switch_db(Nutrient, db_alias) as Nutrient:
            with switch_db(Food, db_alias) as Food:

                for row in csv.DictReader(
                    data, fieldnames=(
                        'ndb_no', 'footnt_no', 'footnt_typ', 'nutr_no', 'footnt_txt'
                    ),
                    delimiter='^', quotechar='~'
                ):
                    created = False

                    # SR22 definition indicates that `footnt_no` and `footnt_typ` are required,
                    # but on occasion, either on is blank.  To compensate for this, we assume
                    # a blank `footnt_no` is '1' and a blank`footnt_typ` is 'N'.
                    if row['footnt_no'] == '':
                        row['footnt_no'] = 1
                    if row['footnt_typ'] not in (FOOTNOTE_DESC, FOOTNOTE_MEAS, FOOTNOTE_NUTR):
                        row['footnt_typ'] = FOOTNOTE_NUTR

                    if row.get('nutr_no'):
                        nutrient = Nutrient.objects.get(number=int(row['nutr_no']))
                    else:
                        nutrient = None

                    try:
                        footnote = Footnote.objects.get(
                            food=Food.objects.get(ndb_number=int(row['ndb_no'])),
                            number=int(row['footnt_no']),
                            nutrient=nutrient
                        )
                        total_updated += 1
                    except Footnote.DoesNotExist:
                        footnote = Footnote(
                            food=Food.objects.get(ndb_number=int(row['ndb_no'])),
                            number=int(row['footnt_no']),
                            nutrient=nutrient
                        )
                        total_created += 1
                        created = True

                    footnote.type = row['footnt_typ']
                    footnote.text = row['footnt_txt']
                    footnote.save()

                    if created:
                        logging.debug('Created %s' % footnote)
                    else:
                        logging.debug('Updated %s' % footnote)

    logging.info('Created %d new footnotes' % total_created)
    logging.info('Updated %d footnotes' % total_updated)

    def create_document(row, test=False):
        # Creates and returns a Genotype document from the values in the row
        db_alias = TEST_DB_ALIAS if test else 'default'
        build_dic = {}
        for key in row:
            if 'date' == key[-4:] or key == 'dtt':
                # Assumes values ending in 'date' are for date fields
                build_dic[key] = datetime.strptime(row[key], "%Y-%m-%d %H:%M:%S.%f")

            # Searches through the other collections for the reference field values
            elif 'datasource' in key:
                with switch_db(DataSource, db_alias) as TestDat:
                    datasource, created = fetch_or_save(
                        TestDat, db_alias=db_alias, name=row[key]
                    )
                build_dic['datasource'] = datasource
            elif 'study' in key:
                with switch_db(Experiment, db_alias) as TestEx:
                    study, created = fetch_or_save(
                        TestEx, db_alias=db_alias, name=row[key]
                    )
                build_dic['study'] = study

            elif key == 'obs':
                # Extracts the dictionary from the obs field
                build_dic[key] = ast.literal_eval(row[key])
            else:
                build_dic[key] = row[key]

        with switch_db(Genotype, db_alias) as TestGen:
            gen, created = fetch_or_save(TestGen, db_alias=db_alias, **build_dic)
        return gen

 def document_compare(self, doc1, doc2):
     """
     Asserts the two given mongoengine.documents are equal. Ignores metadata fields and fields such
     as time stamps that default to datetime.now()
     """
     for key in doc1._fields_ordered:
         # ignores metadata fields and datetime fields that default to datetime.now()
         if key != 'id' and key[0] != '_' and key != 'dtt' and key != 'lastupdateddate':
             with self.subTest(key=key):
                 val = doc1[key]
                 if isinstance(doc1[key], dict):
                     self.assertDictEqual(doc1[key], doc2[key])
                 elif isinstance(val, DBRef):
                     if key == 'study':
                         with switch_db(Experiment, TEST_DB_ALIAS) as TestEx:
                             study = TestEx.objects.get(id=val.id)
                             self.document_compare(study, doc2[key])
                     elif key == 'datasource':
                         with switch_db(DataSource, TEST_DB_ALIAS) as TestDs:
                             ds = TestDs.objects.get(id=val.id)
                             self.document_compare(ds, doc2[key])
                     else:
                         self.fail("Unexpected reference field: " + key)
                 else:
                     self.assertEqual(doc1[key], doc2[key])

 def tearDown(self):
     """
     Clears the test database
     """
     with switch_db(Experiment, TEST_DB_ALIAS) as TestEx:
         TestEx.objects.all().delete()
     with switch_db(DataSource, TEST_DB_ALIAS) as TestDs:
         TestDs.objects.all().delete()
     with switch_db(Genotype, TEST_DB_ALIAS) as TestGen:
         TestGen.objects.all().delete()

    def load(self):
        logger.debug(u"iniciando metodo load() (uuid: %s)" % self.metadata['uuid'])

        logger.debug(u"salvando modelo %s no opac (_id: %s)" % (
            self.opac_model_name, self.opac_model_instance._id))

        with \
            switch_db(OpacCollection, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacJournal, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacIssue, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacArticle, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacSponsor, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacNews, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacPages, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacPressRelease, OPAC_WEBAPP_DB_NAME):

                self.opac_model_instance.switch_db(OPAC_WEBAPP_DB_NAME)
                self.opac_model_instance.save()
                self.opac_model_instance.reload()

        # atualizamos os dados do registro LOAD
        with switch_db(self.load_model_class, OPAC_PROC_DB_NAME):
            self.metadata['process_finish_at'] = datetime.now()
            self.metadata['process_completed'] = True
            self.metadata['must_reprocess'] = False
            json_opac_data = self.opac_model_instance.to_json()
            cleaned_json_opac_data = json_opac_data.replace('$', '')  # retiramos o $
            self.metadata['loaded_data'] = json.loads(cleaned_json_opac_data)
            self.load_model_instance.update(**self.metadata)
            self.load_model_instance.save()
            self.load_model_instance.reload()
            logger.debug(u"modelo %s no opac_proc (uuid: %s) foi atualizado" % (
                self.load_model_name, self.metadata['uuid']))

        logger.debug(u"finalizando metodo load() (uuid: %s)" % self.metadata['uuid'])

def query_genotype_by_experiment(experi_id):
    db_alias = TEST_DB_ALIAS if testing else 'default'
    # Make query
    try:
        with switch_db(Experiment, db_alias) as Exper:
            ex = Exper.objects.get(id=experi_id)
    except Experiment.DoesNotExist:
        raise Http404("Experiment does not exist")
    with switch_db(Genotype, db_alias) as Gen:
        genotype = Gen.objects(study=ex)
    return genotype, ex

    def prepare(self):
        logger.debug(u"iniciando metodo prepare (uuid: %s)" % self.metadata['uuid'])
        obj_dict = self.transform_model_instance_to_python()
        obj_dict['_id'] = self._uuid_str

        with \
            switch_db(OpacCollection, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacJournal, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacIssue, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacArticle, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacSponsor, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacNews, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacPages, OPAC_WEBAPP_DB_NAME), \
            switch_db(OpacPressRelease, OPAC_WEBAPP_DB_NAME):

                if self.opac_model_instance is None:
                    # crio uma nova instância
                    self.opac_model_instance = self.opac_model_class(**obj_dict)
                    self.opac_model_instance.switch_db(OPAC_WEBAPP_DB_NAME)
                else:  # já tenho uma instância no banco
                    self.opac_model_instance.switch_db(OPAC_WEBAPP_DB_NAME)
                    for k, v in obj_dict.iteritems():
                        self.opac_model_instance[k] = v
                    self.opac_model_instance.save()

        logger.debug(u"modelo opac (_id: %s) encontrado. atualizando registro" % obj_dict['_id'])

        logger.debug(u"finalizando metodo prepare(uuid: %s)" % self.metadata['uuid'])
        logger.debug(u'opac_model_instance SALVO: %s' % self.opac_model_instance.to_json())
        return self.opac_model_instance

def create_update_weights(data, encoding, db_alias):
    total_created = 0
    total_updated = 0

    logging.info('Processing %d weights' % len(data))

    from usda_mongo.models import Weight, Food

    with switch_db(Weight, db_alias) as Weight:
        with switch_db(Food, db_alias) as Food:

            for row in UnicodeDictReader(
                data, fieldnames=(
                    'ndb_no', 'seq', 'amount', 'msre_desc', 'gm_wgt', 'num_data_pts', 'std_dev'
                ),
                delimiter='^', quotechar='~',
                encoding=encoding
            ):
                created = False

                try:
                    weight = Weight.objects.get(
                        food=Food.objects.get(ndb_number=int(row['ndb_no'])),
                        sequence=int(row['seq'])
                    )
                    total_updated += 1
                except Weight.DoesNotExist:
                    weight = Weight(
                        food=Food.objects.get(ndb_number=int(row['ndb_no'])),
                        sequence=int(row['seq'])
                    )
                    total_created += 1
                    created = True

                weight.amount = float(row.get('amount'))
                weight.description = row.get('msre_desc')
                weight.gram_weight = float(row.get('gm_wgt'))
                if row.get('num_data_pts'):
                    weight.number_of_data_points = float(row['num_data_pts'])
                if row.get('std_dev'):
                    weight.standard_deviation = float(row['std_dev'])
                weight.save()

                if created:
                    logging.debug('Created %s' % weight)
                else:
                    logging.debug('Updated %s' % weight)

    logging.info('Created %d new weights' % total_created)
    logging.info('Updated %d weights' % total_updated)

def get_ref_fields(gen, testing):
    ref_fields = {"study": gen.study, "datasource": gen.datasource}
    if testing:
        # Reference fields would only hold database references instead of the actual
        # document, so need to query the test database for the documents
        study_son = ref_fields["study"].as_doc()
        ds_son = ref_fields["datasource"].as_doc()
        study_id = study_son.get("$id")
        with switch_db(Experiment, TEST_DB_ALIAS) as Exper:
            ref_fields["study"] = Exper.objects.get(id=study_id)
        ds_id = ds_son.get("$id")
        with switch_db(DataSource, TEST_DB_ALIAS) as Dat:
            ref_fields["datasource"] = Dat.objects.get(id=ds_id)
    return ref_fields

 def get_document_by_id(self, object_id):
     """
     metodo que retorna o documento procurando pelo object_id.
     """
     register_connections()
     with switch_db(self.model_class, OPAC_PROC_LOGS_DB_NAME):
         return self.model_class.objects.get(id=object_id)

def task_delete_selected_collections(selected_uuids):
    """
        Task para apagar Coleções Carregadas.
        @param:
        - selected_uuids: lista de UUIDs dos documentos a serem removidos

        Se a lista `selected_uuids` for maior a SLICE_SIZE
            A lista será fatiada em listas de tamanho: SLICE_SIZE
        Se a lista `selected_uuids` for < a SLICE_SIZE
            Será feito uma delete direto no queryset
    """

    stage = 'load'
    model = 'collection'
    model_class = LoadCollection
    get_db_connection()
    r_queues = RQueues()
    SLICE_SIZE = 1000

    if len(selected_uuids) > SLICE_SIZE:
        list_of_list_of_uuids = list(chunks(selected_uuids, SLICE_SIZE))
        for list_of_uuids in list_of_list_of_uuids:
            uuid_as_string_list = [str(uuid) for uuid in list_of_uuids]
            r_queues.enqueue(stage, model, task_delete_selected_collections, uuid_as_string_list)
    else:
        # removemos o conjunto de documentos do LoadCollection indicados pelos uuids
        documents_to_delete = model_class.objects.filter(uuid__in=selected_uuids)
        documents_to_delete.delete()

        # convertemos os uuid para _id e filtramos esses documentos no OPAC
        register_connections()
        opac_pks = [str(uuid).replace('-', '') for uuid in selected_uuids]
        with switch_db(opac_models.Collection, OPAC_WEBAPP_DB_NAME) as opac_model:
            selected_opac_records = opac_model.objects.filter(pk__in=opac_pks)
            selected_opac_records.delete()

    def switch_db(self, db_alias):
        """
        Temporarily switch the database for a document instance.

        Only really useful for archiving off data and calling `save()`::

            user = User.objects.get(id=user_id)
            user.switch_db('archive-db')
            user.save()

        If you need to read from another database see
        :class:`~mongoengine.context_managers.switch_db`

        :param db_alias: The database alias to use for saving the document
        """
        with switch_db(self.__class__, db_alias) as cls:
            collection = cls._get_collection()
            db = cls._get_db()
        self._get_collection = lambda: collection
        self._get_db = lambda: db
        self._collection = collection
        self._created = True
        self.__objects = self._qs
        self.__objects._collection_obj = collection
        return self

def run(*args):
    global MODE
    if 'override' in args:
        Logger.Warning("OVERRIDE MODE!")
        MODE = "OVERRIDE"

    Logger.Message("Loading process in mode: " + MODE  + "started.")
    global db_alias
    if testing:
        db_alias = TEST_DB_ALIAS
    # for keeping track of documents saved to db by this run of script
    global created_doc_ids
    created_doc_ids = []

    dirs = DataDir.objects.all()
    for d in dirs:
        Logger.Message("Processing data dir: " + d.path)
        path = Path(d.path)
        try:
            look_for_config_dir(path)
        except Exception as e:
            Logger.Error(str(e))
            # 'Cancels' the script, by removing from db all documents saved to db in this script run-through
            for doc_type, doc_id in created_doc_ids:
                with switch_db(doc_type, db_alias) as Col:
                    Col.objects.get(id=doc_id).delete()
            raise e

 def test_yaml_marked_loaded_no_load(self):
     """
     Test script does not load data from a directory where the config.yaml file has been marked as loaded
     """
     yaml_parser = configuration_parser.YamlConfigParser(path_string_yaml_full)
     yaml_parser.mark_loaded()
     load_from_config.load_in_dir(path_string_yaml)
     with switch_db(Experiment, TEST_DB_ALIAS) as TestEx:
         query = TestEx.objects.all()
         self.assertEqual(len(query), 0)
     with switch_db(DataSource, TEST_DB_ALIAS) as TestDs:
         query = TestDs.objects.all()
         self.assertEqual(len(query), 0)
     with switch_db(Genotype, TEST_DB_ALIAS) as TestGen:
         query = TestGen.objects.all()
         self.assertEqual(len(query), 0)