def run_stripe_test(api_key: str):
from snapflow_stripe import module as stripe
if not api_key:
api_key = TEST_API_KEY
storage = get_tmp_sqlite_db_url()
env = Environment(DataspaceCfg(metadata_storage="sqlite://", storages=[storage]))
env.add_module(stripe)
# Initial graph
raw_charges = GraphCfg(
key="import_charges",
function="stripe.import_charges",
params={"api_key": api_key},
)
clean_charges = GraphCfg(
key="clean_charges", function="stripe.clean_charges", input="import_charges"
)
g = GraphCfg(nodes=[raw_charges, clean_charges])
results = env.produce(
clean_charges.key, g, target_storage=storage, execution_timelimit_seconds=1
)
records = results[0].stdout().as_records()
assert len(records) >= 100
assert records[0]["amount"] > 0
def get_env(key="_test", db_url=None):
if db_url is None:
db_url = get_tmp_sqlite_db_url()
env = Environment(key=key, metadata_storage=db_url)
env.add_module(core)
env.add_schema(Customer)
env.add_schema(Metric)
return env
def test_simple_import():
dburl = get_tmp_sqlite_db_url()
env = Environment(metadata_storage=dburl)
g = Graph(env)
env.add_module(core)
df = pd.DataFrame({"a": range(10), "b": range(10)})
g.create_node(key="n1", function="import_dataframe", params={"dataframe": df})
blocks = env.produce("n1", g)
assert_almost_equal(blocks[0].as_dataframe(), df, check_dtype=False)
def produce_function_output_for_static_input(
function: DataFunction,
params: Dict[str, Any] = None,
input: Any = None,
inputs: Any = None,
env: Optional[Environment] = None,
module: Optional[SnapflowModule] = None,
target_storage: Optional[Storage] = None,
upstream: Any = None, # TODO: DEPRECATED
) -> Iterator[List[DataBlock]]:
inputs = input or inputs or upstream
if env is None:
db = get_tmp_sqlite_db_url()
env = Environment(metadata_storage=db)
if module:
env.add_module(module)
with provide_test_storages(function, target_storage) as target_storage:
if target_storage:
target_storage = env.add_storage(target_storage)
with env.md_api.begin():
g = Graph(env)
input_datas = inputs
input_nodes: Dict[str, Node] = {}
pi = function.get_interface()
if not isinstance(inputs, dict):
assert len(pi.get_non_recursive_inputs()) == 1
input_datas = {
pi.get_single_non_recursive_input().name: inputs
}
for inpt in pi.inputs.values():
if inpt.is_self_reference:
continue
assert inpt.name is not None
input_data = input_datas[inpt.name]
if isinstance(input_data, str):
input_data = DataInput(data=input_data)
assert isinstance(input_data, DataInput)
n = g.create_node(
key=f"_input_{inpt.name}",
function="core.import_dataframe",
params={
"dataframe": input_data.as_dataframe(env),
"schema": input_data.get_schema_key(),
},
)
input_nodes[inpt.name] = n
test_node = g.create_node(
key=f"{function.name}",
function=function,
params=params,
inputs=input_nodes,
)
blocks = env.produce(test_node,
to_exhaustion=False,
target_storage=target_storage)
yield blocks
def make_test_env(**kwargs) -> Environment:
if "metadata_storage" not in kwargs:
url = get_tmp_sqlite_db_url()
metadata_storage = Storage.from_url(url)
kwargs["metadata_storage"] = metadata_storage
env = Environment(settings=SnapflowSettings(abort_on_function_error=True), **kwargs)
test_module = SnapflowModule(
"_test",
)
for schema in [TestSchema1, TestSchema2, TestSchema3, TestSchema4]:
env.add_schema(schema)
env.add_module(test_module)
return env
def test_database_handler():
dburl = get_tmp_sqlite_db_url()
s = Storage(dburl)
name = "_test"
handler = get_handler(DatabaseTableFormat, s.storage_engine)
handler().create_empty(name, s, test_records_schema)
s.get_api().bulk_insert_records(name, test_records)
assert list(handler().infer_field_names(name, s)) == list(test_records[0].keys())
assert handler().infer_field_type(name, s, "f1") == Text()
assert handler().infer_field_type(name, s, "f2") == Integer()
assert handler().infer_field_type(name, s, "f3") == DEFAULT_FIELD_TYPE
assert handler().infer_field_type(name, s, "f4") == Date()
assert handler().infer_field_type(name, s, "f5") == DEFAULT_FIELD_TYPE
def test_conform():
from snapflow.modules import core
TestSchemaA = create_quick_schema("TestSchemaA", [("a", "Integer"),
("b", "Integer")],
namespace="core")
TestSchemaB = create_quick_schema(
"TestSchemaB",
[("a", "Integer"), ("c", "Integer"), ("d", "Text")],
implementations=[Implementation("TestSchemaA", {"b": "c"})],
namespace="core",
)
core.add_schema(TestSchemaA)
core.add_schema(TestSchemaB)
input_data = """
a,c,d
1,2,i
1,3,i
1,4,i
2,2,i
"""
expected = """
a,b
1,2
1,3
1,4
2,2
"""
# expected_df = str_as_dataframe(expected, schema=core.schemas.CoreTestSchema)
data_input = DataInput(input_data, schema=TestSchemaB)
s = get_tmp_sqlite_db_url()
for p in [
dataframe_conform_to_schema,
sql_conform_to_schema,
]:
with produce_snap_output_for_static_input(
p,
input=data_input,
target_storage=s,
params={"schema": "TestSchemaA"}) as dbs:
assert len(dbs) == 1
db = dbs[0]
expected_df = DataInput(expected, schema=TestSchemaA).as_dataframe(
db.manager.env)
df = db.as_dataframe()
print(expected_df)
print(df)
def provide_test_storages(
function: DataFunction,
target_storage: Storage) -> Iterator[Optional[Storage]]:
if target_storage:
yield target_storage # TODO
elif function.required_storage_engines:
# TODO: multiple engines -- is it AND or OR?? each entry is AND and inside entry commas delim OR
eng = get_engine_for_scheme(function.required_storage_engines[0])
api_cls = eng.get_api_cls()
if issubclass(api_cls, DatabaseApi):
if not api_cls.dialect_is_supported():
raise TestFeatureNotImplementedError(eng)
with api_cls.temp_local_database() as url:
yield Storage(url)
elif "database" in function.required_storage_classes:
yield Storage(get_tmp_sqlite_db_url())
else:
yield None
def test_app():
db_url = get_tmp_sqlite_db_url()
runner = CliRunner()
result = runner.invoke(app, ["-m", db_url, "logs"])
assert result.exit_code == 0
result = runner.invoke(app, ["-m", db_url, "nodes"])
assert result.exit_code == 0
result = runner.invoke(app, ["-m", db_url, "blocks"])
assert result.exit_code == 0
result = runner.invoke(
app, ["-m", db_url, "generate", "schema"], input='{"f1": 1, "f2": "hi"}'
)
assert result.exit_code == 0
with runner.isolated_filesystem():
result = runner.invoke(app, ["-m", db_url, "init"])
assert result.exit_code == 0
pth = os.path.join(os.getcwd(), SNAPFLOW_PROJECT_FILE_NAME)
assert os.path.exists(pth)
def test_shopify():
from snapflow_shopify import module as shopify
api_key = ensure_api_key()
storage = get_tmp_sqlite_db_url()
env = Environment(
DataspaceCfg(metadata_storage="sqlite://", storages=[storage]))
# Initial graph
orders = GraphCfg(
key="import_orders",
function="shopify.import_orders",
params={"admin_url": api_key},
)
g = GraphCfg(nodes=[orders])
results = env.produce(orders.key,
g,
target_storage=storage,
execution_timelimit_seconds=1)
records = results[0].stdout().as_records()
assert len(records) > 0
def test_multi_env():
db_url = get_tmp_sqlite_db_url()
cfg = EnvironmentConfiguration(
key=f"_test_{rand_str()}",
metadata_storage_url=db_url,
settings=SnapflowSettings(add_core_module=False),
)
env1 = Environment.from_config(cfg)
with env1.md_api.begin():
env1.md_api.add(DataBlockMetadata(realized_schema_key="Any"))
env1.md_api.flush()
assert env1.md_api.count(select(DataBlockMetadata)) == 1
cfg = EnvironmentConfiguration(
key=f"_test_{rand_str()}",
metadata_storage_url=db_url,
settings=SnapflowSettings(add_core_module=False),
)
env2 = Environment.from_config(cfg)
with env2.md_api.begin():
assert env2.md_api.count(select(DataBlockMetadata)) == 0
env2.md_api.add(DataBlockMetadata(realized_schema_key="Any"))
env2.md_api.flush()
assert env2.md_api.count(select(DataBlockMetadata)) == 1
def test_dedupe():
from snapflow.modules import core
input_data = """
k1,k2,f1,f2,f3,f4
1,2,abc,1.1,1,2012-01-01
1,2,def,1.1,{"1":2},2012-01-02
1,3,abc,1.1,2,2012-01-01
1,4,,,"[1,2,3]",2012-01-01
2,2,1.0,2.1,"[1,2,3]",2012-01-01
"""
expected = """
k1,k2,f1,f2,f3,f4
1,2,def,1.1,{"1":2},2012-01-02
1,3,abc,1.1,2,2012-01-01
1,4,,,"[1,2,3]",2012-01-01
2,2,1.0,2.1,"[1,2,3]",2012-01-01
"""
# expected_df = str_as_dataframe(expected, schema=core.schemas.CoreTestSchema)
data_input = DataInput(input_data, schema="CoreTestSchema", module=core)
s = get_tmp_sqlite_db_url()
for p in [
# sql_dedupe_unique_keep_newest_row,
dataframe_dedupe_unique_keep_newest_row,
]:
with produce_snap_output_for_static_input(p,
input=data_input,
target_storage=s) as dbs:
assert len(dbs) == 1
db = dbs[0]
expected_df = DataInput(expected,
schema="CoreTestSchema",
module=core).as_dataframe(db.manager.env)
df = db.as_dataframe()
assert_dataframes_are_almost_equal(
df, expected_df, schema=core.schemas.CoreTestSchema)
def temp_local_database(cls) -> Iterator[str]:
db_url = get_tmp_sqlite_db_url("__test_dcp_sqlite")
yield db_url
