def test_get_lookup_key_field__by_key_field(self):
class FakeModel:
foo = mock.MagicMock()
lookup = MappingLookup(table="contact",
key_field="foo",
name="AccountId")
assert lookup.get_lookup_key_field(FakeModel()) == "foo"
def test_get_lookup_key_field__by_key_field_wrong_case(self):
class FakeModel:
account_id = mock.MagicMock()
# we can correct mismatched mapping files if the mistake is just
# old-fashioned SQL with new Mapping File
lookup = MappingLookup(table="contact", key_field="AccountId", name="AccountId")
assert lookup.get_lookup_key_field(FakeModel()) == "account_id"
def test_get_lookup_key_field__mismatched_name(self):
class FakeModel:
account_id = mock.MagicMock()
# some mistakes can't be fixed.
lookup = MappingLookup(table="contact", key_field="Foo", name="Foo")
with pytest.raises(KeyError):
lookup.get_lookup_key_field(FakeModel())
def test_extract_respects_key_field(self, log_mock):
task = _make_task(
ExtractData,
{"options": {"database_url": "sqlite://", "mapping": "mapping.yml"}},
)
mapping = MappingStep(
sf_object="Opportunity",
table="Opportunity",
fields={"Id": "sf_id", "Name": "Name"},
lookups={
"AccountId": MappingLookup(
table="Account", key_field="account_id", name="AccountId"
)
},
)
step = mock.Mock()
task.session = mock.Mock()
task._sql_bulk_insert_from_records = mock.Mock()
task._import_results(mapping, step)
task.session.connection.assert_called_once_with()
step.get_results.assert_called_once_with()
task._sql_bulk_insert_from_records.assert_called_once_with(
connection=task.session.connection.return_value,
table="Opportunity",
columns=["sf_id", "Name", "account_id"],
record_iterable=log_mock.return_value,
)
def test_convert_lookups_to_id(self):
task = _make_task(
ExtractData, {"options": {"database_url": "sqlite:///", "mapping": ""}}
)
task.session = mock.Mock()
task.models = {
"Account": mock.Mock(),
"Account_sf_ids": mock.Mock(),
"Opportunity": mock.Mock(),
"Opportunity_sf_ids": mock.Mock(),
}
task.mapping = {
"Account": MappingStep(sf_object="Account"),
"Opportunity": MappingStep(sf_object="Opportunity"),
}
task._convert_lookups_to_id(
MappingStep(
sf_object="Opportunity",
lookups={"AccountId": MappingLookup(table="Account", name="AccountId")},
),
["AccountId"],
)
task.session.query.return_value.filter.return_value.update.assert_called_once_with(
{task.models["Opportunity"].AccountId: task.models["Account_sf_ids"].id},
synchronize_session=False,
)
task.session.commit.assert_called_once_with()
def test_convert_lookups_to_id__sqlite(self):
task = _make_task(
ExtractData,
{"options": {
"database_url": "sqlite:///",
"mapping": ""
}})
task.session = mock.Mock()
task.models = {
"Account": mock.Mock(),
"Account_sf_ids": mock.Mock(),
"Opportunity": mock.Mock(),
"Opportunity_sf_ids": mock.Mock(),
}
task.mapping = {
"Account": {
"table": "Account",
"sf_id_table": "Account_sf_ids"
},
"Opportunity": {
"table": "Opportunity",
"sf_id_table": "Opportunity_sf_ids",
},
}
task.session.query.return_value.filter.return_value.update.side_effect = (
NotImplementedError)
item = mock.Mock()
task.session.query.return_value.join.return_value = [(item, "1")]
task._convert_lookups_to_id(
{
"sf_object": "Opportunity",
"table": "Opportunity",
"sf_id_table": "Opportunity_sf_ids",
"lookups": {
"AccountId": MappingLookup(table="Account",
name="AccountId")
},
},
["AccountId"],
)
task.session.bulk_update_mappings.assert_called_once_with(
task.models["Opportunity"], [{
"id": item.id,
"AccountId": "1"
}])
task.session.commit.assert_called_once_with()
def test_import_results__autopk(self, csv_mock):
task = _make_task(
ExtractData,
{
"options": {
"database_url": "sqlite://",
"mapping": "mapping.yml"
}
},
)
mapping = {
"sf_object": "Opportunity",
"table": "Opportunity",
"sf_id_table": "Opportunity_sf_ids",
"oid_as_pk": False,
"fields": {
"Name": "Name"
},
"lookups": {
"AccountId": MappingLookup(table="Account", name="AccountId")
},
}
step = mock.Mock()
step.get_results.return_value = iter(
[["111", "Test Opportunity", "1"],
["222", "Test Opportunity 2", "1"]])
task.session = mock.Mock()
task._sql_bulk_insert_from_records = mock.Mock()
task._convert_lookups_to_id = mock.Mock()
task._import_results(mapping, step)
task.session.connection.assert_called_once_with()
step.get_results.assert_called_once_with()
task._sql_bulk_insert_from_records.assert_has_calls([
mock.call(
connection=task.session.connection.return_value,
table="Opportunity",
columns=["Name", "AccountId"],
record_iterable=csv_mock.return_value,
),
mock.call(
connection=task.session.connection.return_value,
table="Opportunity_sf_ids",
columns=["sf_id"],
record_iterable=csv_mock.return_value,
),
])
task._convert_lookups_to_id.assert_called_once_with(
mapping, ["AccountId"])
def test_query_db__joins_self_lookups(self, aliased):
task = _make_task(
LoadData,
{"options": {
"database_url": "sqlite://",
"mapping": "test.yml"
}})
model = mock.Mock()
task.models = {"accounts": model}
task.metadata = mock.Mock()
task.metadata.tables = {"accounts_sf_ids": mock.Mock()}
task.session = mock.Mock()
model.__table__ = mock.Mock()
model.__table__.primary_key.columns.keys.return_value = ["sf_id"]
columns = {"sf_id": mock.Mock(), "name": mock.Mock()}
model.__table__.columns = columns
mapping = {
"sf_object": "Account",
"table": "accounts",
"action": "update",
"oid_as_pk": True,
"fields": {
"Id": "sf_id",
"Name": "name"
},
"lookups": {
"ParentId":
MappingLookup(table="accounts",
key_field="parent_id",
name="ParentId")
},
}
task._query_db(mapping)
# Validate that the column set is accurate
task.session.query.assert_called_once_with(
model.sf_id,
model.__table__.columns["name"],
aliased.return_value.columns.sf_id,
)
# Validate that we asked for an outer join on the self-lookup
aliased.assert_called_once_with(
task.metadata.tables["accounts_sf_ids"])
task.session.query.return_value.outerjoin.assert_called_once_with(
aliased.return_value, False)
def _expand_mapping(self):
"""Walk the mapping and generate any required 'after' steps
to handle dependent and self-lookups."""
# Expand the mapping to handle dependent lookups
self.after_steps = defaultdict(dict)
for step in self.mapping.values():
step["action"] = step.get("action", "insert")
if step.get("lookups") and any(
[lookup.get("after") for lookup in step["lookups"].values()]):
# We have deferred/dependent lookups.
# Synthesize mapping steps for them.
sobject = step["sf_object"]
after_list = {
lookup["after"]
for lookup in step["lookups"].values()
if lookup.get("after")
}
for after in after_list:
lookups = {
lookup_field: lookup
for lookup_field, lookup in step["lookups"].items()
if lookup.get("after") == after
}
name = f"Update {sobject} Dependencies After {after}"
mapping = {
"sf_object": sobject,
"action": "update",
"table": step["table"],
"lookups": {},
"fields": {},
}
mapping["lookups"]["Id"] = MappingLookup(
name="Id",
table=step["table"],
key_field=self.models[step["table"]].__table__.
primary_key.columns.keys()[0],
)
for lookup in lookups:
mapping["lookups"][lookup] = lookups[lookup].copy()
mapping["lookups"][lookup]["after"] = None
self.after_steps[after][name] = mapping
def test_get_complete_field_map(self):
m = MappingStep(
sf_object="Account",
fields=["Name", "AccountSite"],
lookups={"ParentId": MappingLookup(table="Account")},
)
assert m.get_complete_field_map() == {
"Name": "Name",
"AccountSite": "AccountSite",
"ParentId": "ParentId",
}
assert m.get_complete_field_map(include_id=True) == {
"Id": "sf_id",
"Name": "Name",
"AccountSite": "AccountSite",
"ParentId": "ParentId",
}
def _expand_mapping(self):
"""Walk the mapping and generate any required 'after' steps
to handle dependent and self-lookups."""
# Expand the mapping to handle dependent lookups
self.after_steps = defaultdict(dict)
for step in self.mapping.values():
if any([lookup.after for lookup in step.lookups.values()]):
# We have deferred/dependent lookups.
# Synthesize mapping steps for them.
sobject = step.sf_object
after_list = {
lookup.after
for lookup in step.lookups.values() if lookup.after
}
for after in after_list:
lookups = {
lookup_field: lookup
for lookup_field, lookup in step.lookups.items()
if lookup.after == after
}
name = f"Update {sobject} Dependencies After {after}"
mapping = MappingStep(
sf_object=sobject,
api=step.api,
action="update",
table=step.table,
)
mapping.lookups["Id"] = MappingLookup(
name="Id",
table=step["table"],
key_field=self.models[step["table"]].__table__.
primary_key.columns.keys()[0],
)
for lookup in lookups:
mapping.lookups[lookup] = lookups[lookup].copy()
mapping.lookups[lookup].after = None
self.after_steps[after][name] = mapping
def test_get_lookup_key_field__snake_case_model(self):
class FakeModel:
account_id = mock.MagicMock()
lookup = MappingLookup(table="contact", name="AccountId")
assert lookup.get_lookup_key_field(FakeModel()) == "account_id"
def test_get_lookup_key_field__no_model(self):
lookup = MappingLookup(table="contact", name="AccountId")
assert lookup.get_lookup_key_field() == "AccountId"
def test_expand_mapping_creates_after_steps(self):
base_path = os.path.dirname(__file__)
mapping_path = os.path.join(base_path, "mapping_after.yml")
task = _make_task(
LoadData,
{
"options": {
"database_url": "sqlite://",
"mapping": mapping_path
}
},
)
task._init_mapping()
model = mock.Mock()
model.__table__ = mock.Mock()
model.__table__.primary_key.columns.keys.return_value = ["sf_id"]
task.models = {"accounts": model, "contacts": model}
task._expand_mapping()
self.assertEqual({}, task.after_steps["Insert Opportunities"])
self.assertEqual(
[
"Update Account Dependencies After Insert Contacts",
"Update Contact Dependencies After Insert Contacts",
],
list(task.after_steps["Insert Contacts"].keys()),
)
lookups = {}
lookups["Id"] = MappingLookup(name="Id",
table="accounts",
key_field="sf_id")
lookups["Primary_Contact__c"] = MappingLookup(
table="contacts",
name="Primary_Contact__c",
)
self.assertEqual(
{
"sf_object": "Account",
"action": "update",
"table": "accounts",
"lookups": lookups,
"fields": {},
},
task.after_steps["Insert Contacts"]
["Update Account Dependencies After Insert Contacts"],
)
lookups = {}
lookups["Id"] = MappingLookup(name="Id",
table="contacts",
key_field="sf_id")
lookups["ReportsToId"] = MappingLookup(table="contacts",
name="ReportsToId")
self.assertEqual(
{
"sf_object": "Contact",
"action": "update",
"table": "contacts",
"fields": {},
"lookups": lookups,
},
task.after_steps["Insert Contacts"]
["Update Contact Dependencies After Insert Contacts"],
)
self.assertEqual(
["Update Account Dependencies After Insert Accounts"],
list(task.after_steps["Insert Accounts"].keys()),
)
lookups = {}
lookups["Id"] = MappingLookup(name="Id",
table="accounts",
key_field="sf_id")
lookups["ParentId"] = MappingLookup(table="accounts", name="ParentId")
self.assertEqual(
{
"sf_object": "Account",
"action": "update",
"table": "accounts",
"fields": {},
"lookups": lookups,
},
task.after_steps["Insert Accounts"]
["Update Account Dependencies After Insert Accounts"],
)
