Skip to content

Commit

Permalink
added privatekey checker and ensure backwards compatibility
Browse files Browse the repository at this point in the history
  • Loading branch information
iireland-ii committed Nov 12, 2024
1 parent fdc1634 commit 639a330
Showing 1 changed file with 11 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@ def __init__(
:param warehouse: Name of the warehouse to use.
:param login_timeout: Timeout in seconds for login. By default, 60 seconds.
:param private_key_file: Location of private key
mutually exclusive to password, if key_file is provided this auth method will be used.
:param private_key_file_pwd: Password for private key file
"""

Expand Down Expand Up @@ -292,12 +293,10 @@ def _fetch_data(
if not query:
return df
try:
# Create a new connection with every run
conn = self._snowflake_connector(
connect_params={
# Build up param connection
connect_params={
"user": self.user,
"account": self.account,
"password": self.api_key,
"private_key_file": self.private_key_file,
"private_key_file_pwd": self.private_key_file_pwd,
"database": self.database,
Expand All @@ -306,6 +305,14 @@ def _fetch_data(
"role": self.role,
"login_timeout": self.login_timeout
}

# Check if private key has been provided
if self.private_key_file is None:
connect_params["password"] = self.api_key.resolve_value()

# Create a new connection with every run
conn = self._snowflake_connector(
connect_params=connect_params
)
if conn is None:
return df
Expand Down

0 comments on commit 639a330

Please sign in to comment.