def test_migrate_params_retval_does_not_match_schema(self):
# LoadedModule.migrate_params() may return invalid data: it's up to the
# caller to validate it. In this test, we test that indeed, invalid
# data may be returned.
code = b"def migrate_params(params):\n return {}"
minio.client.put_object(
Bucket=minio.ExternalModulesBucket,
Key="imported/abcdef/imported.py",
Body=code,
ContentLength=len(code),
)
with self.assertLogs("cjwstate.modules.loaded_module"):
lm = LoadedModule.for_module_version(
MockModuleVersion(
"imported",
"abcdef",
ParamDType.Dict({"x": ParamDType.String()}),
"now",
)
)
self.assertEqual(
# should have 'x' key
lm.migrate_params({}),
{},
)
def test_change_parameters_deny_invalid_params(self, load_module):
workflow = Workflow.create_and_init()
wf_module = workflow.tabs.first().wf_modules.create(
order=0,
slug="step-1",
module_id_name="x",
last_relevant_delta_id=workflow.last_delta_id,
params={"x": 1},
)
ModuleVersion.create_or_replace_from_spec({
"id_name":
"x",
"name":
"x",
"category":
"Clean",
"parameters": [{
"id_name": "x",
"type": "integer"
}],
})
load_module.return_value.param_schema = ParamDType.Dict(
{"x": ParamDType.Integer()})
load_module.return_value.migrate_params = lambda x: x
with self.assertRaises(ValueError):
# Now the user requests to change params, giving an invalid param.
self.run_with_async_db(
commands.do(
ChangeParametersCommand,
workflow_id=workflow.id,
wf_module=wf_module,
new_values={"x": "Threeve"},
))
def test_simple(self, load_module):
load_module.return_value.migrate_params.return_value = {"A": "B"}
load_module.return_value.fetch.return_value = FetchResult(self.output_path, [])
result = fetch.fetch_or_wrap_error(
self.ctx,
self.chroot_context,
self.basedir,
WfModule(params={"A": "input"}, secrets={"C": "wrong"}),
MockModuleVersion(
id_name="A", param_schema=ParamDType.Dict({"A": ParamDType.String()})
),
{"C": "D"},
None,
None,
self.output_path,
)
self.assertEqual(result, FetchResult(self.output_path, []))
load_module.return_value.migrate_params.assert_called_with({"A": "input"})
load_module.return_value.fetch.assert_called_with(
chroot_context=self.chroot_context,
basedir=self.basedir,
params=Params({"A": "B"}),
secrets={"C": "D"},
last_fetch_result=None,
input_parquet_filename=None,
output_filename=self.output_path.name,
)
def param_schema(self) -> ParamDType.Dict:
if "param_schema" in self.spec:
# Module author wrote a schema in the YAML, to define storage of 'custom' parameters
json_schema = self.spec["param_schema"]
return ParamDType.parse({
"type": "dict",
"properties": json_schema
})
else:
# Usual case: infer schema from module parameter types
# Use of dict here means schema is not sensitive to parameter ordering, which is good
return ParamDType.Dict(
dict((f.id_name, f.dtype) for f in self.param_fields
if f.dtype is not None))
def _load_tab_flows(workflow: Workflow, delta_id: int) -> List[TabFlow]:
"""
Query `workflow` for each tab's `TabFlow` (ordered by tab position).
Raise `ModuleError` or `ValueError` if migrate_params() fails. Failed
migration means the whole execute can't happen.
"""
ret = []
with workflow.cooperative_lock(): # reloads workflow
if workflow.last_delta_id != delta_id:
raise UnneededExecution
for tab_model in workflow.live_tabs.all():
steps = [
ExecuteStep(
step,
(step.module_version.param_schema if step.module_version
is not None else ParamDType.Dict({})),
# We need to invoke the kernel and migrate _all_ modules'
# params (WfModule.get_params), because we can only check
# for tab cycles after migrating (and before calling any
# render()).
_get_migrated_params(step),
) for step in tab_model.live_wf_modules.all()
]
ret.append(TabFlow(Tab(tab_model.slug, tab_model.name), steps))
return ret
def test_load_dynamic_ignore_test_py(self):
code = b"def render(table, params):\n return table * 2"
minio.client.put_object(
Bucket=minio.ExternalModulesBucket,
Key="imported/abcdef/imported.py",
Body=code,
ContentLength=len(code),
)
# write other .py files that aren't module code and should be ignored
minio.client.put_object(
Bucket=minio.ExternalModulesBucket,
Key="imported/abcdef/setup.py",
Body=b"",
ContentLength=0,
)
minio.client.put_object(
Bucket=minio.ExternalModulesBucket,
Key="imported/abcdef/test_imported.py",
Body=b"",
ContentLength=0,
)
with self.assertLogs("cjwstate.modules.loaded_module"):
LoadedModule.for_module_version(
MockModuleVersion("imported", "abcdef", ParamDType.Dict({}), "now")
)
def test_migrate_params_crash(self):
code = b"def migrate_params(params):\n raise RuntimeError('xxx')"
minio.client.put_object(
Bucket=minio.ExternalModulesBucket,
Key="imported/abcdef/imported.py",
Body=code,
ContentLength=len(code),
)
with self.assertLogs("cjwstate.modules.loaded_module"):
lm = LoadedModule.for_module_version(
MockModuleVersion(
"imported",
"abcdef",
ParamDType.Dict({"x": ParamDType.String()}),
"now",
)
)
with self.assertRaisesRegex(ModuleExitedError, "xxx"):
# should have 'x' key
lm.migrate_params({})
def test_load_dynamic_is_cached(self):
code = b"def render(table, params):\n return table * 2"
minio.client.put_object(
Bucket=minio.ExternalModulesBucket,
Key="imported/abcdef/imported.py",
Body=code,
ContentLength=len(code),
)
with self.assertLogs("cjwstate.modules.loaded_module"):
lm = LoadedModule.for_module_version(
MockModuleVersion("imported", "abcdef", ParamDType.Dict({}), "now")
)
with patch("importlib.util.module_from_spec", None):
lm2 = LoadedModule.for_module_version(
MockModuleVersion("imported", "abcdef", ParamDType.Dict({}), "now")
)
self.assertIs(lm.compiled_module, lm2.compiled_module)
def test_load_dynamic(self):
code = b"def render(table, params):\n return table * 2"
minio.client.put_object(
Bucket=minio.ExternalModulesBucket,
Key="imported/abcdef/imported.py",
Body=code,
ContentLength=len(code),
)
with self.assertLogs("cjwstate.modules.loaded_module"):
lm = LoadedModule.for_module_version(
MockModuleVersion("imported", "abcdef", ParamDType.Dict({}), "now")
)
self.assertEqual(lm.name, "imported:abcdef")
# This ends up being kinda an integration test.
with ExitStack() as ctx:
basedir = Path(ctx.enter_context(tempdir_context(prefix="test-basedir-")))
basedir.chmod(0o755)
input_table = ctx.enter_context(
arrow_table_context({"A": [1]}, dir=basedir)
)
input_table.path.chmod(0o644)
output_tf = ctx.enter_context(tempfile.NamedTemporaryFile(dir=basedir))
ctx.enter_context(self.assertLogs("cjwstate.modules.loaded_module"))
result = lm.render(
basedir=basedir,
input_table=input_table,
params=Params({"col": "A"}),
tab=Tab("tab-1", "Tab 1"),
fetch_result=None,
output_filename=Path(output_tf.name).name,
)
assert_render_result_equals(result, RenderResult(arrow_table({"A": [2]})))
def test_change_parameters_across_module_versions(self, load_module):
workflow = Workflow.create_and_init()
# Initialize a WfModule that used module 'x' version '1' (which we
# don't need to write in code -- after all, that version might be long
# gone when ChangeParametersCommand is called.
wf_module = workflow.tabs.first().wf_modules.create(
order=0,
slug="step-1",
module_id_name="x",
last_relevant_delta_id=workflow.last_delta_id,
params={
"version": "v1",
"x": 1
}, # version-'1' params
)
# Now install version '2' of module 'x'.
#
# Version '2''s migrate_params() could do anything; in this test, it
# simply changes 'version' from 'v1' to 'v2'
ModuleVersion.create_or_replace_from_spec(
{
"id_name":
"x",
"name":
"x",
"category":
"Clean",
"parameters": [
{
"id_name": "version",
"type": "string"
},
{
"id_name": "x",
"type": "integer"
},
],
},
source_version_hash="2",
)
load_module.return_value.param_dtype = ParamDType.Dict({
"version":
ParamDType.String(),
"x":
ParamDType.Integer()
})
load_module.return_value.migrate_params = lambda params: {
**params,
"version": "v2",
}
# Now the user requests to change params.
#
# The user was _viewing_ version '2' of module 'x', though
# `wf_module.params` was at version 1. (Workbench ran
# `migrate_params()` without saving the result when it
# presented `params` to the user.) So the changes should apply atop
# _migrated_ params.
cmd = self.run_with_async_db(
commands.do(
ChangeParametersCommand,
workflow_id=workflow.id,
wf_module=wf_module,
new_values={"x": 2},
))
self.assertEqual(
wf_module.params,
{
"version": "v2", # migrate_params() ran
"x": 2, # and we applied changes on top of its output
},
)
self.run_with_async_db(commands.undo(cmd))
self.assertEqual(
wf_module.params,
{
"version": "v1",
"x": 1
} # exactly what we had before
)
def test_load_static(self):
# Test with a _real_ static module
lm = LoadedModule.for_module_version(
MockModuleVersion("pastecsv", "internal", ParamDType.Dict({}), "now")
)
self.assertEqual(lm.name, "pastecsv:internal")