cumulusci-plus 5.0.0__py3-none-any.whl

cumulusci/tasks/bulkdata/upsert_utils.py

@@ -0,0 +1,130 @@
+import logging
+import typing as T
+
+from simple_salesforce import Salesforce
+from sqlalchemy import Column, MetaData, Table, Unicode, UniqueConstraint, and_
+from sqlalchemy.engine.base import Connection
+from sqlalchemy.exc import IntegrityError
+
+from cumulusci.core.exceptions import BulkDataException
+from cumulusci.tasks.bulkdata.mapping_parser import CaseInsensitiveDict, MappingStep
+from cumulusci.tasks.bulkdata.query_transformers import LoadQueryExtender
+from cumulusci.tasks.bulkdata.step import DataOperationType
+from cumulusci.tasks.bulkdata.utils import create_table_if_needed, sf_query_to_table
+
+logger = logging.getLogger(__name__)
+
+
+def _create_empty_upsert_key_table(
+    tablename, metadata, fieldnames: T.List[str]
+) -> Table:
+    """Create an empty table matching the upsert semantics"""
+    fields = [Column(fieldname, Unicode(255)) for fieldname in fieldnames]
+
+    # the data fields must be unique
+    non_id_fields = [fieldname for fieldname in fieldnames if fieldname.lower() != "id"]
+    assert len(fields) == len(non_id_fields) + 1
+    fields.append(UniqueConstraint(*non_id_fields))
+
+    return create_table_if_needed(tablename, metadata, fields)
+
+
+def upsert_tablename_for_obj_and_keys(sf_object, keys):
+    return f"upsert_{sf_object}_{'_'.join(keys)}"
+
+
+def extract_upsert_key_data(
+    sf_object: str,
+    keys: tuple,
+    context,
+    metadata: MetaData,
+    connection: Connection,
+) -> Table:
+    """Create a table with keys and IDs from Salesforce"""
+    tablename = upsert_tablename_for_obj_and_keys(sf_object, keys)
+    if "Id" not in (key.title() for key in keys):
+        keys = ("Id",) + keys
+
+    table = _create_empty_upsert_key_table(tablename, metadata, keys)
+
+    soql_query = f"select {','.join(keys)} from {sf_object}"
+    try:
+        logger.info(f"Extracting records for upsert: `{soql_query}`")
+        table = sf_query_to_table(
+            table=table,
+            sobject=sf_object,
+            fields=keys,
+            api_options={},
+            context=context,
+            query=soql_query,
+            metadata=metadata,
+            connection=connection,
+        )
+    except IntegrityError as e:
+        message = str(e)
+        if "UNIQUE constraint failed" not in message:  # pragma: no cover
+            raise e
+        raise BulkDataException(f"Duplicate values for upsert key:\n {e.params}") from e
+
+    return table
+
+
+def needs_etl_upsert(mapping: MappingStep, sf: Salesforce):
+    """Is this an upsert that Salesforce cannot do natively?"""
+    # is this an upsert and one that Salesforce cannot do by itself?
+    keys = mapping.update_key
+
+    if len(keys) > 1:
+        return True
+
+    describe_data = mapping.describe_data(sf)
+
+    key = keys[0]
+    is_sf_builtin_key = (
+        key == "Id"
+        or describe_data[key]["externalId"]
+        or describe_data[key]["idLookup"]
+    )
+    return not is_sf_builtin_key
+
+
+class AddUpsertsToQuery(LoadQueryExtender):
+    @property
+    def columns_to_add(self):
+        if self.mapping.action == DataOperationType.ETL_UPSERT:
+            # We've retrieved IDs from the org, so include them.
+
+            # If we treat "Id" as an "external_id_name" then it's
+            # allowed to be sparse.
+            tablename = upsert_tablename_for_obj_and_keys(
+                self.mapping.sf_object, self.mapping.update_key
+            )
+            table = self.metadata.tables[tablename]
+            id_column = table.columns["Id"]
+            return [id_column]
+
+    @property
+    def outerjoins_to_add(self):
+        if self.mapping.action == DataOperationType.ETL_UPSERT:
+            tablename = upsert_tablename_for_obj_and_keys(
+                self.mapping.sf_object, self.mapping.update_key
+            )
+            table = self.metadata.tables[tablename]
+            return [self._make_outerjoin_for_upsert(table)]
+
+    def _make_outerjoin_for_upsert(
+        self,
+        upsert_key_table: Table,
+    ):
+        """Make a query that joins user-provided data to Salesforce IDs"""
+        update_data_table = self.metadata.tables[self.mapping.table]
+
+        # normalize columns
+        update_table_columns = CaseInsensitiveDict(update_data_table.c.items())
+
+        join_on_clauses = [
+            (upsert_key_table.columns[column_name] == update_table_columns[column_name])
+            for column_name in self.mapping.update_key
+        ]
+
+        return (upsert_key_table, and_(*join_on_clauses))

