yujunjun/OpenMetadata
1
0
Fork 0
mirror of https://github.com/open-metadata/OpenMetadata.git synced 2025-10-27 16:55:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix #3095: Added support for private key (#3408)

Browse Source 
* ISSUE-3095: Add support for private key in snowflake

rebase main branch

* Fix #3095: Added support for private key

* Fix #3095: Added support for private key and review changes

* Fix #3095: Added support for private key and review changes

* changed direct assignment of dictionary

Co-authored-by: Onkar Ravgan <onkarravgan@Onkars-MacBook-Pro.local>
This commit is contained in:
Onkar Ravgan 2022-03-15 18:19:02 +05:30 committed by GitHub
parent c860177a49
commit ad04136847
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 32 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ingestion/setup.py
View File

@ -109,7 +109,7 @@ plugins: Dict[str, Set[str]] = {
"psycopg2-binary",
"GeoAlchemy2",
},
"snowflake": {"snowflake-sqlalchemy<=1.3.2"},
"snowflake": {"snowflake-sqlalchemy<=1.3.2", "cryptography"},
"snowflake-usage": {"snowflake-sqlalchemy<=1.3.2"},
"sample-entity": {"faker~=8.1.1"},
"superset": {},

17
ingestion/src/metadata/ingestion/source/snowflake.py
View File

@ -9,8 +9,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import os
from typing import Optional
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import dsa, rsa
from snowflake.sqlalchemy.custom_types import VARIANT
from snowflake.sqlalchemy.snowdialect import SnowflakeDialect, ischema_names
from sqlalchemy.engine import reflection
@ -57,6 +61,19 @@ class SnowflakeConfig(SQLConnectionConfig):
class SnowflakeSource(SQLSource):
def __init__(self, config, metadata_config, ctx):
if config.connect_args.get("private_key"):
private_key = config.connect_args["private_key"]
p_key = serialization.load_pem_private_key(
bytes(private_key, "utf-8"),
password=os.environ["SNOWFLAKE_PRIVATE_KEY_PASSPHRASE"].encode(),
backend=default_backend(),
)
pkb = p_key.private_bytes(
encoding=serialization.Encoding.DER,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
)
config.connect_args["private_key"] = pkb
super().__init__(config, metadata_config, ctx)
def fetch_sample_data(self, schema: str, table: str) -> Optional[TableData]:

15
ingestion/src/metadata/ingestion/source/sql_source.py
View File

@ -11,6 +11,7 @@
"""
Generic source to build SQL connectors.
"""
import copy
import json
import logging
import re
@ -71,6 +72,16 @@ def _get_table_description(schema: str, table: str, inspector: Inspector) -> str
return description
def _get_private_key_config(config: SQLConnectionConfig) -> SQLConnectionConfig:
new_config = copy.deepcopy(config)
if new_config.connect_args.get("private_key"):
new_config.connect_args["private_key"] = str(
new_config.connect_args["private_key"]
)
return new_config
return config
class SQLSource(Source[OMetaDatabaseAndTable]):
"""
Source Connector implementation to extract
@ -87,7 +98,9 @@ class SQLSource(Source[OMetaDatabaseAndTable]):
super().__init__(ctx)
self.config = config
self.metadata_config = metadata_config
self.service = get_database_service_or_create(config, metadata_config)
self.service = get_database_service_or_create(
_get_private_key_config(config), metadata_config
)
self.metadata = OpenMetadata(metadata_config)
self.status = SQLSourceStatus()
self.sql_config = self.config