feat: refactor all get_sqla_engine to use contextmanager in codebase (#21943)

This commit is contained in:
Hugh A. Miles II
2022-11-15 13:45:14 -05:00
committed by GitHub
parent 06f87e1467
commit e23efefc46
41 changed files with 635 additions and 595 deletions

View File

@@ -37,28 +37,28 @@ def load_random_time_series_data(
"""Loading random time series data from a zip file in the repo"""
tbl_name = "random_time_series"
database = database_utils.get_example_database()
engine = database.get_sqla_engine()
schema = inspect(engine).default_schema_name
table_exists = database.has_table_by_name(tbl_name)
with database.get_sqla_engine_with_context() as engine:
schema = inspect(engine).default_schema_name
table_exists = database.has_table_by_name(tbl_name)
if not only_metadata and (not table_exists or force):
url = get_example_url("random_time_series.json.gz")
pdf = pd.read_json(url, compression="gzip")
if database.backend == "presto":
pdf.ds = pd.to_datetime(pdf.ds, unit="s")
pdf.ds = pdf.ds.dt.strftime("%Y-%m-%d %H:%M%:%S")
else:
pdf.ds = pd.to_datetime(pdf.ds, unit="s")
if not only_metadata and (not table_exists or force):
url = get_example_url("random_time_series.json.gz")
pdf = pd.read_json(url, compression="gzip")
if database.backend == "presto":
pdf.ds = pd.to_datetime(pdf.ds, unit="s")
pdf.ds = pdf.ds.dt.strftime("%Y-%m-%d %H:%M%:%S")
else:
pdf.ds = pd.to_datetime(pdf.ds, unit="s")
pdf.to_sql(
tbl_name,
engine,
schema=schema,
if_exists="replace",
chunksize=500,
dtype={"ds": DateTime if database.backend != "presto" else String(255)},
index=False,
)
pdf.to_sql(
tbl_name,
engine,
schema=schema,
if_exists="replace",
chunksize=500,
dtype={"ds": DateTime if database.backend != "presto" else String(255)},
index=False,
)
print("Done loading table!")
print("-" * 80)