<commit_before>from collections import defaultdict
import datetime
from sqlalchemy import Column, MetaData, Table, Unicode, create_engine, text
from sqlalchemy.orm import aliased, Session
from sqlalchemy.ext.automap import automap_base

from cumulusci.core.exceptions import BulkDataException, TaskOptionsError
from cumulusci.core.utils import process_bool_arg
from cumulusci.tasks.bulkdata.utils import (
    get_lookup_key_field,
    SqlAlchemyMixin,
    RowErrorChecker,
)
from cumulusci.tasks.bulkdata.step import (
    BulkApiDmlOperation,
    DataOperationStatus,
    DataOperationType,
)
from cumulusci.tasks.salesforce import BaseSalesforceApiTask
from cumulusci.utils import os_friendly_path

from cumulusci.tasks.bulkdata.mapping_parser import parse_from_yaml


class LoadData(BaseSalesforceApiTask, SqlAlchemyMixin):
    """Perform Bulk API operations to load data defined by a mapping from a local store into an org."""

    task_options = {
        "database_url": {
            "description": "The database url to a database containing the test data to load"
        },
        "mapping": {
            "description": "The path to a yaml file containing mappings of the database fields to Salesforce object fields",
            "required": True,
        },
        "start_step": {
            "description": "If specified, skip steps before this one in the mapping",
            "required": False,
        },
        "sql_path": {
            "description": "If specified, a database will be created from an SQL script at the provided path"
        },
        "ignore_row_errors": {
            "description": "If True, allow the load to continue even if individual rows fail to load."
        },
        "reset_oids": {
            "description": "If True (the default), and the _sf_ids tables exist, reset them before continuing.",
            "required": False,
        },
        "bulk_mode": {
            "description": "Set to Serial to force serial mode on all jobs. Parallel is the default."
        },
    }
    row_warning_limit = 10

    def _init_options(self, kwargs):
        super(LoadData, self)._init_options(kwargs)

        self.options["ignore_row_errors"] = process_bool_arg(
            self.options.get("ignore_row_errors", False)
        )
        if self.options.get("database_url"):
            # prefer database_url if it's set
            self.options["sql_path"] = None
        elif self.options.get("sql_path"):
            self.options["sql_path"] = os_friendly_path(self.options["sql_path"])
            self.options["database_url"] = None
        else:
            raise TaskOptionsError(
                "You must set either the database_url or sql_path option."
            )
        self.reset_oids = self.options.get("reset_oids", True)
        self.bulk_mode = (
            self.options.get("bulk_mode") and self.options.get("bulk_mode").title()
        )
        if self.bulk_mode and self.bulk_mode not in ["Serial", "Parallel"]:
            raise TaskOptionsError("bulk_mode must be either Serial or Parallel")

    def _run_task(self):
        self._init_mapping()
        self._init_db()
        self._expand_mapping()

        start_step = self.options.get("start_step")
        started = False
        for name, mapping in self.mapping.items():
            # Skip steps until start_step
            if not started and start_step and name != start_step:
                self.logger.info(f"Skipping step: {name}")
                continue

            started = True

            self.logger.info(f"Running step: {name}")
            result = self._load_mapping(mapping)
            if result.status is DataOperationStatus.JOB_FAILURE:
                raise BulkDataException(
                    f"Step {name} did not complete successfully: {','.join(result.job_errors)}"
                )

            if name in self.after_steps:
                for after_name, after_step in self.after_steps[name].items():
                    self.logger.info(f"Running post-load step: {after_name}")
                    result = self._load_mapping(after_step)
                    if result.status is DataOperationStatus.JOB_FAILURE:
                        raise BulkDataException(
                            f"Step {after_name} did not complete successfully: {','.join(result.job_errors)}"
                        )

    def _load_mapping(self, mapping):
        """Load data for a single step."""

        if mapping.get("fields", {}).get("RecordTypeId"):
            conn = self.session.connection()
            self._load_record_types([mapping["sf_object"]], conn)

        mapping["oid_as_pk"] = bool(mapping.get("fields", {}).get("Id"))

        bulk_mode = mapping.get("bulk_mode") or self.bulk_mode or "Parallel"

        step = BulkApiDmlOperation(
            sobject=mapping["sf_object"],
            operation=(
                DataOperationType.INSERT
                if mapping.get("action") == "insert"
                else DataOperationType.UPDATE
            ),
            api_options={"bulk_mode": bulk_mode},
            context=self,
            fields=self._get_columns(mapping),
        )

        local_ids = []
        step.start()
        step.load_records(self._stream_queried_data(mapping, local_ids))<commit_msg>Fix a bug when loading from SQLite<commit_after>from collections import defaultdict
import datetime
from sqlalchemy import Column, MetaData, Table, Unicode, create_engine, text
from sqlalchemy.orm import aliased, Session
from sqlalchemy.ext.automap import automap_base

