cumulusci-plus 5.0.0__py3-none-any.whl

cumulusci/tasks/bulkdata/extract.py

@@ -0,0 +1,441 @@
+import itertools
+import re
+from contextlib import contextmanager
+
+from sqlalchemy import Column, MetaData, Table, Unicode, create_engine
+from sqlalchemy.orm import create_session, mapper
+
+from cumulusci.core.exceptions import (
+    BulkDataException,
+    ConfigError,
+    CumulusCIException,
+    TaskOptionsError,
+)
+from cumulusci.core.utils import process_bool_arg
+from cumulusci.tasks.bulkdata.dates import adjust_relative_dates
+from cumulusci.tasks.bulkdata.mapping_parser import (
+    parse_from_yaml,
+    validate_and_inject_mapping,
+)
+from cumulusci.tasks.bulkdata.step import (
+    DataOperationStatus,
+    DataOperationType,
+    get_query_operation,
+)
+from cumulusci.tasks.bulkdata.utils import (
+    SqlAlchemyMixin,
+    consume,
+    create_table,
+    sql_bulk_insert_from_records,
+    sql_bulk_insert_from_records_incremental,
+)
+from cumulusci.tasks.salesforce import BaseSalesforceApiTask
+from cumulusci.utils import log_progress
+
+
+class ExtractData(SqlAlchemyMixin, BaseSalesforceApiTask):
+    """Perform Bulk Queries to extract data for a mapping and persist to a SQL file or database."""
+
+    task_options = {
+        "database_url": {
+            "description": "A DATABASE_URL where the query output should be written"
+        },
+        "mapping": {
+            "description": "The path to a yaml file containing mappings of the database fields to Salesforce object fields",
+            "required": True,
+        },
+        "sql_path": {
+            "description": "If set, an SQL script will be generated at the path provided "
+            + "This is useful for keeping data in the repository and allowing diffs."
+        },
+        "inject_namespaces": {
+            "description": "If True, the package namespace prefix will be "
+            "automatically added to (or removed from) objects "
+            "and fields based on the name used in the org. Defaults to True."
+        },
+        "drop_missing_schema": {
+            "description": "Set to True to skip any missing objects or fields instead of stopping with an error."
+        },
+    }
+
+    def _init_options(self, kwargs):
+        super(ExtractData, self)._init_options(kwargs)
+        if self.options.get("database_url"):
+            # prefer database_url if it's set
+            self.options["sql_path"] = None
+        elif not self.options.get("sql_path"):
+            raise TaskOptionsError(
+                "You must set either the database_url or sql_path option."
+            )
+
+        inject_namespaces = self.options.get("inject_namespaces")
+        self.options["inject_namespaces"] = process_bool_arg(
+            True if inject_namespaces is None else inject_namespaces
+        )
+        self.options["drop_missing_schema"] = process_bool_arg(
+            self.options.get("drop_missing_schema") or False
+        )
+        self._id_generators = {}
+
+    def _run_task(self):
+        self._init_mapping()
+        with self._init_db():
+            for mapping in self.mapping.values():
+                soql = self._soql_for_mapping(mapping)
+                self._run_query(soql, mapping)
+
+            self._map_autopks()
+
+            if self.options.get("sql_path"):
+                self._sqlite_dump()
+
+    @contextmanager
+    def _init_db(self):
+        """Initialize the database and automapper."""
+        self.models = {}
+
+        with self._database_url() as database_url:
+
+            # initialize the DB engine
+            parent_engine = create_engine(database_url)
+            with parent_engine.connect() as connection:
+                # initialize DB metadata
+                self.metadata = MetaData()
+                self.metadata.bind = connection
+
+                # Create the tables
+                self._create_tables()
+
+                # initialize session
+                self.session = create_session(bind=connection, autocommit=False)
+
+                yield self.session, self.metadata, connection
+
+    def _init_mapping(self):
+        """Load a YAML mapping file."""
+        mapping_file_path = self.options["mapping"]
+        if not mapping_file_path:
+            raise TaskOptionsError("Mapping file path required")
+        self.logger.info(f"Mapping file: {self.options['mapping']}")
+
+        self.mapping = parse_from_yaml(mapping_file_path)
+
+        validate_and_inject_mapping(
+            mapping=self.mapping,
+            sf=self.sf,
+            namespace=self.project_config.project__package__namespace,
+            data_operation=DataOperationType.QUERY,
+            inject_namespaces=self.options["inject_namespaces"],
+            drop_missing=self.options["drop_missing_schema"],
+            org_has_person_accounts_enabled=self.org_config.is_person_accounts_enabled,
+        )
+
+    def _soql_for_mapping(self, mapping):
+        """Return a SOQL query suitable for extracting data for this mapping."""
+        sf_object = mapping.sf_object
+        fields = mapping.get_extract_field_list()
+        soql = f"SELECT {', '.join(fields)} FROM {sf_object}"
+
+        if mapping.record_type:
+            soql += f" WHERE RecordType.DeveloperName = '{mapping.record_type}'"
+
+        if mapping.soql_filter is not None:
+            soql = self.append_filter_clause(
+                soql=soql, filter_clause=mapping.soql_filter
+            )
+
+        return soql
+
+    def _run_query(self, soql, mapping):
+        """Execute a Bulk or REST API query job and store the results."""
+
+        step = get_query_operation(
+            sobject=mapping.sf_object,
+            api=mapping.api,
+            fields=list(mapping.get_extract_field_list()),
+            api_options={},
+            context=self,
+            query=soql,
+        )
+
+        self.logger.info(f"Extracting data for sObject {mapping['sf_object']}")
+        step.query()
+
+        if step.job_result.status is DataOperationStatus.SUCCESS:
+            if step.job_result.records_processed:
+                self.logger.info("Downloading and importing records")
+                self._import_results(mapping, step)
+            else:
+                self.logger.info(f"No records found for sObject {mapping['sf_object']}")
+        else:
+            raise BulkDataException(
+                f"Unable to execute query: {','.join(step.job_result.job_errors)}"
+            )
+
+    def _import_results(self, mapping, step):
+        """Ingest results from the Bulk API query."""
+        conn = self.session.connection()
+
+        # Map SF field names to local db column names
+        field_map = mapping.get_complete_field_map(include_id=True)
+        columns = [field_map[f] for f in field_map]  # Get values in insertion order.
+
+        record_type = mapping.record_type
+        if record_type:
+            columns.append("record_type")
+
+        # TODO: log_progress needs to know our batch size, when made configurable.
+        record_iterator = log_progress(step.get_results(), self.logger)
+        if record_type:
+            record_iterator = (record + [record_type] for record in record_iterator)
+
+        # Convert relative dates to stable dates.
+        if mapping.anchor_date:
+            date_context = mapping.get_relative_date_context(
+                list(field_map.keys()), self.sf
+            )
+            if date_context[0] or date_context[1]:
+                record_iterator = (
+                    adjust_relative_dates(
+                        mapping, date_context, record, DataOperationType.QUERY
+                    )
+                    for record in record_iterator
+                )
+
+        # Set Name field as blank for Person Account "Account" records.
+        if (
+            mapping.sf_object == "Account"
+            and "Name" in field_map
+            and self.org_config.is_person_accounts_enabled
+        ):
+            # Bump indices by one since record's ID is the first column.
+            Name_index = columns.index(mapping.fields["Name"])
+            IsPersonAccount_index = columns.index(mapping.fields["IsPersonAccount"])
+
+            def strip_name_field(record):
+                nonlocal Name_index, IsPersonAccount_index
+                if record[IsPersonAccount_index].lower() == "true":
+                    record[Name_index] = ""
+                return record
+
+            record_iterator = (strip_name_field(record) for record in record_iterator)
+
+        if mapping.get_oid_as_pk():
+            sql_bulk_insert_from_records(
+                connection=conn,
+                table=self.metadata.tables[mapping.table],
+                columns=columns,
+                record_iterable=record_iterator,
+            )
+        else:
+            # If using the autogenerated id field, split out the returned records
+            # into two separate streams and load into the main table and the sf_id_table
+            values, ids = itertools.tee(record_iterator)
+
+            # Generate naming format (<sobject>-<number>)
+            id_source_sobj, id_source_global = itertools.tee(
+                self._id_generator_for_object(mapping.sf_object)
+            )
+
+            f_ids = ((row[0], next(id_source_global)) for row in ids)
+            f_values = ([next(id_source_sobj)] + row[1:] for row in values)
+
+            values_chunks = sql_bulk_insert_from_records_incremental(
+                connection=conn,
+                table=self.metadata.tables[mapping.table],
+                columns=["id"] + columns[1:],
+                record_iterable=f_values,
+            )
+            ids_chunks = sql_bulk_insert_from_records_incremental(
+                connection=conn,
+                table=self.metadata.tables[mapping.get_sf_id_table()],
+                columns=["sf_id", "id"],
+                record_iterable=f_ids,
+            )
+
+            # do the inserts one chunk at a time based on all of the
+            # generators nested previously.
+            consume(zip(values_chunks, ids_chunks))
+
+        if "RecordTypeId" in mapping.fields:
+            self._extract_record_types(
+                mapping.sf_object,
+                mapping.get_source_record_type_table(),
+                conn,
+                self.org_config.is_person_accounts_enabled,
+            )
+
+        self.session.commit()
+
+    def _id_generator_for_object(self, sobject: str):
+        """Generates strings for local ids in format {sobject}-{counter}
+        (example: Account-2)"""
+        if sobject not in self._id_generators:
+
+            def _generate_ids():
+                counter = 1
+                while True:
+                    yield f"{sobject}-{counter}"
+                    counter += 1
+
+            self._id_generators[sobject] = _generate_ids()
+
+        return self._id_generators[sobject]
+
+    def _map_autopks(self):
+        # Convert Salesforce Ids to autopks
+        for m in self.mapping.values():
+            lookup_keys = list(m.lookups.keys())
+            if not m.get_oid_as_pk():
+                if lookup_keys:
+                    self._convert_lookups_to_id(m, lookup_keys)
+
+        # Drop sf_id tables
+        for m in self.mapping.values():
+            if not m.get_oid_as_pk():
+                self.metadata.tables[m.get_sf_id_table()].drop()
+
+    def _get_mapping_for_table(self, table):
+        """Return the first mapping for a table name"""
+
+        # Get all mappings for lookups
+        mappings = [
+            mapping
+            for mapping in self.mapping.values()
+            if (isinstance(table, str) and mapping["table"] == table)
+            or (isinstance(table, list) and mapping["table"] in table)
+        ]
+
+        # For polymorphic lookups, raise exception if missing mappings
+        if isinstance(table, list) and len(mappings) != len(table):
+            missing_tables = set(table) - set(mapping["table"] for mapping in mappings)
+            raise CumulusCIException(
+                f"The following tables are missing in the mapping file: {missing_tables}"
+            )
+
+        return mappings
+
+    def _convert_lookups_to_id(self, mapping, lookup_keys):
+        """Rewrite persisted Salesforce Ids to refer to auto-PKs."""
+
+        def throw(string):  # pragma: no cover
+            raise BulkDataException(string)
+
+        for lookup_key in lookup_keys:
+            lookup_info = mapping.lookups.get(lookup_key) or throw(
+                f"Cannot find lookup info {lookup_key}"
+            )
+            model = self.models.get(mapping.table)
+
+            lookup_mappings = self._get_mapping_for_table(lookup_info.table) or throw(
+                f"Cannot find lookup mapping for {lookup_info.table}"
+            )
+
+            key_field = lookup_info.get_lookup_key_field()
+
+            key_attr = getattr(model, key_field, None) or throw(
+                f"key_field {key_field} not found in table {mapping.table}"
+            )
+
+            # Keep track of total mapping operations
+            total_mapping_operations = 0
+
+            for lookup_mapping in lookup_mappings:
+                lookup_model = self.models.get(lookup_mapping.get_sf_id_table())
+                try:
+                    update_query = (
+                        self.session.query(model)
+                        .filter(key_attr.isnot(None), key_attr == lookup_model.sf_id)
+                        .update({key_attr: lookup_model.id}, synchronize_session=False)
+                    )
+                    total_mapping_operations += update_query.rowcount
+                except NotImplementedError:
+                    # Some databases such as sqlite don't support multitable update
+                    mappings = []
+                    for row, lookup_id in self.session.query(
+                        model, lookup_model.id
+                    ).join(lookup_model, key_attr == lookup_model.sf_id):
+                        mappings.append({"id": row.id, key_field: lookup_id})
+                    total_mapping_operations += len(mappings)
+                    self.session.bulk_update_mappings(model, mappings)
+            # Count the total number of rows excluding those with no entry for that field
+            total_rows = (
+                self.session.query(model)
+                .filter(
+                    key_attr.isnot(None),  # Ensure key_attr is not None
+                    key_attr.isnot(""),  # Ensure key_attr is not an empty string
+                )
+                .count()
+            )
+
+            if total_mapping_operations != total_rows:
+                raise ConfigError(
+                    f"Total mapping operations ({total_mapping_operations}) do not match total non-empty rows ({total_rows}) for lookup_key: {lookup_key}. Mention all related tables for lookup: {lookup_key}"
+                )
+        self.session.commit()
+
+    def _create_tables(self):
+        """Create a table for each mapping step."""
+        for mapping in self.mapping.values():
+            self._create_table(mapping)
+        self.metadata.create_all()
+
+    def _create_table(self, mapping):
+        """Create a table for the given mapping."""
+        model_name = f"{mapping.table}Model"
+        mapper_kwargs = {}
+        self.models[mapping.table] = type(model_name, (object,), {})
+
+        t = create_table(mapping, self.metadata)
+
+        if "RecordTypeId" in mapping.fields:
+            # We're using Record Type Mapping support.
+            # If multiple mappings point to the same table, don't recreate the table
+            if mapping.get_source_record_type_table() not in self.models:
+                self._create_record_type_table(mapping.get_source_record_type_table())
+
+        if not mapping.get_oid_as_pk():
+            # If multiple mappings point to the same table, don't recreate the table
+            if mapping.get_sf_id_table() not in self.models:
+                sf_id_model_name = f"{mapping.get_sf_id_table()}Model"
+                self.models[mapping.get_sf_id_table()] = type(
+                    sf_id_model_name, (object,), {}
+                )
+                sf_id_fields = [
+                    Column("id", Unicode(255), primary_key=True),
+                    Column("sf_id", Unicode(24)),
+                ]
+                id_t = Table(mapping.get_sf_id_table(), self.metadata, *sf_id_fields)
+                mapper(self.models[mapping.get_sf_id_table()], id_t)
+
+        mapper(self.models[mapping.table], t, **mapper_kwargs)
+
+    def _sqlite_dump(self):
+        """Write a SQLite script output file."""
+        path = self.options["sql_path"]
+        with open(path, "w", encoding="utf-8") as f:
+            for line in self.session.connection().connection.iterdump():
+                f.write(line + "\n")
+
+    def append_filter_clause(self, soql, filter_clause):
+        """Function that applies filter clause to soql if it is defined in mapping yml file"""
+
+        if not filter_clause:
+            return soql
+
+        # If WHERE keyword is specified in the maping file replace it with empty string.
+        # match WHERE keyword only at the start of the string and whitespace after it.
+        filter_clause = re.sub(
+            pattern=r"^WHERE\s+",
+            repl="",
+            string=filter_clause.strip(),
+            flags=re.IGNORECASE,
+        )
+
+        # If WHERE keyword is already in soql query(because of record type filter) add AND clause
+        if " WHERE " in soql:
+            soql = f"{soql} AND {filter_clause}"
+        else:
+            soql = f"{soql} WHERE {filter_clause}"
+
+        return soql

