Feat/integrate connections toml

schemachange/cli.py

@@ -9,7 +9,7 @@
 from schemachange.config.get_merged_config import get_merged_config
 from schemachange.deploy import deploy
 from schemachange.redact_config_secrets import redact_config_secrets
-from schemachange.session.SnowflakeSession import get_session_from_config
+from schemachange.session.SnowflakeSession import SnowflakeSession
 
 # region Global Variables
 # metadata
@@ -63,11 +63,11 @@ def main():
         )
     else:
         config.check_for_deploy_args()
-        session = get_session_from_config(
-            config=config,
+        session = SnowflakeSession(
             schemachange_version=SCHEMACHANGE_VERSION,
-            snowflake_application_name=SNOWFLAKE_APPLICATION_NAME,
+            application=SNOWFLAKE_APPLICATION_NAME,
             logger=logger,
+            **config.get_session_kwargs(),
         )
         deploy(config=config, session=session)
 

schemachange/config/DeployConfig.py

@@ -6,7 +6,11 @@
 
 from schemachange.config.BaseConfig import BaseConfig
 from schemachange.config.ChangeHistoryTable import ChangeHistoryTable
-from schemachange.config.utils import get_snowflake_identifier_string
+from schemachange.config.utils import (
+    get_snowflake_identifier_string,
+    validate_file_path,
+    get_oauth_token,
+)
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
@@ -18,6 +22,12 @@ class DeployConfig(BaseConfig):
     snowflake_warehouse: str | None = None
     snowflake_database: str | None = None
     snowflake_schema: str | None = None
+    snowflake_authenticator: str | None = "snowflake"
+    snowflake_password: str | None = None
+    snowflake_oauth_token: str | None = None
+    snowflake_private_key_path: Path | None = None
+    connections_file_path: Path | None = None
+    connection_name: str | None = None
     # TODO: Turn change_history_table into three arguments. There's no need to parse it from a string
     change_history_table: ChangeHistoryTable | None = dataclasses.field(
         default_factory=ChangeHistoryTable
@@ -26,41 +36,62 @@ class DeployConfig(BaseConfig):
     autocommit: bool = False
     dry_run: bool = False
     query_tag: str | None = None
-    oauth_config: dict | None = None
 
     @classmethod
     def factory(
         cls,
         config_file_path: Path,
-        snowflake_role: str | None = None,
-        snowflake_warehouse: str | None = None,
-        snowflake_database: str | None = None,
-        snowflake_schema: str | None = None,
         change_history_table: str | None = None,
         **kwargs,
     ):
         if "subcommand" in kwargs:
             kwargs.pop("subcommand")
 
+        for sf_input in [
+            "snowflake_role",
+            "snowflake_warehouse",
+            "snowflake_database",
+            "snowflake_schema",
+        ]:
+            if sf_input in kwargs and kwargs[sf_input] is not None:
+                kwargs[sf_input] = get_snowflake_identifier_string(
+                    kwargs[sf_input], sf_input
+                )
+
+        for sf_path_input in [
+            "snowflake_private_key_path",
+            "snowflake_token_path",
+        ]:
+            if sf_path_input in kwargs and kwargs[sf_path_input] is not None:
+                kwargs[sf_path_input] = validate_file_path(
+                    file_path=kwargs[sf_path_input]
+                )
+
+        # If set by an environment variable, pop snowflake_token_path from kwargs
+        if "snowflake_oauth_token" in kwargs:
+            kwargs.pop("snowflake_token_path", None)
+            kwargs.pop("oauth_config", None)
+        # Load it from a file, if provided
+        elif "snowflake_token_path" in kwargs:
+            kwargs.pop("oauth_config", None)
+            oauth_token_path = kwargs.pop("snowflake_token_path")
+            with open(oauth_token_path) as f:
+                kwargs["snowflake_oauth_token"] = f.read()
+        # Make the oauth call if authenticator == "oauth"
+
+        elif "oauth_config" in kwargs:
+            oauth_config = kwargs.pop("oauth_config")
+            authenticator = kwargs.get("snowflake_authenticator")
+            if authenticator is not None and authenticator.lower() == "oauth":
+                kwargs["snowflake_oauth_token"] = get_oauth_token(oauth_config)
+
         change_history_table = ChangeHistoryTable.from_str(
             table_str=change_history_table
         )
 
         return super().factory(
             subcommand="deploy",
             config_file_path=config_file_path,
-            snowflake_role=get_snowflake_identifier_string(
-                snowflake_role, "snowflake_role"
-            ),
-            snowflake_warehouse=get_snowflake_identifier_string(
-                snowflake_warehouse, "snowflake_warehouse"
-            ),
-            snowflake_database=get_snowflake_identifier_string(
-                snowflake_database, "snowflake_database"
-            ),
-            snowflake_schema=get_snowflake_identifier_string(
-                snowflake_schema, "snowflake_schema"
-            ),
             change_history_table=change_history_table,
             **kwargs,
         )
@@ -74,6 +105,31 @@ def check_for_deploy_args(self) -> None:
             "snowflake_role": self.snowflake_role,
             "snowflake_warehouse": self.snowflake_warehouse,
         }
+
+        # OAuth based authentication
+        if self.snowflake_authenticator.lower() == "oauth":
+            req_args["snowflake_oauth_token"] = self.snowflake_oauth_token
+
+        # External Browser based SSO
+        elif self.snowflake_authenticator.lower() == "externalbrowser":
+            pass
+
+        # IDP based Authentication, limited to Okta
+        elif self.snowflake_authenticator.lower()[:8] == "https://":
+            req_args["snowflake_password"] = self.snowflake_password
+
+        elif self.snowflake_authenticator.lower() == "snowflake_jwt":
+            req_args["snowflake_private_key_path"] = self.snowflake_private_key_path
+
+        elif self.snowflake_authenticator.lower() == "snowflake":
+            req_args["snowflake_password"] = self.snowflake_password
+
+        else:
+            raise ValueError(
+                f"{self.snowflake_authenticator} is not supported authenticator option. "
+                "Choose from snowflake, snowflake_jwt, externalbrowser, oauth, https://<subdomain>.okta.com."
+            )
+
         missing_args = [key for key, value in req_args.items() if value is None]
 
         if len(missing_args) == 0:
@@ -83,3 +139,23 @@ def check_for_deploy_args(self) -> None:
         raise ValueError(
             f"Missing config values. The following config values are required: {missing_args}"
         )
+
+    def get_session_kwargs(self) -> dict:
+        session_kwargs = {
+            "account": self.snowflake_account,
+            "user": self.snowflake_user,
+            "role": self.snowflake_role,
+            "warehouse": self.snowflake_warehouse,
+            "database": self.snowflake_database,
+            "schema": self.snowflake_schema,
+            "authenticator": self.snowflake_authenticator,
+            "password": self.snowflake_password,
+            "oauth_token": self.snowflake_oauth_token,
+            "private_key_path": self.snowflake_private_key_path,
+            "connections_file_path": self.connections_file_path,
+            "connection_name": self.connection_name,
+            "change_history_table": self.change_history_table,
+            "autocommit": self.autocommit,
+            "query_tag": self.query_tag,
+        }
+        return {k: v for k, v in session_kwargs.items() if v is not None}

schemachange/config/get_merged_config.py

@@ -6,7 +6,13 @@
 from schemachange.config.DeployConfig import DeployConfig
 from schemachange.config.RenderConfig import RenderConfig
 from schemachange.config.parse_cli_args import parse_cli_args
-from schemachange.config.utils import load_yaml_config, validate_directory
+from schemachange.config.utils import (
+    load_yaml_config,
+    validate_directory,
+    get_env_kwargs,
+    get_connection_kwargs,
+    validate_file_path,
+)
 
 
 def get_yaml_config_kwargs(config_file_path: Optional[Path]) -> dict:
@@ -20,28 +26,30 @@ def get_yaml_config_kwargs(config_file_path: Optional[Path]) -> dict:
     }
 
     if "verbose" in kwargs:
