def get_data_with_annotation():
password = getenv('password')
user = getenv('user')
database = getenv("database")
host = getenv("host")
port = getenv("port")
output_dir = request.json.get("output_dir")
driver = PostgresDBDriver(password=password,
user=user,
database_name=database,
host=host,
port=port,
tables=[])
data_fetcher = DataFetcher(db_driver=driver)
tables = [PythonFreeConnexTable.load_from_json(t) for t in json.loads(request.json.get("tables"))]
try:
tables = data_fetcher.store_data(output_dir=output_dir, tables=tables)
return jsonify(
[t.to_json() for t in tables]
)
except Exception as e:
traceback.print_exc()
return Response(str(e), status=500)
def setUp(self) -> None:
in_docker = getenv("in_docker")
password = getenv('password')
user = getenv('user')
database = getenv("database")
host = getenv("host")
port = getenv("port")
driver = PostgresDBDriver(password=password,
user=user,
database_name=database,
host=host,
port=port,
tables=[])
if in_docker:
driver.create_db_with_columns(data=SET_UP_SCRIPT)
def generate_python_result():
"""
Perform a code generation and return the actual results
:return:
"""
try:
data: dict = request.json
tables = [PythonFreeConnexTable.load_from_json(t) for t in json.loads(data['table'])]
sql = data['sql']
password = getenv('password')
user = getenv('user')
database = data.get("database", None) if data.get("database", None) else getenv("database")
host = getenv("host")
port = getenv("port")
annotation_name = data['annotation_name']
num_of_rows = data.get("num_of_rows", 100)
print("Generating python result")
driver = PostgresDBDriver(password=password,
user=user,
database_name=database,
host=host,
port=port,
tables=tables)
client_queue = Queue()
server_queue = Queue()
client = Process(target=run_query,
args=(sql, E_role.CLIENT, client_queue, tables, driver, annotation_name, num_of_rows))
server = Process(target=run_query,
args=(sql, E_role.SERVER, server_queue, tables, driver, annotation_name, num_of_rows))
client.start()
server.start()
client.join()
server.join()
client_result = client_queue.get()
server_result = server_queue.get()
if client_result[0]:
raise RuntimeError(client_result[1])
if server_result[0]:
raise RuntimeError(server_result[1])
print(f"Client results: {len(client_result[1])}, server results: {len(server_result[1])}")
return jsonify(
{"client_result": client_result[1], "server_result": server_result[1], "joinGraph": client_result[2],
"isFreeConnex": client_result[3], "errorTables": client_result[4]})
except Exception as e:
traceback.print_exc()
return Response(str(e), status=500)
def test_simple_join(self):
db_plan = PostgresDBPlan.from_json(simple_plan[0]['Plan'], tables=self.tables)
db_driver = PostgresDBDriver(host="", database_name="", user="", password="", port="", tables=self.tables)
codegen = CodeGenDB(sql=SIMPLE_PLAN_SQL, db_driver=db_driver, tables=self.tables, annotation_name="demo")
codegen.parse(query_plan=db_plan)
codegen.to_output()
root = codegen.root_table
self.assertEqual("orders", root.variable_table_name)
self.assertTrue(codegen.is_free_connex())
def create_db():
password = getenv('password')
user = getenv('user')
database = getenv("database")
host = getenv("host")
port = getenv("port")
driver = PostgresDBDriver(password=password,
user=user,
database_name=database,
host=host,
port=port,
tables=[])
try:
driver.create_db_with_columns(data=request.json['data'])
return Response(status=201)
except Exception as e:
traceback.print_exc()
return Response(str(e), status=500)
def setUp(self):
super().setUp()
password = getenv('password')
user = getenv('user')
database = getenv("database")
host = getenv("host")
port = getenv("port")
self.tables = [FreeConnexTable.load_from_json(t) for t in TEST_CONFIG]
self.driver = PostgresDBDriver(password=password,
user=user,
database_name=database,
host=host,
port=port,
tables=self.tables)
def run_query(sql, role: E_role, queue: Queue, tables):
init_global_party(address="0.0.0.0", port=7766, role=role)
db_plan = PostgresDBPlan.from_json(simple_plan[0]['Plan'], tables=tables)
db_driver = PostgresDBDriver(host="",
database_name="",
user="",
password="",
port="",
tables=tables)
codegen = CodeGenPython(sql=sql,
db_driver=db_driver,
tables=tables,
annotation_name="q3_annot")
codegen.parse(db_plan)
results = codegen.to_output(output_table_name="orders", limit_size=3)
queue.put(results)
def generate_code_by_db():
"""
Perform a code generation by using db execution plan
:return:
"""
try:
data: dict = request.json
tables = [FreeConnexTable.load_from_json(t) for t in json.loads(data['table'])]
sql = data['sql']
password = getenv('password')
user = getenv('user')
database = data.get("database", None) if data.get("database", None) else getenv("database")
host = getenv("host")
port = getenv("port")
annotation_name = data['annotation_name']
driver = PostgresDBDriver(password=password,
user=user,
database_name=database,
host=host,
port=port,
tables=tables)
parser = CodeGenDB(db_driver=driver, sql=sql, tables=tables, annotation_name=annotation_name)
if "plan" in data and data["plan"] != "":
plan = PostgresDBPlan.from_json(data["plan"], tables=tables)
parser.parse(query_plan=plan)
else:
parser.parse()
output = parser.to_output(function_name=data.get("functionName"))
graph = parser.root_table.to_json_graph(
output_attrs=parser.get_output_attributes()) if parser.root_table else {}
is_free_connex, error_tables = parser.is_free_connex()
error_tables = [e.variable_table_name for e in error_tables]
return jsonify(
{"code": output, "joinGraph": graph, "isFreeConnex": is_free_connex, "errorTables": error_tables})
except Exception as e:
traceback.print_exc()
return Response(str(e), status=500)
def get_table_schema():
"""
Get tables' schema.
:return:
"""
password = getenv('password')
user = getenv('user')
database = getenv("database")
host = getenv("host")
port = getenv("port")
driver = PostgresDBDriver(password=password,
user=user,
database_name=database,
host=host,
port=port,
tables=[])
schema_fetcher = SchemaFetcher(db_driver=driver)
try:
tables = schema_fetcher.get_schema()
return jsonify([t.to_json() for t in tables])
except Exception as e:
traceback.print_exc()
return Response(str(e), status=500)
LINEITEM
where
c_custkey = o_custkey
and l_orderkey = o_orderkey
and o_orderdate >= date '1993-08-01'
and o_orderdate < date '1993-08-01' + interval '3' month
and l_returnflag = 'R'
group by
c_custkey,
c_name,
c_nationkey
order by
revenue desc
limit
20;
"""
if __name__ == '__main__':
password = getenv('password')
user = getenv('user')
database = "tpch"
host = "localhost"
port = "5432"
with open("./table_config.json", 'r') as f:
tables = [Table.load_from_json(t) for t in json.load(f)]
driver = PostgresDBDriver(password=password, user=user, database_name=database, host=host, port=port,
tables=tables)
codegen = CodeGenDB(sql=sql, db_driver=driver, tables=tables)
codegen.parse().to_file("db-test.cpp")