cumulusci-plus 5.0.0__py3-none-any.whl

cumulusci/tasks/bulkdata/generate.py

@@ -0,0 +1,4 @@
+# for backwards compatibility
+from .generate_mapping import GenerateMapping
+
+__all__ = ["GenerateMapping"]

cumulusci/tasks/bulkdata/generate_and_load_data.py

@@ -0,0 +1,232 @@
+import os
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Optional, Union
+
+from sqlalchemy import MetaData, create_engine
+
+from cumulusci.core.config import TaskConfig
+from cumulusci.core.exceptions import TaskOptionsError
+from cumulusci.core.utils import import_global
+from cumulusci.tasks.bulkdata import LoadData
+from cumulusci.tasks.bulkdata.utils import generate_batches
+from cumulusci.tasks.salesforce import BaseSalesforceApiTask
+
+
+class GenerateAndLoadData(BaseSalesforceApiTask):
+    """Orchestrate creating tempfiles, generating data, loading data, cleaning up tempfiles and batching."""
+
+    task_docs = """
+    Orchestrate creating tempfiles, generating data, loading data, cleaning up tempfiles and batching.
+
+    CCI has features for generating data and for loading them into orgs. This task pulls them
+    together to give some useful additional features, such as storing the intermediate data in
+    a tempfile (the default behavior) and generating the data in batches instead of all at
+    once (controlled by the `batch_size` option).
+
+    The simplest possible usage is to specify the number of records you'd like generated, a
+    mapping file that defines the schema and a data generation task written in Python to actually
+    generate the data.
+
+    Use the `num_records` option to specify how many records to generate.
+    Use the `mapping` option to specify a mapping file.
+    Use 'data_generation_task' to specify what Python class to use to generate the data.'
+    Use 'batch_size' to specify how many records to generate and upload in every batch.
+
+    By default it creates the data in a temporary file and then cleans it up later. Specify database_url if you
+    need more control than that. Existing data tables will be emptied before being refilled.
+    Your database will be completely deleted!
+
+    If you use database_url and batch_size together, latter batches will overwrite
+    earlier batches in the database and the first batch will replace tables if they exist.
+
+    A table mapping IDs to SFIds will persist across batches and will grow monotonically.
+
+    If your generator class makes heavy use of Faker, you might be interested in this patch
+    which frequently speeds Faker up. Adding that code to the bottom of your generator file may
+    help accelerate it.
+
+    https://sfdc.co/bwKxDD
+    """
+
+    task_options = {
+        "num_records": {
+            "description": "How many records to generate. Precise calcuation depends on the generator.",
+            "required": True,
+        },
+        "num_records_tablename": {
+            "description": "Which table to count records in.",
+            "required": False,
+        },
+        "batch_size": {
+            "description": "How many records to create and load at a time.",
+            "required": False,
+        },
+        "data_generation_task": {
+            "description": "Fully qualified class path of a task to generate the data. Look at cumulusci.tasks.bulkdata.tests.dummy_data_factory to learn how to write them.",
+            "required": True,
+        },
+        "data_generation_options": {
+            "description": "Options to pass to the data generator.",
+            "required": False,
+        },
+        "vars": {
+            "description": "Variables that the generate or load tasks might need.",
+        },
+        "replace_database": {
+            "description": "Confirmation that it is okay to delete the data in database_url",
+        },
+        "working_directory": {
+            "description": "Store temporary files in working_directory for easier debugging."
+        },
+        **LoadData.task_options,
+    }
+    task_options["mapping"]["required"] = False
+
+    def _init_options(self, kwargs):
+        super()._init_options(kwargs)
+        mapping_file = self.options.get("mapping", None)
+        if mapping_file:
+            self.mapping_file = os.path.abspath(mapping_file)
+            if not os.path.exists(self.mapping_file):
+                raise TaskOptionsError(f"{self.mapping_file} cannot be found.")
+        else:
+            self.mapping_file = None
+
+        self.database_url = self.options.get("database_url")
+        num_records = self.options.get("num_records")
+
+        self.num_records = int(num_records) if num_records is not None else None
+
+        batch_size = self.options.get("batch_size", self.num_records)
+        if batch_size is not None:
+            self.batch_size = int(batch_size)
+
+            if self.batch_size <= 0:
+                raise TaskOptionsError("Batch size should be greater than zero")
+        else:
+            self.batch_size = None
+        class_path = self.options.get("data_generation_task")
+        if class_path:
+            self.data_generation_task = import_global(class_path)
+        else:
+            raise TaskOptionsError("No data generation task specified")
+
+        self.working_directory = self.options.get("working_directory", None)
+        self.database_url = self.options.get("database_url")
+
+        if self.database_url:
+            engine, metadata = self._setup_engine(self.database_url)
+            tables = metadata.tables
+
+            if len(list(tables)) and not self.options.get("replace_database"):
+                raise TaskOptionsError(
+                    f"Database {self.database_url} has tables "
+                    f"({list(tables)}) "
+                    "but `replace_database` was not specified"
+                )
+
+    def _run_task(self):
+        with TemporaryDirectory() as tempdir:
+            working_directory = self.options.get("working_directory")
+            if working_directory:
+                tempdir = Path(working_directory)
+                tempdir.mkdir(exist_ok=True)
+            if self.batch_size:
+                batches = generate_batches(self.num_records, self.batch_size)
+            else:
+                batches = [(None, 0, 1)]
+            results = []
+            for current_batch_size, index, total_batches in batches:
+                if total_batches > 1:
+                    self.logger.info(
+                        f"Generating a data batch, batch_size={current_batch_size} "
+                        f"index={index} total_records={self.num_records}"
+                    )
+                res = self._generate_batch(
+                    database_url=self.database_url,
+                    tempdir=self.working_directory or tempdir,
+                    mapping_file=self.mapping_file,
+                    batch_size=current_batch_size,
+                    index=index,
+                    total_batches=total_batches,
+                )
+                results.append(res)
+        self.return_values = {"load_results": results}
+        return self.return_values
+
+    def _datagen(self, subtask_options):
+        task_config = TaskConfig({"options": subtask_options})
+        data_gen_task = self.data_generation_task(
+            self.project_config, task_config, org_config=self.org_config
+        )
+        data_gen_task()
+
+    def _dataload(self, subtask_options) -> dict:
+        subtask_config = TaskConfig({"options": subtask_options})
+        subtask = LoadData(
+            project_config=self.project_config,
+            task_config=subtask_config,
+            org_config=self.org_config,
+            flow=self.flow,
+            name=self.name,
+            stepnum=self.stepnum,
+        )
+        subtask()
+        return subtask.return_values
+
+    def _generate_batch(
+        self,
+        *,
+        database_url: Optional[str],
+        tempdir: Union[Path, str, None],
+        mapping_file: Union[Path, str, None],
+        batch_size: Optional[int],
+        index: int,
+        total_batches: int,
+    ) -> dict:
+        """Generate a batch in database_url or a tempfile if it isn't specified."""
+        if not database_url:
+            sqlite_path = Path(tempdir) / "generated_data.db"
+            database_url = f"sqlite:///{sqlite_path}"
+
+        self._cleanup_object_tables(*self._setup_engine(database_url))
+
+        subtask_options = {
+            **self.options,
+            "mapping": mapping_file,
+            "reset_oids": False,
+            "database_url": database_url,
+            "num_records": batch_size,
+            "current_batch_number": index,
+            "working_directory": tempdir,
+        }
+
+        # some generator tasks can generate the mapping file instead of reading it
+        if not subtask_options.get("mapping"):
+            temp_mapping = Path(tempdir) / "temp_mapping.yml"
+            mapping_file = self.options.get("generate_mapping_file", temp_mapping)
+            subtask_options["generate_mapping_file"] = mapping_file
+        self._datagen(subtask_options)
+        if not subtask_options.get("mapping"):
+            subtask_options["mapping"] = mapping_file
+        return self._dataload(subtask_options)
+
+    def _setup_engine(self, database_url):
+        """Set up the database engine"""
+        engine = create_engine(database_url)
+
+        metadata = MetaData(engine)
+        metadata.reflect()
+        return engine, metadata
+
+    def _cleanup_object_tables(self, engine, metadata):
+        """Delete all tables that do not relate to id->OID mapping"""
+        tables = metadata.tables
+        tables_to_drop = [
+            table
+            for tablename, table in tables.items()
+            if not tablename.endswith("sf_ids")
+        ]
+        if tables_to_drop:
+            metadata.drop_all(tables=tables_to_drop)

