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()