[GitHub] [airflow] dstandish commented on a change in pull request #11350: Allow execution of multiple sql statements in SnowflakeHook

[GitBox <gi...@apache.org>]

dstandish commented on a change in pull request #11350:
URL: https://github.com/apache/airflow/pull/11350#discussion_r550005187



##########
File path: airflow/providers/snowflake/hooks/snowflake.py
##########
@@ -142,3 +141,18 @@ def set_autocommit(self, conn, autocommit: Any) -> None:
 
     def get_autocommit(self, conn):
         return getattr(conn, 'autocommit_mode', False)
+
+    def run(self, sql, autocommit=False, parameters=None):
+        """
+        Snowflake-connector doesn't allow natively the execution of multiple SQL statements in the same
+        call. So for allowing to pass files or strings with several queries this method is coded,
+        that relies on run from DBApiHook
+        """
+        if isinstance(sql, str):
+            with closing(self.get_conn()) as conn:

Review comment:
       you do not need to wrap `SnowflakeConnection` with `closing` because it already has a good `__exit__`

##########
File path: airflow/providers/snowflake/hooks/snowflake.py
##########
@@ -142,3 +141,18 @@ def set_autocommit(self, conn, autocommit: Any) -> None:
 
     def get_autocommit(self, conn):
         return getattr(conn, 'autocommit_mode', False)
+
+    def run(self, sql, autocommit=False, parameters=None):
+        """
+        Snowflake-connector doesn't allow natively the execution of multiple SQL statements in the same
+        call. So for allowing to pass files or strings with several queries this method is coded,
+        that relies on run from DBApiHook
+        """
+        if isinstance(sql, str):

Review comment:
       ok so you want to run with execute_string if `str` and use dbapi.run otherwise  (e.g. if it is `List[str]`)
   
   it seems reasonable. but another option would be to use snowflake's `split_statements` function
   
   https://github.com/snowflakedb/snowflake-connector-python/blob/master/src/snowflake/connector/util_text.py#L32
   
   what you could do is always split and then you can always dbapi run.  this would save you the type checking and autocommit logic.
   
   but i think you can reasonably go either way.
   
   one thought though... at my prev company we liked to execute sql files.  what we did was first check if `sql` was a file that exists with `Path(sql).exists()` and if so then `sql = Path(sql).read_text()`.  then execute as normal.  you can do with templating but i find this easier.  anyway it's an idea for your consideration if you want to make this a little more generic --- `Union[Path, List[str], str]` where str can be either single statement, multiple statement, or path to sql file.

##########
File path: airflow/providers/snowflake/hooks/snowflake.py
##########
@@ -142,3 +141,18 @@ def set_autocommit(self, conn, autocommit: Any) -> None:
 
     def get_autocommit(self, conn):
         return getattr(conn, 'autocommit_mode', False)
+
+    def run(self, sql, autocommit=False, parameters=None):
+        """
+        Snowflake-connector doesn't allow natively the execution of multiple SQL statements in the same
+        call. So for allowing to pass files or strings with several queries this method is coded,
+        that relies on run from DBApiHook
+        """
+        if isinstance(sql, str):
+            with closing(self.get_conn()) as conn:
+                if self.supports_autocommit:

Review comment:
       this generic `if self.supports_autocommit` logic makes sense in DbApiHook because it's a generic base class that means to support many different databases -- some which support autocommit and some which don't.  in your case here, you are implementing for one specific database and you _know_ whether it supports autocommit or not -- so you don't need to check --- just do the right thing.

##########
File path: airflow/providers/snowflake/hooks/snowflake.py
##########
@@ -142,3 +141,18 @@ def set_autocommit(self, conn, autocommit: Any) -> None:
 
     def get_autocommit(self, conn):
         return getattr(conn, 'autocommit_mode', False)
+
+    def run(self, sql, autocommit=False, parameters=None):
+        """
+        Snowflake-connector doesn't allow natively the execution of multiple SQL statements in the same
+        call. So for allowing to pass files or strings with several queries this method is coded,
+        that relies on run from DBApiHook
+        """
+        if isinstance(sql, str):
+            with closing(self.get_conn()) as conn:
+                if self.supports_autocommit:
+                    self.set_autocommit(conn, autocommit)
+
+                    conn.execute_string(sql, parameters)
+        else:
+            super().run(sql, autocommit, parameters)

Review comment:
       @JavierLopezT what you want to do is use mock.patch to verify that in the `run` method, the right underlying method is called given the right input.  google for basic examples of mock patch.
   
   in this case you would want to mock the `get_conn` method so that its `return_value` is a mock object.  then after `run` finishes, you can verify that the mock object was called with method `execute_string` and your value for `sql`
   
   there are lots of examples in the repo just try to find something similar




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org