<commit_before>import datetime
import logging

from corehq.apps.es import CaseSearchES
from django.core.management import BaseCommand
from casexml.apps.case.mock import CaseStructure, CaseFactory

logger = logging.getLogger('enikshay_2b_reason_for_test_fix')


class Command(BaseCommand):

    def add_arguments(self, parser):
        parser.add_argument(
            'domain',
            help="The domain to migrate."
        )
        parser.add_argument(
            "log_path",
            help="Path to write the log to"
        )
        parser.add_argument(
            '--commit',
            action='store_true',
            help="actually modifies the cases. Without this flag, it's a dry run."
        )

    def handle(self, domain, log_path, **options):
        commit = options['commit']
        factory = CaseFactory(domain)

        logger.info("Starting {} migration on {} at {}".format(
            "real" if commit else "fake", domain, datetime.datetime.utcnow()
        ))

        cases = (CaseSearchES()
                 .domain(domain)
                 .case_type("episode")
                 .case_property_query("updated_by_migration", "enikshay_2b_case_properties", "must")
                 .run().hits)

        with open(log_path, "w") as f:
            for case in cases:
                case_props = {prop['key']: prop['value'] for prop in case['case_properties']}
                treatment_status = None
                treatment_initiated = case_props.get('treatment_initiated')
                diagnosing_facility_id = case_props.get('diagnosing_facility_id')
                treatment_initiating_facility_id = case_props.get('treatment_initiating_facility_id')

                if treatment_initiated == 'yes_phi' and \
                        diagnosing_facility_id and treatment_initiating_facility_id and \
                        diagnosing_facility_id != treatment_initiating_facility_id:
                    treatment_status = 'initiated_outside_facility'
                elif treatment_initiated == 'yes_phi' and \
                        diagnosing_facility_id and treatment_initiating_facility_id:
                    treatment_status = 'initiated_first_line_treatment'
                elif treatment_initiated == 'yes_private':
                    treatment_status = 'initiated_outside_rntcp'

                if treatment_status:
                    case_id = case['_id']
                    f.write(case_id + "\n")
                    logger.info(case_id)

                    case_structure = CaseStructure(
                        case_id=case_id,
                        walk_related=False,
                        attrs={
                            "create": False,
                            "update": {
                                "treatment_status": treatment_status,
                                "updated_by_migration": "enikshay_2b_treatment_status_fix",
                            },
                        },
                    )

                    if commit:
                        factory.create_or_update_case(case_structure)
        logger.info("Migration finished at {}".format(datetime.datetime.utcnow()))
<commit_msg>scroll instead of fetch all<commit_after>import datetime
import logging

from corehq.apps.es import CaseSearchES
from django.core.management import BaseCommand
from casexml.apps.case.mock import CaseStructure, CaseFactory

logger = logging.getLogger('enikshay_2b_reason_for_test_fix')


class Command(BaseCommand):

    def add_arguments(self, parser):
        parser.add_argument(
            'domain',
            help="The domain to migrate."
        )
        parser.add_argument(
            "log_path",
            help="Path to write the log to"
        )
        parser.add_argument(
            '--commit',
            action='store_true',
            help="actually modifies the cases. Without this flag, it's a dry run."
        )

    def handle(self, domain, log_path, **options):
        commit = options['commit']
        factory = CaseFactory(domain)

        logger.info("Starting {} migration on {} at {}".format(
            "real" if commit else "fake", domain, datetime.datetime.utcnow()
        ))

        cases = (CaseSearchES()
                 .domain(domain)
                 .case_type("episode")
                 .case_property_query("updated_by_migration", "enikshay_2b_case_properties", "must")
                 .scroll())

        with open(log_path, "w") as f:
            for case in cases:
                case_props = {prop['key']: prop['value'] for prop in case['case_properties']}
                treatment_status = None
                treatment_initiated = case_props.get('treatment_initiated')
                diagnosing_facility_id = case_props.get('diagnosing_facility_id')
                treatment_initiating_facility_id = case_props.get('treatment_initiating_facility_id')

                if treatment_initiated == 'yes_phi' and \
                        diagnosing_facility_id and treatment_initiating_facility_id and \
                        diagnosing_facility_id != treatment_initiating_facility_id:
                    treatment_status = 'initiated_outside_facility'
                elif treatment_initiated == 'yes_phi' and \
                        diagnosing_facility_id and treatment_initiating_facility_id:
                    treatment_status = 'initiated_first_line_treatment'
                elif treatment_initiated == 'yes_private':
                    treatment_status = 'initiated_outside_rntcp'

                if treatment_status:
                    case_id = case['_id']
                    f.write(case_id + "\n")
                    logger.info(case_id)

                    case_structure = CaseStructure(
                        case_id=case_id,
                        walk_related=False,
                        attrs={
                            "create": False,
                            "update": {
                                "treatment_status": treatment_status,
                                "updated_by_migration": "enikshay_2b_treatment_status_fix",
                            },
                        },
                    )

                    if commit:
                        factory.create_or_update_case(case_structure)
        logger.info("Migration finished at {}".format(datetime.datetime.utcnow()))

{'fraction_non_alphanumeric': 0.08291692692067458, 'fraction_numerical': 0.0009369144284821987, 'mean_word_length': 1.6799163179916319, 'pattern_counts': {'":': 8, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 2}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5632692', 'n_tokens_mistral': 1702, 'n_tokens_neox': 1577, 'n_words': 318}