cumulusci/tasks/bulkdata/utils.py

@@ -0,0 +1,249 @@
+import collections
+import logging
+import tempfile
+import typing as T
+from contextlib import contextmanager, nullcontext
+from pathlib import Path
+
+from requests.structures import CaseInsensitiveDict as RequestsCaseInsensitiveDict
+from simple_salesforce import Salesforce
+from sqlalchemy import Boolean, Column, MetaData, Table, Unicode, inspect
+from sqlalchemy.engine.base import Connection
+from sqlalchemy.orm import Session, mapper
+
+from cumulusci.core.enums import StrEnum
+from cumulusci.core.exceptions import BulkDataException
+from cumulusci.utils.iterators import iterate_in_chunks
+
+
+class DataApi(StrEnum):
+    """Enum defining requested Salesforce data API for an operation."""
+
+    BULK = "bulk"
+    REST = "rest"
+    SMART = "smart"
+
+
+class CaseInsensitiveDict(RequestsCaseInsensitiveDict):
+    def __init__(self, *args, **kwargs):
+        self._canonical_keys = {}
+        super().__init__(*args, **kwargs)
+
+    def canonical_key(self, name):
+        return self._canonical_keys[name.lower()]
+
+    def __setitem__(self, key, value):
+        super().__setitem__(key, value)
+        self._canonical_keys[key.lower()] = key
+
+
+class SqlAlchemyMixin:
+    logger: logging.Logger
+    metadata: MetaData
+    models: dict
+    options: dict
+    session: Session
+    sf: Salesforce
+
+    def _create_record_type_table(self, table_name):
+        """Create a table to store mapping between Record Type Ids and Developer Names."""
+        rt_map_model_name = f"{table_name}Model"
+        self.models[table_name] = type(rt_map_model_name, (object,), {})
+        rt_map_fields = [
+            Column("record_type_id", Unicode(18), primary_key=True),
+            Column("developer_name", Unicode(255)),
+            Column("is_person_type", Boolean),
+        ]
+        rt_map_table = Table(table_name, self.metadata, *rt_map_fields)
+        mapper(self.models[table_name], rt_map_table)
+
+    def _extract_record_types(
+        self, sobject, tablename: str, conn, is_person_accounts_enabled: bool
+    ):
+        """Query for Record Type information and persist it in the database."""
+        self.logger.info(f"Extracting Record Types for {sobject}")
+        if is_person_accounts_enabled:
+            query = f"SELECT Id, DeveloperName, IsPersonType FROM RecordType WHERE SObjectType='{sobject}'"
+        else:
+            query = f"SELECT Id, DeveloperName FROM RecordType WHERE SObjectType='{sobject}'"
+
+        result = self.sf.query(query)
+
+        if result["totalSize"]:
+            sql_bulk_insert_from_records(
+                connection=conn,
+                table=self.metadata.tables[tablename],
+                columns=["record_type_id", "developer_name", "is_person_type"],
+                record_iterable=(
+                    [rt["Id"], rt["DeveloperName"], rt.get("IsPersonType", False)]
+                    for rt in result["records"]
+                ),
+            )
+
+    @contextmanager
+    def _temp_database_url(self):
+        with tempfile.TemporaryDirectory() as t:
+            tempdb = Path(t) / "temp_db.db"
+
+            self.logger.debug(f"Using temporary database {tempdb}")
+            database_url = f"sqlite:///{tempdb}"
+            yield database_url
+
+    def _database_url(self):
+        database_url = self.options.get("database_url")
+        if database_url:
+            return nullcontext(enter_result=database_url)
+        else:
+            return self._temp_database_url()
+
+
+def _handle_primary_key(mapping, fields):
+    """Provide support for legacy mappings which used the OID as the pk but
+    default to using an autoincrementing int pk and a separate sf_id column"""
+
+    id_column = "id"
+    if mapping.get_oid_as_pk():
+        # Get Id column from mapping
+        id_column = mapping.fields["Id"]
+
+    fields.append(Column(id_column, Unicode(255), primary_key=True))
+
+
+def create_table(mapping, metadata) -> Table:
+    """Given a mapping data structure (from mapping.yml) and SQLAlchemy
+    metadata, create a table matching the mapping.
+
+    Mapping should be a MappingStep instance"""
+
+    fields = []
+    _handle_primary_key(mapping, fields)
+
+    # make a field list to create
+    for field, db in mapping.get_complete_field_map().items():
+        if field == "Id":
+            continue
+
+        fields.append(Column(db, Unicode(255)))
+
+    if mapping.record_type:
+        fields.append(Column("record_type", Unicode(255)))
+    return create_table_if_needed(mapping.table, metadata, fields)
+
+
+def create_table_if_needed(tablename, metadata, fields: T.List[Column]) -> Table:
+    t = Table(tablename, metadata, *fields)
+    inspector = inspect(metadata.bind)
+    if inspector.has_table(tablename):
+        raise BulkDataException(f"Table already exists: {tablename}")
+    t.create(metadata.bind)
+    return t
+
+
+def generate_batches(num_records, batch_size):
+    """Generate batch size list for splitting a number of tasks into batch jobs.
+
+    Given a number of records to split up, and a batch size, generate a
+    stream of batchsize, index pairs"""
+    num_batches, left_over = divmod(num_records, batch_size)
+    if left_over:
+        num_batches += 1  # need an extra batch for the clean-up
+    for i in range(0, num_batches):
+        is_last_batch = i == num_batches - 1
+        if is_last_batch and left_over:
+            batch_size = left_over
+        if batch_size > 0:
+            yield batch_size, i, num_batches
+
+
+class RowErrorChecker:
+    def __init__(self, logger, ignore_row_errors, row_warning_limit):
+        self.logger = logger
+        self.ignore_row_errors = ignore_row_errors
+        self.row_warning_limit = row_warning_limit
+        self.row_error_count = 0
+
+    def check_for_row_error(self, result, row_id):
+        if not result.success:
+            msg = f"Error on record with id {row_id}: {result.error}"
+            if self.ignore_row_errors:
+                if self.row_error_count < self.row_warning_limit:
+                    self.logger.warning(msg)
+                elif self.row_error_count == self.row_warning_limit:
+                    self.logger.warning("Further warnings suppressed")
+                self.row_error_count += 1
+                return self.row_error_count
+            else:
+                raise BulkDataException(msg)
+
+
+def consume(iterator):
+    """Consume an iterator for its side effects.
+
+    Simplified from the function in https://docs.python.org/3/library/itertools.html
+    """
+    collections.deque(iterator, maxlen=0)
+
+
+def sql_bulk_insert_from_records(
+    *,
+    connection: Connection,
+    table: Table,
+    columns: T.Tuple[str],
+    record_iterable: T.Iterable,
+) -> None:
+    """Persist records from the given generator into the local database."""
+    consume(
+        sql_bulk_insert_from_records_incremental(
+            connection=connection,
+            table=table,
+            columns=columns,
+            record_iterable=record_iterable,
+        )
+    )
+
+
+def sql_bulk_insert_from_records_incremental(
+    *,
+    connection: Connection,
+    table: Table,
+    columns: T.Tuple[str],
+    record_iterable: T.Iterable,
+):
+    """Generator that persists batches of records from the given generator into the local database
+
+    Yields after every batch."""
+    dict_iterable = (dict(zip(columns, row)) for row in record_iterable)
+    for group in iterate_in_chunks(10000, dict_iterable):
+        with connection.begin():
+            connection.execute(table.insert(), group)
+        # self.session.flush()  -- Did this line do anything?
+        yield
+
+
+def sf_query_to_table(
+    *,
+    table: Table,
+    metadata: MetaData,
+    connection: Connection,
+    sobject: str,
+    fields: T.List[str],
+    **queryargs,
+) -> Table:
+    """Cache data from Salesforce in a SQL Alchemy Table."""
+    from cumulusci.tasks.bulkdata.step import DataOperationStatus, get_query_operation
+
+    qs = get_query_operation(sobject=sobject, fields=fields, **queryargs)
+    qs.query()
+    if qs.job_result.status is not DataOperationStatus.SUCCESS:  # pragma: no cover
+        raise BulkDataException(
+            f"Unable to query records for {sobject}: {','.join(qs.job_result.job_errors)}"
+        )
+    results = qs.get_results()
+
+    sql_bulk_insert_from_records(
+        connection=connection,
+        table=table,
+        columns=tuple(fields),
+        record_iterable=results,
+    )
+    return table

