diff --git a/railib/api.py b/railib/api.py index 5ffcc53..df4bfc2 100644 --- a/railib/api.py +++ b/railib/api.py @@ -715,7 +715,7 @@ def list_models(ctx: Context, database: str, engine: str) -> list: # Generate a rel literal relation for the given dict. def _gen_literal_dict(items: dict) -> str: result = [] - for k, v in items: + for k, v in items.items(): result.append(f"{_gen_literal(k)},{_gen_literal(v)}") return "{" + ";".join(result) + "}" @@ -761,6 +761,14 @@ def _gen_syntax_config(syntax: dict = {}) -> str: return result +# Generate list of config schema options for `load_csv` +def _gen_schema_config(schema: dict = {}) -> str: + result = "" + for k, v in schema.items(): + result += f"def config:schema{k} = \"{v}\"\n" + return result + + # `syntax`: # * header: a map from col number to name (base 1) # * header_row: row number of header, 0 means no header (default: 1) @@ -777,6 +785,7 @@ def load_csv( relation: str, data: str or io.TextIOBase, syntax: dict = {}, + schema: dict = {}, ) -> dict: if isinstance(data, str): pass # ok @@ -786,8 +795,9 @@ def load_csv( raise TypeError(f"bad type for arg 'data': {data.__class__.__name__}") inputs = {"data": data} command = _gen_syntax_config(syntax) + command += _gen_schema_config(schema) command += "def config:data = data\n" "def insert:%s = load_csv[config]" % relation - return exec_v1(ctx, database, engine, command, inputs=inputs, readonly=False) + return exec(ctx, database, engine, command, inputs=inputs, readonly=False) def load_json( @@ -805,7 +815,7 @@ def load_json( raise TypeError(f"bad type for arg 'data': {data.__class__.__name__}") inputs = {"data": data} command = "def config:data = data\n" "def insert:%s = load_json[config]" % relation - return exec_v1(ctx, database, engine, command, inputs=inputs, readonly=False) + return exec(ctx, database, engine, command, inputs=inputs, readonly=False) def exec_v1( diff --git a/tests/integration.py b/tests/integration.py index a882b38..7a73b93 100644 --- a/tests/integration.py +++ b/tests/integration.py @@ -32,7 +32,6 @@ def create_engine_wait(ctx: api.Context, engine: str): custom_headers = json.loads(os.getenv('CUSTOM_HEADERS', '{}')) if client_id is None: - print("using config from path") cfg = config.read() else: file = tempfile.NamedTemporaryFile(mode="w") @@ -51,19 +50,19 @@ def create_engine_wait(ctx: api.Context, engine: str): ctx = api.Context(**cfg) -suffix = uuid.uuid4() -engine = f"python-sdk-{suffix}" -dbname = f"python-sdk-{suffix}" - class TestTransactionAsync(unittest.TestCase): def setUp(self): - create_engine_wait(ctx, engine) - api.create_database(ctx, dbname) + suffix = uuid.uuid4() + self.engine = f"python-sdk-{suffix}" + self.dbname = f"python-sdk-{suffix}" + + create_engine_wait(ctx, self.engine) + api.create_database(ctx, self.dbname) def test_v2_exec(self): cmd = "x, x^2, x^3, x^4 from x in {1; 2; 3; 4; 5}" - rsp = api.exec(ctx, dbname, engine, cmd) + rsp = api.exec(ctx, self.dbname, self.engine, cmd) # transaction self.assertEqual("COMPLETED", rsp.transaction["state"]) @@ -89,8 +88,83 @@ def test_v2_exec(self): 1, 16, 81, 256, 625]}, rsp.results[0]["table"].to_pydict()) def tearDown(self): - api.delete_engine(ctx, engine) - api.delete_database(ctx, dbname) + api.delete_engine(ctx, self.engine) + api.delete_database(ctx, self.dbname) + + +class TestDataload(unittest.TestCase): + def setUp(self): + suffix = uuid.uuid4() + self.engine = f"python-sdk-{suffix}" + self.dbname = f"python-sdk-{suffix}" + + create_engine_wait(ctx, self.engine) + api.create_database(ctx, self.dbname) + + def test_load_json(self): + json = '{ "test" : 123 }' + resp = api.load_json(ctx, self.dbname, self.engine, 'test_relation', json) + self.assertEqual("COMPLETED", resp.transaction["state"]) + + resp = api.exec(ctx, self.dbname, self.engine, 'def output = test_relation') + self.assertEqual("COMPLETED", resp.transaction["state"]) + self.assertEqual({'v1': [123]}, resp.results[0]["table"].to_pydict()) + + def test_load_csv(self): + csv = 'foo,bar\n1,2' + resp = api.load_csv(ctx, self.dbname, self.engine, 'test_relation', csv) + self.assertEqual("COMPLETED", resp.transaction["state"]) + + resp = api.exec(ctx, self.dbname, self.engine, 'def output = test_relation') + self.assertEqual("COMPLETED", resp.transaction["state"]) + self.assertEqual({'v1': [2], 'v2': ['2']}, resp.results[0]["table"].to_pydict()) + self.assertEqual({'v1': [2], 'v2': ['1']}, resp.results[1]["table"].to_pydict()) + + def test_load_csv_with_syntax(self): + csv = 'foo|bar\n1,2' + resp = api.load_csv( + ctx, + self.dbname, + self.engine, + 'test_relation', + csv, + { + 'header': {1: 'foo', 2: 'bar'}, + 'delim': '|', + 'quotechar': "'", + 'header_row': 0, + 'escapechar': ']' + } + ) + self.assertEqual("COMPLETED", resp.transaction["state"]) + + resp = api.exec(ctx, self.dbname, self.engine, 'def output = test_relation') + self.assertEqual("COMPLETED", resp.transaction["state"]) + self.assertEqual({'v1': [2], 'v2': [2], 'v3': ['1,2']}, resp.results[0]["table"].to_pydict()) + + def test_load_csv_with_schema(self): + csv = 'foo,bar\n1,test' + resp = api.load_csv( + ctx, + self.dbname, + self.engine, + 'test_relation', + csv, + schema={ + ':foo': 'int', + ':bar': 'string' + } + ) + self.assertEqual("COMPLETED", resp.transaction["state"]) + + resp = api.exec(ctx, self.dbname, self.engine, 'def output = test_relation') + self.assertEqual("COMPLETED", resp.transaction["state"]) + self.assertEqual({'v1': [2], 'v2': ['test']}, resp.results[0]["table"].to_pydict()) + self.assertEqual({'v1': [2], 'v2': [1]}, resp.results[1]["table"].to_pydict()) + + def tearDown(self): + api.delete_engine(ctx, self.engine) + api.delete_database(ctx, self.dbname) if __name__ == '__main__':