def exporter_db_conn():
"""Fixture for db connection."""
def _handler(postgresql):
"""Init DB with data."""
conn = psycopg2.connect(**postgresql.dsn())
cursor = conn.cursor()
with open("../database/vmaas_user_create_postgresql.sql",
"r") as psql_user:
cursor.execute(psql_user.read())
with open("../database/vmaas_db_postgresql.sql", "r") as vmaas_db:
cursor.execute(vmaas_db.read())
with open("test_data/exporter/exporter_test_data.sql",
"r") as test_data:
cursor.execute(test_data.read())
cursor.close()
conn.commit()
conn.close()
# Create temporary posgresql server
# pylint: disable=invalid-name
Postgresql = testing.postgresql.PostgresqlFactory(
cache_initialized_db=True, on_initialized=_handler)
postgresql = Postgresql()
os.environ["POSTGRESQL_PORT"] = str(postgresql.dsn()["port"])
DatabaseHandler.close_connection()
init_db()
conn = psycopg2.connect(**postgresql.dsn())
yield conn
# teardown - close connection, stop postgresql
conn.close()
postgresql.stop()
def fake_database():
with testing.postgresql.Postgresql() as postgresql:
db = psycopg2.connect(**postgresql.dsn())
cur = db.cursor()
create_database(cur)
db.commit()
yield postgresql.dsn()
db.close()
def handler(postgresql):
conn = psycopg2.connect(**postgresql.dsn(password="******"))
cursor = conn.cursor()
cursor.execute("ALTER USER postgres PASSWORD '{}'".format(PASSWORD))
cursor.close()
conn.commit()
conn.close()
def handler(postgresql):
conn = psycopg2.connect(**postgresql.dsn())
cursor = conn.cursor()
cursor.execute(prod_create_sql)
cursor.close()
conn.commit()
conn.close()
def create_pg(vmaas_user=VMAAS_USER, vmaas_pg=VMAAS_PG):
""" Create postgres instance with given SQL scripts. """
def _handler(postgresql):
""" Handler which intializes scheme in instance. """
conn = psycopg2.connect(**postgresql.dsn())
with conn.cursor() as cursor:
cursor.execute("DROP SCHEMA public CASCADE")
cursor.execute("CREATE SCHEMA public")
cursor.execute("GRANT ALL ON SCHEMA public TO postgres")
cursor.execute("GRANT ALL ON SCHEMA public TO public")
with conn.cursor() as cursor:
cursor.execute(vmaas_user.read_text(encoding="utf-8"))
cursor.execute(vmaas_pg.read_text(encoding="utf-8"))
conn.commit()
conn.close()
# pylint: disable=invalid-name
Postgresql = testing.postgresql.PostgresqlFactory(
cache_initialized_db=True,
on_initialized=_handler,
)
postgresql = Postgresql()
os.environ["POSTGRESQL_PORT"] = str(postgresql.dsn()["port"])
return postgresql
def test_importing_a_text_file(self):
"""
Test imprting file to database
"""
with testing.postgresql.Postgresql() as postgresql:
create_engine(postgresql.url())
dsn = postgresql.dsn()
database = etl.database.Database()
database.connection = psycopg2.connect(**dsn)
cursor = database.connection.cursor()
cursor.execute('create table test_table (col1 varchar(20), col2 varchar(20));')
# create test directory in tmp folder
root_path = '/tmp/'
directory_name = str(uuid.uuid4())
directory_path = os.path.join(root_path, directory_name)
os.mkdir(directory_path)
# create test file in test directory
test_file_path = os.path.join(directory_path, 'test.file')
f = open(test_file_path, 'w')
f.write('col1|col2\n')
f.write('test1|test2')
f.close()
# test the function
database.import_file(test_file_path, 'test_table')
cursor.execute('select * from test_table;')
row = cursor.fetchone()
# clean up
shutil.rmtree(directory_path)
self.assertEqual(row[0], 'test1')
self.assertEqual(row[1], 'test2')
def test_execute_bad_query():
q = 'SELECT blarg'
with testing.postgresql.Postgresql() as postgresql:
url = postgresql.url()
dsn = postgresql.dsn()
conn = psycopg2.connect(dbname=dsn['database'],
user=dsn['user'],
password='',
host=dsn['host'],
port=dsn['port'])
cur = conn.cursor()
try:
cur.execute(q)
except Exception:
pass
subsegment = xray_recorder.current_segment().subsegments[0]
assert subsegment.name == 'execute'
sql = subsegment.sql
assert sql['database_type'] == 'PostgreSQL'
assert sql['user'] == dsn['user']
assert sql['url'] == url
assert sql['database_version']
exception = subsegment.cause['exceptions'][0]
assert exception.type == 'ProgrammingError'
def connect_pg(postgresql, **kwargs):
"""PostgreSQL cursor context manager."""
conn = psycopg2.connect(**postgresql.dsn(**kwargs))
cursor = conn.cursor()
yield cursor
cursor.close()
conn.commit()
conn.close()
def handler(postgresql):
conn = psycopg2.connect(**postgresql.dsn())
cursor = conn.cursor()
cursor.execute("create table foo(bar text)")
cursor.execute("INSERT INTO foo VALUES ('hello'), ('fixture'), ('baz')")
cursor.close()
conn.commit()
conn.close()
def handler(postgresql):
conn = psycopg2.connect(**postgresql.dsn())
cursor = conn.cursor()
cursor.execute(
"CREATE TABLE messages(id SERIAL PRIMARY KEY, content TEXT)")
cursor.execute("INSERT INTO messages (content) VALUES ('hello'), ('ciao')")
cursor.close()
conn.commit()
conn.close()
def postgresql(postgresql_factory):
postgresql = postgresql_factory()
dsn = postgresql.dsn()
os.environ["PGDATABASE"] = dsn["database"]
os.environ["PGHOST"] = dsn["host"]
os.environ["PGPORT"] = str(dsn["port"])
os.environ["PGUSER"] = dsn["user"]
yield postgresql
postgresql.stop()
def db_manager(request):
postgresql = testing.postgresql.Postgresql()
db_manager = DatabaseManager(create_tables=True, **postgresql.dsn())
def fin():
postgresql.stop()
request.addfinalizer(fin)
return db_manager
def db_execute_extra(cnn, query, extra):
global conn
try:
cursor = cnn.cursor()
cursor.execute(query, extra)
except:
conn = psy.connect(**postgresql.dsn())
cursor = conn.cursor()
cursor.execute(query, extra)
return cursor
def db_conn():
"""Fixture for db connection."""
postgresql = create_pg()
init_db()
conn = psycopg2.connect(**postgresql.dsn())
yield conn
# teardown - close connection, stop postgresql
conn.close()
postgresql.stop()
def test_set_connection(self):
"""
Test set connection
"""
with testing.postgresql.Postgresql() as postgresql:
create_engine(postgresql.url())
dsn = postgresql.dsn()
database = etl.database.Database()
connection = database.set_connection(dsn['host'], dsn['port'], dsn['database'], dsn['user'], None)
self.assertEqual(str(type(connection)), '<class \'psycopg2.extensions.connection\'>')
def test_register_extensions():
with testing.postgresql.Postgresql() as postgresql:
url = postgresql.url()
dsn = postgresql.dsn()
conn = psycopg2.connect('dbname=' + dsn['database'] +
' password=mypassword' + ' host=' +
dsn['host'] + ' port=' + str(dsn['port']) +
' user='******'user'])
assert psycopg2.extras.register_uuid(None, conn)
assert psycopg2.extras.register_uuid(None, conn.cursor())
def test_query_as_string():
with testing.postgresql.Postgresql() as postgresql:
url = postgresql.url()
dsn = postgresql.dsn()
conn = psycopg2.connect('dbname=' + dsn['database'] +
' password=mypassword' + ' host=' +
dsn['host'] + ' port=' + str(dsn['port']) +
' user='******'user'])
test_sql = psycopg2.sql.Identifier('test')
assert test_sql.as_string(conn)
assert test_sql.as_string(conn.cursor())
def _handler(postgresql):
"""Init DB with data."""
conn = psycopg2.connect(**postgresql.dsn())
cursor = conn.cursor()
with open("../database/vmaas_user_create_postgresql.sql",
"r") as psql_user:
cursor.execute(psql_user.read())
with open("../database/vmaas_db_postgresql.sql", "r") as vmaas_db:
cursor.execute(vmaas_db.read())
cursor.close()
conn.commit()
conn.close()
def test_db():
db = DB()
with testing.postgresql.Postgresql() as postgresql:
db.conn = psycopg2.connect(**postgresql.dsn())
db.cur = db.conn.cursor()
db.cur.execute(
"CREATE TABLE django_session (session_key varchar(255),session_data varchar(255),expire_date varchar(255))"
)
db.cur.execute(
"INSERT INTO django_session (session_key, session_data, expire_date) VALUES (%s, %s, %s)",
('123', "321", datetime.datetime.now()))
assert db.get_session_data('123') == '321'
def initialize_test_database(postgresql):
db_port = postgresql.dsn()['port']
db_host = postgresql.dsn()['host']
db_user = postgresql.dsn()['user']
database = postgresql.dsn()['database']
os.environ['DB_NAME'] = database
os.environ['DB_CONNECTION'] = ('postgresql://' + db_user + '@' + db_host +
':' + str(db_port) + '/' + database)
models.engine = create_engine(os.environ['DB_CONNECTION'])
models.Session = sessionmaker(bind=models.engine)
alembicArgs = [
'--raiseerr',
'upgrade',
'head',
]
alembic.config.main(argv=alembicArgs)
def _handler(postgresql):
"""Initializes PostgreSQL with data."""
with connect_pg(postgresql) as cursor:
pg_env = postgresql.dsn()
cursor.execute("CREATE USER ve_db_admin WITH CREATEROLE")
cursor.execute(
"ALTER USER ve_db_admin WITH PASSWORD 've_db_admin_pwd'")
cursor.execute(
f"ALTER DATABASE {pg_env['database']} OWNER TO ve_db_admin")
with connect_pg(postgresql, user="******") as cursor:
cursor.execute(ve_user.read_text())
cursor.execute(ve_pg.read_text())
cursor.execute(ve_data.read_text())
def handler(postgresql):
conn = psycopg2.connect(**postgresql.dsn())
curr = conn.cursor()
curr.execute(CREATE_LAST_UPDATE_TABLE)
curr.execute(CREATE_TABLE_HISTORIQUE_TAXON)
curr.execute(CREATE_TABLE_TAXON)
curr.execute(CREATE_TABLE_TAXREF_UPDATE)
curr.execute(CREATE_TYPE_TAXREF_DATA)
curr.execute(CREATE_FUNCTION_GET_TAXON)
curr.execute(CREATE_FUNCTION_GET_ALL_TAXON)
curr.close()
conn.commit()
conn.close()
def handler(postgresql):
conn = psycopg2.connect(**postgresql.dsn())
cursor = conn.cursor()
cursor.execute(cust_create_sql)
cursor.execute(insert_cust_sql,
(101, 'Spooky', 'Dogg', 'M', 10, '2008-04-02', '*****@*****.**', '11/25', 'UK'))
cursor.execute(insert_cust_sql,
(102, 'Charlie', 'Bone', 'M', 10, '2008-04-02', '*****@*****.**', '11/25', 'UK'))
cursor.execute(insert_cust_sql,
(103, 'Lulu', 'Catz', 'M', 13, '2005-01-23', '*****@*****.**', '11/20', 'USA'))
cursor.close()
conn.commit()
conn.close()
def run_tests():
with testing.postgresql.Postgresql() as postgresql:
# Minimum settings required for the app's tests.
settings_dict = {
'INSTALLED_APPS': ('tests', ),
'DATABASES': {
'default': {
'ENGINE': 'django.db.backends.postgresql',
'NAME': postgresql.dsn()['database'],
'USER': postgresql.dsn()['user'],
'HOST': postgresql.dsn()['host'],
'PORT': postgresql.dsn()['port'],
},
},
}
# Making Django run this way is a two-step process. First, call
# settings.configure() to give Django settings to work with:
from django.conf import settings
settings.configure(**settings_dict)
# Then, call django.setup() to initialize the application cache
# and other bits:
import django
if hasattr(django, 'setup'):
django.setup()
# Now we instantiate a test runner...
from django.test.utils import get_runner
TestRunner = get_runner(settings)
_monkeypatch_create_btree_gist()
# And then we run tests and return the results.
test_runner = TestRunner(verbosity=1, interactive=True)
failures = test_runner.run_tests(['tests'])
sys.exit(bool(failures))
def set_env(postgresql):
"""Sets PostgreSQL related environment."""
pg_env = postgresql.dsn()
os.environ["POSTGRESQL_DATABASE"] = pg_env["database"]
os.environ["POSTGRESQL_HOST"] = pg_env["host"]
os.environ["POSTGRESQL_PORT"] = str(pg_env["port"])
# update module variables to make DatabaseHandler pool work
database_handler.DB_NAME = pg_env["database"]
database_handler.DB_HOST = pg_env["host"]
database_handler.DB_PORT = int(pg_env["port"])
database_handler.DB_USER = "******"
database_handler.DB_PASS = "******"
database_handler.pg_testing = postgresql
return pg_env
def _handler(postgresql):
""" Handler which intializes scheme in instance. """
conn = psycopg2.connect(**postgresql.dsn())
with conn.cursor() as cursor:
cursor.execute("DROP SCHEMA public CASCADE")
cursor.execute("CREATE SCHEMA public")
cursor.execute("GRANT ALL ON SCHEMA public TO postgres")
cursor.execute("GRANT ALL ON SCHEMA public TO public")
with conn.cursor() as cursor:
pg_env = postgresql.dsn()
cursor.execute("CREATE USER vmaas_db_admin WITH CREATEROLE")
cursor.execute(
"ALTER USER vmaas_db_admin WITH PASSWORD 'vmaas_db_admin_pwd'")
cursor.execute(
f"ALTER DATABASE {pg_env['database']} OWNER TO vmaas_db_admin")
with conn.cursor() as cursor:
cursor.execute(vmaas_user.read_text())
cursor.execute(vmaas_pg.read_text())
conn.commit()
def setup_module(request):
"""
Setup module-level fixtures.
"""
# Create temporary database for testing.
postgresql = testing.postgresql.Postgresql()
connection_params = postgresql.dsn()
# Create connection string from params.
conn_str = 'postgresql://{user}@{host}:{port}/{database}'.format(
user=connection_params['user'],
host=connection_params['host'],
port=connection_params['port'],
database=connection_params['database'],
)
# Create `metabase` and `data` schemata.
engine = sqlalchemy.create_engine(conn_str)
engine.execute(sqlalchemy.schema.CreateSchema('metabase'))
engine.execute(sqlalchemy.schema.CreateSchema('data'))
# Create metabase tables with alembic scripts.
alembic_cfg = Config()
alembic_cfg.set_main_option('script_location', 'alembic')
alembic_cfg.set_main_option('sqlalchemy.url', conn_str)
alembic.command.upgrade(alembic_cfg, 'head')
# Mock settings to connect to testing database. Use this database for
# both the metabase and data schemata.
mock_params = MagicMock()
mock_params.metabase_connection_string = conn_str
mock_params.data_connection_string = conn_str
def teardown_module():
"""
Delete the temporary database.
"""
postgresql.stop()
request.addfinalizer(teardown_module)
return_db = collections.namedtuple('db',
['postgresql', 'engine', 'mock_params'])
return return_db(postgresql=postgresql,
engine=engine,
mock_params=mock_params)
def test_user() -> Tuple[str, str]:
created_at = datetime.datetime.now(datetime.timezone.utc)
test_email = f"test-{created_at.utcnow()}"
test_password = ''.join(
random.choice(string.ascii_letters) for _ in range(24))
pwd_hash = security.generate_password_hash(test_password)
# Initialize a testing database if env vars not defined
if not env.POSTGRES_CONFIG:
postgresql = testing.postgresql.Postgresql()
env.POSTGRES_CONFIG = postgresql.dsn()
db.init_db()
conn = psycopg2.connect(**env.POSTGRES_CONFIG)
with conn.cursor() as cur:
cur.execute(
"""
INSERT INTO users (email, pwhash, created_at)
VALUES (%s, %s, %s)
ON CONFLICT DO NOTHING;
""", (test_email, pwd_hash, created_at))
conn.commit()
yield test_email, test_password
# Clean up the database
with conn.cursor() as cur:
cur.execute(
"""
DELETE FROM samples
WHERE dataset_id IN (
SELECT datasets.id
FROM datasets
WHERE datasets.name ILIKE %s);
""", ('test%', ))
cur.execute(
"""
DELETE FROM datasets
WHERE datasets.name ILIKE %s;
""", ('test%', ))
cur.execute(
"""
DELETE FROM users
WHERE email = %s;
""", (test_email, ))
conn.commit()
def handler(postgresql):
conn = psycopg2.connect(**postgresql.dsn())
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE vault(
id SERIAL PRIMARY KEY,
account CHAR(20) NOT NULL,
password BYTEA NOT NULL
)
''')
query = '''INSERT INTO vault (account, password) VALUES %s'''
data = [('twitter', b'pass'), ('fb', b'fb')]
execute_values(cursor, query, data)
cursor.close()
conn.commit()
conn.close()
def _handler(postgresql):
""" Handler which intializes scheme in instance. """
conn = psycopg2.connect(**postgresql.dsn())
with conn.cursor() as cursor:
cursor.execute("DROP SCHEMA public CASCADE")
cursor.execute("CREATE SCHEMA public")
cursor.execute("GRANT ALL ON SCHEMA public TO postgres")
cursor.execute("GRANT ALL ON SCHEMA public TO public")
with conn.cursor() as cursor:
cursor.execute(vmaas_user.read_text())
cursor.execute(vmaas_pg.read_text())
conn.commit()
conn.close()
def _postgres_impl(mocked_config):
"""Implementation of fixture to initialise a temporary PostgreSQL instance with a clean DB schema."""
# The system needs to be set to the C locale other than en_US.UTF8 to assume that,
# in collation order uppercase will come before lowercase.
postgresql = testing.postgresql.Postgresql(initdb_args='-U postgres -A trust --lc-collate=C.UTF-8 '
'--lc-ctype=C.UTF-8')
dsn = postgresql.dsn()
# Monkey-patch Postgres config to use temp postgres instance
for setting in ['database', 'host', 'port', 'user', 'password']:
setattr(mocked_config.db_config, setting, dsn.get(setting, None))
# Run dirbs-db install_roles using db args from the temp postgres instance
runner = CliRunner()
result = runner.invoke(dirbs_db_cli, ['install_roles'], obj={'APP_CONFIG': mocked_config})
assert result.exit_code == 0
with create_db_connection(mocked_config.db_config) as conn, conn.cursor() as cursor:
cursor.execute('CREATE SCHEMA hll;')
cursor.execute('GRANT USAGE ON SCHEMA hll TO dirbs_core_base;')
cursor.execute('CREATE EXTENSION hll SCHEMA hll;')
cursor.execute(sql.SQL('ALTER DATABASE {0} OWNER TO dirbs_core_power_user')
.format(sql.Identifier(dsn.get('database'))))
# Run dirbs-db install using db args from the temp postgres instance
result = runner.invoke(dirbs_db_cli, ['install'], catch_exceptions=False, obj={'APP_CONFIG': mocked_config})
assert result.exit_code == 0
# Create the necessary roles for security tests
with create_db_connection(mocked_config.db_config) as conn, conn.cursor() as cursor:
cursor.execute('CREATE ROLE dirbs_import_operator_user IN ROLE dirbs_core_import_operator LOGIN')
cursor.execute('CREATE ROLE dirbs_import_gsma_user IN ROLE dirbs_core_import_gsma LOGIN')
cursor.execute('CREATE ROLE dirbs_import_pairing_list_user IN ROLE dirbs_core_import_pairing_list LOGIN')
cursor.execute('CREATE ROLE dirbs_import_stolen_list_user IN ROLE dirbs_core_import_stolen_list LOGIN')
cursor.execute('CREATE ROLE dirbs_import_registration_list_user '
'IN ROLE dirbs_core_import_registration_list LOGIN')
cursor.execute('CREATE ROLE dirbs_import_golden_list_user IN ROLE dirbs_core_import_golden_list LOGIN')
cursor.execute('CREATE ROLE dirbs_classify_user IN ROLE dirbs_core_classify LOGIN')
cursor.execute('CREATE ROLE dirbs_listgen_user IN ROLE dirbs_core_listgen LOGIN')
cursor.execute('CREATE ROLE dirbs_report_user IN ROLE dirbs_core_report LOGIN')
cursor.execute('CREATE ROLE dirbs_api_user IN ROLE dirbs_core_api LOGIN')
cursor.execute('CREATE ROLE dirbs_catalog_user IN ROLE dirbs_core_catalog LOGIN')
cursor.execute('CREATE ROLE dirbs_poweruser_login IN ROLE dirbs_core_power_user LOGIN')
cursor.execute('CREATE ROLE unknown_user LOGIN')
yield postgresql
postgresql.stop()