cumulusci/tasks/bulkdata/generate_and_load_data_from_yaml.py

@@ -0,0 +1,19 @@
+from cumulusci.tasks.bulkdata import GenerateAndLoadData
+from cumulusci.tasks.bulkdata.generate_from_yaml import GenerateDataFromYaml
+
+bulkgen_task = "cumulusci.tasks.bulkdata.generate_from_yaml.GenerateDataFromYaml"
+
+
+# Class to glue together the task_options from GenerateAndLoadData
+# and GenerateDataFromYaml
+class GenerateAndLoadDataFromYaml(GenerateAndLoadData):
+    """Generate and load data from Snowfakery in as many batches as necessary"""
+
+    task_options = {
+        **GenerateAndLoadData.task_options,
+        **GenerateDataFromYaml.task_options,
+    }
+
+    def _init_options(self, kwargs):
+        args = {"data_generation_task": bulkgen_task, **kwargs}
+        super()._init_options(args)

cumulusci/tasks/bulkdata/generate_from_yaml.py

@@ -0,0 +1,183 @@
+import os
+import shutil
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Optional
+
+from snowfakery import generate_data
+
+from cumulusci.core.exceptions import TaskOptionsError
+from cumulusci.core.utils import process_list_arg, process_list_of_pairs_dict_arg
+from cumulusci.tasks.bulkdata.base_generate_data_task import BaseGenerateDataTask
+
+
+class GenerateDataFromYaml(BaseGenerateDataTask):
+    """Generate sample data from a YAML template file."""
+
+    task_docs = """
+    Generate YAML from a Snowfakery recipe.
+    """
+
+    task_options = {
+        **BaseGenerateDataTask.task_options,
+        "generator_yaml": {
+            "description": "A Snowfakery recipe file to use",
+            "required": True,
+        },
+        "vars": {
+            "description": "Pass values to override options in the format VAR1:foo,VAR2:bar"
+        },
+        "plugin_options": {
+            "description": "Pass values to override plugin options in the format VAR1:foo,VAR2:bar"
+        },
+        "generate_mapping_file": {
+            "description": "A path to put a mapping file inferred from the generator_yaml",
+            "required": False,
+        },
+        "num_records_tablename": {
+            "description": "A string representing which table determines when the recipe execution is done.",
+            "required": False,
+        },
+        "continuation_file": {
+            "description": "YAML file generated by Snowfakery representing next steps for data generation"
+        },
+        "generate_continuation_file": {
+            "description": "Path for Snowfakery to put its next continuation file"
+        },
+        "working_directory": {
+            "description": "Default path for temporary / working files"
+        },
+        "loading_rules": {
+            "description": "Path to .load.yml file containing rules to use to load the file. Defaults to <recipename>.load.yml . Multiple files can be comma separated."
+        },
+        "num_records": {
+            "description": (
+                "Target number of records. "
+                "You will get at least this many records, but may get more. "
+                "The recipe will always execute to completion, so if it creates "
+                "3 records per execution and you ask for 5, you will get 6."
+            ),
+            "required": False,
+        },
+    }
+    stopping_criteria = None
+
+    def __init__(self, *args, **kwargs):
+        self.vars = {}
+        self.plugin_options = {}
+        super().__init__(*args, **kwargs)
+
+    def _init_options(self, kwargs):
+        super()._init_options(kwargs)
+        self.yaml_file = os.path.abspath(self.options["generator_yaml"])
+        if not os.path.exists(self.yaml_file):
+            raise TaskOptionsError(f"Cannot find {self.yaml_file}")
+        if "vars" in self.options:
+            self.vars = process_list_of_pairs_dict_arg(self.options["vars"])
+        if "plugin_options" in self.options:
+            self.plugin_options = process_list_of_pairs_dict_arg(
+                self.options["plugin_options"]
+            )
+        self.generate_mapping_file = self.options.get("generate_mapping_file")
+        if self.generate_mapping_file:
+            self.generate_mapping_file = os.path.abspath(self.generate_mapping_file)
+        num_records = self.options.get("num_records")
+        if num_records is not None:
+            num_records = int(num_records)
+            num_records_tablename = self.options.get("num_records_tablename")
+
+            if not num_records_tablename:
+                raise TaskOptionsError(
+                    "Cannot specify num_records without num_records_tablename."
+                )
+
+            self.stopping_criteria = (num_records, num_records_tablename)
+        self.working_directory = self.options.get("working_directory")
+        loading_rules = process_list_arg(self.options.get("loading_rules")) or []
+        self.loading_rules = [Path(path) for path in loading_rules if path]
+
+    def _generate_data(self, db_url, mapping_file_path, num_records, current_batch_num):
+        """Generate all of the data"""
+        if num_records is not None:  # num_records is None means execute Snowfakery once
+            self.logger.info(f"Generating batch {current_batch_num} with {num_records}")
+        self.generate_data(db_url, num_records, current_batch_num)
+        self.logger.info("Generated batch")
+
+    def default_continuation_file_path(self):
+        return Path(self.working_directory) / "continuation.yml"
+
+    def get_old_continuation_file(self) -> Optional[Path]:
+        """Use a continuation file if specified or look for one in the working directory
+
+        Return None if no file can be found.
+
+        If this code is used within a GenerateAndLoad loop, the continuation files will go
+        into the working directory specified by the GenerateAndLoad caller.
+        """
+        old_continuation_file = self.options.get("continuation_file")
+
+        if old_continuation_file:
+            old_continuation_file = Path(old_continuation_file)
+            if not old_continuation_file.exists():
+                raise TaskOptionsError(f"{old_continuation_file} does not exist")
+        elif self.working_directory:
+            path = self.default_continuation_file_path()
+            if path.exists():
+                old_continuation_file = path
+
+        return old_continuation_file
+
+    @contextmanager
+    def open_new_continuation_file(self):
+        """Create a continuation file based on config or working directory
+
+        Return None if there is no config nor working directory.
+
+        If this code is used within a GenerateAndLoad loop, the continuation files will go
+        into the working directory specified by the GenerateAndLoad caller.
+        """
+        if self.options.get("generate_continuation_file"):
+            continuation_path = self.options["generate_continuation_file"]
+
+        elif self.working_directory:
+            continuation_path = Path(self.working_directory) / "continuation_next.yml"
+        else:
+            continuation_path = None
+
+        if continuation_path:
+            with open(continuation_path, "w+") as new_continuation_file:
+                yield new_continuation_file
+        else:
+            yield None
+
+    def generate_data(self, dburl, num_records, current_batch_num):
+        old_continuation_file = self.get_old_continuation_file()
+        if old_continuation_file:
+            # reopen to ensure file pointer is at starting point
+            old_continuation_file = open(old_continuation_file, "r")
+        with self.open_new_continuation_file() as new_continuation_file:
+            generate_data(
+                yaml_file=self.yaml_file,
+                user_options=self.vars,
+                target_number=self.stopping_criteria,
+                continuation_file=old_continuation_file,
+                generate_continuation_file=new_continuation_file,
+                generate_cci_mapping_file=self.generate_mapping_file,
+                dburl=dburl,
+                load_declarations=self.loading_rules,
+                should_create_cci_record_type_tables=True,
+                plugin_options={
+                    "org_config": self.org_config,
+                    "project_config": self.project_config,
+                    **self.plugin_options,
+                },
+            )
+
+        if (
+            new_continuation_file
+            and Path(new_continuation_file.name).exists()
+            and self.working_directory
+        ):
+            shutil.move(
+                new_continuation_file.name, self.default_continuation_file_path()
+            )