cumulusci/tasks/bulkdata/extract_dataset_utils/calculate_dependencies.py

@@ -0,0 +1,117 @@
+import typing as T
+
+from cumulusci.salesforce_api.filterable_objects import NOT_EXTRACTABLE
+from cumulusci.salesforce_api.org_schema import Schema
+
+from .synthesize_extract_declarations import (
+    SimplifiedExtractDeclaration,
+    synthesize_declaration_for_sobject,
+)
+
+
+class SObjDependency(T.NamedTuple):
+    table_name_from: str
+    table_names_to: T.Union[str, T.Tuple[str, ...]]
+    field_name: str
+    priority: bool = False
+
+
+def _calculate_dependencies_for_declarations(
+    decls: T.Sequence[SimplifiedExtractDeclaration], schema: Schema
+) -> T.Dict[str, T.List[SObjDependency]]:
+    """Ensure that required lookups are fulfilled for list of SimplifiedExtractDeclarations
+
+    Do this by adding their referent tables (in full) to the extract.
+    Module-internal helper function.
+    """
+    dependencies = {}
+    for decl in decls:
+        assert isinstance(decl.sf_object, str)
+        new_dependencies = _collect_dependencies_for_sobject(
+            decl.sf_object, decl.fields, schema, only_required_fields=False
+        )
+        dependencies.update(new_dependencies)
+    return dependencies
+
+
+def _collect_dependencies_for_sobject(
+    source_sfobject: str,
+    fields: T.List[str],
+    schema: Schema,
+    only_required_fields: bool,
+) -> T.Dict[str, T.List[SObjDependency]]:
+    """Ensure that required lookups are fulfilled for a single SObject
+
+    Do this by adding its referent tables (in full) to the extract.
+    Module-internal helper function.
+    """
+    dependencies = {}
+    for field_name in fields:
+        field_info = schema[source_sfobject].fields[field_name]
+        if not field_info.createable:  # pragma: no cover
+            continue
+        references = field_info.referenceTo
+        if references:
+            # Remove RecordType from references
+            if "RecordType" in references:
+                references.remove("RecordType")
+                if not references:
+                    continue
+
+            targets = tuple(
+                target for target in references if target not in NOT_EXTRACTABLE
+            )
+            field_disallowed = not targets or not field_info.createable
+            field_allowed = not (only_required_fields or field_disallowed)
+            if field_info.requiredOnCreate or field_allowed:
+                dependencies.setdefault(source_sfobject, []).append(
+                    SObjDependency(
+                        source_sfobject,
+                        targets,
+                        field_name,
+                        field_info.requiredOnCreate,
+                    )
+                )
+
+    return dependencies
+
+
+def extend_declarations_to_include_referenced_tables(
+    decl_list: T.Sequence[SimplifiedExtractDeclaration], schema: Schema
+) -> T.Sequence[SimplifiedExtractDeclaration]:
+    """Extend the declarations to complete required or requested lookups recursively"""
+    decls = {decl.sf_object: decl for decl in decl_list}
+    dependencies = _calculate_dependencies_for_declarations(decl_list, schema)
+    to_process = list(decls.keys())
+
+    while to_process:
+        sf_object = to_process.pop()
+        assert isinstance(sf_object, str)
+        my_dependencies = dependencies.get(sf_object, ())
+        for dep in my_dependencies:
+            target_tables = dep.table_names_to
+            for target_table in target_tables:
+                sobj = schema.get(target_table)
+                target_extractable = (
+                    target_table not in NOT_EXTRACTABLE and sobj and sobj.extractable
+                )
+                if target_table not in decls and target_extractable:
+                    required_fields = [
+                        field.name
+                        for field in schema[target_table].fields.values()
+                        if field.requiredOnCreate
+                    ]
+                    decls[target_table] = synthesize_declaration_for_sobject(
+                        target_table, required_fields, schema[target_table].fields
+                    )
+
+                    new_dependencies = _collect_dependencies_for_sobject(
+                        target_table,
+                        decls[target_table].fields,
+                        schema,
+                        only_required_fields=True,
+                    )
+                    dependencies.update(new_dependencies)
+                    to_process.append(target_table)
+
+    return list(decls.values())

