put back logging_utils to enable custom logging for spark applications (

#10)

src/cehrbert_data/apps/generate_training_data.py

@@ -23,6 +23,7 @@
     get_measurement_table,
     validate_table_names,
 )
+from cehrbert_data.utils.logging_utils import add_console_logging
 
 
 def main(
@@ -333,6 +334,9 @@ def main(
 
     ARGS = parser.parse_args()
 
+    # Enable logging
+    add_console_logging()
+
     main(
         ARGS.input_folder,
         ARGS.output_folder,

src/cehrbert_data/cohorts/spark_app_base.py

@@ -22,6 +22,7 @@
     get_descendant_concept_ids,
     preprocess_domain_table,
 )
+from cehrbert_data.utils.logging_utils import add_console_logging
 
 from ..cohorts.query_builder import ENTRY_COHORT, NEGATIVE_COHORT, QueryBuilder
 
@@ -677,6 +678,10 @@ def create_prediction_cohort(
     :param ehr_table_list:
     :return:
     """
+    # Add logging to spark application output when enable_logging is set to True
+    if spark_args.enable_logging:
+        add_console_logging()
+
     # Toggle the prior/post observation_period depending on the is_window_post_index flag
     prior_observation_period = (
         0 if spark_args.is_window_post_index else spark_args.observation_window + spark_args.hold_off_window

src/cehrbert_data/utils/logging_utils.py

@@ -0,0 +1,12 @@
+import logging
+import sys
+
+
+def add_console_logging():
+    root = logging.getLogger()
+    root.setLevel(logging.INFO)
+
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setLevel(logging.DEBUG)
+    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    handler.setFormatter(formatter)

src/cehrbert_data/utils/spark_parse_args.py

@@ -82,6 +82,12 @@ def create_spark_args():
         --single_contribution: Whether patients should contribute only once to the training data.
     """
     parser = argparse.ArgumentParser(description="Arguments for spark applications for generating cohort definitions")
+    parser.add_argument(
+        "--enable_logging",
+        dest="enable_logging",
+        action="store_true",
+        help="Indicate whether we should log the info to the spark-submit output",
+    )
     parser.add_argument(
         "-c",
         "--cohort_name",