def to_python_value(
self,
ctx: FlyteContext,
lv: Literal,
expected_python_type: typing.Type[DoltTable],
) -> DoltTable:
if not (lv and lv.scalar and lv.scalar.generic and "config" in lv.scalar.generic):
raise ValueError("DoltTable requires DoltConfig to load python value")
conf_dict = MessageToDict(lv.scalar.generic["config"])
conf = DoltConfig(**conf_dict)
db = dolt.Dolt(conf.db_path)
with tempfile.NamedTemporaryFile() as f:
dolt_int.load(
db=db,
tablename=conf.tablename,
sql=conf.sql,
filename=f.name,
branch_conf=conf.branch_conf,
meta_conf=conf.meta_conf,
remote_conf=conf.remote_conf,
load_args=conf.io_args,
)
df = pandas.read_csv(f)
lv.data = df
return lv
def to_python_value(
self,
ctx: FlyteContext,
lv: Literal,
expected_python_type: typing.Type[DoltTable],
) -> DoltTable:
if not (lv and lv.scalar and lv.scalar.generic
and lv.scalar.generic["config"]):
return pandas.DataFrame()
conf = DoltConfig(**lv.scalar.generic["config"])
db = dolt.Dolt(conf.db_path)
with tempfile.NamedTemporaryFile() as f:
dolt_int.load(
db=db,
tablename=conf.tablename,
sql=conf.sql,
filename=f.name,
branch_conf=conf.branch_conf,
meta_conf=conf.meta_conf,
remote_conf=conf.remote_conf,
load_args=conf.io_args,
)
df = pandas.read_csv(f)
lv.data = df
return lv
