def handle(self, *args, **options): elasticsearch_client = instantiate_elasticsearch_client() config = process_cli_parameters(options, elasticsearch_client) start = perf_counter() printf({"msg": f"Starting script\n{'=' * 56}"}) start_msg = "target index: {index_name} | FY(s): {fiscal_years} | Starting from: {starting_date}" printf({"msg": start_msg.format(**config)}) if config["load_type"] == "transactions": ensure_view_exists(settings.ES_TRANSACTIONS_ETL_VIEW_NAME) elif config["load_type"] == "awards": ensure_view_exists(settings.ES_AWARDS_ETL_VIEW_NAME) loader = Rapidloader(config, elasticsearch_client) loader.run_load_steps() loader.complete_process() printf({ "msg": "---------------------------------------------------------------" }) printf({"msg": f"Script completed in {perf_counter() - start:.2f}s"}) printf({ "msg": "---------------------------------------------------------------" })
def test_incremental_load_into_award_index(award_data_fixture, elasticsearch_award_index, monkeypatch): """Test the ``elasticsearch_loader`` django management command to incrementally load updated data into the awards ES index from the DB, overwriting the doc that was already there """ original_db_awards_count = Award.objects.count() elasticsearch_award_index.update_index() client = elasticsearch_award_index.client # type: Elasticsearch assert client.indices.exists(elasticsearch_award_index.index_name) es_award_docs = client.count( index=elasticsearch_award_index.index_name)["count"] assert es_award_docs == original_db_awards_count # Inject ETL arg into config for this run, to suppress processing deletes. Test incremental load only elasticsearch_award_index.etl_config["process_deletes"] = False elasticsearch_award_index.etl_config["start_datetime"] = datetime.now( timezone.utc) es_etl_config = _process_es_etl_test_config(client, elasticsearch_award_index) # Now modify one of the DB objects awd = Award.objects.first() # type: Award awd.total_obligation = 9999 awd.save() # Must use mock sql function to share test DB conn+transaction in ETL code # Patching on the module into which it is imported, not the module where it is defined monkeypatch.setattr( "usaspending_api.etl.elasticsearch_loader_helpers.extract_data.execute_sql_statement", mock_execute_sql) # Also override SQL function listed in config object with the mock one es_etl_config["execute_sql_func"] = mock_execute_sql ensure_view_exists(es_etl_config["sql_view"], force=True) loader = Controller(es_etl_config) assert loader.__class__.__name__ == "Controller" loader.prepare_for_etl() loader.dispatch_tasks() client.indices.refresh(elasticsearch_award_index.index_name) assert client.indices.exists(elasticsearch_award_index.index_name) es_award_docs = client.count( index=elasticsearch_award_index.index_name)["count"] assert es_award_docs == original_db_awards_count es_awards = client.search(index=elasticsearch_award_index.index_name) updated_award = [ a for a in es_awards["hits"]["hits"] if a["_source"]["award_id"] == awd.id ][0] assert int(updated_award["_source"]["total_obligation"]) == 9999
def handle(self, *args, **options): elasticsearch_client = instantiate_elasticsearch_client() config = parse_cli_args(options, elasticsearch_client) start = perf_counter() logger.info(format_log(f"Starting script\n{'=' * 56}")) start_msg = "target index: {index_name} | Starting from: {starting_date}" logger.info(format_log(start_msg.format(**config))) ensure_view_exists(config["sql_view"], force=True) error_addition = "" loader = Controller(config) if config["is_incremental_load"]: toggle_refresh_off(elasticsearch_client, config["index_name"]) # Turned back on at end. try: if config["process_deletes"]: loader.run_deletes() if not config["deletes_only"]: loader.prepare_for_etl() loader.dispatch_tasks() except Exception as e: logger.error(f"{str(e)}") error_addition = "before encountering a problem during execution.... " raise SystemExit(1) else: loader.complete_process() if config["drop_db_view"]: logger.info( format_log(f"Dropping SQL view '{config['sql_view']}'")) drop_etl_view(config["sql_view"], True) finally: msg = f"Script duration was {perf_counter() - start:.2f}s {error_addition}|" headers = f"{'-' * (len(msg) - 2)} |" logger.info(format_log(headers)) logger.info(format_log(msg)) logger.info(format_log(headers)) # Used to help pipeline determine when job passed but needs attention if config["raise_status_code_3"]: raise SystemExit(3)
def django_db_setup( request, django_test_environment, django_db_blocker, django_db_use_migrations, django_db_keepdb, django_db_createdb, django_db_modify_db_settings, ): """This is an override of the original implementation in the https://github.com/pytest-dev/pytest-django plugin from file /pytest-django/fixtures.py. Because this "hides" the original implementation, it may get out-of-date as that plugin is upgraded. This override takes the implementation from pytest-django Release 3.5.1, and extends it in order to execute materialized views as part of database setup. If requirements.txt shows a different version than the one this is based on: compare, update, and test. More work could be put into trying to patch, replace, or wrap implementation of ``django.test.utils.setup_databases``, which is the actual method that needs to be wrapped and extended. """ from pytest_django.compat import setup_databases, teardown_databases from pytest_django.fixtures import _disable_native_migrations setup_databases_args = {} if not django_db_use_migrations: _disable_native_migrations() if django_db_keepdb and not django_db_createdb: setup_databases_args["keepdb"] = True with django_db_blocker.unblock(): db_cfg = setup_databases(verbosity=request.config.option.verbose, interactive=False, **setup_databases_args) # If migrations are skipped, assume matviews and views are not to be (re)created either # Other scenarios (such as reuse or keep DB) may still lead to creation of a non-existent DB, so they must be # (re)created under those conditions if not django_db_use_migrations: logger.warning( "Skipping generation of materialized views or other views in this test run because migrations are also " "being skipped. ") else: generate_matviews(materialized_views_as_traditional_views=True) ensure_view_exists(settings.ES_TRANSACTIONS_ETL_VIEW_NAME) ensure_view_exists(settings.ES_AWARDS_ETL_VIEW_NAME) ensure_business_categories_functions_exist() call_command("load_broker_static_data") def teardown_database(): with django_db_blocker.unblock(): try: teardown_databases(db_cfg, verbosity=request.config.option.verbose) except Exception as exc: request.node.warn( pytest.PytestWarning( "Error when trying to teardown test databases: %r" % exc)) if not django_db_keepdb: request.addfinalizer(teardown_database)