cumulusci/tasks/command.py

@@ -0,0 +1,178 @@
+""" Tasks for running a command in a subprocess
+
+Command - run a command with optional environment variables
+SalesforceCommand - run a command with credentials passed
+SalesforceBrowserTest - a task designed to wrap browser testing that could run locally or remotely
+"""
+
+import json
+import os
+import sys
+
+import sarge
+
+from cumulusci.core.exceptions import CommandException
+from cumulusci.core.tasks import BaseTask
+from cumulusci.core.utils import process_bool_arg
+
+
+class Command(BaseTask):
+    """Execute a shell command in a subprocess"""
+
+    task_docs = """
+        **Example Command-line Usage:**
+        ``cci task run command -o command "echo 'Hello command task!'"``
+
+        **Example Task to Run Command:**
+
+        .. code-block:: yaml
+
+            hello_world:
+                description: Says hello world
+                class_path: cumulusci.tasks.command.Command
+                options:
+                    command: echo 'Hello World!'
+"""
+
+    task_options = {
+        "command": {"description": "The command to execute", "required": True},
+        "dir": {
+            "description": "If provided, the directory where the command "
+            "should be run from."
+        },
+        "env": {
+            "description": "Environment variables to set for command. Must "
+            "be flat dict, either as python dict from YAML or "
+            "as JSON string."
+        },
+        "pass_env": {
+            "description": "If False, the current environment variables "
+            "will not be passed to the child process. "
+            "Defaults to True",
+            "required": True,
+        },
+        "interactive": {
+            "description": "If True, the command will use stderr, stdout, "
+            "and stdin of the main process."
+            "Defaults to False."
+        },
+    }
+
+    def _init_options(self, kwargs):
+        super(Command, self)._init_options(kwargs)
+        if "pass_env" not in self.options:
+            self.options["pass_env"] = True
+        if "dir" not in self.options or not self.options["dir"]:
+            self.options["dir"] = "."
+        if "interactive" not in self.options:
+            self.options["interactive"] = False
+        if "env" not in self.options:
+            self.options["env"] = {}
+        else:
+            try:
+                self.options["env"] = json.loads(self.options["env"])
+            except TypeError:
+                # assume env is already dict
+                pass
+
+    def _run_task(self):
+        env = self._get_env()
+        self._run_command(env)
+
+    def _get_env(self):
+        if process_bool_arg(self.options["pass_env"]):
+            env = os.environ.copy()
+        else:
+            env = {}
+
+        env.update(self.options["env"])
+        return env
+
+    def _process_output(self, line):
+        self.logger.info(line.decode("utf-8").rstrip())
+
+    def _handle_returncode(self, returncode, stderr):
+        if returncode:
+            message = "Return code: {}".format(returncode)
+            if stderr:
+                message += "\nstderr: {}".format(stderr.read().decode("utf-8"))
+            self.logger.error(message)
+            raise CommandException(message)
+
+    def _get_command(self):
+        return self.options["command"]
+
+    def _run_command(
+        self, env, command=None, output_handler=None, return_code_handler=None
+    ):
+        if not command:
+            command = self._get_command()
+
+        interactive_mode = process_bool_arg(self.options["interactive"])
+
+        self.logger.info(f"Running command: {command}")
+
+        p = sarge.Command(
+            command,
+            stdout=sys.stdout if interactive_mode else sarge.Capture(buffer_size=-1),
+            stderr=sys.stderr if interactive_mode else sarge.Capture(buffer_size=-1),
+            shell=True,
+            env=env,
+            cwd=self.options.get("dir"),
+        )
+        if interactive_mode:
+            p.run(input=sys.stdin)
+        else:
+            p.run(async_=True)
+            # Handle output lines
+            if not output_handler:
+                output_handler = self._process_output
+            while True:
+                line = p.stdout.readline(timeout=1.0)
+                if line:
+                    output_handler(line)
+                elif p.poll() is not None:
+                    break
+            p.wait()
+        self.return_values = {"returncode": p.returncode}
+
+        # Handle return code
+        if not return_code_handler:
+            return_code_handler = self._handle_returncode
+        return_code_handler(p.returncode, None if interactive_mode else p.stderr)
+
+
+class SalesforceCommand(Command):
+    """Execute a Command with SF credentials provided on the environment.
+
+    Provides:
+     * SF_ORG_INSTANCE_URL
+     * SF_ACCESS_TOKEN
+    """
+
+    salesforce_task = True
+
+    def _update_credentials(self):
+        self.org_config.refresh_oauth_token(self.project_config.keychain)
+
+    def _get_env(self):
+        env = super(SalesforceCommand, self)._get_env()
+        env["SF_ACCESS_TOKEN"] = self.org_config.access_token
+        env["SF_ORG_INSTANCE_URL"] = self.org_config.instance_url
+        return env
+
+
+task_options = Command.task_options.copy()
+task_options["extra"] = {
+    "description": "If provided, will be appended to the end of the "
+    "command.  Use to pass extra args to the command.",
+    "required": False,
+}
+task_options["use_saucelabs"] = {
+    "description": "If True, use SauceLabs to run the tests. The "
+    "SauceLabs credentials will be fetched from the "
+    "saucelabs service in the keychain and passed as "
+    "environment variables to the command.  Defaults to "
+    "False to run tests in the local browser.",
+    "required": True,
+}