yujunjun/OpenMetadata
1
0
Fork 0
mirror of https://github.com/open-metadata/OpenMetadata.git synced 2025-09-08 00:18:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Pipeline code refactored

Browse Source
This commit is contained in:
Ayush Shah 2021-08-26 01:17:01 +05:30
parent bd40ae45ce
commit 2d1ad58e1b
1 changed files with 6 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
ingestion/src/metadata/ingestion/source/bigquery_usage.py
View File

@ -19,9 +19,9 @@ from metadata.ingestion.models.table_queries import TableQuery
from sqlalchemy.engine import create_engine from sqlalchemy.engine import create_engine
from sqlalchemy.inspection import inspect from sqlalchemy.inspection import inspect
from metadata.ingestion.ometa.auth_provider import MetadataServerConfig from metadata.ingestion.ometa.auth_provider import MetadataServerConfig
from metadata.ingestion.source.sql_alchemy_helper import SQLSourceStatus
from metadata.ingestion.api.source import Source, SourceStatus from metadata.ingestion.api.source import Source, SourceStatus
from typing import Iterator, Union, Dict, Any, Iterable from typing import Dict, Any, Iterable
from metadata.ingestion.source.sql_alchemy_helper import SQLSourceStatus
from metadata.utils.helpers import get_start_and_end from metadata.utils.helpers import get_start_and_end
from metadata.ingestion.source.bigquery import BigQueryConfig from metadata.ingestion.source.bigquery import BigQueryConfig
@ -76,8 +76,8 @@ class BigqueryUsageSource(Source):
super().__init__(ctx) super().__init__(ctx)
start, end = get_start_and_end(config.duration) start, end = get_start_and_end(config.duration)
self.project_id = config.project_id self.project_id = config.project_id
self.alchemy_helper = create_engine(self.get_connection_url(), **config.options) self.engine = create_engine(self.get_connection_url(), **config.options).connect()
inspector = inspect(self.alchemy_helper) inspector = inspect(self.engine)
for schema in inspector.get_schema_names(): for schema in inspector.get_schema_names():
dataset_name = schema dataset_name = schema
self.sql_stmt = BigqueryUsageSource.SQL_STATEMENT.format( self.sql_stmt = BigqueryUsageSource.SQL_STATEMENT.format(
@ -86,8 +86,6 @@ class BigqueryUsageSource(Source):
end_date=end, end_date=end,
dataset=dataset_name dataset=dataset_name
) )
self._extract_iter: Union[None, Iterator] = None
self._database = 'bigquery'
self.status = SQLSourceStatus() self.status = SQLSourceStatus()
def get_connection_url(self): def get_connection_url(self):
@ -111,7 +109,7 @@ class BigqueryUsageSource(Source):
:return: :return:
""" """
try: try:
rows = self.alchemy_helper.connect().execute(self.sql_stmt).fetchall() rows = self.engine.execute(self.sql_stmt).fetchall()
for row in rows: for row in rows:
yield row yield row
except Exception as err: except Exception as err:
@ -125,7 +123,7 @@ class BigqueryUsageSource(Source):
yield tq yield tq
def close(self): def close(self):
self.alchemy_helper.connect().close() self.engine.close()
def get_status(self) -> SourceStatus: def get_status(self) -> SourceStatus:
return self.status return self.status