def second_session(connection: sa.engine.Connection, sessionmaker: orm.sessionmaker): transaction = connection.begin() sess = sessionmaker(bind=connection) yield sess transaction.rollback() sess.close()
def session(connection: sa.engine.Connection, sessionmaker: orm.sessionmaker): # pylint: disable=redefined-outer-name """ yields temporalized session -- per test """ transaction = connection.begin() sess = sessionmaker(bind=connection) yield sess transaction.rollback() sess.close()
def scrape_geojson(conn: sa.engine.Connection, dataset_id: str) -> None: log = logger.bind(dataset_id=dataset_id, method="scrape_geojson") params = {"method": "export", "format": "GeoJSON"} url = f"{BASE}/geospatial/{dataset_id}" with utils.download_file(url, params=params) as fname: try: df = gpd.read_file(fname) except ValueError as e: raise exceptions.SocrataParseError from e for column in df.columns: if column == "geometry": continue # Bad type inference try: df[column] = df[column].astype(int) continue except (ValueError, TypeError): pass try: df[column] = df[column].astype(float) continue except (ValueError, TypeError): pass try: df[column] = pd.to_datetime(df[column]) continue except (ValueError, TypeError): pass log.info("Inserting") del df["geometry"] trans = conn.begin() try: conn.execute(f'DROP TABLE IF EXISTS "{RAW_SCHEMA}.{dataset_id}"') df.to_sql( f"{dataset_id}", conn, if_exists="replace", index=False, schema=RAW_SCHEMA, ) except Exception: trans.rollback() raise trans.commit()