Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

adding privatekey auth param #1190

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open

adding privatekey auth param #1190

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
25a7992
adding privatekey auth param
iireland-ii Nov 14, 2024
f2c15d1
Merge branch 'main' into snowflake-privatekey-update
iireland-ii Nov 14, 2024
cd2658f
removed whitespace
iireland-ii Nov 14, 2024
adbfef5
updating linting
iireland-ii Nov 14, 2024
63a2eeb
indenting
iireland-ii Nov 14, 2024
8de2b06
final lint updates
iireland-ii Nov 14, 2024
3f7a045
Merge branch 'main' into snowflake-privatekey-update
iireland-ii Nov 14, 2024
0492dc4
docstring update - key-pair info
iireland-ii Nov 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 24 additions & 11 deletions ...ke/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,8 @@ def __init__(
user: str,
account: str,
api_key: Secret = Secret.from_env_var("SNOWFLAKE_API_KEY"), # noqa: B008
private_key_file: Optional[str] = None,
private_key_file_pwd: Optional[str] = None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this should be using the Secret - check above how it was done for api_key

database: Optional[str] = None,
db_schema: Optional[str] = None,
warehouse: Optional[str] = None,
Expand All @@ -82,11 +84,16 @@ def __init__(
:param db_schema: Name of the schema to use.
: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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this docstring is very confusing - what's key_file? can you maybe rephrase it?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I chose to match the snowflake connector params to keep it consistent with the snowflake documentation, see sample code here:

I've added some further comments to the docstring to include a note to the following link:

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can't see your changes - also my comment was not about the params names themselves, rather the docstring explaining what they are used for.

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
"""

self.user = user
self.account = account
self.api_key = api_key
self.private_key_file = private_key_file
self.private_key_file_pwd = private_key_file_pwd
self.database = database
self.db_schema = db_schema
self.warehouse = warehouse
Expand All @@ -104,6 +111,8 @@ def to_dict(self) -> Dict[str, Any]:
user=self.user,
account=self.account,
api_key=self.api_key.to_dict(),
private_key_file=self.private_key_file,
private_key_file_pwd=self.private_key_file_pwd,
database=self.database,
db_schema=self.db_schema,
warehouse=self.warehouse,
Expand Down Expand Up @@ -275,18 +284,22 @@ def _fetch_data(
if not query:
return df
try:
# Build up param connection
connect_params = {
"user": self.user,
"account": self.account,
"private_key_file": self.private_key_file,
"private_key_file_pwd": self.private_key_file_pwd,
"database": self.database,
"schema": self.db_schema,
"warehouse": self.warehouse,
"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={
"user": self.user,
"account": self.account,
"password": self.api_key.resolve_value(),
"database": self.database,
"schema": self.db_schema,
"warehouse": self.warehouse,
"login_timeout": self.login_timeout,
}
)
conn = self._snowflake_connector(connect_params=connect_params)
if conn is None:
return df
except (ForbiddenError, ProgrammingError) as e:
Expand Down
6 changes: 6 additions & 0 deletions integrations/snowflake/tests/test_snowflake_table_retriever.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -473,6 +473,8 @@ def test_to_dict_default(self, monkeypatch: MagicMock) -> None:
db_schema="test_schema",
warehouse="test_warehouse",
login_timeout=30,
private_key_file=None,
private_key_file_pwd=None,
)

data = component.to_dict()
Expand All @@ -494,6 +496,8 @@ def test_to_dict_default(self, monkeypatch: MagicMock) -> None:
"db_schema": "test_schema",
"warehouse": "test_warehouse",
"login_timeout": 30,
"private_key_file": None,
"private_key_file_pwd": None,
},
}

Expand Down Expand Up @@ -529,6 +533,8 @@ def test_to_dict_with_parameters(self, monkeypatch: MagicMock) -> None:
"db_schema": "SMALL_TOWNS",
"warehouse": "COMPUTE_WH",
"login_timeout": 30,
"private_key_file": None,
"private_key_file_pwd": None,
},
}

Expand Down