private key to okta (#216)

Author: Ttommeke

src/koheesio/__about__.py

@@ -12,6 +12,7 @@
 
 LICENSE_INFO = "Licensed as Apache 2.0"
 SOURCE = "https://github.com/Nike-Inc/koheesio"
+
 __version__ = "0.10.6"
 __logo__ = (
     75,

src/koheesio/integrations/snowflake/__init__.py

@@ -18,8 +18,9 @@
     Login name for the Snowflake user.
     Alias for `sfUser`.
 password : SecretStr
-    Password for the Snowflake user.
-    Alias for `sfPassword`.
+    The password for the snowflake user. Alias for `sfPassword`. Either this or private_key must be provided.
+private_key : SecretStr
+    The private_key for the snowflake user. Either this or password must be provided.
 database : str
     The database to use for the session after connecting.
     Alias for `sfDatabase`.
@@ -172,8 +173,9 @@ class SnowflakeBaseModel(BaseModel, ExtraParamsMixin, ABC):  # type: ignore[misc
         Login name for the Snowflake user.
         Alias for `sfUser`.
     password : SecretStr
-        Password for the Snowflake user.
-        Alias for `sfPassword`.
+        The password for the snowflake user. Alias for `sfPassword`. Either this or private_key must be provided.
+    private_key : SecretStr
+        The private_key for the snowflake user. Either this or password must be provided.
     role : str
         The default security role to use for the session after connecting.
         Alias for `sfRole`.
@@ -199,7 +201,8 @@ class SnowflakeBaseModel(BaseModel, ExtraParamsMixin, ABC):  # type: ignore[misc
         examples=["example.snowflakecomputing.com"],
     )
     user: str = Field(default=..., alias="sfUser", description="Login name for the Snowflake user")
-    password: SecretStr = Field(default=..., alias="sfPassword", description="Password for the Snowflake user")
+    password: Optional[SecretStr] = Field(default=None, alias="sfPassword", description="Password for the Snowflake user")
+    private_key: Optional[SecretStr] = Field(default=None, alias="pem_private_key", description="PEM private key for the Snowflake user")
     role: str = Field(
         default=..., alias="sfRole", description="The default security role to use for the session after connecting"
     )
@@ -227,6 +230,17 @@ class SnowflakeBaseModel(BaseModel, ExtraParamsMixin, ABC):  # type: ignore[misc
         examples=[{"sfCompress": "on", "continue_on_error": "off"}],
     )
 
+    @model_validator(mode="after")
+    def check_authentication_method(self) -> "SnowflakeBaseModel":
+        """Ensure at least one of password or private_key is provided."""
+        if not self.password and not self.private_key:
+            raise ValueError("You must provide either 'password' or 'private_key'.")
+        if self.password and self.private_key:
+            raise ValueError(
+                "You must provide either 'password' or 'private_key', not both."
+            )
+        return self
+
     @property
     def options(self) -> Dict[str, Any]:
         """Shorthand for accessing self.params provided for backwards compatibility"""
@@ -260,6 +274,7 @@ def get_options(self, by_alias: bool = True, include: Optional[Set[str]] = None)
             # schema and password have to be handled separately
             "sfSchema",
             "password",
+            "private_key",
         } - (include or set())
 
         fields = self.model_dump(
@@ -268,13 +283,22 @@ def get_options(self, by_alias: bool = True, include: Optional[Set[str]] = None)
             exclude=exclude_set,
         )
 
+        fields.update({ "sfSchema" if by_alias else "schema": self.sfSchema })
+
         # handle schema and password
-        fields.update(
-            {
-                "sfSchema" if by_alias else "schema": self.sfSchema,
-                "sfPassword" if by_alias else "password": self.password.get_secret_value(),
-            }
-        )
+        if self.password:
+            fields.update(
+                {
+                    "sfPassword" if by_alias else "password": self.password.get_secret_value(),
+                }
+            )
+        elif self.private_key:
+            fields.update(
+                {
+                    "private_key": self.private_key.get_secret_value(),
+                    "pem_private_key": self.private_key.get_secret_value(),
+                }
+            )
 
         # handle include
         if include:
@@ -379,6 +403,7 @@ def get_options(self, by_alias: bool = False, include: Optional[Set[str]] = None
                 "database",
                 "schema",
                 "password",
+                "private_key",
             }
         return super().get_options(by_alias=by_alias, include=include)
 
@@ -389,6 +414,8 @@ def conn(self) -> Generator:
             raise RuntimeError("Snowflake connector is not installed. Please install `snowflake-connector-python`.")
 
         sf_options = self.get_options()
+
+
         _conn = self._snowflake_connector.connect(**sf_options)
         self.log.info(f"Connected to Snowflake account: {sf_options['account']}")
 
@@ -442,7 +469,9 @@ class GrantPrivilegesOnObject(SnowflakeRunQueryPython):
     user : str
         The username. Alias for `sfUser`
     password : SecretStr
-        The password. Alias for `sfPassword`
+        The password for the snowflake user. Alias for `sfPassword`. Either this or private_key must be provided.
+    private_key : SecretStr
+        The private_key for the snowflake user. Either this or password must be provided.
     role : str
         The role name
     object : str

src/koheesio/integrations/spark/snowflake.py

@@ -640,7 +640,10 @@ class SnowflakeWriter(SnowflakeBaseModel, Writer):
     def execute(self) -> SnowflakeWriter.Output:
         """Write to Snowflake"""
         self.log.debug(f"writing to {self.table} with mode {self.insert_type}")
-        self.df.write.format(self.format).options(**self.get_options()).option("dbtable", self.table).mode(
+        
+        options = self.get_options()
+        
+        self.df.write.format(self.format).options(**options).option("dbtable", self.table).mode(
             self.insert_type
         ).save()
 

src/koheesio/spark/readers/jdbc.py

@@ -75,7 +75,8 @@ class JdbcReader(Reader, ExtraParamsMixin):
         "the hostname of the server.",
     )
     user: str = Field(default=..., description="User to authenticate to the server")
-    password: SecretStr = Field(default=..., description="Password belonging to the username")
+    password: Optional[SecretStr] = Field(default=None, description="Password belonging to the username")
+    private_key: Optional[SecretStr] = Field(default=None, alias="pem_private_key", description="Private key for authentication")
     dbtable: Optional[str] = Field(
         default=None,
         description="Database table name, also include schema name",
@@ -114,6 +115,14 @@ def check_dbtable_or_query(self) -> "JdbcReader":
         if self.query and self.dbtable:
             self.log.warning("Query is filled in, dbtable will be ignored!")
 
+        """Ensure at least one of password or private_key is provided."""
+        if not self.password and not self.private_key:
+            raise ValueError("You must provide either 'password' or 'private_key'.")
+        if self.password and self.private_key:
+            raise ValueError(
+                "You must provide either 'password' or 'private_key', not both."
+            )
+
         return self
 
     @property
@@ -124,8 +133,11 @@ def options(self) -> Dict[str, Any]:
     def execute(self) -> "JdbcReader.Output":
         """Wrapper around Spark's jdbc read format"""
         options = self.get_options()
-
+        
         if pw := self.password:
             options["password"] = pw.get_secret_value()
+        if pk := self.private_key:
+            options["pem_private_key"] = pk.get_secret_value()
+            options["private_key"] = pk.get_secret_value()
 
         self.output.df = self.spark.read.format(self.format).options(**options).load()