Snowflake key pair authentication support (#158)

raster_loader/cli/snowflake.py

@@ -4,7 +4,7 @@
 import click
 from functools import wraps, partial
 
-from raster_loader.utils import get_default_table_name
+from raster_loader.utils import get_default_table_name, check_private_key
 from raster_loader.io.snowflake import SnowflakeConnection
 
 
@@ -40,7 +40,20 @@ def snowflake(args=None):
     required=False,
     default=None,
 )
+@click.option(
+    "--private-key-path",
+    help="The path to the private key file. (PEM format)",
+    required=False,
+    default=None,
+)
+@click.option(
+    "--private-key-passphrase",
+    help="The passphrase for the private key.",
+    required=False,
+    default=None,
+)
 @click.option("--role", help="The role to use for the file upload.", default=None)
+@click.option("--warehouse", help="Name of the default warehouse to use.", default=None)
 @click.option(
     "--file_path", help="The path to the raster file.", required=False, default=None
 )
@@ -110,7 +123,10 @@ def upload(
     username,
     password,
     token,
+    private_key_path,
+    private_key_passphrase,
     role,
+    warehouse,
     file_path,
     file_url,
     database,
@@ -132,13 +148,31 @@ def upload(
         get_block_dims,
     )
 
-    if (token is None and (username is None or password is None)) or all(
-        v is not None for v in [token, username, password]
+    if not (
+        (token is not None and username is None)
+        or (
+            token is None
+            and username is not None
+            and password is not None
+            and private_key_path is None
+        )
+        or (
+            token is None
+            and username is not None
+            and password is None
+            and private_key_path is not None
+        )
     ):
         raise ValueError(
-            "Either --token or --username and --password must be provided."
+            "Either (--token) or (--username and --private-key-path) or"
+            " (--username and --password) must be provided."
         )
 
+    if private_key_path is not None:
+        check_private_key(private_key_path, private_key_passphrase)
+        if username is None:
+            raise ValueError("--username must be provided when using a private key.")
+
     if file_path is None and file_url is None:
         raise ValueError("Either --file_path or --file_url must be provided.")
 
@@ -167,11 +201,14 @@ def upload(
     connector = SnowflakeConnection(
         username=username,
         password=password,
+        private_key_path=private_key_path,
+        private_key_passphrase=private_key_passphrase,
         token=token,
         account=account,
         database=database,
         schema=schema,
         role=role,
+        warehouse=warehouse,
     )
 
     source = file_path if is_local_file else file_url
@@ -226,29 +263,76 @@ def upload(
     required=False,
     default=None,
 )
+@click.option(
+    "--private-key-path",
+    help="The path to the private key file. (PEM format)",
+    required=False,
+    default=None,
+)
+@click.option(
+    "--private-key-passphrase",
+    help="The passphrase for the private key.",
+    required=False,
+    default=None,
+)
 @click.option("--role", help="The role to use for the file upload.", default=None)
+@click.option("--warehouse", help="Name of the default warehouse to use.", default=None)
 @click.option("--database", help="The name of the database.", required=True)
 @click.option("--schema", help="The name of the schema.", required=True)
 @click.option("--table", help="The name of the table.", required=True)
 @click.option("--limit", help="Limit number of rows returned", default=10)
-def describe(account, username, password, token, role, database, schema, table, limit):
+def describe(
+    account,
+    username,
+    password,
+    token,
+    private_key_path,
+    private_key_passphrase,
+    role,
+    warehouse,
+    database,
+    schema,
+    table,
+    limit,
+):
 
-    if (token is None and (username is None or password is None)) or all(
-        v is not None for v in [token, username, password]
+    if not (
+        (token is not None and username is None)
+        or (
+            token is None
+            and username is not None
+            and password is not None
+            and private_key_path is None
+        )
+        or (
+            token is None
+            and username is not None
+            and password is None
+            and private_key_path is not None
+        )
     ):
         raise ValueError(
-            "Either --token or --username and --password must be provided."
+            "Either (--token) or (--username and --private-key-path) or"
+            " (--username and --password) must be provided."
         )
 
+    if private_key_path is not None:
+        check_private_key(private_key_path, private_key_passphrase)
+        if username is None:
+            raise ValueError("--username must be provided when using a private key.")
+
     fqn = f"{database}.{schema}.{table}"
     connector = SnowflakeConnection(
         username=username,
         password=password,
+        private_key_path=private_key_path,
+        private_key_passphrase=private_key_passphrase,
         token=token,
         account=account,
         database=database,
         schema=schema,
         role=role,
+        warehouse=warehouse,
     )
     df = connector.get_records(fqn, limit)
     print(f"Table: {fqn}")

raster_loader/io/snowflake.py

@@ -33,28 +33,54 @@
 
 
 class SnowflakeConnection(DataWarehouseConnection):
-    def __init__(self, username, password, account, database, schema, token, role):
+    def __init__(
+        self,
+        username,
+        password,
+        account,
+        database,
+        schema,
+        token,
+        private_key_path,
+        private_key_passphrase,
+        role,
+        warehouse,
+    ):
         if not _has_snowflake:
             import_error_snowflake()
 
         # TODO: Write a proper static factory for this
-        if token is None:
+        if token is not None:
             self.client = snowflake.connector.connect(
+                authenticator="oauth",
+                token=token,
+                account=account,
+                database=database.upper(),
+                schema=schema.upper(),
+                role=role.upper() if role is not None else None,
+                warehouse=warehouse,
+            )
+        elif private_key_path is not None:
+            self.client = snowflake.connector.connect(
+                authenticator="snowflake_jwt",
                 user=username,
-                password=password,
+                private_key_file=private_key_path,
+                private_key_file_pwd=private_key_passphrase,
                 account=account,
                 database=database.upper(),
                 schema=schema.upper(),
                 role=role.upper() if role is not None else None,
+                warehouse=warehouse,
             )
         else:
             self.client = snowflake.connector.connect(
-                authenticator="oauth",
-                token=token,
+                user=username,
+                password=password,
                 account=account,
                 database=database.upper(),
                 schema=schema.upper(),
                 role=role.upper() if role is not None else None,
+                warehouse=warehouse,
             )
 
     def band_rename_function(self, band_name: str):

raster_loader/tests/snowflake/test_cli.py

@@ -181,7 +181,8 @@ def test_snowflake_credentials_validation(*args, **kwargs):
     )
     assert result.exit_code == 1
     assert (
-        "Either --token or --username and --password must be provided." in result.output
+        "Either (--token) or (--username and --private-key-path) or"
+        " (--username and --password) must be provided." in result.output
     )
 
     result = runner.invoke(
@@ -213,7 +214,8 @@ def test_snowflake_credentials_validation(*args, **kwargs):
     )
     assert result.exit_code == 1
     assert (
-        "Either --token or --username and --password must be provided." in result.output
+        "Either (--token) or (--username and --private-key-path) or"
+        " (--username and --password) must be provided." in result.output
     )
 
 

raster_loader/utils.py

@@ -37,6 +37,22 @@ def get_default_table_name(base_path: str, band):
     return re.sub(r"[^a-zA-Z0-9_-]", "_", table)
 
 
+def check_private_key(private_key_path: str, private_key_passphrase: str):
+    # Check that the private key file exists
+    if not os.path.exists(private_key_path):
+        raise ValueError(f"Private key file {private_key_path} not found")
+
+    with open(private_key_path, "r") as f:
+        private_key = f.read()
+        if (
+            private_key.startswith("-----BEGIN ENCRYPTED PRIVATE KEY-----")
+            and private_key_passphrase is None
+        ):
+            raise ValueError(
+                "The private key file is encrypted. Please provide a passphrase."
+            )
+
+
 # Modify the __init__ so that self.line = "" instead of None
 def new_init(
     self, message, category, filename, lineno, file=None, line=None, source=None