class SchemaListCreateDeleteShare(TestCase): """ Tests that items reach the execute_sql method in pg.py. Does not test execute_sql itself. """ def setUp(self): # some words to test out self.good_nouns = ["good", "good_noun", "good-noun"] # some words that shoudl throw validation errors self.bad_nouns = ["_foo", "foo_", "-foo", "foo-", "foo bar", "injection;attack", ";injection", "injection;"] self.username = "******" self.password = "******" # mock the execute_sql function self.mock_execute_sql = self.create_patch("core.db.backend.pg.PGBackend.execute_sql") self.mock_execute_sql.return_value = True # mock the is_valid_noun_name, which checks for injection attacks self.mock_check_for_injections = self.create_patch("core.db.backend.pg.PGBackend._check_for_injections") # mock open connection, or else it will try to # create a real db connection self.mock_open_connection = self.create_patch("core.db.backend.pg.PGBackend.__open_connection__") # mock the psycopg2.extensions.AsIs - many of the pg.py methods use it # Its return value (side effect) is the call value self.mock_as_is = self.create_patch("core.db.backend.pg.AsIs") self.mock_as_is.side_effect = lambda x: x # create an instance of PGBackend self.backend = PGBackend(self.username, self.password, repo_base=self.username) def create_patch(self, name): # helper method for creating patches patcher = patch(name) thing = patcher.start() self.addCleanup(patcher.stop) return thing def reset_mocks(self): # clears the mock call arguments and sets their call counts to 0 self.mock_as_is.reset_mock() self.mock_execute_sql.reset_mock() self.mock_check_for_injections.reset_mock() # testing externally called methods in PGBackend def test_create_repo(self): create_repo_sql = "CREATE SCHEMA IF NOT EXISTS %s AUTHORIZATION %s" reponame = "reponame" self.mock_execute_sql.return_value = {"status": True, "row_count": -1, "tuples": [], "fields": []} res = self.backend.create_repo(reponame) self.assertEqual(self.mock_execute_sql.call_args[0][0], create_repo_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1][0], reponame) self.assertEqual(self.mock_execute_sql.call_args[0][1][1], self.username) self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections.called) self.assertEqual(res, True) def test_list_repos(self): # the user is already logged in, so there's not much to be tested here # except that the arguments are passed correctly list_repo_sql = ( "SELECT schema_name AS repo_name " "FROM information_schema.schemata " "WHERE schema_owner != %s" ) mock_settings = self.create_patch("core.db.backend.pg.settings") mock_settings.DATABASES = {"default": {"USER": "******"}} self.mock_execute_sql.return_value = { "status": True, "row_count": 1, "tuples": [("test_table",)], "fields": [{"type": 1043, "name": "table_name"}], } params = (mock_settings.DATABASES["default"]["USER"],) res = self.backend.list_repos() self.assertEqual(self.mock_execute_sql.call_args[0][0], list_repo_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(res, ["test_table"]) def test_delete_repo_happy_path_cascade(self): drop_schema_sql = "DROP SCHEMA %s %s" repo_name = "repo_name" self.mock_execute_sql.return_value = {"status": True, "row_count": -1, "tuples": [], "fields": []} res = self.backend.delete_repo(repo=repo_name, force=True) self.assertEqual(self.mock_execute_sql.call_args[0][0], drop_schema_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1][0], repo_name) self.assertEqual(self.mock_execute_sql.call_args[0][1][1], "CASCADE") self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections) self.assertEqual(res, True) def test_delete_repo_no_cascade(self): drop_schema_sql = "DROP SCHEMA %s %s" repo_name = "repo_name" self.mock_execute_sql.return_value = {"status": True, "row_count": -1, "tuples": [], "fields": []} res = self.backend.delete_repo(repo=repo_name, force=False) self.assertEqual(self.mock_execute_sql.call_args[0][0], drop_schema_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1][0], repo_name) self.assertEqual(self.mock_execute_sql.call_args[0][1][1], None) self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections.called) self.assertEqual(res, True) def test_add_collaborator(self): privileges = [ "SELECT", "INSERT", "UPDATE", "DELETE", "TRUNCATE", "REFERENCES", "TRIGGER", "CREATE", "CONNECT", "TEMPORARY", "EXECUTE", "USAGE", ] add_collab_query = ( "BEGIN;" "GRANT USAGE ON SCHEMA %s TO %s;" "GRANT %s ON ALL TABLES IN SCHEMA %s TO %s;" "ALTER DEFAULT PRIVILEGES IN SCHEMA %s " "GRANT %s ON TABLES TO %s;" "COMMIT;" ) self.mock_execute_sql.return_value = {"status": True, "row_count": -1, "tuples": [], "fields": []} product = itertools.product(self.good_nouns, self.good_nouns, privileges) # test every combo here. For now, don't test combined privileges for repo, receiver, privilege in product: params = (repo, receiver, privilege, repo, receiver, repo, privilege, receiver) res = self.backend.add_collaborator(repo=repo, collaborator=receiver, privileges=[privilege]) self.assertEqual(self.mock_execute_sql.call_args[0][0], add_collab_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 3) self.assertEqual(res, True) self.reset_mocks() def test_add_collaborator_concatinates_privileges(self): privileges = ["SELECT", "USAGE"] repo = "repo" receiver = "receiver" self.mock_execute_sql.return_value = {"status": True, "row_count": -1, "tuples": [], "fields": []} self.backend.add_collaborator(repo=repo, collaborator=receiver, privileges=privileges) # make sure that the privileges are passed as a string in params self.assertTrue("SELECT, USAGE" in self.mock_execute_sql.call_args[0][1]) def test_delete_collaborator(self): delete_collab_sql = ( "BEGIN;" "REVOKE ALL ON ALL TABLES IN SCHEMA %s " "FROM %s CASCADE;" "REVOKE ALL ON SCHEMA %s FROM %s CASCADE;" "ALTER DEFAULT PRIVILEGES IN SCHEMA %s " "REVOKE ALL ON TABLES FROM %s;" "COMMIT;" ) self.mock_execute_sql.return_value = {"status": True, "row_count": -1, "tuples": [], "fields": []} repo = "repo_name" username = "******" params = (repo, username, repo, username, repo, username) res = self.backend.delete_collaborator(repo=repo, collaborator=username) self.assertEqual(self.mock_execute_sql.call_args[0][0], delete_collab_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 2) self.assertEqual(res, True) def test_list_tables(self): repo = "repo" list_tables_query = ( "SELECT table_name FROM " "information_schema.tables " "WHERE table_schema = %s " "AND table_type = 'BASE TABLE';" ) params = (repo,) # execute sql should return this: self.mock_execute_sql.return_value = { "status": True, "row_count": 1, "tuples": [("test_table",)], "fields": [{"type": 1043, "name": "table_name"}], } # mocking out execute_sql's complicated return JSON mock_list_repos = self.create_patch("core.db.backend.pg.PGBackend.list_repos") mock_list_repos.return_value = [repo] res = self.backend.list_tables(repo) self.assertEqual(self.mock_execute_sql.call_args[0][0], list_tables_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(res, ["test_table"]) def test_list_views(self): repo = "repo" list_views_query = ( "SELECT table_name FROM information_schema.tables " "WHERE table_schema = %s " "AND table_type = 'VIEW';" ) params = (repo,) # mocking out execute_sql's complicated return JSON self.mock_execute_sql.return_value = { "status": True, "row_count": 1, "tuples": [("test_view",)], "fields": [{"type": 1043, "name": "view_name"}], } # list_views depends on list_repos, which is being mocked out mock_list_repos = self.create_patch("core.db.backend.pg.PGBackend.list_repos") mock_list_repos.return_value = [repo] res = self.backend.list_views(repo) self.assertEqual(self.mock_execute_sql.call_args[0][0], list_views_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(res, ["test_view"]) def test_get_schema(self): self.mock_execute_sql.return_value = { "status": True, "row_count": 2, "tuples": [(u"id", u"integer"), (u"words", u"text")], "fields": [{"type": 1043, "name": "column_name"}, {"type": 1043, "name": "data_type"}], } repo = "repo" table = "table" get_schema_query = ( "SELECT column_name, data_type " "FROM information_schema.columns " "WHERE table_name = %s " "AND table_schema = %s;" ) params = ("table", "repo") self.backend.get_schema(repo, table) self.assertEqual(self.mock_execute_sql.call_args[0][0], get_schema_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 2) def test_create_user_no_create_db(self): create_user_query = "CREATE ROLE %s WITH LOGIN " "NOCREATEDB NOCREATEROLE NOCREATEUSER " "PASSWORD %s" username = "******" password = "******" self.backend.create_user(username, password, create_db=False) params = (username, password) mock_create_user_database = self.create_patch("core.db.backend.pg.PGBackend.create_user_database") self.assertEqual(self.mock_execute_sql.call_args[0][0], create_user_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 1) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertFalse(mock_create_user_database.called) def test_create_user_calls_create_db(self): username = "******" password = "******" mock_create_user_database = self.create_patch("core.db.backend.pg.PGBackend.create_user_database") self.backend.create_user(username=username, password=password, create_db=True) self.assertTrue(mock_create_user_database.called) def test_create_user_db(self): create_db_query_1 = "CREATE DATABASE %s; " create_db_query_2 = "ALTER DATABASE %s OWNER TO %s; " username = "******" self.backend.create_user_database(username) params_1 = (username,) params_2 = (username, username) call_args_1 = self.mock_execute_sql.call_args_list[0][0] self.assertEqual(call_args_1[0], create_db_query_1) self.assertEqual(call_args_1[1], params_1) call_args_2 = self.mock_execute_sql.call_args_list[1][0] self.assertEqual(call_args_2[0], create_db_query_2) self.assertEqual(call_args_2[1], params_2) self.assertEqual(self.mock_as_is.call_count, len(params_1 + params_2)) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_remove_user(self): query = "DROP ROLE %s;" username = "******" params = (username,) self.backend.remove_user(username) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_remove_database(self): # mock out list_all_users mock_list_all_users = self.create_patch("core.db.backend.pg.PGBackend.list_all_users") mock_list_all_users.return_value = ["tweedledee", "tweedledum"] self.backend.remove_database(self.username) # revoke statement stuff revoke_query = "REVOKE ALL ON DATABASE %s FROM %s;" revoke_params_1 = (self.username, "tweedledee") revoke_params_2 = (self.username, "tweedledum") self.assertEqual(self.mock_execute_sql.call_args_list[0][0][0], revoke_query) self.assertEqual(self.mock_execute_sql.call_args_list[0][0][1], revoke_params_1) self.assertEqual(self.mock_execute_sql.call_args_list[1][0][0], revoke_query) self.assertEqual(self.mock_execute_sql.call_args_list[1][0][1], revoke_params_2) # drop statement stuff drop_query = "DROP DATABASE %s;" drop_params = (self.username,) self.assertEqual(self.mock_execute_sql.call_args_list[2][0][0], drop_query) self.assertEqual(self.mock_execute_sql.call_args_list[2][0][1], drop_params) self.assertEqual(self.mock_as_is.call_count, 5) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_change_password(self): query = "ALTER ROLE %s WITH PASSWORD %s;" params = (self.username, self.password) self.backend.change_password(self.username, self.password) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 1) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_list_collaborators(self): query = "SELECT unnest(nspacl) FROM pg_namespace WHERE nspname=%s;" repo = "repo_name" params = (repo,) self.mock_execute_sql.return_value = { "status": True, "row_count": 2, "tuples": [("al_carter=UC/al_carter",), ("foo_bar=U/al_carter",)], "fields": [{"type": 1033, "name": "unnest"}], } res = self.backend.list_collaborators(repo) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertFalse(self.mock_as_is.called) self.assertEqual(res, ["al_carter", "foo_bar"]) def test_list_all_users(self): query = "SELECT usename FROM pg_catalog.pg_user WHERE usename != %s" params = (self.username,) self.mock_execute_sql.return_value = { "status": True, "row_count": 2, "tuples": [(u"delete_me_alpha_user",), (u"delete_me_beta_user",)], "fields": [{"type": 19, "name": "usename"}], } res = self.backend.list_all_users() self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertFalse(self.mock_as_is.called) self.assertEqual(res, ["delete_me_alpha_user", "delete_me_beta_user"]) def test_has_base_privilege(self): query = "SELECT has_database_privilege(%s, %s);" privilege = "CONNECT" params = (self.username, privilege) self.mock_execute_sql.return_value = {"status": True, "row_count": -1, "tuples": [[True]], "fields": []} res = self.backend.has_base_privilege(login=self.username, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_repo_privilege(self): query = "SELECT has_schema_privilege(%s, %s, %s);" repo = "repo" privilege = "CONNECT" params = (self.username, repo, privilege) self.mock_execute_sql.return_value = {"status": True, "row_count": -1, "tuples": [[True]], "fields": []} res = self.backend.has_repo_privilege(login=self.username, repo=repo, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_table_privilege(self): query = "SELECT has_table_privilege(%s, %s, %s);" table = "table" privilege = "CONNECT" params = (self.username, table, privilege) self.mock_execute_sql.return_value = {"status": True, "row_count": -1, "tuples": [[True]], "fields": []} res = self.backend.has_table_privilege(login=self.username, table=table, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_column_privilege(self): query = "SELECT has_column_privilege(%s, %s, %s, %s);" table = "table" column = "column" privilege = "CONNECT" params = (self.username, table, column, privilege) self.mock_execute_sql.return_value = {"status": True, "row_count": -1, "tuples": [[True]], "fields": []} res = self.backend.has_column_privilege(login=self.username, table=table, column=column, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_export_table_with_header(self): query = "COPY %s TO %s WITH %s %s DELIMITER %s;" table_name = "user_name.repo_name.table_name" file_path = "file_path" file_format = "file_format" delimiter = "," header = True params = (table_name, file_path, file_format, "HEADER", delimiter) self.backend.export_table(table_name, file_path, file_format, delimiter, header) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 4) def test_export_table_with_no_header(self): table_name = "table_name" file_path = "file_path" file_format = "file_format" delimiter = "," header = False params = (table_name, file_path, file_format, "", delimiter) self.backend.export_table(table_name, file_path, file_format, delimiter, header) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_export_query_with_header(self): query = "COPY (%s) TO %s WITH %s %s DELIMITER %s;" passed_query = "myquery" file_path = "file_path" file_format = "CSV" delimiter = "," header = True params = (passed_query, file_path, file_format, "HEADER", delimiter) self.backend.export_query(passed_query, file_path, file_format, delimiter, header) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 2) def test_export_query_with_no_header(self): passed_query = "myquery" file_path = "file_path" file_format = "CSV" delimiter = "," header = False params = (passed_query, file_path, file_format, "", delimiter) self.backend.export_query(passed_query, file_path, file_format, delimiter, header) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_export_query_only_executes_text_before_semicolon(self): passed_query = " text before semicolon; text after; " file_path = "file_path" file_format = "CSV" delimiter = "," header = False passed_query_cleaned = " text before semicolon" params = (passed_query_cleaned, file_path, file_format, "", delimiter) self.backend.export_query(passed_query, file_path, file_format, delimiter, header) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_import_file_with_header(self): query = "COPY %s FROM %s WITH %s %s DELIMITER %s ENCODING %s QUOTE %s;" table_name = "user_name.repo_name.table_name" file_path = "file_path" file_format = "file_format" delimiter = "," header = True encoding = "ISO-8859-1" quote_character = '"' params = (table_name, file_path, file_format, "HEADER", delimiter, encoding, quote_character) self.backend.import_file(table_name, file_path, file_format, delimiter, header, encoding, quote_character) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 5) def test_import_table_with_no_header(self): table_name = "table_name" file_path = "file_path" file_format = "file_format" delimiter = "," header = False encoding = "ISO-8859-1" quote_character = '"' params = (table_name, file_path, file_format, "", delimiter, encoding, quote_character) self.backend.import_file(table_name, file_path, file_format, delimiter, header, encoding, quote_character) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_import_file_w_dbtruck(self): # DBTruck is not tested for safety/security... At all. # The method does so little # that it doesn't even make much sense to test it. pass
class SchemaListCreateDeleteShare(MockingMixin, TestCase): """ Tests that items reach the execute_sql method in pg.py. Does not test execute_sql itself. """ def setUp(self): # some words to test out self.good_nouns = ['good', 'good_noun', 'good-noun'] # some words that shoudl throw validation errors self.bad_nouns = [ '_foo', 'foo_', '-foo', 'foo-', 'foo bar', 'injection;attack', ';injection', 'injection;', ] self.username = "******" self.password = "******" # mock the execute_sql function self.mock_execute_sql = self.create_patch( 'core.db.backend.pg.PGBackend.execute_sql') self.mock_execute_sql.return_value = True # mock the mock_check_for_injections, which checks for injection # attacks self.mock_check_for_injections = self.create_patch( 'core.db.backend.pg.PGBackend._check_for_injections') self.mock_validate_table_name = self.create_patch( 'core.db.backend.pg.PGBackend._validate_table_name') # mock open connection, or else it will try to # create a real db connection self.mock_open_connection = self.create_patch( 'core.db.backend.pg.PGBackend.__open_connection__') # mock the psycopg2.extensions.AsIs - many of the pg.py methods use it # Its return value (side effect) is the call value self.mock_as_is = self.create_patch('core.db.backend.pg.AsIs') self.mock_as_is.side_effect = lambda x: x # create an instance of PGBackend self.backend = PGBackend(self.username, self.password, repo_base=self.username) def reset_mocks(self): # clears the mock call arguments and sets their call counts to 0 self.mock_as_is.reset_mock() self.mock_execute_sql.reset_mock() self.mock_check_for_injections.reset_mock() # testing externally called methods in PGBackend def test_create_repo(self): create_repo_sql = 'CREATE SCHEMA IF NOT EXISTS %s AUTHORIZATION %s' reponame = 'reponame' self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } res = self.backend.create_repo(reponame) self.assertEqual(self.mock_execute_sql.call_args[0][0], create_repo_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1][0], reponame) self.assertEqual(self.mock_execute_sql.call_args[0][1][1], self.username) self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections.called) self.assertEqual(res, True) def test_list_repos(self): # the user is already logged in, so there's not much to be tested here # except that the arguments are passed correctly list_repo_sql = ('SELECT schema_name AS repo_name ' 'FROM information_schema.schemata ' 'WHERE schema_owner != %s') mock_settings = self.create_patch("core.db.backend.pg.settings") mock_settings.DATABASES = {'default': {'USER': '******'}} self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [('test_table', )], 'fields': [{ 'type': 1043, 'name': 'table_name' }] } params = (mock_settings.DATABASES['default']['USER'], ) res = self.backend.list_repos() self.assertEqual(self.mock_execute_sql.call_args[0][0], list_repo_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(res, ['test_table']) def test_rename_repo(self): alter_repo_sql = 'ALTER SCHEMA %s RENAME TO %s' self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [('test_table', )], 'fields': [{ 'type': 1043, 'name': 'table_name' }] } params = ('old_name', 'new_name') res = self.backend.rename_repo('old_name', 'new_name') self.assertEqual(res, True) self.assertEqual(self.mock_execute_sql.call_args[0][0], alter_repo_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertTrue(self.mock_execute_sql.called) self.assertEqual(self.mock_check_for_injections.call_count, 2) def test_delete_repo_happy_path_cascade(self): drop_schema_sql = 'DROP SCHEMA %s %s' repo_name = 'repo_name' self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } res = self.backend.delete_repo(repo=repo_name, force=True) self.assertEqual(self.mock_execute_sql.call_args[0][0], drop_schema_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1][0], repo_name) self.assertEqual(self.mock_execute_sql.call_args[0][1][1], 'CASCADE') self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections) self.assertEqual(res, True) def test_delete_repo_no_cascade(self): drop_schema_sql = 'DROP SCHEMA %s %s' repo_name = 'repo_name' self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } res = self.backend.delete_repo(repo=repo_name, force=False) self.assertEqual(self.mock_execute_sql.call_args[0][0], drop_schema_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1][0], repo_name) self.assertEqual(self.mock_execute_sql.call_args[0][1][1], '') self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections.called) self.assertEqual(res, True) def test_add_collaborator(self): privileges = [ 'SELECT', 'INSERT', 'UPDATE', 'DELETE', 'TRUNCATE', 'REFERENCES', 'TRIGGER', 'CREATE', 'CONNECT', 'TEMPORARY', 'EXECUTE', 'USAGE' ] add_collab_query = ('BEGIN;' 'GRANT USAGE ON SCHEMA %s TO %s;' 'GRANT %s ON ALL TABLES IN SCHEMA %s TO %s;' 'ALTER DEFAULT PRIVILEGES IN SCHEMA %s ' 'GRANT %s ON TABLES TO %s;' 'COMMIT;') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } product = itertools.product(self.good_nouns, self.good_nouns, privileges) # test every combo here. For now, don't test combined privileges for repo, receiver, privilege in product: params = (repo, receiver, privilege, repo, receiver, repo, privilege, receiver) res = self.backend.add_collaborator(repo=repo, collaborator=receiver, db_privileges=[privilege]) self.assertEqual(self.mock_execute_sql.call_args[0][0], add_collab_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 3) self.assertEqual(res, True) self.reset_mocks() def test_add_collaborator_concatenates_privileges(self): privileges = ['SELECT', 'USAGE'] repo = 'repo' receiver = 'receiver' self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } self.backend.add_collaborator(repo=repo, collaborator=receiver, db_privileges=privileges) # make sure that the privileges are passed as a string in params self.assertTrue( 'SELECT, USAGE' in self.mock_execute_sql.call_args[0][1]) def test_delete_collaborator(self): delete_collab_sql = ('BEGIN;' 'REVOKE ALL ON ALL TABLES IN SCHEMA %s ' 'FROM %s CASCADE;' 'REVOKE ALL ON SCHEMA %s FROM %s CASCADE;' 'ALTER DEFAULT PRIVILEGES IN SCHEMA %s ' 'REVOKE ALL ON TABLES FROM %s;' 'COMMIT;') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } repo = 'repo_name' username = '******' params = (repo, username, repo, username, repo, username) res = self.backend.delete_collaborator(repo=repo, collaborator=username) self.assertEqual(self.mock_execute_sql.call_args[0][0], delete_collab_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 2) self.assertEqual(res, True) def test_create_table(self): repo = 'repo' table = 'table' params = [{ 'column_name': 'id', 'data_type': 'integer' }, { 'column_name': 'words', 'data_type': 'text' }] expected_params = ('repo', 'table', 'id integer, words text') create_table_query = ('CREATE TABLE %s.%s (%s)') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } res = self.backend.create_table(repo, table, params) # checks repo, table, and all param values for injections self.assertEqual(self.mock_check_for_injections.call_count, 5) self.assertEqual(self.mock_validate_table_name.call_count, 1) final_query = self.mock_execute_sql.call_args[0][0] final_params = self.mock_execute_sql.call_args[0][1] self.assertEqual(final_query, create_table_query) self.assertEqual(final_params, expected_params) self.assertEqual(res, True) # create table test_repo.test_table (id integer, words text) def test_list_tables(self): repo = 'repo' list_tables_query = ('SELECT table_name FROM ' 'information_schema.tables ' 'WHERE table_schema = %s ' 'AND table_type = \'BASE TABLE\';') params = (repo, ) # execute sql should return this: self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [('test_table', )], 'fields': [{ 'type': 1043, 'name': 'table_name' }] } # mocking out execute_sql's complicated return JSON mock_list_repos = self.create_patch( 'core.db.backend.pg.PGBackend.list_repos') mock_list_repos.return_value = [repo] res = self.backend.list_tables(repo) self.assertEqual(self.mock_execute_sql.call_args[0][0], list_tables_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(res, ['test_table']) def test_describe_table_without_detail(self): repo = 'repo' table = 'table' detail = False query = ("SELECT %s " "FROM information_schema.columns " "WHERE table_schema = %s and table_name = %s;") params = ('column_name, data_type', repo, table) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'id', u'integer'), (u'words', u'text')], 'fields': [{ 'type': 1043, 'name': 'column_name' }, { 'type': 1043, 'name': 'data_type' }] } res = self.backend.describe_table(repo, table, detail) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(res, [(u'id', u'integer'), (u'words', u'text')]) def test_describe_table_query_in_detail(self): repo = 'repo' table = 'table' detail = True query = ("SELECT %s " "FROM information_schema.columns " "WHERE table_schema = %s and table_name = %s;") params = ('*', repo, table) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'id', u'integer'), (u'words', u'text'), ('foo', 'bar')], 'fields': [{ 'type': 1043, 'name': 'column_name' }, { 'type': 1043, 'name': 'data_type' }] } res = self.backend.describe_table(repo, table, detail) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(res, [(u'id', u'integer'), (u'words', u'text'), ('foo', 'bar')]) def test_list_table_permissions(self): repo = 'repo' table = 'table' query = ("select privilege_type from " "information_schema.role_table_grants where table_schema=%s " "and table_name=%s and grantee=%s") params = ('repo', 'table', self.username) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'SELECT'), (u'UPDATE')], 'fields': [{ 'type': 1043, 'name': 'privilege_type' }] } self.backend.list_table_permissions(repo, table) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_delete_table(self): repo = 'repo_name' table = 'table_name' force = False expected_query = ('DROP TABLE %s.%s.%s %s') expected_params = ('username', 'repo_name', 'table_name', 'RESTRICT') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } res = self.backend.delete_table(repo, table, force) final_query = self.mock_execute_sql.call_args[0][0] final_params = self.mock_execute_sql.call_args[0][1] self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(self.mock_validate_table_name.call_count, 1) self.assertEqual(final_query, expected_query) self.assertEqual(final_params, expected_params) self.assertEqual(res, True) def test_list_views(self): repo = 'repo' list_views_query = ('SELECT table_name FROM information_schema.tables ' 'WHERE table_schema = %s ' 'AND table_type = \'VIEW\';') params = (repo, ) # mocking out execute_sql's complicated return JSON self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [('test_view', )], 'fields': [{ 'type': 1043, 'name': 'view_name' }] } # list_views depends on list_repos, which is being mocked out mock_list_repos = self.create_patch( 'core.db.backend.pg.PGBackend.list_repos') mock_list_repos.return_value = [repo] res = self.backend.list_views(repo) self.assertEqual(self.mock_execute_sql.call_args[0][0], list_views_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(res, ['test_view']) def test_delete_view(self): repo = 'repo_name' view = 'view_name' force = False expected_query = ('DROP VIEW %s.%s.%s %s') expected_params = ('username', 'repo_name', 'view_name', 'RESTRICT') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } res = self.backend.delete_view(repo, view, force) final_query = self.mock_execute_sql.call_args[0][0] final_params = self.mock_execute_sql.call_args[0][1] self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(self.mock_validate_table_name.call_count, 1) self.assertEqual(final_query, expected_query) self.assertEqual(final_params, expected_params) self.assertEqual(res, True) def test_create_view(self): repo = 'repo_name' view = 'view_name' sql = 'SELECT * FROM table' expected_params = ('repo_name', 'view_name', 'SELECT * FROM table') create_view_query = ('CREATE VIEW %s.%s AS (%s)') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } res = self.backend.create_view(repo, view, sql) # checks repo and view for injections self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(self.mock_validate_table_name.call_count, 1) final_query = self.mock_execute_sql.call_args[0][0] final_params = self.mock_execute_sql.call_args[0][1] self.assertEqual(final_query, create_view_query) self.assertEqual(final_params, expected_params) self.assertEqual(res, True) def test_describe_view_without_detail(self): repo = 'repo_name' view = 'view_name' detail = False query = ("SELECT %s " "FROM information_schema.columns " "WHERE table_schema = %s and table_name = %s;") params = ('column_name, data_type', repo, view) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'id', u'integer'), (u'words', u'text')], 'fields': [{ 'type': 1043, 'name': 'column_name' }, { 'type': 1043, 'name': 'data_type' }] } res = self.backend.describe_view(repo, view, detail) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(res, [(u'id', u'integer'), (u'words', u'text')]) def test_get_schema(self): self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'id', u'integer'), (u'words', u'text')], 'fields': [{ 'type': 1043, 'name': 'column_name' }, { 'type': 1043, 'name': 'data_type' }] } repo = 'repo' table = 'table' get_schema_query = ('SELECT column_name, data_type ' 'FROM information_schema.columns ' 'WHERE table_name = %s ' 'AND table_schema = %s;') params = ('table', 'repo') self.backend.get_schema(repo, table) self.assertEqual(self.mock_execute_sql.call_args[0][0], get_schema_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(self.mock_validate_table_name.call_count, 1) def test_create_public_user_no_create_db(self): create_user_query = ('CREATE ROLE %s WITH LOGIN ' 'NOCREATEDB NOCREATEROLE NOCREATEUSER ' 'PASSWORD %s') username = PUBLIC_ROLE password = '******' self.backend.create_user(username, password, create_db=False) params = (username, password) mock_create_user_database = self.create_patch( 'core.db.backend.pg.PGBackend.create_user_database') # import pdb; pdb.set_trace() self.assertEqual(self.mock_execute_sql.call_args[0][0], create_user_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 1) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertFalse(mock_create_user_database.called) def test_create_normal_user_no_create_db(self): create_user_query = 'GRANT %s to %s' username = '******' password = '******' self.backend.create_user(username, password, create_db=False) params = (PUBLIC_ROLE, username) mock_create_user_database = self.create_patch( 'core.db.backend.pg.PGBackend.create_user_database') # import pdb; pdb.set_trace() self.assertEqual(self.mock_execute_sql.call_args[0][0], create_user_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertFalse(mock_create_user_database.called) def test_create_user_calls_create_db(self): username = '******' password = '******' mock_create_user_database = self.create_patch( 'core.db.backend.pg.PGBackend.create_user_database') self.backend.create_user(username=username, password=password, create_db=True) self.assertTrue(mock_create_user_database.called) def test_create_user_db(self): create_db_query_1 = 'CREATE DATABASE %s; ' create_db_query_2 = 'ALTER DATABASE %s OWNER TO %s; ' username = '******' self.backend.create_user_database(username) params_1 = (username, ) params_2 = (username, username) call_args_1 = self.mock_execute_sql.call_args_list[0][0] self.assertEqual(call_args_1[0], create_db_query_1) self.assertEqual(call_args_1[1], params_1) call_args_2 = self.mock_execute_sql.call_args_list[1][0] self.assertEqual(call_args_2[0], create_db_query_2) self.assertEqual(call_args_2[1], params_2) self.assertEqual(self.mock_as_is.call_count, len(params_1 + params_2)) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_remove_user(self): query = 'DROP ROLE %s;' username = "******" params = (username, ) self.backend.remove_user(username) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_remove_database(self): # mock out list_all_users mock_list_all_users = self.create_patch( 'core.db.backend.pg.PGBackend.list_all_users') mock_list_all_users.return_value = ['tweedledee', 'tweedledum'] self.backend.remove_database(self.username) # revoke statement stuff revoke_query = 'REVOKE ALL ON DATABASE %s FROM %s;' revoke_params_1 = (self.username, 'tweedledee') revoke_params_2 = (self.username, 'tweedledum') self.assertEqual(self.mock_execute_sql.call_args_list[0][0][0], revoke_query) self.assertEqual(self.mock_execute_sql.call_args_list[0][0][1], revoke_params_1) self.assertEqual(self.mock_execute_sql.call_args_list[1][0][0], revoke_query) self.assertEqual(self.mock_execute_sql.call_args_list[1][0][1], revoke_params_2) # drop statement stuff drop_query = 'DROP DATABASE %s;' drop_params = (self.username, ) self.assertEqual(self.mock_execute_sql.call_args_list[2][0][0], drop_query) self.assertEqual(self.mock_execute_sql.call_args_list[2][0][1], drop_params) self.assertEqual(self.mock_as_is.call_count, 5) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_change_password(self): query = 'ALTER ROLE %s WITH PASSWORD %s;' params = (self.username, self.password) self.backend.change_password(self.username, self.password) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 1) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_list_collaborators(self): query = 'SELECT unnest(nspacl) FROM pg_namespace WHERE nspname=%s;' repo = 'repo_name' params = (repo, ) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [('al_carter=UC/al_carter', ), ('foo_bar=U/al_carter', )], 'fields': [{ 'type': 1033, 'name': 'unnest' }] } expected_result = [{ 'username': '******', 'db_permissions': 'UC' }, { 'username': '******', 'db_permissions': 'U' }] res = self.backend.list_collaborators(repo) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertFalse(self.mock_as_is.called) self.assertEqual(res, expected_result) def test_list_all_users(self): query = 'SELECT usename FROM pg_catalog.pg_user WHERE usename != %s' params = (self.username, ) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'delete_me_alpha_user', ), (u'delete_me_beta_user', )], 'fields': [{ 'type': 19, 'name': 'usename' }] } res = self.backend.list_all_users() self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertFalse(self.mock_as_is.called) self.assertEqual(res, ['delete_me_alpha_user', 'delete_me_beta_user']) def test_has_base_privilege(self): query = 'SELECT has_database_privilege(%s, %s);' privilege = 'CONNECT' params = (self.username, privilege) self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': [] } res = self.backend.has_base_privilege(login=self.username, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_repo_db_privilege(self): query = 'SELECT has_schema_privilege(%s, %s, %s);' repo = 'repo' privilege = 'CONNECT' params = (self.username, repo, privilege) self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': [] } res = self.backend.has_repo_db_privilege(login=self.username, repo=repo, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_table_privilege(self): query = 'SELECT has_table_privilege(%s, %s, %s);' table = 'table' privilege = 'CONNECT' params = (self.username, table, privilege) self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': [] } res = self.backend.has_table_privilege(login=self.username, table=table, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_column_privilege(self): query = 'SELECT has_column_privilege(%s, %s, %s, %s);' table = 'table' column = 'column' privilege = 'CONNECT' params = (self.username, table, column, privilege) self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': [] } res = self.backend.has_column_privilege(login=self.username, table=table, column=column, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) @patch('core.db.backend.pg.os.makedirs') @patch('core.db.backend.pg.os.remove') @patch('core.db.backend.pg.shutil.move') def test_export_table_with_header(self, *args): table_name = 'repo_name.table_name' table_name_prepped = 'SELECT * FROM %s' % table_name file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = True query = ('COPY (SELECT * FROM repo_name.table_name) ' 'TO STDOUT WITH CSV HEADER DELIMITER \',\';') self.backend.connection = Mock() mock_connection = self.backend.connection mock_mogrify = mock_connection.cursor.return_value.mogrify mock_mogrify.return_value = query mock_copy_expert = mock_connection.cursor.return_value.copy_expert with patch("__builtin__.open", mock_open()): self.backend.export_table(table_name, file_path, file_format, delimiter, header) mock_mogrify.assert_called_once_with( 'COPY (%s) TO STDOUT WITH %s %s DELIMITER %s;', (table_name_prepped, file_format, 'HEADER' if header else '', delimiter)) # Kind of a meaningless check since we have to mock the return value # of mogrify, but at least it ensures the result of mogrify is passed # into copy_expert as is. self.assertEqual(mock_copy_expert.call_args[0][0], query) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 4) self.assertEqual(self.mock_validate_table_name.call_count, 1) @patch('core.db.backend.pg.os.makedirs') @patch('core.db.backend.pg.os.remove') @patch('core.db.backend.pg.shutil.move') def test_export_table_with_no_header(self, *args): table_name = 'repo_name.table_name' table_name_prepped = 'SELECT * FROM %s' % table_name file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = False query = ('COPY (SELECT * FROM repo_name.table_name) ' 'TO STDOUT WITH CSV DELIMITER \',\';') self.backend.connection = Mock() mock_connection = self.backend.connection mock_mogrify = mock_connection.cursor.return_value.mogrify mock_mogrify.return_value = query mock_copy_expert = mock_connection.cursor.return_value.copy_expert with patch("__builtin__.open", mock_open()): self.backend.export_table(table_name, file_path, file_format, delimiter, header) mock_mogrify.assert_called_once_with( 'COPY (%s) TO STDOUT WITH %s %s DELIMITER %s;', (table_name_prepped, file_format, 'HEADER' if header else '', delimiter)) # Kind of a meaningless check since we have to mock the return value # of mogrify, but at least it ensures the result of mogrify is passed # into copy_expert as is. self.assertEqual(mock_copy_expert.call_args[0][0], query) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 4) self.assertEqual(self.mock_validate_table_name.call_count, 1) @patch('core.db.backend.pg.os.makedirs') @patch('core.db.backend.pg.os.remove') @patch('core.db.backend.pg.shutil.move') def test_export_view(self, *args): view_name = 'repo_name.view_name' view_name_prepped = 'SELECT * FROM %s' % view_name file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = True query = ('COPY (SELECT * FROM repo_name.view_name) ' 'TO STDOUT WITH CSV HEADER DELIMITER \',\';') self.backend.connection = Mock() mock_connection = self.backend.connection mock_mogrify = mock_connection.cursor.return_value.mogrify mock_mogrify.return_value = query mock_copy_expert = mock_connection.cursor.return_value.copy_expert with patch("__builtin__.open", mock_open()): self.backend.export_view(view_name, file_path, file_format, delimiter, header) mock_mogrify.assert_called_once_with( 'COPY (%s) TO STDOUT WITH %s %s DELIMITER %s;', (view_name_prepped, file_format, 'HEADER' if header else '', delimiter)) # Kind of a meaningless check since we have to mock the return value # of mogrify, but at least it ensures the result of mogrify is passed # into copy_expert as is. self.assertEqual(mock_copy_expert.call_args[0][0], query) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 4) self.assertEqual(self.mock_validate_table_name.call_count, 1) @patch('core.db.backend.pg.os.makedirs') @patch('core.db.backend.pg.os.remove') @patch('core.db.backend.pg.shutil.move') def _export_query_test_helper(self, *args, **kwargs): passed_query = kwargs['passed_query'] passed_query_cleaned = kwargs.get('passed_query_cleaned', passed_query) file_path = kwargs['file_path'] file_format = kwargs['file_format'] delimiter = kwargs['delimiter'] header = kwargs['header'] query = kwargs['query'] self.backend.connection = Mock() mock_connection = self.backend.connection mock_mogrify = mock_connection.cursor.return_value.mogrify mock_mogrify.return_value = query mock_copy_expert = mock_connection.cursor.return_value.copy_expert with patch("__builtin__.open", mock_open()): self.backend.export_query(passed_query, file_path, file_format, delimiter, header) mock_mogrify.assert_called_once_with( 'COPY (%s) TO STDOUT WITH %s %s DELIMITER %s;', (passed_query_cleaned, file_format, 'HEADER' if header else '', delimiter)) # Kind of a meaningless check since we have to mock the return value # of mogrify, but at least it ensures the result of mogrify is passed # into copy_expert as is. self.assertEqual(mock_copy_expert.call_args[0][0], query) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 2) def test_export_query_with_header(self, *args): self._export_query_test_helper( passed_query='myquery', file_path='file_path', file_format='CSV', delimiter=',', header=True, query=('COPY (myquery) ' 'TO STDOUT WITH CSV HEADER DELIMITER \',\';')) def test_export_query_with_no_header(self, *args): self._export_query_test_helper( passed_query='myquery', file_path='file_path', file_format='CSV', delimiter=',', header=False, query=('COPY (myquery) ' 'TO STDOUT WITH CSV DELIMITER \',\';')) def test_export_query_only_executes_text_before_semicolon(self, *args): self._export_query_test_helper( passed_query=' text before semicolon; text after; ', passed_query_cleaned='text before semicolon', file_path='file_path', file_format='CSV', delimiter=',', header=False, query=('COPY (text before semicolon) ' 'TO STDOUT WITH CSV DELIMITER \',\';')) def test_import_file_with_header(self): query = 'COPY %s FROM %s WITH %s %s DELIMITER %s ENCODING %s QUOTE %s;' table_name = 'user_name.repo_name.table_name' file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = True encoding = 'ISO-8859-1' quote_character = '"' params = (table_name, file_path, file_format, 'HEADER', delimiter, encoding, quote_character) self.backend.import_file(table_name, file_path, file_format, delimiter, header, encoding, quote_character) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 4) self.assertEqual(self.mock_validate_table_name.call_count, 1) def test_import_table_with_no_header(self): table_name = 'table_name' file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = False encoding = 'ISO-8859-1' quote_character = '"' params = (table_name, file_path, file_format, '', delimiter, encoding, quote_character) self.backend.import_file(table_name, file_path, file_format, delimiter, header, encoding, quote_character) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_import_file_w_dbtruck(self): # DBTruck is not tested for safety/security... At all. # The method does so little # that it doesn't even make much sense to test it. pass def test_can_user_access_rls_table(self): mock_settings = self.create_patch("core.db.backend.pg.settings") mock_settings.POLICY_SCHEMA = 'SCHEMA' mock_settings.POLICY_TABLE = 'TABLE' self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [ (True, ), ] } username = '******' permissions = ['select', 'update'] expected_query = ( "SELECT exists(" "SELECT * FROM %s.%s where grantee=lower(%s) and (" "lower(policy_type)=lower(%s) or lower(policy_type)=lower(%s)" "))") expected_params = ('SCHEMA', 'TABLE', 'delete_me_user', 'select', 'update') self.backend.can_user_access_rls_table(username, permissions) self.assertEqual(self.mock_execute_sql.call_args[0][0], expected_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], expected_params)
class SchemaListCreateDeleteShare(MockingMixin, TestCase): """ Tests that items reach the execute_sql method in pg.py. Does not test execute_sql itself. """ def setUp(self): # some words to test out self.good_nouns = ['good', 'good_noun', 'good-noun'] # some words that shoudl throw validation errors self.bad_nouns = ['_foo', 'foo_', '-foo', 'foo-', 'foo bar', 'injection;attack', ';injection', 'injection;', ] self.username = "******" self.password = "******" # mock the execute_sql function self.mock_execute_sql = self.create_patch( 'core.db.backend.pg.PGBackend.execute_sql') self.mock_execute_sql.return_value = True # mock the mock_check_for_injections, which checks for injection # attacks self.mock_check_for_injections = self.create_patch( 'core.db.backend.pg.PGBackend._check_for_injections') self.mock_validate_table_name = self.create_patch( 'core.db.backend.pg.PGBackend._validate_table_name') # mock open connection, or else it will try to # create a real db connection self.mock_open_connection = self.create_patch( 'core.db.backend.pg.PGBackend.__open_connection__') # mock the psycopg2.extensions.AsIs - many of the pg.py methods use it # Its return value (side effect) is the call value self.mock_as_is = self.create_patch('core.db.backend.pg.AsIs') self.mock_as_is.side_effect = lambda x: x # create an instance of PGBackend self.backend = PGBackend(self.username, self.password, repo_base=self.username) def reset_mocks(self): # clears the mock call arguments and sets their call counts to 0 self.mock_as_is.reset_mock() self.mock_execute_sql.reset_mock() self.mock_check_for_injections.reset_mock() # testing externally called methods in PGBackend def test_create_repo(self): create_repo_sql = 'CREATE SCHEMA IF NOT EXISTS %s AUTHORIZATION %s' reponame = 'reponame' self.mock_execute_sql.return_value = {'status': True, 'row_count': -1, 'tuples': [], 'fields': []} res = self.backend.create_repo(reponame) self.assertEqual( self.mock_execute_sql.call_args[0][0], create_repo_sql) self.assertEqual( self.mock_execute_sql.call_args[0][1][0], reponame) self.assertEqual( self.mock_execute_sql.call_args[0][1][1], self.username) self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections.called) self.assertEqual(res, True) def test_list_repos(self): # the user is already logged in, so there's not much to be tested here # except that the arguments are passed correctly list_repo_sql = ('SELECT schema_name AS repo_name ' 'FROM information_schema.schemata ' 'WHERE schema_owner != %s') mock_settings = self.create_patch("core.db.backend.pg.settings") mock_settings.DATABASES = {'default': {'USER': '******'}} self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [ ('test_table',)], 'fields': [{'type': 1043, 'name': 'table_name'}]} params = (mock_settings.DATABASES['default']['USER'],) res = self.backend.list_repos() self.assertEqual( self.mock_execute_sql.call_args[0][0], list_repo_sql) self.assertEqual( self.mock_execute_sql.call_args[0][1], params) self.assertEqual(res, ['test_table']) def test_rename_repo(self): alter_repo_sql = 'ALTER SCHEMA %s RENAME TO %s' self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [ ('test_table',)], 'fields': [{'type': 1043, 'name': 'table_name'}]} params = ('old_name', 'new_name') res = self.backend.rename_repo('old_name', 'new_name') self.assertEqual(res, True) self.assertEqual( self.mock_execute_sql.call_args[0][0], alter_repo_sql) self.assertEqual( self.mock_execute_sql.call_args[0][1], params) self.assertTrue(self.mock_execute_sql.called) self.assertEqual(self.mock_check_for_injections.call_count, 2) def test_delete_repo_happy_path_cascade(self): drop_schema_sql = 'DROP SCHEMA %s %s' repo_name = 'repo_name' self.mock_execute_sql.return_value = {'status': True, 'row_count': -1, 'tuples': [], 'fields': []} res = self.backend.delete_repo(repo=repo_name, force=True) self.assertEqual( self.mock_execute_sql.call_args[0][0], drop_schema_sql) self.assertEqual( self.mock_execute_sql.call_args[0][1][0], repo_name) self.assertEqual( self.mock_execute_sql.call_args[0][1][1], 'CASCADE') self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections) self.assertEqual(res, True) def test_delete_repo_no_cascade(self): drop_schema_sql = 'DROP SCHEMA %s %s' repo_name = 'repo_name' self.mock_execute_sql.return_value = {'status': True, 'row_count': -1, 'tuples': [], 'fields': []} res = self.backend.delete_repo(repo=repo_name, force=False) self.assertEqual( self.mock_execute_sql.call_args[0][0], drop_schema_sql) self.assertEqual( self.mock_execute_sql.call_args[0][1][0], repo_name) self.assertEqual( self.mock_execute_sql.call_args[0][1][1], '') self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections.called) self.assertEqual(res, True) def test_add_collaborator(self): privileges = ['SELECT', 'INSERT', 'UPDATE', 'DELETE', 'TRUNCATE', 'REFERENCES', 'TRIGGER', 'CREATE', 'CONNECT', 'TEMPORARY', 'EXECUTE', 'USAGE'] add_collab_query = ('BEGIN;' 'GRANT USAGE ON SCHEMA %s TO %s;' 'GRANT %s ON ALL TABLES IN SCHEMA %s TO %s;' 'ALTER DEFAULT PRIVILEGES IN SCHEMA %s ' 'GRANT %s ON TABLES TO %s;' 'COMMIT;' ) self.mock_execute_sql.return_value = {'status': True, 'row_count': -1, 'tuples': [], 'fields': []} product = itertools.product(self.good_nouns, self.good_nouns, privileges) # test every combo here. For now, don't test combined privileges for repo, receiver, privilege in product: params = (repo, receiver, privilege, repo, receiver, repo, privilege, receiver) res = self.backend.add_collaborator( repo=repo, collaborator=receiver, db_privileges=[privilege]) self.assertEqual( self.mock_execute_sql.call_args[0][0], add_collab_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 3) self.assertEqual(res, True) self.reset_mocks() def test_add_collaborator_concatenates_privileges(self): privileges = ['SELECT', 'USAGE'] repo = 'repo' receiver = 'receiver' self.mock_execute_sql.return_value = {'status': True, 'row_count': -1, 'tuples': [], 'fields': []} self.backend.add_collaborator(repo=repo, collaborator=receiver, db_privileges=privileges) # make sure that the privileges are passed as a string in params self.assertTrue( 'SELECT, USAGE' in self.mock_execute_sql.call_args[0][1]) def test_delete_collaborator(self): delete_collab_sql = ('BEGIN;' 'REVOKE ALL ON ALL TABLES IN SCHEMA %s ' 'FROM %s CASCADE;' 'REVOKE ALL ON SCHEMA %s FROM %s CASCADE;' 'ALTER DEFAULT PRIVILEGES IN SCHEMA %s ' 'REVOKE ALL ON TABLES FROM %s;' 'COMMIT;' ) self.mock_execute_sql.return_value = {'status': True, 'row_count': -1, 'tuples': [], 'fields': []} repo = 'repo_name' username = '******' params = (repo, username, repo, username, repo, username) res = self.backend.delete_collaborator( repo=repo, collaborator=username) self.assertEqual( self.mock_execute_sql.call_args[0][0], delete_collab_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 2) self.assertEqual(res, True) def test_create_table(self): repo = 'repo' table = 'table' params = [ {'column_name': 'id', 'data_type': 'integer'}, {'column_name': 'words', 'data_type': 'text'} ] expected_params = ('repo', 'table', 'id integer, words text') create_table_query = ('CREATE TABLE %s.%s (%s)') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': []} res = self.backend.create_table(repo, table, params) # checks repo, table, and all param values for injections self.assertEqual(self.mock_check_for_injections.call_count, 5) self.assertEqual(self.mock_validate_table_name.call_count, 1) final_query = self.mock_execute_sql.call_args[0][0] final_params = self.mock_execute_sql.call_args[0][1] self.assertEqual(final_query, create_table_query) self.assertEqual(final_params, expected_params) self.assertEqual(res, True) # create table test_repo.test_table (id integer, words text) def test_list_tables(self): repo = 'repo' list_tables_query = ('SELECT table_name FROM ' 'information_schema.tables ' 'WHERE table_schema = %s ' 'AND table_type = \'BASE TABLE\';') params = (repo,) # execute sql should return this: self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [ ('test_table',)], 'fields': [{'type': 1043, 'name': 'table_name'}]} # mocking out execute_sql's complicated return JSON mock_list_repos = self.create_patch( 'core.db.backend.pg.PGBackend.list_repos') mock_list_repos.return_value = [repo] res = self.backend.list_tables(repo) self.assertEqual( self.mock_execute_sql.call_args[0][0], list_tables_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(res, ['test_table']) def test_describe_table_without_detail(self): repo = 'repo' table = 'table' detail = False query = ("SELECT %s " "FROM information_schema.columns " "WHERE table_schema = %s and table_name = %s;") params = ('column_name, data_type', repo, table) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'id', u'integer'), (u'words', u'text')], 'fields': [ {'type': 1043, 'name': 'column_name'}, {'type': 1043, 'name': 'data_type'} ] } res = self.backend.describe_table(repo, table, detail) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(res, [(u'id', u'integer'), (u'words', u'text')]) def test_describe_table_query_in_detail(self): repo = 'repo' table = 'table' detail = True query = ("SELECT %s " "FROM information_schema.columns " "WHERE table_schema = %s and table_name = %s;") params = ('*', repo, table) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [ (u'id', u'integer'), (u'words', u'text'), ('foo', 'bar') ], 'fields': [ {'type': 1043, 'name': 'column_name'}, {'type': 1043, 'name': 'data_type'} ] } res = self.backend.describe_table(repo, table, detail) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual( res, [(u'id', u'integer'), (u'words', u'text'), ('foo', 'bar')]) def test_list_table_permissions(self): repo = 'repo' table = 'table' query = ("select privilege_type from " "information_schema.role_table_grants where table_schema=%s " "and table_name=%s and grantee=%s") params = ('repo', 'table', self.username) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [ (u'SELECT'), (u'UPDATE')], 'fields': [{'type': 1043, 'name': 'privilege_type'}] } self.backend.list_table_permissions(repo, table) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_delete_table(self): repo = 'repo_name' table = 'table_name' force = False expected_query = ('DROP TABLE %s.%s.%s %s') expected_params = ('username', 'repo_name', 'table_name', 'RESTRICT') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': []} res = self.backend.delete_table(repo, table, force) final_query = self.mock_execute_sql.call_args[0][0] final_params = self.mock_execute_sql.call_args[0][1] self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(self.mock_validate_table_name.call_count, 1) self.assertEqual(final_query, expected_query) self.assertEqual(final_params, expected_params) self.assertEqual(res, True) def test_list_views(self): repo = 'repo' list_views_query = ('SELECT table_name FROM information_schema.tables ' 'WHERE table_schema = %s ' 'AND table_type = \'VIEW\';') params = (repo,) # mocking out execute_sql's complicated return JSON self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [ ('test_view',)], 'fields': [{'type': 1043, 'name': 'view_name'}]} # list_views depends on list_repos, which is being mocked out mock_list_repos = self.create_patch( 'core.db.backend.pg.PGBackend.list_repos') mock_list_repos.return_value = [repo] res = self.backend.list_views(repo) self.assertEqual( self.mock_execute_sql.call_args[0][0], list_views_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(res, ['test_view']) def test_delete_view(self): repo = 'repo_name' view = 'view_name' force = False expected_query = ('DROP VIEW %s.%s.%s %s') expected_params = ('username', 'repo_name', 'view_name', 'RESTRICT') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': []} res = self.backend.delete_view(repo, view, force) final_query = self.mock_execute_sql.call_args[0][0] final_params = self.mock_execute_sql.call_args[0][1] self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(self.mock_validate_table_name.call_count, 1) self.assertEqual(final_query, expected_query) self.assertEqual(final_params, expected_params) self.assertEqual(res, True) def test_create_view(self): repo = 'repo_name' view = 'view_name' sql = 'SELECT * FROM table' expected_params = ('repo_name', 'view_name', 'SELECT * FROM table') create_view_query = ('CREATE VIEW %s.%s AS (%s)') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': []} res = self.backend.create_view(repo, view, sql) # checks repo and view for injections self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(self.mock_validate_table_name.call_count, 1) final_query = self.mock_execute_sql.call_args[0][0] final_params = self.mock_execute_sql.call_args[0][1] self.assertEqual(final_query, create_view_query) self.assertEqual(final_params, expected_params) self.assertEqual(res, True) def test_describe_view_without_detail(self): repo = 'repo_name' view = 'view_name' detail = False query = ("SELECT %s " "FROM information_schema.columns " "WHERE table_schema = %s and table_name = %s;") params = ('column_name, data_type', repo, view) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'id', u'integer'), (u'words', u'text')], 'fields': [ {'type': 1043, 'name': 'column_name'}, {'type': 1043, 'name': 'data_type'} ] } res = self.backend.describe_view(repo, view, detail) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(res, [(u'id', u'integer'), (u'words', u'text')]) def test_get_schema(self): self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'id', u'integer'), (u'words', u'text')], 'fields': [ {'type': 1043, 'name': 'column_name'}, {'type': 1043, 'name': 'data_type'} ] } repo = 'repo' table = 'table' get_schema_query = ('SELECT column_name, data_type ' 'FROM information_schema.columns ' 'WHERE table_name = %s ' 'AND table_schema = %s;' ) params = ('table', 'repo') self.backend.get_schema(repo, table) self.assertEqual( self.mock_execute_sql.call_args[0][0], get_schema_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(self.mock_validate_table_name.call_count, 1) def test_create_public_user_no_create_db(self): create_user_query = ('CREATE ROLE %s WITH LOGIN ' 'NOCREATEDB NOCREATEROLE NOCREATEUSER ' 'PASSWORD %s') username = PUBLIC_ROLE password = '******' self.backend.create_user(username, password, create_db=False) params = (username, password) mock_create_user_database = self.create_patch( 'core.db.backend.pg.PGBackend.create_user_database') # import pdb; pdb.set_trace() self.assertEqual( self.mock_execute_sql.call_args[0][0], create_user_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 1) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertFalse(mock_create_user_database.called) def test_create_normal_user_no_create_db(self): create_user_query = 'GRANT %s to %s' username = '******' password = '******' self.backend.create_user(username, password, create_db=False) params = (PUBLIC_ROLE, username) mock_create_user_database = self.create_patch( 'core.db.backend.pg.PGBackend.create_user_database') # import pdb; pdb.set_trace() self.assertEqual( self.mock_execute_sql.call_args[0][0], create_user_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertFalse(mock_create_user_database.called) def test_create_user_calls_create_db(self): username = '******' password = '******' mock_create_user_database = self.create_patch( 'core.db.backend.pg.PGBackend.create_user_database') self.backend.create_user( username=username, password=password, create_db=True) self.assertTrue(mock_create_user_database.called) def test_create_user_db(self): create_db_query_1 = 'CREATE DATABASE %s; ' create_db_query_2 = 'ALTER DATABASE %s OWNER TO %s; ' username = '******' self.backend.create_user_database(username) params_1 = (username,) params_2 = (username, username) call_args_1 = self.mock_execute_sql.call_args_list[0][0] self.assertEqual(call_args_1[0], create_db_query_1) self.assertEqual(call_args_1[1], params_1) call_args_2 = self.mock_execute_sql.call_args_list[1][0] self.assertEqual(call_args_2[0], create_db_query_2) self.assertEqual(call_args_2[1], params_2) self.assertEqual(self.mock_as_is.call_count, len(params_1 + params_2)) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_remove_user(self): query = 'DROP ROLE %s;' username = "******" params = (username,) self.backend.remove_user(username) self.assertEqual( self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_remove_database(self): # mock out list_all_users mock_list_all_users = self.create_patch( 'core.db.backend.pg.PGBackend.list_all_users') mock_list_all_users.return_value = ['tweedledee', 'tweedledum'] self.backend.remove_database(self.username) # revoke statement stuff revoke_query = 'REVOKE ALL ON DATABASE %s FROM %s;' revoke_params_1 = (self.username, 'tweedledee') revoke_params_2 = (self.username, 'tweedledum') self.assertEqual( self.mock_execute_sql.call_args_list[0][0][0], revoke_query) self.assertEqual( self.mock_execute_sql.call_args_list[0][0][1], revoke_params_1) self.assertEqual( self.mock_execute_sql.call_args_list[1][0][0], revoke_query) self.assertEqual( self.mock_execute_sql.call_args_list[1][0][1], revoke_params_2) # drop statement stuff drop_query = 'DROP DATABASE %s;' drop_params = (self.username,) self.assertEqual( self.mock_execute_sql.call_args_list[2][0][0], drop_query) self.assertEqual( self.mock_execute_sql.call_args_list[2][0][1], drop_params) self.assertEqual(self.mock_as_is.call_count, 5) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_change_password(self): query = 'ALTER ROLE %s WITH PASSWORD %s;' params = (self.username, self.password) self.backend.change_password(self.username, self.password) self.assertEqual( self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 1) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_list_collaborators(self): query = 'SELECT unnest(nspacl) FROM pg_namespace WHERE nspname=%s;' repo = 'repo_name' params = (repo, ) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [ ('al_carter=UC/al_carter',), ('foo_bar=U/al_carter',) ], 'fields': [{'type': 1033, 'name': 'unnest'}]} expected_result = [ {'username': '******', 'db_permissions': 'UC'}, {'username': '******', 'db_permissions': 'U'} ] res = self.backend.list_collaborators(repo) self.assertEqual( self.mock_execute_sql.call_args[0][0], query) self.assertEqual( self.mock_execute_sql.call_args[0][1], params) self.assertFalse(self.mock_as_is.called) self.assertEqual(res, expected_result) def test_list_all_users(self): query = 'SELECT usename FROM pg_catalog.pg_user WHERE usename != %s' params = (self.username,) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'delete_me_alpha_user',), (u'delete_me_beta_user',)], 'fields': [{'type': 19, 'name': 'usename'}] } res = self.backend.list_all_users() self.assertEqual( self.mock_execute_sql.call_args[0][0], query) self.assertEqual( self.mock_execute_sql.call_args[0][1], params) self.assertFalse(self.mock_as_is.called) self.assertEqual(res, ['delete_me_alpha_user', 'delete_me_beta_user']) def test_has_base_privilege(self): query = 'SELECT has_database_privilege(%s, %s);' privilege = 'CONNECT' params = (self.username, privilege) self.mock_execute_sql.return_value = {'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': []} res = self.backend.has_base_privilege( login=self.username, privilege=privilege) self.assertEqual( self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_repo_db_privilege(self): query = 'SELECT has_schema_privilege(%s, %s, %s);' repo = 'repo' privilege = 'CONNECT' params = (self.username, repo, privilege) self.mock_execute_sql.return_value = {'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': []} res = self.backend.has_repo_db_privilege( login=self.username, repo=repo, privilege=privilege) self.assertEqual( self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_table_privilege(self): query = 'SELECT has_table_privilege(%s, %s, %s);' table = 'table' privilege = 'CONNECT' params = (self.username, table, privilege) self.mock_execute_sql.return_value = {'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': []} res = self.backend.has_table_privilege( login=self.username, table=table, privilege=privilege) self.assertEqual( self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_column_privilege(self): query = 'SELECT has_column_privilege(%s, %s, %s, %s);' table = 'table' column = 'column' privilege = 'CONNECT' params = (self.username, table, column, privilege) self.mock_execute_sql.return_value = {'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': []} res = self.backend.has_column_privilege( login=self.username, table=table, column=column, privilege=privilege) self.assertEqual( self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) @patch('core.db.backend.pg.os.makedirs') @patch('core.db.backend.pg.os.remove') @patch('core.db.backend.pg.shutil.move') def test_export_table_with_header(self, *args): table_name = 'repo_name.table_name' table_name_prepped = 'SELECT * FROM %s' % table_name file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = True query = ('COPY (SELECT * FROM repo_name.table_name) ' 'TO STDOUT WITH CSV HEADER DELIMITER \',\';') self.backend.connection = Mock() mock_connection = self.backend.connection mock_mogrify = mock_connection.cursor.return_value.mogrify mock_mogrify.return_value = query mock_copy_expert = mock_connection.cursor.return_value.copy_expert with patch("__builtin__.open", mock_open()): self.backend.export_table(table_name, file_path, file_format, delimiter, header) mock_mogrify.assert_called_once_with( 'COPY (%s) TO STDOUT WITH %s %s DELIMITER %s;', (table_name_prepped, file_format, 'HEADER' if header else '', delimiter)) # Kind of a meaningless check since we have to mock the return value # of mogrify, but at least it ensures the result of mogrify is passed # into copy_expert as is. self.assertEqual(mock_copy_expert.call_args[0][0], query) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 4) self.assertEqual(self.mock_validate_table_name.call_count, 1) @patch('core.db.backend.pg.os.makedirs') @patch('core.db.backend.pg.os.remove') @patch('core.db.backend.pg.shutil.move') def test_export_table_with_no_header(self, *args): table_name = 'repo_name.table_name' table_name_prepped = 'SELECT * FROM %s' % table_name file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = False query = ('COPY (SELECT * FROM repo_name.table_name) ' 'TO STDOUT WITH CSV DELIMITER \',\';') self.backend.connection = Mock() mock_connection = self.backend.connection mock_mogrify = mock_connection.cursor.return_value.mogrify mock_mogrify.return_value = query mock_copy_expert = mock_connection.cursor.return_value.copy_expert with patch("__builtin__.open", mock_open()): self.backend.export_table(table_name, file_path, file_format, delimiter, header) mock_mogrify.assert_called_once_with( 'COPY (%s) TO STDOUT WITH %s %s DELIMITER %s;', (table_name_prepped, file_format, 'HEADER' if header else '', delimiter)) # Kind of a meaningless check since we have to mock the return value # of mogrify, but at least it ensures the result of mogrify is passed # into copy_expert as is. self.assertEqual(mock_copy_expert.call_args[0][0], query) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 4) self.assertEqual(self.mock_validate_table_name.call_count, 1) @patch('core.db.backend.pg.os.makedirs') @patch('core.db.backend.pg.os.remove') @patch('core.db.backend.pg.shutil.move') def test_export_view(self, *args): view_name = 'repo_name.view_name' view_name_prepped = 'SELECT * FROM %s' % view_name file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = True query = ('COPY (SELECT * FROM repo_name.view_name) ' 'TO STDOUT WITH CSV HEADER DELIMITER \',\';') self.backend.connection = Mock() mock_connection = self.backend.connection mock_mogrify = mock_connection.cursor.return_value.mogrify mock_mogrify.return_value = query mock_copy_expert = mock_connection.cursor.return_value.copy_expert with patch("__builtin__.open", mock_open()): self.backend.export_view(view_name, file_path, file_format, delimiter, header) mock_mogrify.assert_called_once_with( 'COPY (%s) TO STDOUT WITH %s %s DELIMITER %s;', (view_name_prepped, file_format, 'HEADER' if header else '', delimiter)) # Kind of a meaningless check since we have to mock the return value # of mogrify, but at least it ensures the result of mogrify is passed # into copy_expert as is. self.assertEqual(mock_copy_expert.call_args[0][0], query) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 4) self.assertEqual(self.mock_validate_table_name.call_count, 1) @patch('core.db.backend.pg.os.makedirs') @patch('core.db.backend.pg.os.remove') @patch('core.db.backend.pg.shutil.move') def _export_query_test_helper(self, *args, **kwargs): passed_query = kwargs['passed_query'] passed_query_cleaned = kwargs.get( 'passed_query_cleaned', passed_query) file_path = kwargs['file_path'] file_format = kwargs['file_format'] delimiter = kwargs['delimiter'] header = kwargs['header'] query = kwargs['query'] self.backend.connection = Mock() mock_connection = self.backend.connection mock_mogrify = mock_connection.cursor.return_value.mogrify mock_mogrify.return_value = query mock_copy_expert = mock_connection.cursor.return_value.copy_expert with patch("__builtin__.open", mock_open()): self.backend.export_query(passed_query, file_path, file_format, delimiter, header) mock_mogrify.assert_called_once_with( 'COPY (%s) TO STDOUT WITH %s %s DELIMITER %s;', (passed_query_cleaned, file_format, 'HEADER' if header else '', delimiter)) # Kind of a meaningless check since we have to mock the return value # of mogrify, but at least it ensures the result of mogrify is passed # into copy_expert as is. self.assertEqual(mock_copy_expert.call_args[0][0], query) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 2) def test_export_query_with_header(self, *args): self._export_query_test_helper( passed_query='myquery', file_path='file_path', file_format='CSV', delimiter=',', header=True, query=('COPY (myquery) ' 'TO STDOUT WITH CSV HEADER DELIMITER \',\';')) def test_export_query_with_no_header(self, *args): self._export_query_test_helper( passed_query='myquery', file_path='file_path', file_format='CSV', delimiter=',', header=False, query=('COPY (myquery) ' 'TO STDOUT WITH CSV DELIMITER \',\';')) def test_export_query_only_executes_text_before_semicolon(self, *args): self._export_query_test_helper( passed_query=' text before semicolon; text after; ', passed_query_cleaned='text before semicolon', file_path='file_path', file_format='CSV', delimiter=',', header=False, query=('COPY (text before semicolon) ' 'TO STDOUT WITH CSV DELIMITER \',\';')) def test_import_file_with_header(self): query = 'COPY %s FROM %s WITH %s %s DELIMITER %s ENCODING %s QUOTE %s;' table_name = 'user_name.repo_name.table_name' file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = True encoding = 'ISO-8859-1' quote_character = '"' params = (table_name, file_path, file_format, 'HEADER', delimiter, encoding, quote_character) self.backend.import_file(table_name, file_path, file_format, delimiter, header, encoding, quote_character) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections. call_count, 3) self.assertEqual(self.mock_validate_table_name.call_count, 1) def test_import_table_with_no_header(self): table_name = 'table_name' file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = False encoding = 'ISO-8859-1' quote_character = '"' params = (table_name, file_path, file_format, '', delimiter, encoding, quote_character) self.backend.import_file(table_name, file_path, file_format, delimiter, header, encoding, quote_character) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_import_file_w_dbtruck(self): # DBTruck is not tested for safety/security... At all. # The method does so little # that it doesn't even make much sense to test it. pass def test_can_user_access_rls_table(self): mock_settings = self.create_patch("core.db.backend.pg.settings") mock_settings.POLICY_SCHEMA = 'SCHEMA' mock_settings.POLICY_TABLE = 'TABLE' self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [ (True,), ]} username = '******' permissions = ['select', 'update'] expected_query = ( "SELECT exists(" "SELECT * FROM %s.%s where grantee=lower(%s) and (" "lower(policy_type)=lower(%s) or lower(policy_type)=lower(%s)" "))") expected_params = ('SCHEMA', 'TABLE', 'delete_me_user', 'select', 'update') self.backend.can_user_access_rls_table(username, permissions) self.assertEqual(self.mock_execute_sql.call_args[0][0], expected_query) self.assertEqual( self.mock_execute_sql.call_args[0][1], expected_params)
class SchemaListCreateDeleteShare(TestCase): """ Tests that items reach the execute_sql method in pg.py. Does not test execute_sql itself. """ def setUp(self): # some words to test out self.good_nouns = ['good', 'good_noun', 'good-noun'] # some words that shoudl throw validation errors self.bad_nouns = [ '_foo', 'foo_', '-foo', 'foo-', 'foo bar', 'injection;attack', ';injection', 'injection;', ] self.username = "******" self.password = "******" # mock the execute_sql function self.mock_execute_sql = self.create_patch( 'core.db.backend.pg.PGBackend.execute_sql') self.mock_execute_sql.return_value = True # mock the is_valid_noun_name, which checks for injection attacks self.mock_check_for_injections = self.create_patch( 'core.db.backend.pg.PGBackend._check_for_injections') # mock open connection, or else it will try to # create a real db connection self.mock_open_connection = self.create_patch( 'core.db.backend.pg.PGBackend.__open_connection__') # mock the psycopg2.extensions.AsIs - many of the pg.py methods use it # Its return value (side effect) is the call value self.mock_as_is = self.create_patch('core.db.backend.pg.AsIs') self.mock_as_is.side_effect = lambda x: x # create an instance of PGBackend self.backend = PGBackend(self.username, self.password, repo_base=self.username) def create_patch(self, name): # helper method for creating patches patcher = patch(name) thing = patcher.start() self.addCleanup(patcher.stop) return thing def reset_mocks(self): # clears the mock call arguments and sets their call counts to 0 self.mock_as_is.reset_mock() self.mock_execute_sql.reset_mock() self.mock_check_for_injections.reset_mock() # testing externally called methods in PGBackend def test_create_repo(self): create_repo_sql = 'CREATE SCHEMA IF NOT EXISTS %s AUTHORIZATION %s' reponame = 'reponame' self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } res = self.backend.create_repo(reponame) self.assertEqual(self.mock_execute_sql.call_args[0][0], create_repo_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1][0], reponame) self.assertEqual(self.mock_execute_sql.call_args[0][1][1], self.username) self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections.called) self.assertEqual(res, True) def test_list_repos(self): # the user is already logged in, so there's not much to be tested here # except that the arguments are passed correctly list_repo_sql = ('SELECT schema_name AS repo_name ' 'FROM information_schema.schemata ' 'WHERE schema_owner != %s') mock_settings = self.create_patch("core.db.backend.pg.settings") mock_settings.DATABASES = {'default': {'USER': '******'}} self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [('test_table', )], 'fields': [{ 'type': 1043, 'name': 'table_name' }] } params = (mock_settings.DATABASES['default']['USER'], ) res = self.backend.list_repos() self.assertEqual(self.mock_execute_sql.call_args[0][0], list_repo_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(res, ['test_table']) def test_delete_repo_happy_path_cascade(self): drop_schema_sql = 'DROP SCHEMA %s %s' repo_name = 'repo_name' self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } res = self.backend.delete_repo(repo=repo_name, force=True) self.assertEqual(self.mock_execute_sql.call_args[0][0], drop_schema_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1][0], repo_name) self.assertEqual(self.mock_execute_sql.call_args[0][1][1], 'CASCADE') self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections) self.assertEqual(res, True) def test_delete_repo_no_cascade(self): drop_schema_sql = 'DROP SCHEMA %s %s' repo_name = 'repo_name' self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } res = self.backend.delete_repo(repo=repo_name, force=False) self.assertEqual(self.mock_execute_sql.call_args[0][0], drop_schema_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1][0], repo_name) self.assertEqual(self.mock_execute_sql.call_args[0][1][1], None) self.assertTrue(self.mock_as_is.called) self.assertTrue(self.mock_check_for_injections.called) self.assertEqual(res, True) def test_add_collaborator(self): privileges = [ 'SELECT', 'INSERT', 'UPDATE', 'DELETE', 'TRUNCATE', 'REFERENCES', 'TRIGGER', 'CREATE', 'CONNECT', 'TEMPORARY', 'EXECUTE', 'USAGE' ] add_collab_query = ('BEGIN;' 'GRANT USAGE ON SCHEMA %s TO %s;' 'GRANT %s ON ALL TABLES IN SCHEMA %s TO %s;' 'ALTER DEFAULT PRIVILEGES IN SCHEMA %s ' 'GRANT %s ON TABLES TO %s;' 'COMMIT;') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } product = itertools.product(self.good_nouns, self.good_nouns, privileges) # test every combo here. For now, don't test combined privileges for repo, receiver, privilege in product: params = (repo, receiver, privilege, repo, receiver, repo, privilege, receiver) res = self.backend.add_collaborator(repo=repo, collaborator=receiver, privileges=[privilege]) self.assertEqual(self.mock_execute_sql.call_args[0][0], add_collab_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 3) self.assertEqual(res, True) self.reset_mocks() def test_add_collaborator_concatinates_privileges(self): privileges = ['SELECT', 'USAGE'] repo = 'repo' receiver = 'receiver' self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } self.backend.add_collaborator(repo=repo, collaborator=receiver, privileges=privileges) # make sure that the privileges are passed as a string in params self.assertTrue( 'SELECT, USAGE' in self.mock_execute_sql.call_args[0][1]) def test_delete_collaborator(self): delete_collab_sql = ('BEGIN;' 'REVOKE ALL ON ALL TABLES IN SCHEMA %s ' 'FROM %s CASCADE;' 'REVOKE ALL ON SCHEMA %s FROM %s CASCADE;' 'ALTER DEFAULT PRIVILEGES IN SCHEMA %s ' 'REVOKE ALL ON TABLES FROM %s;' 'COMMIT;') self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [], 'fields': [] } repo = 'repo_name' username = '******' params = (repo, username, repo, username, repo, username) res = self.backend.delete_collaborator(repo=repo, collaborator=username) self.assertEqual(self.mock_execute_sql.call_args[0][0], delete_collab_sql) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 2) self.assertEqual(res, True) def test_list_tables(self): repo = 'repo' list_tables_query = ('SELECT table_name FROM ' 'information_schema.tables ' 'WHERE table_schema = %s ' 'AND table_type = \'BASE TABLE\';') params = (repo, ) # execute sql should return this: self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [('test_table', )], 'fields': [{ 'type': 1043, 'name': 'table_name' }] } # mocking out execute_sql's complicated return JSON mock_list_repos = self.create_patch( 'core.db.backend.pg.PGBackend.list_repos') mock_list_repos.return_value = [repo] res = self.backend.list_tables(repo) self.assertEqual(self.mock_execute_sql.call_args[0][0], list_tables_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(res, ['test_table']) def test_list_views(self): repo = 'repo' list_views_query = ('SELECT table_name FROM information_schema.tables ' 'WHERE table_schema = %s ' 'AND table_type = \'VIEW\';') params = (repo, ) # mocking out execute_sql's complicated return JSON self.mock_execute_sql.return_value = { 'status': True, 'row_count': 1, 'tuples': [('test_view', )], 'fields': [{ 'type': 1043, 'name': 'view_name' }] } # list_views depends on list_repos, which is being mocked out mock_list_repos = self.create_patch( 'core.db.backend.pg.PGBackend.list_repos') mock_list_repos.return_value = [repo] res = self.backend.list_views(repo) self.assertEqual(self.mock_execute_sql.call_args[0][0], list_views_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertEqual(res, ['test_view']) def test_get_schema(self): self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'id', u'integer'), (u'words', u'text')], 'fields': [{ 'type': 1043, 'name': 'column_name' }, { 'type': 1043, 'name': 'data_type' }] } repo = 'repo' table = 'table' get_schema_query = ('SELECT column_name, data_type ' 'FROM information_schema.columns ' 'WHERE table_name = %s ' 'AND table_schema = %s;') params = ('table', 'repo') self.backend.get_schema(repo, table) self.assertEqual(self.mock_execute_sql.call_args[0][0], get_schema_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_check_for_injections.call_count, 2) def test_create_user_no_create_db(self): create_user_query = ('CREATE ROLE %s WITH LOGIN ' 'NOCREATEDB NOCREATEROLE NOCREATEUSER ' 'PASSWORD %s') username = '******' password = '******' self.backend.create_user(username, password, create_db=False) params = (username, password) mock_create_user_database = self.create_patch( 'core.db.backend.pg.PGBackend.create_user_database') self.assertEqual(self.mock_execute_sql.call_args[0][0], create_user_query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 1) self.assertEqual(self.mock_check_for_injections.call_count, 1) self.assertFalse(mock_create_user_database.called) def test_create_user_calls_create_db(self): username = '******' password = '******' mock_create_user_database = self.create_patch( 'core.db.backend.pg.PGBackend.create_user_database') self.backend.create_user(username=username, password=password, create_db=True) self.assertTrue(mock_create_user_database.called) def test_create_user_db(self): create_db_query_1 = 'CREATE DATABASE %s; ' create_db_query_2 = 'ALTER DATABASE %s OWNER TO %s; ' username = '******' self.backend.create_user_database(username) params_1 = (username, ) params_2 = (username, username) call_args_1 = self.mock_execute_sql.call_args_list[0][0] self.assertEqual(call_args_1[0], create_db_query_1) self.assertEqual(call_args_1[1], params_1) call_args_2 = self.mock_execute_sql.call_args_list[1][0] self.assertEqual(call_args_2[0], create_db_query_2) self.assertEqual(call_args_2[1], params_2) self.assertEqual(self.mock_as_is.call_count, len(params_1 + params_2)) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_remove_user(self): query = 'DROP ROLE %s;' username = "******" params = (username, ) self.backend.remove_user(username) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, len(params)) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_remove_database(self): # mock out list_all_users mock_list_all_users = self.create_patch( 'core.db.backend.pg.PGBackend.list_all_users') mock_list_all_users.return_value = ['tweedledee', 'tweedledum'] self.backend.remove_database(self.username) # revoke statement stuff revoke_query = 'REVOKE ALL ON DATABASE %s FROM %s;' revoke_params_1 = (self.username, 'tweedledee') revoke_params_2 = (self.username, 'tweedledum') self.assertEqual(self.mock_execute_sql.call_args_list[0][0][0], revoke_query) self.assertEqual(self.mock_execute_sql.call_args_list[0][0][1], revoke_params_1) self.assertEqual(self.mock_execute_sql.call_args_list[1][0][0], revoke_query) self.assertEqual(self.mock_execute_sql.call_args_list[1][0][1], revoke_params_2) # drop statement stuff drop_query = 'DROP DATABASE %s;' drop_params = (self.username, ) self.assertEqual(self.mock_execute_sql.call_args_list[2][0][0], drop_query) self.assertEqual(self.mock_execute_sql.call_args_list[2][0][1], drop_params) self.assertEqual(self.mock_as_is.call_count, 5) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_change_password(self): query = 'ALTER ROLE %s WITH PASSWORD %s;' params = (self.username, self.password) self.backend.change_password(self.username, self.password) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 1) self.assertEqual(self.mock_check_for_injections.call_count, 1) def test_list_collaborators(self): query = 'SELECT unnest(nspacl) FROM pg_namespace WHERE nspname=%s;' repo = 'repo_name' params = (repo, ) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [('al_carter=UC/al_carter', ), ('foo_bar=U/al_carter', )], 'fields': [{ 'type': 1033, 'name': 'unnest' }] } res = self.backend.list_collaborators(repo) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertFalse(self.mock_as_is.called) self.assertEqual(res, ['al_carter', 'foo_bar']) def test_list_all_users(self): query = 'SELECT usename FROM pg_catalog.pg_user WHERE usename != %s' params = (self.username, ) self.mock_execute_sql.return_value = { 'status': True, 'row_count': 2, 'tuples': [(u'delete_me_alpha_user', ), (u'delete_me_beta_user', )], 'fields': [{ 'type': 19, 'name': 'usename' }] } res = self.backend.list_all_users() self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertFalse(self.mock_as_is.called) self.assertEqual(res, ['delete_me_alpha_user', 'delete_me_beta_user']) def test_has_base_privilege(self): query = 'SELECT has_database_privilege(%s, %s);' privilege = 'CONNECT' params = (self.username, privilege) self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': [] } res = self.backend.has_base_privilege(login=self.username, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_repo_privilege(self): query = 'SELECT has_schema_privilege(%s, %s, %s);' repo = 'repo' privilege = 'CONNECT' params = (self.username, repo, privilege) self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': [] } res = self.backend.has_repo_privilege(login=self.username, repo=repo, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_table_privilege(self): query = 'SELECT has_table_privilege(%s, %s, %s);' table = 'table' privilege = 'CONNECT' params = (self.username, table, privilege) self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': [] } res = self.backend.has_table_privilege(login=self.username, table=table, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_has_column_privilege(self): query = 'SELECT has_column_privilege(%s, %s, %s, %s);' table = 'table' column = 'column' privilege = 'CONNECT' params = (self.username, table, column, privilege) self.mock_execute_sql.return_value = { 'status': True, 'row_count': -1, 'tuples': [[True]], 'fields': [] } res = self.backend.has_column_privilege(login=self.username, table=table, column=column, privilege=privilege) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 0) self.assertEqual(res, True) def test_export_table_with_header(self): query = 'COPY %s TO %s WITH %s %s DELIMITER %s;' table_name = 'user_name.repo_name.table_name' file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = True params = (table_name, file_path, file_format, 'HEADER', delimiter) self.backend.export_table(table_name, file_path, file_format, delimiter, header) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 4) def test_export_table_with_no_header(self): table_name = 'table_name' file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = False params = (table_name, file_path, file_format, '', delimiter) self.backend.export_table(table_name, file_path, file_format, delimiter, header) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_export_query_with_header(self): query = 'COPY (%s) TO %s WITH %s %s DELIMITER %s;' passed_query = 'myquery' file_path = 'file_path' file_format = 'CSV' delimiter = ',' header = True params = (passed_query, file_path, file_format, 'HEADER', delimiter) self.backend.export_query(passed_query, file_path, file_format, delimiter, header) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 2) def test_export_query_with_no_header(self): passed_query = 'myquery' file_path = 'file_path' file_format = 'CSV' delimiter = ',' header = False params = (passed_query, file_path, file_format, '', delimiter) self.backend.export_query(passed_query, file_path, file_format, delimiter, header) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_export_query_only_executes_text_before_semicolon(self): passed_query = ' text before semicolon; text after; ' file_path = 'file_path' file_format = 'CSV' delimiter = ',' header = False passed_query_cleaned = ' text before semicolon' params = (passed_query_cleaned, file_path, file_format, '', delimiter) self.backend.export_query(passed_query, file_path, file_format, delimiter, header) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_import_file_with_header(self): query = 'COPY %s FROM %s WITH %s %s DELIMITER %s ENCODING %s QUOTE %s;' table_name = 'user_name.repo_name.table_name' file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = True encoding = 'ISO-8859-1' quote_character = '"' params = (table_name, file_path, file_format, 'HEADER', delimiter, encoding, quote_character) self.backend.import_file(table_name, file_path, file_format, delimiter, header, encoding, quote_character) self.assertEqual(self.mock_execute_sql.call_args[0][0], query) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) self.assertEqual(self.mock_as_is.call_count, 3) self.assertEqual(self.mock_check_for_injections.call_count, 5) def test_import_table_with_no_header(self): table_name = 'table_name' file_path = 'file_path' file_format = 'file_format' delimiter = ',' header = False encoding = 'ISO-8859-1' quote_character = '"' params = (table_name, file_path, file_format, '', delimiter, encoding, quote_character) self.backend.import_file(table_name, file_path, file_format, delimiter, header, encoding, quote_character) self.assertEqual(self.mock_execute_sql.call_args[0][1], params) def test_import_file_w_dbtruck(self): # DBTruck is not tested for safety/security... At all. # The method does so little # that it doesn't even make much sense to test it. pass