cumulusci/tasks/bulkdata/extract_dataset_utils/extract_yml.py

@@ -0,0 +1,123 @@
+import re
+import typing as T
+from pathlib import Path
+
+from pydantic import Field, validator
+
+from cumulusci.core.enums import StrEnum
+from cumulusci.tasks.bulkdata.utils import DataApi
+from cumulusci.utils.yaml.model_parser import CCIDictModel, HashableBaseModel
+
+object_decl = re.compile(r"objects\((\w+)\)", re.IGNORECASE)
+field_decl = re.compile(r"fields\((\w+)\)", re.IGNORECASE)
+
+
+class SFObjectGroupTypes(StrEnum):
+    all = "all"
+    custom = "custom"
+    standard = "standard"
+
+
+class SFFieldGroupTypes(StrEnum):
+    all = "all"
+    custom = "custom"
+    standard = "standard"
+    required = "required"
+
+
+class ExtractDeclaration(HashableBaseModel):
+    where: T.Optional[str] = None
+    fields_: T.Union[T.List[str], str] = Field(["FIELDS(ALL)"], alias="fields")
+    api: DataApi = DataApi.SMART
+    sf_object: T.Optional[str] = None  # injected, not declared explicitly
+
+    @staticmethod
+    def parse_field_complex_type(fieldspec):
+        """If it's something like FIELDS(...), parse it out"""
+        if match := field_decl.match(fieldspec):
+            matching_group = match.groups()[0].lower()
+            field_group_type = getattr(SFFieldGroupTypes, matching_group, None)
+            return field_group_type
+        else:
+            return None
+
+    def assert_sf_object_fits_pattern(self):
+        if self.is_group:
+            assert (
+                self.group_type in SFObjectGroupTypes
+            ), f"Expected OBJECTS(ALL), OBJECTS(CUSTOM) or OBJECTS(STANDARD), not `{self.group_type.upper()}`"
+        else:
+            assert self.sf_object.isidentifier(), (
+                "Value should start with OBJECTS( or be a simple alphanumeric field name"
+                f" (underscores allowed) not {self.sf_object}"
+            )
+        return self.sf_object
+
+    def assert_check_where_against_complex(self):
+        """Check that a where clause was not used with a group declaration."""
+        assert not (
+            self.where and self.is_group
+        ), "Cannot specify a `where` clause on a declaration for multiple kinds of objects."
+
+    @validator("fields_")
+    def normalize_fields(cls, vals):
+        if isinstance(vals, str):
+            vals = [vals]
+        for val in vals:
+            cls.validate_field(val)
+        return vals
+
+    @classmethod
+    def validate_field(cls, val):
+        assert cls.parse_field_complex_type(val) or val.isidentifier(), val
+        if group_type := cls.parse_field_complex_type(val):
+            assert group_type, (
+                "group_type in OBJECT(group_type) should be one of "
+                f"{tuple(SFFieldGroupTypes.__members__.keys())}, "
+                f"not {val}"
+            )
+
+        return val
+
+    @property
+    def is_group(self):
+        return bool(self.group_type)
+
+    @property
+    def group_type(self):
+        """If it's something like OBJECT(...), parse it out"""
+        val = self._parse_group_type(self.sf_object)
+        if val:
+            group_type = getattr(SFObjectGroupTypes, val, None)
+            assert group_type, (
+                "group_type in OBJECT(group_type) should be one of "
+                f"{tuple(SFObjectGroupTypes.__members__.keys())}, "
+                f"not {val}"
+            )
+            return group_type
+
+    @staticmethod
+    def _parse_group_type(val):
+        if "(" in val:
+            return object_decl.match(val)[1].lower()
+        else:
+            return None
+
+
+class ExtractRulesFile(CCIDictModel):
+    version: int = 1
+    extract: T.Dict[str, ExtractDeclaration]
+
+    @validator("extract")
+    def inject_sf_object_name(cls, val):
+        for sf_object, decl in val.items():
+            decl.sf_object = sf_object
+            decl.assert_sf_object_fits_pattern()
+            decl.assert_check_where_against_complex()
+
+        return val
+
+    @classmethod
+    def parse_extract(cls, source: T.Union[str, Path, T.IO]):
+        """Return the extract key after parsing an extract file."""
+        return super().parse_from_yaml(source).extract