Merge pull request #32 from jojoduquartier/develop

snowflake connection slightly tested and functional
jojoduquartier · Aug 24, 2020 · 03586a5 · 03586a5
2 parents 2db3440 + 4a8ed0c
commit 03586a5
Show file tree

Hide file tree

Showing 3 changed files with 114 additions and 64 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -26,4 +26,9 @@ the manifest file has been added and a wheel file created
 ## [Version 1.0.4]
 - sqlalchemy_engine is read only
 - host, credential and key attibutes of configurer are read only
-- user can now create subsets. In the root folder where host is created, you will have a `subsets` folder. the goal is to be able to create different subsets for each project. Say I have 10 databases, 5 oracle and 5 mysql but I only need 2 each for a specific project. You can create a subset with a totally different encryption key that you can use on the server you want the project to run.
+- user can now create subsets. In the root folder where host is created, you will have a `subsets` folder. the goal is to be able to create different subsets for each project. Say I have 10 databases, 5 oracle and 5 mysql but I only need 2 each for a specific project. You can create a subset with a totally different encryption key that you can use on the server you want the project to run.
+
+## [Version 1.0.5]
+- no breaking changes
+- snowflake connection is rather different so changes were made so that user can enter all snowflake specific connection parameters
+- If users want to use the snowflake connection object, they can use a `raw_connection` argument in `dsdbmanager.snowflake_.Snowflake.create_engine`
diff --git a/dsdbmanager/configuring.py b/dsdbmanager/configuring.py
@@ -176,6 +176,25 @@ def write_credentials(
 
         return None
 
+    def _snowflake_params(self):
+        account: str = click.prompt(
+            "Account for database - probably the beging of what ends with snowflakecomputing.com",
+            type=str
+        )
+
+        database: str = click.prompt("Database to connect to", type=str)
+        schema: str = click.prompt("Schema for database", type=str)
+        warehouse: str = click.prompt("Warehouse for the database", type=str)
+        role: str = click.prompt("User role", default='', type=str)
+
+        return dict(
+            account=account,
+            database=database,
+            schema=schema,
+            warehouse=warehouse,
+            role=role
+        )
+
     def add_new_database_info(self):
         """
         Add new database
@@ -197,61 +216,77 @@ def add_new_database_info(self):
             host_file[flavor] = {}
         current_flavor_info = host_file.get(flavor, {})
 
-        # database name
-        name: str = click.prompt(
-            "Name for database - for some flavors like mysql database name is used to create engine",
-            type=str
-        )
-        name = name.strip()
-        if name in current_flavor_info:
-            confirmation = click.confirm(
-                f"There exists a {name} database in your {flavor}. Do you wish to replace it?")
-            if not confirmation:
-                return None
-
-        # additional_infos
-        host = click.prompt(
-            "Host/Database Address or Snowflake Account",
-            type=str
-        )
-        schema = click.prompt("Schema - Enter if none", default='', type=str)
-        sid = click.prompt("SID - Enter if none", default='', type=str)
-        service_name = click.prompt(
-            "Service Name - Enter if none",
-            default='',
-            type=str
-        )
-        port = click.prompt(
-            "Port Number - Enter if none",
-            default=-1,
-            type=int
-        )
+        if flavor == 'snowflake':
+            snowflake_dict = self._snowflake_params()
+            snowflake_dict = {k: v.strip() for k, v in snowflake_dict.items()}
+            name = snowflake_dict['database']
+            if name in current_flavor_info:
+                confirmation = click.confirm(
+                    f"There exists a {name} database in your {flavor}. Do you wish to replace it?"
+                )
+                if not confirmation:
+                    return None
+
+            # set data
+            host_file[flavor][name] = snowflake_dict
+
+        else:
+            # database name
+            name: str = click.prompt(
+                "Name for database - for some flavors like mysql database name is used to create engine",
+                type=str
+            )
+            name = name.strip()
+            if name in current_flavor_info:
+                confirmation = click.confirm(
+                    f"There exists a {name} database in your {flavor}. Do you wish to replace it?")
+                if not confirmation:
+                    return None
+
+            # additional_infos
+            host = click.prompt(
+                "Host/Database Address or Snowflake Account",
+                type=str
+            )
+            schema = click.prompt(
+                "Schema - Enter if none", default='', type=str)
+            sid = click.prompt("SID - Enter if none", default='', type=str)
+            service_name = click.prompt(
+                "Service Name - Enter if none",
+                default='',
+                type=str
+            )
+            port = click.prompt(
+                "Port Number - Enter if none",
+                default=-1,
+                type=int
+            )
 
-        host_dict = dict(
-            name=name,
-            host=host,
-            schema=schema,
-            port=port,
-            service_name=service_name,
-            sid=sid
-        )
+            host_dict = dict(
+                name=name,
+                host=host,
+                schema=schema,
+                port=port,
+                service_name=service_name,
+                sid=sid
+            )
 
-        # we don't want to store schema, service name or port if
-        # they are not required
-        if not schema:
-            _ = host_dict.pop('schema')
+            # we don't want to store schema, service name or port if
+            # they are not required
+            if not schema:
+                _ = host_dict.pop('schema')
 
-        if not service_name:
-            _ = host_dict.pop('service_name')
+            if not service_name:
+                _ = host_dict.pop('service_name')
 
-        if not sid:
-            _ = host_dict.pop('sid')
+            if not sid:
+                _ = host_dict.pop('sid')
 
-        if port == -1:
-            _ = host_dict.pop('port')
+            if port == -1:
+                _ = host_dict.pop('port')
 
-        # set data
-        host_file[flavor][name] = host_dict
+            # set data
+            host_file[flavor][name] = host_dict
 
         try:
             with self.host_location.open('w') as f:

diff --git a/dsdbmanager/snowflake_.py b/dsdbmanager/snowflake_.py
@@ -14,39 +14,49 @@ def __init__(self, db_name: str, host_dict: host_type = None):
         :param host_dict: optional database info with host, ports etc
         """
         self.db_name = db_name
-        self.host_dict: host_type = ConfigFilesManager().get_hosts() if not host_dict else host_dict
+        self.host_dict: host_type = ConfigFilesManager(
+        ).get_hosts() if not host_dict else host_dict
 
         if not self.host_dict or 'snowflake' not in self.host_dict:
             raise MissingFlavor("No databases available for snowflake", None)
 
         self.host_dict = self.host_dict.get('snowflake').get(self.db_name, {})
 
         if not self.host_dict:
-            raise MissingDatabase(f"{self.db_name} has not been added for snowflake", None)
-
-    def create_engine(self, user: str = None, pwd: str = None, **kwargs):
+            raise MissingDatabase(
+                f"{self.db_name} has not been added for snowflake", None)
+
+    def create_engine(
+        self,
+        user: str = None,
+        pwd: str = None,
+        raw_connection=False,
+        **kwargs
+    ):
         """
 
         :param user: username
         :param pwd: password
+        :param raw_connection: true or false
         :param kwargs: for compatibility/additional sqlalchemy create_engine kwargs or things like role, warehouse etc.
-        :return: sqlalchemy engine
+        :return: sqlalchemy engine if raw_connection is false or snowflake connector when true
         """
         try:
             from snowflake.sqlalchemy import URL
+            from snowflake.connector import connect
         except ImportError as e:
-            raise MissingPackage("You need the snowflake-sqlalchemy package to initiate connection", e)
-
-        host = self.host_dict.get('host')
+            raise MissingPackage(
+                "You need the snowflake-sqlalchemy and snowflake packages to initiate connection", e)
 
-        url = URL(
-            account=host,
+        url = dict(
+            self.host_dict.items(),
             user=user,
-            password=pwd,
-            database=self.db_name,
-            **kwargs
+            password=pwd
         )
 
         # TODO - find a way to identify kwrags consumed by URL and pull them out of kwargs and pass the rest below
+        if raw_connection:
+            # this could be obtained from the raw_connection method from sqlalchemy
+            return connect(**url)
 
-        return sa.create_engine(url)
+        return sa.create_engine(URL(**url), **kwargs)