from cumulusci.core.exceptions import BulkDataException, TaskOptionsError
from cumulusci.core.utils import process_bool_arg
from cumulusci.tasks.bulkdata.utils import (
    get_lookup_key_field,
    SqlAlchemyMixin,
    RowErrorChecker,
)
from cumulusci.tasks.bulkdata.step import (
    BulkApiDmlOperation,
    DataOperationStatus,
    DataOperationType,
)
from cumulusci.tasks.salesforce import BaseSalesforceApiTask
from cumulusci.utils import os_friendly_path

from cumulusci.tasks.bulkdata.mapping_parser import parse_from_yaml


class LoadData(BaseSalesforceApiTask, SqlAlchemyMixin):
    """Perform Bulk API operations to load data defined by a mapping from a local store into an org."""

    task_options = {
        "database_url": {
            "description": "The database url to a database containing the test data to load"
        },
        "mapping": {
            "description": "The path to a yaml file containing mappings of the database fields to Salesforce object fields",
            "required": True,
        },
        "start_step": {
            "description": "If specified, skip steps before this one in the mapping",
            "required": False,
        },
        "sql_path": {
            "description": "If specified, a database will be created from an SQL script at the provided path"
        },
        "ignore_row_errors": {
            "description": "If True, allow the load to continue even if individual rows fail to load."
        },
        "reset_oids": {
            "description": "If True (the default), and the _sf_ids tables exist, reset them before continuing.",
            "required": False,
        },
        "bulk_mode": {
            "description": "Set to Serial to force serial mode on all jobs. Parallel is the default."
        },
    }
    row_warning_limit = 10

    def _init_options(self, kwargs):
        super(LoadData, self)._init_options(kwargs)

        self.options["ignore_row_errors"] = process_bool_arg(
            self.options.get("ignore_row_errors", False)
        )
        if self.options.get("database_url"):
            # prefer database_url if it's set
            self.options["sql_path"] = None
        elif self.options.get("sql_path"):
            self.options["sql_path"] = os_friendly_path(self.options["sql_path"])
            self.options["database_url"] = None
        else:
            raise TaskOptionsError(
                "You must set either the database_url or sql_path option."
            )
        self.reset_oids = self.options.get("reset_oids", True)
        self.bulk_mode = (
            self.options.get("bulk_mode") and self.options.get("bulk_mode").title()
        )
        if self.bulk_mode and self.bulk_mode not in ["Serial", "Parallel"]:
            raise TaskOptionsError("bulk_mode must be either Serial or Parallel")

    def _run_task(self):
        self._init_mapping()
        self._init_db()
        self._expand_mapping()

        start_step = self.options.get("start_step")
        started = False
        for name, mapping in self.mapping.items():
            # Skip steps until start_step
            if not started and start_step and name != start_step:
                self.logger.info(f"Skipping step: {name}")
                continue

            started = True

            self.logger.info(f"Running step: {name}")
            result = self._load_mapping(mapping)
            if result.status is DataOperationStatus.JOB_FAILURE:
                raise BulkDataException(
                    f"Step {name} did not complete successfully: {','.join(result.job_errors)}"
                )

            if name in self.after_steps:
                for after_name, after_step in self.after_steps[name].items():
                    self.logger.info(f"Running post-load step: {after_name}")
                    result = self._load_mapping(after_step)
                    if result.status is DataOperationStatus.JOB_FAILURE:
                        raise BulkDataException(
                            f"Step {after_name} did not complete successfully: {','.join(result.job_errors)}"
                        )

    def _load_mapping(self, mapping):
        """Load data for a single step."""

        if mapping.get("fields", {}).get("RecordTypeId"):
            conn = self.session.connection()
            self._load_record_types([mapping["sf_object"]], conn)
            self.session.commit()

        mapping["oid_as_pk"] = bool(mapping.get("fields", {}).get("Id"))

        bulk_mode = mapping.get("bulk_mode") or self.bulk_mode or "Parallel"

        step = BulkApiDmlOperation(
            sobject=mapping["sf_object"],
            operation=(
                DataOperationType.INSERT
                if mapping.get("action") == "insert"
                else DataOperationType.UPDATE
            ),
            api_options={"bulk_mode": bulk_mode},
            context=self,
            fields=self._get_columns(mapping),
        )

        local_ids = []
        step.start()
        step.load_records(self._stream_queried_data(mapping, local_ids))
{'fraction_non_alphanumeric': 0.09553361542078044, 'fraction_numerical': 0.00037611659614480487, 'mean_word_length': 2.2867737948084055, 'pattern_counts': {'":': 36, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, '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': '949771', 'n_tokens_mistral': 2902, 'n_tokens_neox': 2841, 'n_words': 756}