-        kwargs["log_level"] = logging.DEBUG
+        if kwargs["verbose"]:
+            kwargs["log_level"] = logging.DEBUG
         kwargs.pop("verbose")
 
     if "vars" in kwargs:
         kwargs["config_vars"] = kwargs.pop("vars")
 
-    return kwargs
+    return {k: v for k, v in kwargs.items() if v is not None}
 
 
 def get_merged_config() -> Union[DeployConfig, RenderConfig]:
-    cli_kwargs = parse_cli_args(sys.argv[1:])
+    env_kwargs: dict[str, str] = get_env_kwargs()
 
-    if "verbose" in cli_kwargs and cli_kwargs["verbose"]:
-        cli_kwargs["log_level"] = logging.DEBUG
-        cli_kwargs.pop("verbose")
+    cli_kwargs = parse_cli_args(sys.argv[1:])
 
-    cli_config_vars = cli_kwargs.pop("config_vars", None)
-    if cli_config_vars is None:
-        cli_config_vars = {}
+    cli_config_vars = cli_kwargs.pop("config_vars")
 
+    connections_file_path = validate_file_path(
+        file_path=cli_kwargs.pop("connections_file_path", None)
+    )
+    connection_name = cli_kwargs.pop("connection_name", None)
     config_folder = validate_directory(path=cli_kwargs.pop("config_folder", "."))
