From 25a79926fb525f77ef8cd5797b95f98b28223578 Mon Sep 17 00:00:00 2001 From: Isaac Ireland Date: Thu, 14 Nov 2024 10:01:19 +0000 Subject: [PATCH 1/6] adding privatekey auth param --- .../snowflake/snowflake_table_retriever.py | 27 +++++++++++++++---- .../tests/test_snowflake_table_retriever.py | 6 +++++ 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py index aa6f5ff4d..083060ec3 100644 --- a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py +++ b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py @@ -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, database: Optional[str] = None, db_schema: Optional[str] = None, warehouse: Optional[str] = None, @@ -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 + 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 @@ -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, @@ -275,17 +284,25 @@ 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.resolve_value(), + "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, + "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 diff --git a/integrations/snowflake/tests/test_snowflake_table_retriever.py b/integrations/snowflake/tests/test_snowflake_table_retriever.py index f5b8fee37..4aa990a6c 100644 --- a/integrations/snowflake/tests/test_snowflake_table_retriever.py +++ b/integrations/snowflake/tests/test_snowflake_table_retriever.py @@ -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() @@ -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 }, } @@ -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 }, } From cd2658fa20fee494f7cba35904ec871505c224fd Mon Sep 17 00:00:00 2001 From: Isaac Ireland Date: Thu, 14 Nov 2024 10:20:55 +0000 Subject: [PATCH 2/6] removed whitespace --- .../retrievers/snowflake/snowflake_table_retriever.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py index 083060ec3..09a575ced 100644 --- a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py +++ b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py @@ -295,11 +295,9 @@ def _fetch_data( "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=connect_params From adbfef57d73d9ae37f7e9a41c32dbf03d0399b83 Mon Sep 17 00:00:00 2001 From: Isaac Ireland Date: Thu, 14 Nov 2024 10:25:11 +0000 Subject: [PATCH 3/6] updating linting --- .../retrievers/snowflake/snowflake_table_retriever.py | 4 +--- .../snowflake/tests/test_snowflake_table_retriever.py | 6 +++--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py index 09a575ced..d2b7744fb 100644 --- a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py +++ b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py @@ -299,9 +299,7 @@ def _fetch_data( 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 - ) + conn = self._snowflake_connector(connect_params=connect_params) if conn is None: return df except (ForbiddenError, ProgrammingError) as e: diff --git a/integrations/snowflake/tests/test_snowflake_table_retriever.py b/integrations/snowflake/tests/test_snowflake_table_retriever.py index 4aa990a6c..5991ba9da 100644 --- a/integrations/snowflake/tests/test_snowflake_table_retriever.py +++ b/integrations/snowflake/tests/test_snowflake_table_retriever.py @@ -474,7 +474,7 @@ def test_to_dict_default(self, monkeypatch: MagicMock) -> None: warehouse="test_warehouse", login_timeout=30, private_key_file=None, - private_key_file_pwd=None + private_key_file_pwd=None, ) data = component.to_dict() @@ -497,7 +497,7 @@ def test_to_dict_default(self, monkeypatch: MagicMock) -> None: "warehouse": "test_warehouse", "login_timeout": 30, "private_key_file": None, - "private_key_file_pwd": None + "private_key_file_pwd": None, }, } @@ -534,7 +534,7 @@ def test_to_dict_with_parameters(self, monkeypatch: MagicMock) -> None: "warehouse": "COMPUTE_WH", "login_timeout": 30, "private_key_file": None, - "private_key_file_pwd": None + "private_key_file_pwd": None, }, } From 63a2eeb15205ee5469ad736548707bdaae847084 Mon Sep 17 00:00:00 2001 From: Isaac Ireland Date: Thu, 14 Nov 2024 10:27:07 +0000 Subject: [PATCH 4/6] indenting --- .../snowflake/snowflake_table_retriever.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py index d2b7744fb..703f63fe9 100644 --- a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py +++ b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py @@ -286,15 +286,15 @@ def _fetch_data( 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 - } + "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() From 8de2b0629cbfd89401b1aa01a20a1c477edc4e18 Mon Sep 17 00:00:00 2001 From: Isaac Ireland Date: Thu, 14 Nov 2024 10:29:29 +0000 Subject: [PATCH 5/6] final lint updates --- .../retrievers/snowflake/snowflake_table_retriever.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py index 703f63fe9..54099a626 100644 --- a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py +++ b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py @@ -285,7 +285,7 @@ def _fetch_data( return df try: # Build up param connection - connect_params={ + connect_params = { "user": self.user, "account": self.account, "private_key_file": self.private_key_file, @@ -293,7 +293,7 @@ def _fetch_data( "database": self.database, "schema": self.db_schema, "warehouse": self.warehouse, - "login_timeout": self.login_timeout + "login_timeout": self.login_timeout, } # Check if private key has been provided if self.private_key_file is None: From 0492dc4ca2c4209c05c3d4e5a07a10e3f31d5617 Mon Sep 17 00:00:00 2001 From: Isaac Ireland Date: Tue, 19 Nov 2024 16:32:03 +0000 Subject: [PATCH 6/6] docstring update - key-pair info --- .../retrievers/snowflake/snowflake_table_retriever.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py index 54099a626..f824ac265 100644 --- a/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py +++ b/integrations/snowflake/src/haystack_integrations/components/retrievers/snowflake/snowflake_table_retriever.py @@ -84,8 +84,10 @@ 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 - mutually exclusive to password, if key_file is provided this auth method will be used. + :param private_key_file: Location of private key-pair file. + This is mutually exclusive to password, if key_file is provided this auth method will be used. + See Snowflake documentation on Key-Pair authentication for further information: + https://docs.snowflake.com/en/user-guide/key-pair-auth :param private_key_file_pwd: Password for private key file """