def create_tables_if_not_existing(engine: Engine):
metadata = MetaData()
metadata.bind = engine
tablename_jira_issue_created = 'JiraIssueCreated'
tablename_jira_issue_updated = 'JiraIssueUpdated'
try:
if not (engine.has_table(tablename_jira_issue_created)
and engine.has_table(tablename_jira_issue_updated)):
issue_created = Table(
tablename_jira_issue_created, metadata,
Column('issue', String(32), primary_key=True, nullable=False),
Column('created', TIMESTAMP, nullable=False))
issue_updated = Table(
tablename_jira_issue_updated, metadata,
Column('id',
BIGINT(),
Sequence('id', start=1, increment=1),
primary_key=True),
Column('issue',
String(32),
ForeignKey(issue_created.c.issue),
nullable=False),
Column('updated', TIMESTAMP, nullable=False),
Column('issue_status', String(32)),
Column('customer', String(200)))
metadata.create_all()
finally:
engine.dispose()
def iterate_csv_dir(db_engine: engine.Engine = None):
current_dir = os.getcwd()
db_config = config.DB()
for _, services_dir, _ in os.walk(f'{current_dir}/csv'):
for service_db in services_dir:
db_engine = create_engine(db_config.print())
db_connection = db_engine.raw_connection()
try:
db_cursor = db_connection.cursor()
db_cursor.execute(
f"SELECT 1 FROM pg_catalog.pg_database WHERE datname = '{service_db}'"
)
db_exists = db_cursor.fetchone()
if not db_exists:
db_cursor.execute(f'CREATE DATABASE {service_db}')
db_cursor.close()
db_connection.commit()
finally:
db_connection.close()
db_engine.dispose()
db_engine = create_engine(db_config.print(service_db))
for filename in os.listdir(f'{current_dir}/csv/{service_db}'):
filename = os.path.splitext(filename)[0]
from_csv_to_db(service_db, filename, db_engine)
db_engine.dispose()
def test_connection(
ctx: object,
metadata: MetaData,
engine: Engine,
real_transaction: bool = False,
ctx_connection_attribute_name: str = '_test_fx_connection',
) -> typing.Generator:
"""Joining a SQLAlchemy session into an external transaction for test suit.
:param object ctx: Context object to inject test connection into attribute
:param MetaData metadata: SQLAlchemy schema metadata
:param bool real_transaction: (Optional) Whether to use engine as connection directly
or make separate connection. Default: `False`
:param str ctx_connection_attribute_name: (Optional) Attribute name for injecting
test connection to the context object
Default: `'_test_fx_connection'`
.. seealso::
Documentation of the SQLAlchemy session used in test suites.
<http://docs.sqlalchemy.org/en/latest/orm/session_transaction.html#joining-a-session-into-an-external-transaction-such-as-for-test-suites>
""" # noqa
if real_transaction:
metadata.create_all(engine)
try:
yield engine
finally:
metadata.drop_all(engine, checkfirst=True)
return
connection = engine.connect()
try:
metadata.drop_all(connection, checkfirst=True)
transaction = connection.begin()
try:
metadata.create_all(bind=connection)
setattr(ctx, ctx_connection_attribute_name, connection)
try:
yield connection
finally:
delattr(ctx, ctx_connection_attribute_name)
finally:
transaction.rollback()
finally:
connection.close()
engine.dispose()