-    config_file_path = Path(config_folder) / "schemachange-config.yml"
+    config_file_name = cli_kwargs.pop("config_file_name")
+    config_file_path = Path(config_folder) / config_file_name
 
     yaml_kwargs = get_yaml_config_kwargs(
         config_file_path=config_file_path,
@@ -50,6 +58,21 @@ def get_merged_config() -> Union[DeployConfig, RenderConfig]:
     if yaml_config_vars is None:
         yaml_config_vars = {}
 
+    if connections_file_path is None:
+        connections_file_path = yaml_kwargs.pop("connections_file_path", None)
+        if config_folder is not None and connections_file_path is not None:
+            connections_file_path = config_folder / connections_file_path
+
+        connections_file_path = validate_file_path(file_path=connections_file_path)
+
+    if connection_name is None:
+        connection_name = yaml_kwargs.pop("connection_name", None)
+
+    connection_kwargs: dict[str, str] = get_connection_kwargs(
+        connections_file_path=connections_file_path,
+        connection_name=connection_name,
+    )
+
     config_vars = {
         **yaml_config_vars,
         **cli_config_vars,
@@ -59,9 +82,15 @@ def get_merged_config() -> Union[DeployConfig, RenderConfig]:
     kwargs = {
         "config_file_path": config_file_path,
         "config_vars": config_vars,
+        **{k: v for k, v in connection_kwargs.items() if v is not None},
         **{k: v for k, v in yaml_kwargs.items() if v is not None},
         **{k: v for k, v in cli_kwargs.items() if v is not None},
+        **{k: v for k, v in env_kwargs.items() if v is not None},
     }
+    if connections_file_path is not None:
+        kwargs["connections_file_path"] = connections_file_path
+    if connection_name is not None:
+        kwargs["connection_name"] = connection_name
 
     if cli_kwargs["subcommand"] == "deploy":
         return DeployConfig.factory(**kwargs)

schemachange/config/parse_cli_args.py

@@ -2,6 +2,7 @@
 
 import argparse
 import json
+import logging
 from enum import Enum
 
 import structlog
@@ -58,6 +59,14 @@ def parse_cli_args(args) -> dict:
         "(the default is the current working directory)",
         required=False,
     )
+    parent_parser.add_argument(
+        "--config-file-name",
+        type=str,
+        default="schemachange-config.yml",
+        help="The schemachange config YAML file name. Must be in the directory supplied as the config-folder "
+        "(Default: schemachange-config.yml)",
+        required=False,
+    )
     parent_parser.add_argument(
         "-f",
         "--root-folder",
@@ -134,6 +143,39 @@ def parse_cli_args(args) -> dict:
         help="The name of the default schema to use. Can be overridden in the change scripts.",
         required=False,
     )
+    parser_deploy.add_argument(
+        "-A",
+        "--snowflake-authenticator",
+        type=str,
+        help="The Snowflake Authenticator to use. One of snowflake, oauth, externalbrowser, or https://<okta_account_name>.okta.com",
+        required=False,
+    )
+    parser_deploy.add_argument(
+        "-k",
+        "--snowflake-private-key-path",
+        type=str,
+        help="Path to file containing private key.",
+        required=False,
+    )
+    parser_deploy.add_argument(
+        "-t",
+        "--snowflake-token-path",
+        type=str,
+        help="Path to the file containing the OAuth token to be used when authenticating with Snowflake.",
+        required=False,
+    )
+    parser_deploy.add_argument(
+        "--connections-file-path",
+        type=str,
+        help="Override the default connections file path at snowflake.connector.constants.CONNECTIONS_FILE (OS specific)",
+        required=False,
+    )
+    parser_deploy.add_argument(
+        "--connection-name",
+        type=str,
+        help="Override the default connection name. Other connection-related values will override these connection values.",
+        required=False,
+    )
     parser_deploy.add_argument(
         "-c",
         "--change-history-table",
@@ -204,7 +246,16 @@ def parse_cli_args(args) -> dict:
     if "log_level" in parsed_kwargs and isinstance(parsed_kwargs["log_level"], Enum):
         parsed_kwargs["log_level"] = parsed_kwargs["log_level"].value
 
+    parsed_kwargs["config_vars"] = {}
     if "vars" in parsed_kwargs:
-        parsed_kwargs["config_vars"] = parsed_kwargs.pop("vars")
+        config_vars = parsed_kwargs.pop("vars")
+        if config_vars is not None:
+            parsed_kwargs["config_vars"] = config_vars
+
+    if "verbose" in parsed_kwargs:
+        parsed_kwargs["log_level"] = (
+            logging.DEBUG if parsed_kwargs["verbose"] else logging.INFO
+        )
+        parsed_kwargs.pop("verbose")
 
-    return parsed_kwargs
+    return {k: v for k, v in parsed_kwargs.items() if v is not None}