Add support to specify role in connections (#19)

* Added error message when refresh failed
*  Fix methods to be able to receive empty context information
* Fix remove running task even though the task was cancelled

entities/sf_data_item.py

@@ -617,6 +617,8 @@ def on_edit_connection_action_triggered(self) -> None:
         sf_connection_string_dialog_window.txtDatabase.setText(
             auth_information["database"]
         )
+        if "role" in auth_information:
+            sf_connection_string_dialog_window.txtRole.setText(auth_information["role"])
 
         sf_connection_string_dialog_window.mAuthSettings.setUsername(
             auth_information["username"]
@@ -676,8 +678,15 @@ def refresh_internal(self) -> None:
         """
         Refreshes the data item.
         """
-        connection_manager: SFConnectionManager = SFConnectionManager.get_instance()
-        if self.item_type != "root":
-            connection_manager.reconnect(self.connection_name)
+        try:
+            connection_manager: SFConnectionManager = SFConnectionManager.get_instance()
+            if self.item_type != "root":
+                connection_manager.reconnect(self.connection_name)
 
-        super().refresh()
+            super().refresh()
+        except Exception as e:
+            QMessageBox.information(
+                None,
+                "Data Item Actions Refresh Error",
+                f"SFDataItem - refresh failed.\n\nExtended error information:\n{str(e)}",
+            )

entities/sf_dynamic_connection_combo_box_widget.py

@@ -9,7 +9,11 @@
 
 from ..helpers.data_base import get_schema_iterator, get_table_iterator
 
-from ..helpers.utils import get_authentification_information, get_qsettings
+from ..helpers.utils import (
+    get_authentification_information,
+    get_connection_child_groups,
+    get_qsettings,
+)
 from processing.gui.wrappers import WidgetWrapper
 
 
@@ -43,7 +47,7 @@ def createWidget(self):
 
     def get_connections_cb_options(self):
         connections_cb_options = []
-        root_groups = self.settings.childGroups()
+        root_groups = get_connection_child_groups()
         for group in root_groups:
             connections_cb_options.append(group)
         connections_cb_options.insert(0, "")

helpers/data_base.py

@@ -180,7 +180,7 @@ def get_cursor_description(
     auth_information: dict,
     query: str,
     connection_name: str,
-    context_information: typing.Dict[str, typing.Union[str, None]],
+    context_information: typing.Dict[str, typing.Union[str, None]] = None,
 ) -> list[snowflake.connector.cursor.ResultMetadata]:
     """
     Executes a query on a Snowflake database and retrieves the cursor description.

helpers/layer_creation.py

@@ -292,7 +292,7 @@ def get_srid_from_table(
     table_information: dict,
     connection_name: str,
     column_name: str,
-    context_information: typing.Dict[str, typing.Union[str, None]],
+    context_information: typing.Dict[str, typing.Union[str, None]] = None,
 ) -> int:
     """
     Retrieves the SRID (Spatial Reference Identifier) for a specified column in a table.

helpers/utils.py

@@ -104,6 +104,9 @@ def get_authentification_information(settings: QSettings, connection_name: str)
     auth_info["username"] = settings.value("username", defaultValue="")
     auth_info["connection_type"] = settings.value("connection_type", defaultValue="")
     auth_info["password"] = settings.value("password", defaultValue="")
+    role = settings.value("role", defaultValue="")
+    if role != "":
+        auth_info["role"] = role
     settings.endGroup()
 
     return auth_info
@@ -184,6 +187,8 @@ def set_connection_settings(connection_settings: dict) -> None:
     settings.setValue("database", connection_settings["database"])
     settings.setValue("username", connection_settings["username"])
     settings.setValue("connection_type", connection_settings["connection_type"])
+    if "role" in connection_settings:
+        settings.setValue("role", connection_settings["role"])
     if connection_settings["connection_type"] == "Default Authentication":
         settings.setValue("password", connection_settings["password"])
     settings.endGroup()
@@ -332,6 +337,9 @@ def get_auth_information(connection_name: str) -> dict:
     auth_info["username"] = settings.value("username", defaultValue="")
     auth_info["connection_type"] = settings.value("connection_type", defaultValue="")
     auth_info["password"] = settings.value("password", defaultValue="")
+    role = settings.value("role", defaultValue="")
+    if role != "":
+        auth_info["role"] = role
     settings.endGroup()
     return auth_info
 

managers/sf_connection_manager.py

@@ -70,6 +70,8 @@ def connect(self, connection_name: str, connection_params: dict) -> None:
                 conn_params["password"] = connection_params["password"]
             elif connection_params["connection_type"] == "Single sign-on (SSO)":
                 conn_params["authenticator"] = "externalbrowser"
+            if "role" in connection_params:
+                conn_params["role"] = connection_params["role"]
 
             self.opened_connections[connection_name] = self.create_snowflake_connection(
                 conn_params

tasks/sf_convert_column_to_layer_task.py

@@ -142,4 +142,4 @@ def finished(self, result: bool) -> None:
                 if isinstance(layer, QgsMapLayer):
                     QgsProject.instance().addMapLayer(layer)
                     QgsProject.instance().layerTreeRoot()
-                    remove_task_from_running_queue(self.path)
+        remove_task_from_running_queue(self.path)

tasks/sf_convert_sql_query_to_layer_task.py

@@ -15,10 +15,10 @@ class SFConvertSQLQueryToLayerTask(QgsTask):
 
     def __init__(
         self,
-        context_information: typing.Dict[str, typing.Union[str, None]],
         query: str,
         geo_column_name: str,
         layer_name: str,
+        context_information: typing.Dict[str, typing.Union[str, None]] = None,
     ):
         try:
             self.query = query

tasks/sf_execute_sql_query_task.py

@@ -12,9 +12,9 @@ class SFExecuteSQLQueryTask(QgsTask):
 
     def __init__(
         self,
-        context_information: typing.Dict[str, typing.Union[str, None]],
         query: str,
         limit: typing.Union[int, None] = None,
+        context_information: typing.Dict[str, typing.Union[str, None]] = None,
     ):
         try:
             self.query = query

ui/sf_connection_string_dialog.py

@@ -92,6 +92,8 @@ def button_box_ok_clicked(self) -> None:
                 "username": self.mAuthSettings.username(),
                 "connection_type": self.cbxConnectionType.currentText(),
             }
+            if self.txtRole.text() != "":
+                conn_settings["role"] = self.txtRole.text()
             if self.cbxConnectionType.currentText() == "Default Authentication":
                 conn_settings["password"] = self.mAuthSettings.password()
             set_connection_settings(conn_settings)
@@ -123,6 +125,8 @@ def test_connection_clicked(self) -> None:
                 "database": self.txtDatabase.text(),
                 "login_timeout": 5,
             }
+            if self.txtRole.text() != "":
+                connection_params["role"] = self.txtRole.text()
             if self.cbxConnectionType.currentText() == "Default Authentication":
                 connection_params["password"] = self.mAuthSettings.password()
                 conn = sf_connection_manager.create_snowflake_connection(

ui/sf_connection_string_dialog.ui

@@ -59,6 +59,19 @@
            <cstring>txtWarehouse</cstring>
           </property>
          </widget>
+        </item>
+         <item row="5" column="0">
+         <widget class="QLabel" name="TextLabel_role">
+          <property name="text">
+           <string>Role</string>
+          </property>
+          <property name="buddy">
+           <cstring>txtRole</cstring>
+          </property>
+         </widget>
+        </item>
+        <item row="5" column="1">
+         <widget class="QLineEdit" name="txtRole"/>
         </item>
         <item row="0" column="0">
          <widget class="QLabel" name="TextLabel_name">

ui/sf_data_source_manager_widget.py

@@ -232,6 +232,9 @@ def on_btn_edit_clicked(self) -> None:
                 another_window.txtAccount.setText(auth_information["account"])
                 another_window.txtDatabase.setText(auth_information["database"])
 
+                if "role" in auth_information:
+                    another_window.txtRole.setText(auth_information["role"])
+
                 another_window.mAuthSettings.setUsername(auth_information["username"])
                 another_window.exec_()
         except Exception as e:

ui/sf_sql_query_dialog.py

@@ -27,7 +27,7 @@ class SFSQLQueryDialog(QDialog, FORM_CLASS_SFCS):
 
     def __init__(
         self,
-        context_information: typing.Dict[str, typing.Union[str, None]],
+        context_information: typing.Dict[str, typing.Union[str, None]] = None,
         parent: typing.Optional[QWidget] = None,
     ) -> None:
         super().__init__(parent)