def load(self):
if self:
if self.__batch:
log.debug('Enabling batch load')
enable_batch_inserting(self._s)
else:
log.debug('Leaving batch mode unchanged')
for type in self._staging:
log.debug(
f'Adding {len(self._staging[type])} instances of {type}')
for instance in self._staging[type]:
self._s.add(instance)
self._s.commit()
self._postload()
else:
log.warning('No data to load')
def session_scope(self):
"""Provide a transactional scope around a series of operations."""
if not hasattr(self, 'Session'):
self.configure()
session = self.Session()
try:
if self._batchmode:
enable_batch_inserting(session)
yield session
session.commit()
except:
session.rollback()
raise
finally:
session.close()
del session
def bind_engine(region: Region,
provider: Provider,
schema_base: DeclarativeMeta):
logger.debug(f'bind async engine: {region}, {provider}, {schema_base}')
# get database engine
engine = get_db_engine(region)
# create table
schema_base.metadata.create_all(engine, checkfirst=True)
# create index
# create_index(region, engine, schema_base)
db_session = sessionmaker(bind=engine)()
enable_batch_inserting(db_session)
return db_session
def get_db_session(region: Region,
provider: Provider,
db_name: str = None,
data_schema: object = None,
force_new: bool = False) -> Session:
if data_schema:
db_name = get_db_name(data_schema=data_schema)
session_key = '{}_{}_{}'.format(region.value, provider.value, db_name)
if force_new:
return get_db_session_factory(region, provider, db_name)()
session = sessions.get(session_key)
if not session:
session = get_db_session_factory(region, provider, db_name)()
enable_batch_inserting(session)
sessions[session_key] = session
return session
def main():
if len(sys.argv) != 3:
usage_and_exit()
if sys.argv[1] == 'yes':
batch = True
elif sys.argv[1] == 'no':
batch = False
else:
usage_and_exit()
try:
count = int(sys.argv[2])
except ValueError:
usage_and_exit()
engine = create_engine("postgresql+psycopg2://postgres@localhost",
executemany_mode="values",
echo=count <= 100)
Base.metadata.create_all(engine)
Session = sessionmaker(bind=engine)
session = Session()
if batch:
enable_batch_inserting(session)
prefix = str(int(time.time()))
for i in range(count):
name = "%s-%d" % (prefix, i)
user = User(name=name)
user.addresses = [
Address(email_address="*****@*****.**" % name),
Address(email_address="*****@*****.**" % name),
]
session.add(user)
session.add_all(user.addresses)
with Timer() as timer:
session.commit()
users = session.query(User).count()
addresses = session.query(Address).count()
print("took", timer.duration, "seconds")
print("have", users, "users")
print("have", addresses, "addresses")
def db_session(): # pragma: no cover
try:
db_hostname = os.environ["sighting_db_hostname_str"]
if db_hostname is None:
db_hostname = "127.0.0.1"
db_protocol = "postgresql+psycopg2://sightings:sightings"
db_address = f"{db_hostname}:5432/sightings"
db_url = f"{db_protocol}@{db_address}"
db = create_engine(
db_url,
executemany_mode="values",
executemany_values_page_size=db_ingest_batch_size,
executemany_batch_page_size=500,
)
SightingTable.__table__.create(bind=db, checkfirst=True)
TechniqueTable.__table__.create(bind=db, checkfirst=True)
base.metadata.create_all(db)
_session = sessionmaker(bind=db)
session = _session()
enable_batch_inserting(session)
return session, db
except exc.DatabaseError as err:
logger.error("{}: {}".format(__file__, err))
exit(-1)