class BigQueryTest(unittest.TestCase): def setUp(self): self.project_id = os.getenv("PROJECT_ID") self.dataset_id = os.getenv("DATASET_ID", "test_dataset") self.table_id = os.getenv("TABLE_ID", "test_table") self.view_id = os.getenv("VIEW_ID", "test_view") if self.project_id is None: print("PROJECT_ID is not defined.") sys.exit(1) self.bq = BigQuery(self.project_id) if self.bq.exists_dataset(self.dataset_id): self.bq.drop_dataset(self.dataset_id, delete_contents=True) self.bq.create_dataset(self.dataset_id) self.bq.dataset_id = self.dataset_id # Set default datasetId def TearDown(self): self.bq.drop_dataset(self.dataset_id, delete_contents=True) def test_error(self): query = "SELECT * FROM " + self.dataset_id + "." + self.table_id with self.assertRaises(NotFoundError): self.bq.create_view(self.view_id, query) schema = [ {"name": "id", "type": "INTEGER", "mode": "REQUIRED"}, {"name": "name", "type": "STRING", "mode": "REQUIRED"}, { "name": "birth", "type": "RECORD", "mode": "NULLABLE", "fields": [ {"name": "year", "type": "INTEGER", "mode": "REQUIRED"}, {"name": "month", "type": "INTEGER", "mode": "REQUIRED"}, {"name": "day", "type": "INTEGER", "mode": "REQUIRED"}, ], }, {"name": "url", "type": "STRING", "mode": "REPEATED"}, ] with self.assertRaises(NotFoundError): self.bq.create_table(self.table_id, schema=schema, dataset_id="not_found_dataset") self.bq.create_table(self.table_id, schema=schema) # "Schema field shouldn't be used as input with a view" with self.assertRaises(Http4xxError): self.bq.create_view(self.view_id, query, schema=schema) def test_normal_with_args(self): print("exists table?") if self.bq.exists_table(project_id=self.project_id, dataset_id=self.dataset_id, table_id=self.table_id): print("exists") print("drop table") res = self.bq.drop_table(project_id=self.project_id, dataset_id=self.dataset_id, table_id=self.table_id) self.assertTrue(bool(res)) else: print("no exists") print("create table") schema = [ {"name": "id", "type": "INTEGER", "mode": "REQUIRED"}, {"name": "name", "type": "STRING", "mode": "REQUIRED"}, { "name": "birth", "type": "RECORD", "mode": "NULLABLE", "fields": [ {"name": "year", "type": "INTEGER", "mode": "REQUIRED"}, {"name": "month", "type": "INTEGER", "mode": "REQUIRED"}, {"name": "day", "type": "INTEGER", "mode": "REQUIRED"}, ], }, {"name": "url", "type": "STRING", "mode": "REPEATED"}, ] expiration_time = str(int(time.time()) + 86400) + "000" res = self.bq.create_table( project_id=self.project_id, dataset_id=self.dataset_id, table_id=self.table_id, schema=schema, description="Description", expiration_time=expiration_time, friendly_name="Friendly Name", ) self.assertTrue(bool(res)) print("info table") res = self.bq.info_table(project_id=self.project_id, dataset_id=self.dataset_id, table_id=self.table_id) self.assertEqual("Description", res["description"]) self.assertEqual(expiration_time, res["expirationTime"]) self.assertEqual("Friendly Name", res["friendlyName"]) pprint(res) print("create view") query = "SELECT * FROM " + self.dataset_id + "." + self.table_id res = self.bq.create_view( project_id=self.project_id, dataset_id=self.dataset_id, table_id=self.view_id, query=query ) self.assertTrue(bool(res)) print("show tables") res = self.bq.show_tables(project_id=self.project_id, dataset_id=self.dataset_id, max_results=1) self.assertIn(self.table_id, res) self.assertIn(self.view_id, res) print("\n".join(res)) print("drop view") res = self.bq.drop_table(self.view_id) self.assertFalse(bool(res)) print("drop table") res = self.bq.drop_table(self.table_id) self.assertFalse(bool(res)) def test_normal(self): print("exists table?") if self.bq.exists_table(self.table_id): print("exists") print("drop table") res = self.bq.drop_table(self.table_id) self.assertTrue(bool(res)) else: print("no exists") print("create table") schema = [ {"name": "id", "type": "INTEGER", "mode": "REQUIRED"}, {"name": "name", "type": "STRING", "mode": "REQUIRED"}, { "name": "birth", "type": "RECORD", "mode": "NULLABLE", "fields": [ {"name": "year", "type": "INTEGER", "mode": "REQUIRED"}, {"name": "month", "type": "INTEGER", "mode": "REQUIRED"}, {"name": "day", "type": "INTEGER", "mode": "REQUIRED"}, ], }, {"name": "url", "type": "STRING", "mode": "REPEATED"}, ] res = self.bq.create_table(self.table_id, schema=schema) self.assertTrue(bool(res)) print("create exists table") res = self.bq.create_table(self.table_id, schema=schema) self.assertFalse(bool(res)) print("info table") res = self.bq.info_table(self.table_id) self.assertTrue(bool(res)) pprint(res) print("create view") query = "SELECT * FROM " + self.dataset_id + "." + self.table_id self.bq.create_view(self.view_id, query) self.assertTrue(bool(res)) print("info view") res = self.bq.info_table(self.view_id) self.assertTrue(bool(res)) pprint(res) print("show tables") res = self.bq.show_tables() self.assertIn(self.table_id, res) self.assertIn(self.view_id, res) print("\n".join(res)) print("drop view") res = self.bq.drop_table(self.view_id) self.assertFalse(bool(res)) print("drop table") res = self.bq.drop_table(self.table_id) self.assertFalse(bool(res)) print("drop no exists table") res = self.bq.drop_table(self.table_id) self.assertFalse(bool(res))