fix pyright issues

Author: MT255026

libraries/dagster-teradata/dagster_teradata/resources.py

@@ -218,7 +218,7 @@ def get_connection(self):
     def set_query_band(self, query_band_text, teradata_conn):
         """Set SESSION Query Band for each connection session."""
         try:
-            query_band_text = _handle_user_query_band_text(query_band_text)
+            query_band_text = _handle_user_query_band_text(self, query_band_text)
             set_query_band_sql = f"SET QUERY_BAND='{query_band_text}' FOR SESSION"
             with teradata_conn.cursor() as cur:
                 cur.execute(set_query_band_sql)

libraries/dagster-teradata/dagster_teradata/ttu/bteq.py

@@ -3,7 +3,7 @@
 import subprocess
 import tempfile
 from contextlib import contextmanager
-from typing import Optional, List, Union, Literal
+from typing import Optional, List, Union, Literal, cast
 
 import paramiko
 from dagster import DagsterError
@@ -201,16 +201,20 @@ def bteq_operator(
             if self.file_path:
                 if not self._setup_ssh_connection(
                     host=self.remote_host,
-                    user=self.remote_user,
+                    user=cast(str, self.remote_user),
                     password=self.remote_remote_password,
                     key_path=self.ssh_key_path,
                     port=self.remote_port,
                 ):
                     raise DagsterError(
                         "Failed to establish SSH connection. Please check the provided credentials."
                     )
-                if self.file_path and is_valid_remote_bteq_script_file(
-                    self.ssh_client, self.file_path
+                if (
+                    self.file_path
+                    and self.ssh_client
+                    and is_valid_remote_bteq_script_file(
+                        self.ssh_client, self.file_path
+                    )
                 ):
                     return self._handle_remote_bteq_file(
                         ssh_client=self.ssh_client,
@@ -230,7 +234,7 @@ def execute_bteq_script(
         bteq_script: str,
         remote_working_dir: str | None,
         bteq_script_encoding: str | None,
-        timeout: int,
+        timeout: Optional[int],  # or: timeout: int | None
         timeout_rc: int | None,
         bteq_session_encoding: str | None,
         bteq_quit_rc: int | list[int] | tuple[int, ...] | None,
@@ -297,7 +301,7 @@ def execute_bteq_script_at_remote(
         bteq_script: str,
         remote_working_dir: str | None,
         bteq_script_encoding: str | None,
-        timeout: int,
+        timeout: Optional[int],  # or: timeout: int | None
         timeout_rc: int | None,
         bteq_session_encoding: str | None,
         bteq_quit_rc: int | list[int] | tuple[int, ...] | None,
@@ -347,8 +351,8 @@ def execute_bteq_script_at_remote(
                 bteq_quit_rc,
                 bteq_session_encoding,
                 tmp_dir,
-                remote_host,
-                remote_user,
+                cast(str, remote_host),
+                cast(str, remote_user),
                 remote_password,
                 ssh_key_path,
                 remote_port,
@@ -359,7 +363,7 @@ def _transfer_to_and_execute_bteq_on_remote(
         file_path: str,
         remote_working_dir: str | None,
         bteq_script_encoding: str | None,
-        timeout: int,
+        timeout: Optional[int],  # or: timeout: int | None
         timeout_rc: int | None,
         bteq_quit_rc: int | list[int] | tuple[int, ...] | None,
         bteq_session_encoding: str | None,
@@ -498,7 +502,7 @@ def execute_bteq_script_at_local(
         self,
         bteq_script: str,
         bteq_script_encoding: str | None,
-        timeout: int,
+        timeout: Optional[int],  # or: timeout: int | None
         timeout_rc: int | None,
         bteq_quit_rc: int | list[int] | tuple[int, ...] | None,
         bteq_session_encoding: str | None,
@@ -547,7 +551,7 @@ def execute_bteq_script_at_local(
         try:
             # https://docs.python.org/3.10/library/subprocess.html#subprocess.Popen.wait  timeout is in seconds
             process.wait(
-                timeout=timeout + 60
+                timeout=(timeout or 0) + 60
             )  # Adding 1 minute extra for BTEQ script timeout
         except subprocess.TimeoutExpired:
             self.on_kill()
@@ -583,10 +587,6 @@ def execute_bteq_script_at_local(
 
         return process.returncode
 
-    def contains_template(parameter_value):
-        """Check if the parameter contains Jinja templating syntax."""
-        return "{{" in parameter_value and "}}" in parameter_value
-
     def on_kill(self):
         """Terminate the subprocess if running."""
         conn = self.connection
@@ -677,7 +677,7 @@ def _handle_remote_bteq_file(
                     remote_user=self.remote_user,
                     remote_password=self.remote_remote_password,
                     ssh_key_path=self.ssh_key_path,
-                    remote_port=self.remote_port,
+                    remote_port=self.remote_port or 22,
                 )
         else:
             raise ValueError(
@@ -718,7 +718,7 @@ def _handle_local_bteq_file(self, file_path: str) -> int | None:
     def _setup_ssh_connection(
         self,
         host: str,
-        user: str,
+        user: Optional[str],
         password: Optional[str],
         key_path: Optional[str],
         port: int,
@@ -753,6 +753,11 @@ def _setup_ssh_connection(
                 self.ssh_client.connect(host, port=port, username=user, pkey=key)
             else:
                 if not password:
+                    if user is None:
+                        raise ValueError(
+                            "Username is required to fetch stored credentials"
+                        )
+                    # Attempt to retrieve stored credentials
                     creds = get_stored_credentials(self, host, user)
                     password = (
                         self.cred_manager.decrypt(creds["password"]) if creds else None

libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py

@@ -3,7 +3,7 @@
 import os
 import shutil
 import stat
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 if TYPE_CHECKING:
     from paramiko import SSHClient
@@ -200,7 +200,7 @@ def _prepare_bteq_command(
 
 
 def prepare_bteq_command_for_remote_execution(
-    timeout: int,
+    timeout: Optional[int],  # or: timeout: int | None
     bteq_script_encoding: str,
     bteq_session_encoding: str,
     timeout_rc: int,
@@ -218,15 +218,15 @@ def prepare_bteq_command_for_remote_execution(
         str: Complete BTEQ command string for remote execution.
     """
     bteq_core_cmd = _prepare_bteq_command(
-        timeout, bteq_script_encoding, bteq_session_encoding, timeout_rc
+        timeout or 600, bteq_script_encoding, bteq_session_encoding, timeout_rc
     )
     bteq_core_cmd.append('"')
     return " ".join(bteq_core_cmd)
 
 
 def prepare_bteq_command_for_local_execution(
     teradata_connection_resource,
-    timeout: int,
+    timeout: Optional[int],  # or: timeout: int | None
     bteq_script_encoding: str,
     bteq_session_encoding: str,
     timeout_rc: int,
@@ -245,7 +245,7 @@ def prepare_bteq_command_for_local_execution(
         str: Complete BTEQ command string for local execution including login details.
     """
     bteq_core_cmd = _prepare_bteq_command(
-        timeout, bteq_script_encoding, bteq_session_encoding, timeout_rc
+        timeout or 600, bteq_script_encoding, bteq_session_encoding, timeout_rc
     )
     host = teradata_connection_resource.host
     login = teradata_connection_resource.user
@@ -311,7 +311,7 @@ def read_file(file_path: str, encoding: str = "UTF-8") -> str:
 
 
 def is_valid_remote_bteq_script_file(
-    ssh_client: SSHClient, remote_file_path: str, logger=None
+    ssh_client: SSHClient, remote_file_path: str | None, logger=None
 ) -> bool:
     """
     Check if a remote file is a valid BTEQ script file.

libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py

@@ -70,6 +70,8 @@ def encrypt(self, data: str) -> str:
             >>> manager = SecureCredentialManager()
             >>> encrypted = manager.encrypt("my_password")
         """
+        if self.key is None:
+            raise ValueError("Encryption key is not initialized")
         f = Fernet(self.key)
         return f.encrypt(data.encode()).decode()
 
@@ -90,6 +92,8 @@ def decrypt(self, encrypted_data: str) -> str:
             >>> manager = SecureCredentialManager()
             >>> decrypted = manager.decrypt(encrypted_string)
         """
+        if self.key is None:
+            raise ValueError("Encryption key is not initialized")
         f = Fernet(self.key)
         return f.decrypt(encrypted_data.encode()).decode()