-
Notifications
You must be signed in to change notification settings - Fork 215
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: SQLAlchemy-based MetaDataProvider (#516)
- Loading branch information
Showing
9 changed files
with
172 additions
and
44 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -127,3 +127,6 @@ yarn-error.log* | |
|
||
# local testing sql files | ||
test*.sql | ||
|
||
# sqlite database generated by test cases | ||
*.db |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,46 @@ | ||
import logging | ||
from typing import List | ||
|
||
from sqlalchemy import MetaData, Table, create_engine, inspect, make_url | ||
from sqlalchemy.exc import NoSuchModuleError, NoSuchTableError, OperationalError | ||
|
||
from sqllineage.core.metadata_provider import MetaDataProvider | ||
from sqllineage.exceptions import MetaDataProviderException | ||
|
||
|
||
logger = logging.getLogger(__name__) | ||
|
||
|
||
class SQLAlchemyMetaDataProvider(MetaDataProvider): | ||
""" | ||
SQLAlchemyMetaDataProvider queries metadata from database using SQLAlchemy | ||
""" | ||
|
||
def __init__(self, url: str): | ||
super().__init__() | ||
self.metadata_obj = MetaData() | ||
try: | ||
self.engine = create_engine(url) | ||
except NoSuchModuleError as e: | ||
u = make_url(url) | ||
raise MetaDataProviderException( | ||
f"SQLAlchemy dialect driver {u.drivername} is not installed correctly" | ||
) from e | ||
try: | ||
self.engine.connect() | ||
except OperationalError as e: | ||
raise MetaDataProviderException(f"Could not connect to {url}") from e | ||
|
||
def _get_table_columns(self, schema: str, table: str, **kwargs) -> List[str]: | ||
columns = [] | ||
if inspect(self.engine).has_schema(schema): | ||
try: | ||
sqlalchemy_table = Table( | ||
table, self.metadata_obj, schema=schema, autoload_with=self.engine | ||
) | ||
columns = [c.name for c in sqlalchemy_table.columns] | ||
except NoSuchTableError: | ||
logger.warning("%s does not exist in database %s", table, schema) | ||
else: | ||
logger.warning("Schema %s does not exist", schema) | ||
return columns |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,26 @@ | ||
import os | ||
|
||
import pytest | ||
|
||
from sqllineage.core.metadata.sqlalchemy import SQLAlchemyMetaDataProvider | ||
from sqllineage.exceptions import MetaDataProviderException | ||
|
||
|
||
def test_sqlalchemy_metadata_provider_connection_fail(): | ||
# connect to a directory as sqlite db, which is not possible. Simulate connection failure | ||
with pytest.raises(MetaDataProviderException): | ||
SQLAlchemyMetaDataProvider(f"sqlite:///{os.path.dirname(__file__)}") | ||
|
||
|
||
def test_sqlalchemy_metadata_provider_driver_not_install(): | ||
# use an unknown driver to connect. Simulate driver not installed | ||
with pytest.raises(MetaDataProviderException): | ||
SQLAlchemyMetaDataProvider("sqlite+unknown_driver:///:memory:") | ||
|
||
|
||
def test_sqlalchemy_metadata_provider_query_fail(): | ||
provider = SQLAlchemyMetaDataProvider("sqlite:///:memory:") | ||
assert ( | ||
provider._get_table_columns("non_existing_schema", "non_existing_table") == [] | ||
) | ||
assert provider._get_table_columns("main", "non_existing_table") == [] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters