def test_scope(self):
config = ConfigFactory.from_dict({})
transformer = BigqueryUsageTransformer()
transformer.init(config)
self.assertEqual(transformer.get_scope(), 'transformer.bigquery_usage')
def test_transform_function(self):
# type: () -> None
config = ConfigFactory.from_dict({})
transformer = BigqueryUsageTransformer()
transformer.init(config)
key = TableColumnUsageTuple(
database=TestBigQueryUsageTransform.DATABASE,
cluster=TestBigQueryUsageTransform.CLUSTER,
schema=TestBigQueryUsageTransform.DATASET,
table=TestBigQueryUsageTransform.TABLE,
column=TestBigQueryUsageTransform.COLUMN,
email=TestBigQueryUsageTransform.EMAIL)
t1 = (key, TestBigQueryUsageTransform.READ_COUNT)
xformed = transformer.transform(t1)
self.assertIsInstance(xformed, TableColumnUsage)
self.assertEqual(len(xformed.col_readers), 1)
col_reader = xformed.col_readers[0]
self.assertEqual(col_reader.cluster,
TestBigQueryUsageTransform.CLUSTER)
self.assertEqual(col_reader.database,
TestBigQueryUsageTransform.DATABASE)
self.assertEqual(col_reader.schema, TestBigQueryUsageTransform.DATASET)
self.assertEqual(col_reader.table, TestBigQueryUsageTransform.TABLE)
self.assertEqual(col_reader.column, TestBigQueryUsageTransform.COLUMN)
self.assertEqual(col_reader.user_email,
TestBigQueryUsageTransform.EMAIL)
self.assertEqual(col_reader.read_count,
TestBigQueryUsageTransform.READ_COUNT)
def test_transform_function(self) -> None:
config = ConfigFactory.from_dict({})
transformer = BigqueryUsageTransformer()
transformer.init(config)
key = TableColumnUsageTuple(database=TestBigQueryUsageTransform.DATABASE,
cluster=TestBigQueryUsageTransform.CLUSTER,
schema=TestBigQueryUsageTransform.DATASET,
table=TestBigQueryUsageTransform.TABLE,
column=TestBigQueryUsageTransform.COLUMN,
email=TestBigQueryUsageTransform.EMAIL)
t1 = (key, TestBigQueryUsageTransform.READ_COUNT)
xformed = transformer.transform(t1)
assert xformed is not None
self.assertIsInstance(xformed, TableColumnUsage)
col_readers = list(xformed.col_readers)
self.assertEqual(len(col_readers), 1)
col_reader = col_readers[0]
self.assertEqual(col_reader.start_label, 'Table')
self.assertEqual(col_reader.start_key, TestBigQueryUsageTransform.TABLE_KEY)
self.assertEqual(col_reader.user_email, TestBigQueryUsageTransform.EMAIL)
self.assertEqual(col_reader.read_count, TestBigQueryUsageTransform.READ_COUNT)
def create_transformer(metadata_type):
if metadata_type == MetadataType.DSL:
return NoopTransformer()
elif metadata_type == MetadataType.USAGE:
return BigqueryUsageTransformer()
else:
raise ValueError('Invalid metadata_type')
def run_bq_tu_job(job_name):
#where_clause_suffix = " "
gcloud_project = "peya-data-pocs"
#label_filter = ""
tmp_folder = '/var/tmp/amundsen/{job_name}'.format(job_name=job_name)
node_files_folder = '{tmp_folder}/nodes'.format(tmp_folder=tmp_folder)
relationship_files_folder = '{tmp_folder}/relationships'.format(tmp_folder=tmp_folder)
bq_usage_extractor = BigQueryTableUsageExtractor()
csv_loader = FsNeo4jCSVLoader()
task = DefaultTask(extractor=bq_usage_extractor,
loader=csv_loader,
transformer=BigqueryUsageTransformer())
job_config = ConfigFactory.from_dict({
'extractor.bigquery_table_usage.{}'.format(BigQueryTableUsageExtractor.PROJECT_ID_KEY):
gcloud_project,
'loader.filesystem_csv_neo4j.{}'.format(FsNeo4jCSVLoader.NODE_DIR_PATH):
node_files_folder,
'loader.filesystem_csv_neo4j.{}'.format(FsNeo4jCSVLoader.RELATION_DIR_PATH):
relationship_files_folder,
'loader.filesystem_csv_neo4j.{}'.format(FsNeo4jCSVLoader.SHOULD_DELETE_CREATED_DIR):
True,
'publisher.neo4j.{}'.format(neo4j_csv_publisher.NODE_FILES_DIR):
node_files_folder,
'publisher.neo4j.{}'.format(neo4j_csv_publisher.RELATION_FILES_DIR):
relationship_files_folder,
'publisher.neo4j.{}'.format(neo4j_csv_publisher.NEO4J_END_POINT_KEY):
neo4j_endpoint,
'publisher.neo4j.{}'.format(neo4j_csv_publisher.NEO4J_USER):
neo4j_user,
'publisher.neo4j.{}'.format(neo4j_csv_publisher.NEO4J_PASSWORD):
neo4j_password,
'publisher.neo4j.{}'.format(neo4j_csv_publisher.JOB_PUBLISH_TAG):
'unique_tag', # should use unique tag here like {ds}
})
job = DefaultJob(conf=job_config,
task=task,
publisher=Neo4jCsvPublisher())
job.launch()
def create_bq_job(metadata_type, gcloud_project):
tmp_folder = f'/var/tmp/amundsen/{metadata_type}'
node_files_folder = f'{tmp_folder}/nodes'
relationship_files_folder = f'{tmp_folder}/relationships'
bq_usage_extractor = BigQueryTableUsageExtractor()
csv_loader = FsNeo4jCSVLoader()
task = DefaultTask(extractor=bq_usage_extractor,
loader=csv_loader,
transformer=BigqueryUsageTransformer())
job_config = ConfigFactory.from_dict({
f'extractor.bigquery_table_usage.{BigQueryTableUsageExtractor.PROJECT_ID_KEY}':
gcloud_project,
f'loader.filesystem_csv_neo4j.{FsNeo4jCSVLoader.NODE_DIR_PATH}':
node_files_folder,
f'loader.filesystem_csv_neo4j.{FsNeo4jCSVLoader.RELATION_DIR_PATH}':
relationship_files_folder,
f'loader.filesystem_csv_neo4j.{FsNeo4jCSVLoader.SHOULD_DELETE_CREATED_DIR}':
True,
f'publisher.neo4j.{neo4j_csv_publisher.NODE_FILES_DIR}':
node_files_folder,
f'publisher.neo4j.{neo4j_csv_publisher.RELATION_FILES_DIR}':
relationship_files_folder,
f'publisher.neo4j.{neo4j_csv_publisher.NEO4J_END_POINT_KEY}':
neo4j_endpoint,
f'publisher.neo4j.{neo4j_csv_publisher.NEO4J_USER}':
neo4j_user,
f'publisher.neo4j.{neo4j_csv_publisher.NEO4J_PASSWORD}':
neo4j_password,
f'publisher.neo4j.{neo4j_csv_publisher.JOB_PUBLISH_TAG}':
'unique_tag', # should use unique tag here like {ds}
})
job = DefaultJob(conf=job_config, task=task, publisher=Neo4jCsvPublisher())
return job