diff --git a/.gitignore b/.gitignore index 51646296..b55e9b56 100644 --- a/.gitignore +++ b/.gitignore @@ -166,3 +166,4 @@ cython_debug/ # NodeJS node_modules/ +/.idea diff --git a/libraries/dagster-teradata/README.md b/libraries/dagster-teradata/README.md index 5684c8ed..d6270a30 100644 --- a/libraries/dagster-teradata/README.md +++ b/libraries/dagster-teradata/README.md @@ -122,6 +122,72 @@ def test_create_teradata_compute_cluster(tmp_path): } ) ``` +## BTEQ Operator + +The bteq_operator method enables execution of SQL statements or BTEQ (Basic Teradata Query) scripts using the Teradata BTEQ utility. +It supports running commands either on the local machine or on a remote machine over SSH — in both cases, the BTEQ utility must be installed on the target system. + +### Key Features + +- Executes SQL provided as a string or from a script file (only one can be used at a time). +- Supports custom encoding for the script or session. +- Configurable timeout and return code handling. +- Remote execution supports authentication using a password or an SSH key. +- Works in both local and remote setups, provided the BTEQ tool is installed on the system where execution takes place. + +> Ensure that the Teradata BTEQ utility is installed on the machine where the SQL statements or scripts will be executed. +> +> This could be: +> * The local machine where Dagster runs the task, for local execution. +> * The remote host accessed via SSH, for remote execution. +> * If executing remotely, also ensure that an SSH server (e.g., sshd) is running and accessible on the remote machine. + +### Parameters + +- `sql`: SQL statement(s) to be executed using BTEQ. (optional, mutually exclusive with `file_path`) +- `file_path`: If provided, this file will be used instead of the `sql` content. This path represents remote file path when executing remotely via SSH, or local file path when executing locally. (optional, mutually exclusive with `sql`) +- `remote_host`: Hostname or IP address for remote execution. If not provided, execution is assumed to be local. *(optional)* +- `remote_user`: Username used for SSH authentication on the remote host. Required if `remote_host` is specified. +- `remote_password`: Password for SSH authentication. Optional, and used as an alternative to `ssh_key_path`. +- `ssh_key_path`: Path to the SSH private key used for authentication. Optional, and used as an alternative to `remote_password`. +- `remote_port`: SSH port number for the remote host. Defaults to `22` if not specified. *(optional)* +- `remote_working_dir`: Temporary directory location on the remote host (via SSH) where the BTEQ script will be transferred and executed. Defaults to `/tmp` if not specified. This is only applicable when `ssh_conn_id` is provided. +- `bteq_script_encoding`: Character encoding for the BTEQ script file. Defaults to ASCII if not specified. +- `bteq_session_encoding`: Character set encoding for the BTEQ session. Defaults to ASCII if not specified. +- `bteq_quit_rc`: Accepts a single integer, list, or tuple of return codes. Specifies which BTEQ return codes should be treated as successful, allowing subsequent tasks to continue execution. +- `timeout`: Timeout (in seconds) for executing the BTEQ command. Default is 600 seconds (10 minutes). +- `timeout_rc`: Return code to use if the BTEQ execution fails due to a timeout. To allow Ops execution to continue after a timeout, include this value in `bteq_quit_rc`. If not specified, a timeout will raise an exception and stop the Ops. + +### Returns + +- Output of the BTEQ execution, or `None` if no output was produced. + +### Raises + +- `ValueError`: For invalid input or configuration +- `DagsterError`: If BTEQ execution fails or times out + +### Notes + +- Either `sql` or `file_path` must be provided, but not both. +- For remote execution, provide either `remote_password` or `ssh_key_path` (not both). +- Encoding and timeout handling are customizable. +- Validates remote port and authentication parameters. + +### Example Usage + +```python +# Local execution with direct SQL +output = bteq_operator(sql="SELECT * FROM table;") + +# Remote execution with file +output = bteq_operator( + file_path="script.sql", + remote_host="example.com", + remote_user="user", + ssh_key_path="/path/to/key.pem" +) +``` ## Development diff --git a/libraries/dagster-teradata/__init__.py b/libraries/dagster-teradata/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/libraries/dagster-teradata/dagster_teradata/__init__.py b/libraries/dagster-teradata/dagster_teradata/__init__.py index 8643c211..7a1098c6 100644 --- a/libraries/dagster-teradata/dagster_teradata/__init__.py +++ b/libraries/dagster-teradata/dagster_teradata/__init__.py @@ -7,6 +7,12 @@ teradata_resource as teradata_resource, ) +from dagster_teradata.teradata_compute_cluster_manager import ( + TeradataComputeClusterManager as TeradataComputeClusterManager, +) + +from dagster_teradata.ttu.bteq import Bteq as Bteq + __version__ = "0.0.3" DagsterLibraryRegistry.register( diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 8949c309..2d91b518 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -1,8 +1,9 @@ from contextlib import closing, contextmanager from datetime import datetime from textwrap import dedent -from typing import Any, List, Mapping, Optional, Sequence, Union, TYPE_CHECKING +from typing import Any, List, Mapping, Optional, Sequence, Union, TYPE_CHECKING, Literal +import re import dagster._check as check import teradatasql from dagster import ( @@ -16,12 +17,51 @@ from dagster._utils.cached_method import cached_method from pydantic import Field -from dagster_teradata import constants +from . import constants +from dagster_teradata.ttu.bteq import Bteq from dagster_teradata.teradata_compute_cluster_manager import ( TeradataComputeClusterManager, ) +def _handle_user_query_band_text(self, query_band_text) -> str: + """Validate given query_band and append if required values missed in query_band.""" + # Ensures 'appname=dagster' and 'org=teradata-internal-telem' are in query_band_text. + if query_band_text is not None: + # checking org doesn't exist in query_band, appending 'org=teradata-internal-telem' + # If it exists, user might have set some value of their own, so doing nothing in that case + pattern = r"org\s*=\s*([^;]*)" + match = re.search(pattern, query_band_text) + if not match: + if not query_band_text.endswith(";"): + query_band_text += ";" + query_band_text += "org=teradata-internal-telem;" + # Making sure appname in query_band contains 'dagster' + pattern = r"appname\s*=\s*([^;]*)" + # Search for the pattern in the query_band_text + match = re.search(pattern, query_band_text) + if match: + appname_value = match.group(1).strip() + # if appname exists and dagster not exists in appname then appending 'dagster' to existing + # appname value + if "dagster" not in appname_value.lower(): + new_appname_value = appname_value + "_dagster" + # Optionally, you can replace the original value in the query_band_text + updated_query_band_text = re.sub( + pattern, f"appname={new_appname_value}", query_band_text + ) + query_band_text = updated_query_band_text + else: + # if appname doesn't exist in query_band, adding 'appname=dagster' + if len(query_band_text.strip()) > 0 and not query_band_text.endswith(";"): + query_band_text += ";" + query_band_text += "appname=dagster;" + else: + query_band_text = "org=teradata-internal-telem;appname=dagster;" + + return query_band_text + + class TeradataResource(ConfigurableResource, IAttachDifferentObjectToOpContext): host: Optional[str] = Field(description="Teradata Database Hostname") user: Optional[str] = Field(description="User login name.") @@ -30,12 +70,25 @@ class TeradataResource(ConfigurableResource, IAttachDifferentObjectToOpContext): default=None, description=("Name of the default database to use."), ) + query_band: Optional[str] = None port: Optional[str] = None tmode: Optional[str] = "ANSI" logmech: Optional[str] = None browser: Optional[str] = None browser_tab_timeout: Optional[int] = None browser_timeout: Optional[int] = None + console_output_encoding: Optional[str] = Field( + default="utf-8", description="Console output encoding." + ) + bteq_session_encoding: Optional[str] = Field( + default="ASCII", description="BTEQ session encoding." + ) + bteq_output_width: Optional[int] = Field( + default=65531, description="BTEQ output width." + ) + bteq_quit_zero: Optional[bool] = Field( + default=False, description="BTEQ quit zero flag." + ) http_proxy: Optional[str] = None http_proxy_user: Optional[str] = None http_proxy_password: Optional[str] = None @@ -63,6 +116,7 @@ def _connection_args(self) -> Mapping[str, Any]: "user", "password", "database", + "query_band", "port", "tmode", "logmech", @@ -158,9 +212,19 @@ def get_connection(self): connection_params["oidc_sslmode"] = self.oidc_sslmode teradata_conn = teradatasql.connect(**connection_params) - + self.set_query_band(self.query_band, teradata_conn) yield teradata_conn + 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(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) + except Exception: + pass + def get_object_to_set_on_execution_context(self) -> Any: # Directly create a TeradataDagsterConnection here for backcompat since the TeradataDagsterConnection # has methods this resource does not have @@ -188,6 +252,7 @@ def __init__( ): self.teradata_connection_resource = teradata_connection_resource self.compute_cluster_manager = TeradataComputeClusterManager(self, log) + self.bteq = Bteq(self, teradata_connection_resource, log) self.log = log @public @@ -283,6 +348,149 @@ def create_fresh_database(teradata: TeradataResource): results = results.append(cursor.fetchall()) # type: ignore return results + def bteq_operator( + self, + sql: Optional[str] = None, + file_path: Optional[str] = None, + remote_host: Optional[str] = None, + remote_user: Optional[str] = None, + remote_password: Optional[str] = None, + ssh_key_path: Optional[str] = None, + remote_port: int = 22, + remote_working_dir: str = "/tmp", + bteq_script_encoding: Optional[str] = "utf-8", + bteq_session_encoding: Optional[str] = "ASCII", + bteq_quit_rc: Union[int, List[int]] = 0, + timeout: int | Literal[600] = 600, # Default to 10 minutes + timeout_rc: int | None = None, + ) -> Optional[int]: + """ + Execute BTEQ commands either locally or on a remote machine via SSH. + + This method provides a unified interface to run BTEQ operations with various configurations: + - Local or remote execution + - Direct SQL input or file-based input + - Custom encoding settings + - Timeout controls + - Authentication via password or SSH key + + Args: + sql (Optional[str]): SQL commands to execute directly. Mutually exclusive with file_path. + file_path (Optional[str]): [Optional] Path to an existing SQL or BTEQ script file. + remote_host (Optional[str]): Hostname or IP for remote execution. None for local execution. + remote_user (Optional[str]): Username for remote authentication. Required if remote_host specified. + remote_password (Optional[str]): Password for remote authentication. Alternative to ssh_key_path. + ssh_key_path (Optional[str]): Path to SSH private key for authentication. Alternative to remote_password. + remote_port (int): SSH port for remote connection. Defaults to 22. + remote_working_dir (str): Working directory on remote machine. Defaults to '/tmp'. + bteq_script_encoding (Optional[str]): Encoding for BTEQ script file. Defaults to 'utf-8'. + bteq_session_encoding (Optional[str]): Encoding for BTEQ session. Defaults to 'ASCII'. + bteq_quit_rc (Union[int, List[int]]): Acceptable return codes for BTEQ execution. Defaults to 0. + timeout (int | Literal[600]): Maximum execution time in seconds. Defaults to 600 (10 minutes). + timeout_rc (int | None): Specific return code to use for timeout cases. Defaults to None. + + Returns: + Optional[str]: The output of the BTEQ execution, or None if no output was produced. + + Raises: + ValueError: If input validation fails, including: + - Missing both sql and file_path + - Providing both sql and file_path + - Missing required remote authentication parameters + - Invalid remote port specification + DagsterError: If BTEQ execution fails or times out + + Note: + - For remote execution, either remote_password or ssh_key_path must be provided (but not both) + - Encoding handling follows these rules: + * ASCII session encoding: Uses default system encoding + * UTF8/UTF16 session encoding: Forces corresponding script encoding + - Timeout handling: + * MAXREQTIME parameter sets maximum execution time per request + * EXITONDELAY forces exit if timeout occurs + * Timeout return code can be customized + + Example: + # Local execution with direct SQL + >>> output = bteq_operator(sql="SELECT * FROM table;") + + # Remote execution with file + >>> output = bteq_operator( + ... file_path="script.sql", + ... remote_host="example.com", + ... remote_user="user", + ... ssh_key_path="/path/to/key.pem" + ... ) + """ + # Validate input parameters + if not sql and not file_path: + raise ValueError( + "BteqOperator requires either the 'sql' or 'file_path' parameter. Both are missing." + ) + + if sum(bool(x) for x in [sql, file_path]) > 1: + raise ValueError( + "BteqOperator requires either the 'sql' or 'file_path' parameter but not both." + ) + + # Validate remote execution parameters if needed + if remote_host: + if not remote_user: + raise ValueError("remote_user must be provided for remote execution") + if not ssh_key_path and not remote_password: + raise ValueError( + "Either ssh_key_path or remote_password must be provided for remote execution" + ) + if remote_password and ssh_key_path: + raise ValueError( + "Cannot specify both remote_password and ssh_key_path for remote execution" + ) + + # Validate network port + if remote_port < 1 or remote_port > 65535: + raise ValueError("remote_port must be a valid port number (1-65535)") + + # Handle encoding configuration + temp_file_read_encoding = "UTF-8" + if not bteq_session_encoding or bteq_session_encoding == "ASCII": + bteq_session_encoding = "" + if bteq_script_encoding == "UTF8": + temp_file_read_encoding = "UTF-8" + elif bteq_script_encoding == "UTF16": + temp_file_read_encoding = "UTF-16" + bteq_script_encoding = "" + elif bteq_session_encoding == "UTF8" and ( + not bteq_script_encoding or bteq_script_encoding == "ASCII" + ): + bteq_script_encoding = "UTF8" + elif bteq_session_encoding == "UTF16": + if not bteq_script_encoding or bteq_script_encoding == "ASCII": + bteq_script_encoding = "UTF8" + + # Map BTEQ encoding to Python file reading encoding + if bteq_script_encoding == "UTF8": + temp_file_read_encoding = "UTF-8" + elif bteq_script_encoding == "UTF16": + temp_file_read_encoding = "UTF-16" + + # Delegate execution to the underlying BTEQ implementation + return self.bteq.bteq_operator( + sql=sql, + file_path=file_path, + remote_host=remote_host, + remote_user=remote_user, + remote_password=remote_password, + ssh_key_path=ssh_key_path, + remote_port=remote_port, + remote_working_dir=remote_working_dir, + bteq_script_encoding=bteq_script_encoding, + bteq_session_encoding=bteq_session_encoding, + bteq_quit_rc=bteq_quit_rc, + timeout=timeout, + timeout_rc=timeout_rc, + temp_file_read_encoding=temp_file_read_encoding, + ) + def drop_database(self, databases: Union[str, Sequence[str]]) -> None: """ Drop one or more databases in Teradata. diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/__init__.py b/libraries/dagster-teradata/dagster_teradata/ttu/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py new file mode 100644 index 00000000..6484a5c1 --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py @@ -0,0 +1,773 @@ +import os +import socket +import subprocess +import tempfile +from contextlib import contextmanager +from typing import Optional, List, Union, Literal, cast + +import paramiko +from dagster import DagsterError +from paramiko.client import SSHClient +from paramiko.ssh_exception import SSHException + +from dagster_teradata.ttu.utils.bteq_util import ( + prepare_bteq_command_for_local_execution, + prepare_bteq_command_for_remote_execution, + prepare_bteq_script_for_local_execution, + prepare_bteq_script_for_remote_execution, + verify_bteq_installed, + verify_bteq_installed_remote, + is_valid_file, + is_valid_remote_bteq_script_file, + transfer_file_sftp, + read_file, + is_valid_encoding, +) +from dagster_teradata.ttu.utils.encryption_utils import ( + SecureCredentialManager, + generate_random_password, + generate_encrypted_file_with_openssl, + decrypt_remote_file_to_string, + get_stored_credentials, +) + + +class Bteq: + """ + Main BTEQ operator class for executing Teradata BTEQ commands either locally or remotely. + + Features: + - Local and remote execution via SSH + - Secure credential handling and encryption + - File-based and direct SQL execution + - Timeout and error handling + - Encoding support (ASCII, UTF-8, UTF-16) + - Temporary file management + + Attributes: + connection: Legacy connection object (maintained for compatibility) + teradata_connection_resource: Contains Teradata connection parameters + log: Logger instance for operation logging + cred_manager: SecureCredentialManager instance for encryption/decryption + ssh_client: Active SSH connection for remote execution (None for local) + """ + + def __init__(self, connection, teradata_connection_resource, log): + """ + Initialize BTEQ operator with connection resources and logger. + + Args: + connection: Legacy connection object (maintained for compatibility) + teradata_connection_resource: Contains Teradata connection parameters + including host, username, password + log: Logger instance for operation logging + """ + self.remote_port = None + self.ssh_key_path = None + self.remote_remote_password = None + self.remote_user = None + self.remote_host = None + self.file_path = None + self.temp_file_read_encoding = None + self.bteq_quit_rc = None + self.bteq_session_encoding = None + self.timeout_rc = None + self.bteq_script_encoding = None + self.sql = None + self.remote_working_dir = None + self.timeout = None + self.connection = connection + self.log = log + self.teradata_connection_resource = teradata_connection_resource + self.cred_manager = SecureCredentialManager() + self.ssh_client = None # Will hold active SSH connection if remote execution + + def bteq_operator( + self, + sql: Optional[str] = None, + file_path: Optional[str] = None, + remote_host: Optional[str] = None, + remote_user: Optional[str] = None, + remote_password: Optional[str] = None, + ssh_key_path: Optional[str] = None, + remote_port: int = 22, + remote_working_dir: str = "/tmp", + bteq_script_encoding: Optional[str] = "utf-8", + bteq_session_encoding: Optional[str] = "ASCII", + bteq_quit_rc: Union[int, List[int]] = 0, + timeout: int | Literal[600] = 600, + timeout_rc: int | None = None, + temp_file_read_encoding: Optional[str] = "UTF-8", + ) -> int | None: + """ + Execute BTEQ commands either locally or remotely. + + Args: + sql: SQL commands to execute directly + file_path: Path to file containing SQL commands + remote_host: Hostname for remote execution (None for local) + remote_user: Username for remote authentication + remote_password: Password for remote authentication + ssh_key_path: Path to SSH private key for authentication + remote_port: SSH port (default: 22) + remote_working_dir: Remote working directory (default: '/tmp') + bteq_script_encoding: Encoding for BTEQ script file + bteq_session_encoding: Encoding for BTEQ session + bteq_quit_rc: Acceptable return codes (default: 0) + timeout: Maximum execution time in seconds (default: 600) + timeout_rc: Specific return code for timeout cases + temp_file_read_encoding: Encoding for reading temporary files + + Returns: + int: Exit status code from BTEQ execution, or None if no execution occurred + + Raises: + ValueError: For invalid input parameters + DagsterError: For execution failures or timeouts + """ + self.sql = sql + self.file_path = file_path + self.remote_host = remote_host + self.remote_user = remote_user + self.remote_remote_password = remote_password + self.ssh_key_path = ssh_key_path + self.remote_port = remote_port + self.remote_working_dir = remote_working_dir + self.bteq_script_encoding = bteq_script_encoding + self.bteq_session_encoding = bteq_session_encoding + self.bteq_quit_rc = bteq_quit_rc + self.timeout = timeout + self.timeout_rc = timeout_rc + self.temp_file_read_encoding = temp_file_read_encoding + + # Local execution + if not self.remote_host: + if self.sql: + bteq_script = prepare_bteq_script_for_local_execution(sql=self.sql) + self.log.debug( + "Executing BTEQ script with SQL content: %s", bteq_script + ) + return self.execute_bteq_script( + bteq_script, + self.remote_working_dir, + self.bteq_script_encoding, + self.timeout, + self.timeout_rc, + self.bteq_session_encoding, + self.bteq_quit_rc, + self.temp_file_read_encoding, + ) + elif self.file_path: + if not is_valid_file(self.file_path): + raise ValueError( + f"The provided file path '{self.file_path}' is invalid or does not exist." + ) + try: + is_valid_encoding( + self.file_path, self.temp_file_read_encoding or "UTF-8" + ) + except UnicodeDecodeError as e: + errmsg = f"The provided file '{self.file_path}' encoding is different from BTEQ I/O encoding i.e.'UTF-8'." + if self.bteq_script_encoding: + errmsg = f"The provided file '{self.file_path}' encoding is different from the specified BTEQ I/O encoding '{self.bteq_script_encoding}'." + raise ValueError(errmsg) from e + return self._handle_local_bteq_file(file_path=self.file_path) + + # Remote execution + elif self.remote_host: + if self.sql: + bteq_script = prepare_bteq_script_for_remote_execution( + teradata_connection_resource=self.teradata_connection_resource, + sql=self.sql, + ) + self.log.debug( + "Executing BTEQ script with SQL content: %s", bteq_script + ) + return self.execute_bteq_script( + bteq_script, + self.remote_working_dir, + self.bteq_script_encoding, + self.timeout, + self.timeout_rc, + self.bteq_session_encoding, + self.bteq_quit_rc, + self.temp_file_read_encoding, + self.remote_host, + self.remote_user, + self.remote_remote_password, + self.ssh_key_path, + self.remote_port, + ) + if self.file_path: + if not self._setup_ssh_connection( + host=self.remote_host, + 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 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, + file_path=self.file_path, + ) + raise ValueError( + f"The provided remote file path '{self.file_path}' is invalid or file does not exist on remote machine at given path." + ) + else: + raise ValueError( + "BteqOperator requires either the 'sql' or 'file_path' parameter. Both are missing." + ) + return None + + def execute_bteq_script( + self, + bteq_script: str, + remote_working_dir: str | None, + bteq_script_encoding: str | None, + 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, + temp_file_read_encoding: str | None, + remote_host: str | None = None, + remote_user: str | None = None, + remote_password: str | None = None, + ssh_key_path: str | None = None, + remote_port: int = 22, + ) -> int | None: + """ + Execute BTEQ script either locally or remotely. + + Args: + bteq_script: The BTEQ script content to execute + remote_working_dir: Working directory for remote execution + bteq_script_encoding: Encoding for BTEQ script file + timeout: Maximum execution time in seconds + timeout_rc: Return code for timeout cases + bteq_session_encoding: Encoding for BTEQ session + bteq_quit_rc: Acceptable return codes + temp_file_read_encoding: Encoding for reading temporary files + remote_host: Remote hostname (None for local) + remote_user: Remote username + remote_password: Remote password + ssh_key_path: Path to SSH private key + remote_port: SSH port + + Returns: + int: Exit status code from BTEQ execution + + Note: + Delegates to execute_bteq_script_at_local or execute_bteq_script_at_remote + based on whether remote_host is specified + """ + if remote_host: + return self.execute_bteq_script_at_remote( + bteq_script=bteq_script, + remote_working_dir=remote_working_dir, + bteq_script_encoding=bteq_script_encoding, + timeout=timeout, + timeout_rc=timeout_rc, + bteq_session_encoding=bteq_session_encoding, + bteq_quit_rc=bteq_quit_rc, + temp_file_read_encoding=temp_file_read_encoding, + remote_host=remote_host, + remote_user=remote_user, + remote_password=remote_password, + ssh_key_path=ssh_key_path, + remote_port=remote_port, + ) + return self.execute_bteq_script_at_local( + bteq_script=bteq_script, + bteq_script_encoding=bteq_script_encoding, + timeout=timeout, + timeout_rc=timeout_rc, + bteq_quit_rc=bteq_quit_rc, + bteq_session_encoding=bteq_session_encoding, + temp_file_read_encoding=temp_file_read_encoding, + ) + + def execute_bteq_script_at_remote( + self, + bteq_script: str, + remote_working_dir: str | None, + bteq_script_encoding: str | None, + 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, + temp_file_read_encoding: str | None, + remote_host: str | None, + remote_user: str | None, + remote_password: str | None, + ssh_key_path: str | None, + remote_port: int = 22, + ) -> int | None: + """ + Execute BTEQ script on a remote machine via SSH. + + Args: + bteq_script: The BTEQ script content to execute + remote_working_dir: Working directory on remote machine + bteq_script_encoding: Encoding for BTEQ script file + timeout: Maximum execution time in seconds + timeout_rc: Return code for timeout cases + bteq_session_encoding: Encoding for BTEQ session + bteq_quit_rc: Acceptable return codes + temp_file_read_encoding: Encoding for reading temporary files + remote_host: Remote hostname + remote_user: Remote username + remote_password: Remote password + ssh_key_path: Path to SSH private key + remote_port: SSH port + + Returns: + int: Exit status code from BTEQ execution + + Raises: + DagsterError: For SSH failures, execution errors, or timeouts + """ + with self.preferred_temp_directory() as tmp_dir: + file_path = os.path.join(tmp_dir, "bteq_script.txt") + with open( + file_path, "w", encoding=str(temp_file_read_encoding or "UTF-8") + ) as f: + f.write(bteq_script) + return self._transfer_to_and_execute_bteq_on_remote( + file_path, + remote_working_dir, + bteq_script_encoding, + timeout, + timeout_rc, + bteq_quit_rc, + bteq_session_encoding, + tmp_dir, + cast(str, remote_host), + cast(str, remote_user), + remote_password, + ssh_key_path, + remote_port, + ) + + def _transfer_to_and_execute_bteq_on_remote( + self, + file_path: str, + remote_working_dir: str | None, + bteq_script_encoding: str | None, + 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, + tmp_dir: str, + remote_host: str, + remote_user: str, + remote_password: str | None = None, + ssh_key_path: str | None = None, + remote_port: int = 22, + ) -> int | None: + """ + Transfer and execute BTEQ script on remote machine with encryption. + + Args: + file_path: Local path to BTEQ script file + remote_working_dir: Remote working directory + bteq_script_encoding: Encoding for BTEQ script file + timeout: Maximum execution time in seconds + timeout_rc: Return code for timeout cases + bteq_quit_rc: Acceptable return codes + bteq_session_encoding: Encoding for BTEQ session + tmp_dir: Local temporary directory + remote_host: Remote hostname + remote_user: Remote username + remote_password: Remote password + ssh_key_path: Path to SSH private key + remote_port: SSH port + + Returns: + int: Exit status code from BTEQ execution + + Note: + - Uses OpenSSL AES-256-CBC encryption for secure file transfer + - Automatically cleans up temporary files + """ + encrypted_file_path = None + remote_encrypted_path = None + try: + if not self._setup_ssh_connection( + host=remote_host, + user=remote_user, + password=remote_password, + key_path=ssh_key_path, + port=remote_port, + ): + raise DagsterError( + "Failed to establish SSH connection. Please check the provided credentials." + ) + if self.ssh_client is None: + raise DagsterError( + "Failed to establish SSH connection. `ssh_client` is None." + ) + verify_bteq_installed_remote(self.ssh_client) + password = generate_random_password() # Encryption/Decryption password + encrypted_file_path = os.path.join(tmp_dir, "bteq_script.enc") + generate_encrypted_file_with_openssl( + file_path, password, encrypted_file_path + ) + remote_encrypted_path = os.path.join( + remote_working_dir or "", "bteq_script.enc" + ) + + transfer_file_sftp( + self.ssh_client, encrypted_file_path, remote_encrypted_path + ) + + bteq_command_str = prepare_bteq_command_for_remote_execution( + timeout=timeout, + bteq_script_encoding=bteq_script_encoding or "", + bteq_session_encoding=bteq_session_encoding or "", + timeout_rc=timeout_rc or -1, + ) + self.log.debug("Executing BTEQ command: %s", bteq_command_str) + + exit_status, stdout, stderr = decrypt_remote_file_to_string( + self.ssh_client, + remote_encrypted_path, + password, + bteq_command_str, + ) + + failure_message = None + self.log.debug("stdout : %s", stdout) + self.log.debug("stderr : %s", stderr) + self.log.debug("exit_status : %s", exit_status) + + if "Failure" in stderr or "Error" in stderr: + failure_message = stderr + # Raising an exception if there is any failure in bteq and also user wants to fail the + # task otherwise just log the error message as warning to not fail the task. + if ( + failure_message + and exit_status != 0 + and exit_status + not in ( + bteq_quit_rc + if isinstance(bteq_quit_rc, (list, tuple)) + else [bteq_quit_rc if bteq_quit_rc is not None else 0] + ) + ): + raise DagsterError(f"BTEQ task failed with error: {failure_message}") + if failure_message: + self.log.warning(failure_message) + return exit_status + # If we get here, everything succeeded + return exit_status # Explicit return instead of else block + except (OSError, socket.gaierror): + raise DagsterError( + "SSH connection timed out. Please check the network or server availability." + ) + except SSHException as e: + raise DagsterError( + f"An unexpected error occurred during SSH connection: {str(e)}" + ) + except DagsterError as e: + raise e + except Exception as e: + raise DagsterError( + f"An unexpected error occurred while executing BTEQ script on remote machine: {str(e)}" + ) + finally: + # Remove the local script file + if encrypted_file_path and os.path.exists(encrypted_file_path): + os.remove(encrypted_file_path) + # Cleanup: Delete the remote temporary file + if encrypted_file_path: + cleanup_en_command = f"rm -f {remote_encrypted_path}" + if self.ssh_client and self.connection: + if self.ssh_client is None: + raise DagsterError( + "Failed to establish SSH connection. `ssh_client` is None." + ) + self.ssh_client.exec_command(cleanup_en_command) + + def execute_bteq_script_at_local( + self, + bteq_script: str, + bteq_script_encoding: str | None, + 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, + temp_file_read_encoding: str | None, + ) -> int | None: + """ + Execute BTEQ script on local machine. + + Args: + bteq_script: The BTEQ script content to execute + bteq_script_encoding: Encoding for BTEQ script file + timeout: Maximum execution time in seconds + timeout_rc: Return code for timeout cases + bteq_quit_rc: Acceptable return codes + bteq_session_encoding: Encoding for BTEQ session + temp_file_read_encoding: Encoding for reading temporary files + + Returns: + int: Exit status code from BTEQ execution + + Raises: + DagsterError: For execution failures or timeouts + """ + verify_bteq_installed() + bteq_command_str = prepare_bteq_command_for_local_execution( + teradata_connection_resource=self.teradata_connection_resource, + timeout=timeout, + bteq_script_encoding=bteq_script_encoding or "", + bteq_session_encoding=bteq_session_encoding or "", + timeout_rc=timeout_rc or -1, + ) + self.log.debug("Executing BTEQ command: %s", bteq_command_str) + + process = subprocess.Popen( + bteq_command_str, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + shell=True, + preexec_fn=os.setsid, + ) + encode_bteq_script = bteq_script.encode(str(temp_file_read_encoding or "UTF-8")) + self.log.debug("encode_bteq_script : %s", encode_bteq_script) + stdout_data, _ = process.communicate(input=encode_bteq_script) + self.log.debug("stdout_data : %s", stdout_data) + try: + # https://docs.python.org/3.10/library/subprocess.html#subprocess.Popen.wait timeout is in seconds + process.wait( + timeout=(timeout or 0) + 60 + ) # Adding 1 minute extra for BTEQ script timeout + except subprocess.TimeoutExpired: + self.on_kill() + raise DagsterError(f"BTEQ command timed out after {timeout} seconds.") + + failure_message = None + if stdout_data is None: + raise DagsterError("Process stdout is None. Unable to read BTEQ output.") + decoded_line = "" + for line in stdout_data.splitlines(): + try: + decoded_line = line.decode("UTF-8").strip() + self.log.debug("decoded_line : %s", decoded_line) + except UnicodeDecodeError: + self.log.warning("Failed to decode line: %s", line) + if "Failure" in decoded_line or "Error" in decoded_line: + failure_message = decoded_line + # Raising an exception if there is any failure in bteq and also user wants to fail the + # task otherwise just log the error message as warning to not fail the task. + if ( + failure_message + and process.returncode != 0 + and process.returncode + not in ( + bteq_quit_rc + if isinstance(bteq_quit_rc, (list, tuple)) + else [bteq_quit_rc if bteq_quit_rc is not None else 0] + ) + ): + raise DagsterError(f"BTEQ task failed with error: {failure_message}") + if failure_message: + self.log.warning(failure_message) + + return process.returncode + + def on_kill(self): + """Terminate the subprocess if running.""" + conn = self.connection + process = conn.get("sp") + if process: + try: + process.terminate() + process.wait(timeout=5) + except subprocess.TimeoutExpired: + self.log.warning( + "Subprocess did not terminate in time. Forcing kill..." + ) + process.kill() + except Exception as e: + self.log.error("Failed to terminate subprocess: %s", str(e)) + + def get_dagster_home_dir(self) -> str: + """Get the DAGSTER_HOME directory from environment variables.""" + return os.environ.get("DAGSTER_HOME", "~/.dagster_home") + + @contextmanager + def preferred_temp_directory(self, prefix="bteq_"): + """ + Context manager for creating a temporary directory. + + Args: + prefix: Prefix for the temporary directory name + + Yields: + str: Path to the created temporary directory + + Note: + Falls back to DAGSTER_HOME if system temp directory is not usable + """ + try: + temp_dir = tempfile.gettempdir() + if not os.path.isdir(temp_dir) or not os.access(temp_dir, os.W_OK): + raise OSError("OS temp dir not usable") + except Exception: + temp_dir = self.get_dagster_home_dir() + + with tempfile.TemporaryDirectory(dir=temp_dir, prefix=prefix) as tmp: + yield tmp + + def _handle_remote_bteq_file( + self, ssh_client: SSHClient, file_path: str | None + ) -> int | None: + """ + Handle execution of a remote BTEQ script file. + + Args: + ssh_client: Active SSH connection + file_path: Path to remote BTEQ script file + + Returns: + int: Exit status code from BTEQ execution + + Raises: + ValueError: For invalid file path + """ + if file_path: + with ssh_client: + sftp = ssh_client.open_sftp() + try: + with sftp.open(file_path, "r") as remote_file: + file_content = remote_file.read().decode( + self.temp_file_read_encoding or "UTF-8" + ) + finally: + sftp.close() + bteq_script = prepare_bteq_script_for_remote_execution( + teradata_connection_resource=self.teradata_connection_resource, + sql=file_content, + ) + self.log.debug( + "Executing BTEQ script with SQL content: %s", bteq_script + ) + return self.execute_bteq_script_at_remote( + bteq_script, + self.remote_working_dir, + self.bteq_script_encoding, + self.timeout, + self.timeout_rc, + self.bteq_session_encoding, + self.bteq_quit_rc, + self.temp_file_read_encoding, + remote_host=self.remote_host, + remote_user=self.remote_user, + remote_password=self.remote_remote_password, + ssh_key_path=self.ssh_key_path, + remote_port=self.remote_port or 22, + ) + else: + raise ValueError( + "Please provide a valid file path for the BTEQ script to be executed on the remote machine." + ) + + def _handle_local_bteq_file(self, file_path: str) -> int | None: + """ + Handle execution of a local BTEQ script file. + + Args: + file_path: Path to local BTEQ script file + + Returns: + int: Exit status code from BTEQ execution, or None if file is invalid + """ + if file_path and is_valid_file(file_path): + file_content = read_file( + file_path, encoding=str(self.temp_file_read_encoding or "UTF-8") + ) + bteq_script = prepare_bteq_script_for_local_execution( + sql=file_content, + ) + self.log.debug("Executing BTEQ script with SQL content: %s", bteq_script) + result = self.execute_bteq_script( + bteq_script, + self.remote_working_dir, + self.bteq_script_encoding, + self.timeout, + self.timeout_rc, + self.bteq_session_encoding, + self.bteq_quit_rc, + self.temp_file_read_encoding, + ) + return result + return None + + def _setup_ssh_connection( + self, + host: str, + user: Optional[str], + password: Optional[str], + key_path: Optional[str], + port: int, + ) -> bool: + """ + Establish SSH connection using either password or key authentication. + + Args: + host: Remote hostname + user: Remote username + password: Remote password (optional if key_path provided) + key_path: Path to SSH private key (optional if password provided) + port: SSH port + + Returns: + bool: True if connection succeeded, False otherwise + + Raises: + DagsterError: If connection fails + + Note: + - Tries stored credentials if no password provided + - Prompts for password if no credentials available + - Stores new credentials if successfully authenticated + """ + try: + self.ssh_client = paramiko.SSHClient() + self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + if key_path: + key = paramiko.RSAKey.from_private_key_file(key_path) + 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 + ) + + self.ssh_client.connect( + host, port=port, username=user, password=password + ) + + self.log.info(f"SSH connected to {user}@{host}") + return True + except Exception as e: + raise DagsterError(f"SSH connection failed: {e}") diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py new file mode 100644 index 00000000..74c1d0af --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py @@ -0,0 +1,343 @@ +from __future__ import annotations + +import os +import shutil +import stat +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from paramiko import SSHClient + +from dagster import DagsterError + + +def identify_os(ssh_client: SSHClient) -> str: + """ + Identify the operating system of a remote machine via SSH. + + Args: + ssh_client (SSHClient): An active SSH connection to the remote machine. + + Returns: + str: The name of the operating system in lowercase (e.g., 'linux', 'windows'). + + Note: + Uses 'uname' command for Unix-like systems and 'ver' for Windows systems. + """ + stdin, stdout, stderr = ssh_client.exec_command("uname || ver") + return stdout.read().decode().lower() + + +def verify_bteq_installed(): + """ + Verify if BTEQ is installed and available in the local system's PATH. + + Raises: + DagsterError: If BTEQ executable is not found in the system PATH. + """ + if shutil.which("bteq") is None: + raise DagsterError( + "BTEQ is not installed or not available in the system's PATH." + ) + + +def verify_bteq_installed_remote(ssh_client: SSHClient): + """ + Verify if BTEQ is installed on a remote machine accessible via SSH. + + Args: + ssh_client (SSHClient): An active SSH connection to the remote machine. + + Raises: + DagsterError: If BTEQ executable is not found on the remote machine. + + Note: + Detects the remote OS and uses appropriate commands to check for BTEQ: + - Windows: 'where bteq' + - macOS: Checks via zsh if available, falls back to 'which bteq' + - Other Unix-like: 'which bteq' + """ + # Detect OS + os_info = identify_os(ssh_client) + + if "windows" in os_info: + check_cmd = "where bteq" + elif "darwin" in os_info: + # Check if zsh exists first for remote shell compatibility + stdin, stdout, stderr = ssh_client.exec_command("command -v zsh") + zsh_path = stdout.read().strip() + if zsh_path: + # If zsh is available, use it to check for bteq in the user's environment + check_cmd = 'zsh -l -c "which bteq"' + else: + # Fallback to default shell if zsh is not available + check_cmd = "which bteq" + else: + # Default command for other Unix-like systems + check_cmd = "which bteq" + + # Execute the command to check for bteq + stdin, stdout, stderr = ssh_client.exec_command(check_cmd) + exit_status = stdout.channel.recv_exit_status() + output = stdout.read().strip() + error = stderr.read().strip() + + # Raise an error if bteq is not found in the PATH + if exit_status != 0 or not output: + raise DagsterError( + "BTEQ is not installed or not available on the remote machine. (%s)", error + ) + + +def transfer_file_sftp(ssh_client: SSHClient, local_path: str, remote_path: str): + """ + Transfer a file from local to remote machine using SFTP. + + Args: + ssh_client (SSHClient): An active SSH connection to the remote machine. + local_path (str): Path to the local file to be transferred. + remote_path (str): Destination path on the remote machine. + + Note: + Currently opens and immediately closes SFTP connection without transfer. + This appears to be a placeholder implementation. + """ + sftp = ssh_client.open_sftp() + sftp.put(local_path, remote_path) + sftp.close() + + +def prepare_bteq_script_for_remote_execution( + teradata_connection_resource, sql: str +) -> str: + """ + Build a BTEQ script with connection details for remote execution. + + Args: + teradata_connection_resource: Resource containing Teradata connection details + (host, user, password). + sql (str): SQL commands to be executed. + + Returns: + str: Complete BTEQ script including login commands and SQL. + + Note: + The .LOGON command is included in the script for remote execution to avoid + exposing credentials in command line arguments. + """ + script_lines = [] + host = teradata_connection_resource.host + login = teradata_connection_resource.user + password = teradata_connection_resource.password + script_lines.append(f".LOGON {host}/{login},{password}") + return _prepare_bteq_script(script_lines, sql) + + +def prepare_bteq_script_for_local_execution(sql: str) -> str: + """ + Build a BTEQ script for local execution. + + Args: + sql (str): SQL commands to be executed. + + Returns: + str: Complete BTEQ script including SQL commands. + + Note: + For local execution, connection details are typically passed as command line + arguments rather than in the script. + """ + script_lines: list[str] = [] + return _prepare_bteq_script(script_lines, sql) + + +def _prepare_bteq_script(script_lines: list[str], sql: str) -> str: + """ + Internal helper function to assemble BTEQ script components. + + Args: + script_lines (list[str]): List of BTEQ commands to include before the SQL. + sql (str): SQL commands to be executed. + + Returns: + str: Complete BTEQ script with all commands and SQL. + """ + script_lines.append(sql.strip()) + script_lines.append(".EXIT") + return "\n".join(script_lines) + + +def _prepare_bteq_command( + timeout: int, + bteq_script_encoding: str, + bteq_session_encoding: str, + timeout_rc: int, +) -> list[str]: + """ + Prepare the core BTEQ command with common parameters. + + Args: + timeout (int): Maximum execution time in seconds before timeout. + bteq_script_encoding (str): Character encoding for the script file. + bteq_session_encoding (str): Character encoding for the session. + timeout_rc (int): Return code to use when timeout occurs. + + Returns: + list[str]: List of command components that can be joined into a full command. + """ + bteq_core_cmd = ["bteq"] + if bteq_session_encoding: + bteq_core_cmd.append(f" -e {bteq_script_encoding}") + bteq_core_cmd.append(f" -c {bteq_session_encoding}") + bteq_core_cmd.append('"') + bteq_core_cmd.append(f".SET EXITONDELAY ON MAXREQTIME {timeout}") + if timeout_rc is not None and timeout_rc >= 0: + bteq_core_cmd.append(f" RC {timeout_rc}") + bteq_core_cmd.append(";") + # Dagster doesn't display the script of BTEQ in UI but only in log so WIDTH is 500 enough + bteq_core_cmd.append(" .SET WIDTH 500;") + return bteq_core_cmd + + +def prepare_bteq_command_for_remote_execution( + timeout: Optional[int], # or: timeout: int | None + bteq_script_encoding: str, + bteq_session_encoding: str, + timeout_rc: int, +) -> str: + """ + Prepare the BTEQ command string for remote execution. + + Args: + timeout (int): Maximum execution time in seconds before timeout. + bteq_script_encoding (str): Character encoding for the script file. + bteq_session_encoding (str): Character encoding for the session. + timeout_rc (int): Return code to use when timeout occurs. + + Returns: + str: Complete BTEQ command string for remote execution. + """ + bteq_core_cmd = _prepare_bteq_command( + 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: Optional[int], # or: timeout: int | None + bteq_script_encoding: str, + bteq_session_encoding: str, + timeout_rc: int, +) -> str: + """ + Prepare the BTEQ command string for local execution. + + Args: + teradata_connection_resource: Resource containing Teradata connection details. + timeout (int): Maximum execution time in seconds before timeout. + bteq_script_encoding (str): Character encoding for the script file. + bteq_session_encoding (str): Character encoding for the session. + timeout_rc (int): Return code to use when timeout occurs. + + Returns: + str: Complete BTEQ command string for local execution including login details. + """ + bteq_core_cmd = _prepare_bteq_command( + timeout or 600, bteq_script_encoding, bteq_session_encoding, timeout_rc + ) + host = teradata_connection_resource.host + login = teradata_connection_resource.user + password = teradata_connection_resource.password + bteq_core_cmd.append(f" .LOGON {host}/{login},{password}") + bteq_core_cmd.append('"') + bteq_command_str = " ".join(bteq_core_cmd) + return bteq_command_str + + +def is_valid_file(file_path: str) -> bool: + """ + Check if a file exists at the given path. + + Args: + file_path (str): Path to the file to check. + + Returns: + bool: True if the file exists, False otherwise. + """ + return os.path.isfile(file_path) + + +def is_valid_encoding(file_path: str, encoding: str = "UTF-8") -> bool: + """ + Check if a file can be read with the specified encoding. + + Args: + file_path (str): Path to the file to be checked. + encoding (str, optional): Encoding to test. Defaults to "UTF-8". + + Returns: + bool: True if the file can be read with the specified encoding, False otherwise. + + Raises: + UnicodeDecodeError: If the file cannot be read with the specified encoding. + """ + with open(file_path, encoding=encoding) as f: + f.read() + return True + + +def read_file(file_path: str, encoding: str = "UTF-8") -> str: + """ + Read the content of a file with the specified encoding. + + Args: + file_path (str): Path to the file to be read. + encoding (str, optional): Encoding to use for reading. Defaults to "UTF-8". + + Returns: + str: Content of the file as a string. + + Raises: + FileNotFoundError: If the file does not exist. + UnicodeDecodeError: If the file cannot be read with the specified encoding. + """ + if not os.path.isfile(file_path): + raise FileNotFoundError(f"The file {file_path} does not exist.") + + with open(file_path, encoding=encoding) as f: + return f.read() + + +def is_valid_remote_bteq_script_file( + ssh_client: SSHClient, remote_file_path: str | None, logger=None +) -> bool: + """ + Check if a remote file is a valid BTEQ script file. + + Args: + ssh_client (SSHClient): An active SSH connection to the remote machine. + remote_file_path (str): Path to the remote file to check. + logger (optional): Logger instance for error logging. + + Returns: + bool: True if the file exists and is a regular file, False otherwise. + """ + if remote_file_path: + sftp_client = ssh_client.open_sftp() + try: + # Get file metadata + file_stat = sftp_client.stat(remote_file_path) + if file_stat.st_mode: + is_regular_file = stat.S_ISREG(file_stat.st_mode) + return is_regular_file + return False + except FileNotFoundError: + if logger: + logger.error("File does not exist on remote at : %s", remote_file_path) + return False + finally: + sftp_client.close() + else: + return False diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py new file mode 100644 index 00000000..3b329371 --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py @@ -0,0 +1,318 @@ +from __future__ import annotations + +import json +import os +import secrets +import string +import subprocess +from datetime import datetime +from typing import Optional + +from cryptography.fernet import Fernet +from dagster import DagsterError + + +class SecureCredentialManager: + """ + A secure credential manager that handles encryption and decryption of sensitive data + using Fernet symmetric encryption. + + Attributes: + key (bytes): The encryption key used for cryptographic operations. + key_file (str): Path to the file where the encryption key is stored. + """ + + def __init__(self, key_file: Optional[str] = None): + """ + Initialize the credential manager with an optional custom key file path. + + Args: + key_file (Optional[str]): Path to the encryption key file. If not provided, + defaults to '~/.ssh/bteq_cred_key'. + + Note: + If the key file doesn't exist, a new encryption key will be generated and stored. + """ + self.key = None + self.key_file = key_file or os.path.expanduser("~/.ssh/bteq_cred_key") + self._load_or_generate_key() + + def _load_or_generate_key(self): + """ + Load an existing encryption key or generate a new one with secure permissions. + + Raises: + DagsterError: If there's any issue in loading or generating the key. + """ + try: + if os.path.exists(self.key_file): + with open(self.key_file, "rb") as f: + self.key = f.read() + else: + self.key = Fernet.generate_key() + key_dir = os.path.dirname(self.key_file) + os.makedirs(key_dir, exist_ok=True) + + with open(self.key_file, "wb") as f: + f.write(self.key) + os.chmod(self.key_file, 0o600) # Restrict to owner-only permissions + except Exception as e: + raise DagsterError(f"Encryption key handling failed: {e}") + + def encrypt(self, data: str) -> str: + """ + Encrypt sensitive string data using Fernet symmetric encryption. + + Args: + data (str): The plaintext string to encrypt. + + Returns: + str: The encrypted string in URL-safe base64 format. + + Example: + >>> 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() + + def decrypt(self, encrypted_data: str) -> str: + """ + Decrypt an encrypted string back to plaintext. + + Args: + encrypted_data (str): The encrypted string in URL-safe base64 format. + + Returns: + str: The decrypted plaintext string. + + Raises: + cryptography.fernet.InvalidToken: If the encrypted data is invalid or corrupted. + + Example: + >>> 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() + + +def generate_random_password(length: int = 12) -> str: + """ + Generate a cryptographically secure random password. + + Args: + length (int, optional): Length of the password to generate. Defaults to 12. + + Returns: + str: A randomly generated password containing letters, digits, and special characters. + + Note: + Uses Python's secrets module for cryptographically secure random generation. + The password includes: + - Uppercase and lowercase letters + - Digits + - Special punctuation characters + + Example: + >>> password = generate_random_password(16) + """ + # Define the character set: letters, digits, and special characters + characters = string.ascii_letters + string.digits + string.punctuation + # Generate a random password using cryptographically secure random selection + password = "".join(secrets.choice(characters) for _ in range(length)) + return password + + +def generate_encrypted_file_with_openssl( + file_path: str, password: str, out_file: str +) -> None: + """ + Encrypt a file using OpenSSL with AES-256-CBC encryption. + + Args: + file_path (str): Path to the input file to encrypt. + password (str): Password to use for encryption. + out_file (str): Path where the encrypted file will be saved. + + Raises: + subprocess.CalledProcessError: If the OpenSSL command fails. + FileNotFoundError: If the input file doesn't exist. + + Note: + Uses the following OpenSSL parameters: + - AES-256-CBC encryption algorithm + - Salt for additional security + - PBKDF2 for key derivation + - Password-based encryption + + Example: + >>> generate_encrypted_file_with_openssl("plain.txt", "secret", "encrypted.enc") + """ + # Run openssl enc with AES-256-CBC, pbkdf2, salt + cmd = [ + "openssl", + "enc", + "-aes-256-cbc", + "-salt", + "-pbkdf2", + "-pass", + f"pass:{password}", + "-in", + file_path, + "-out", + out_file, + ] + subprocess.run(cmd, check=True) + + +def decrypt_remote_file_to_string( + ssh_client, remote_enc_file: str, password: str, bteq_command_str: str +) -> tuple[int, str, str]: + """ + Decrypt a remote file and pipe the output to a BTEQ command. + + Args: + ssh_client: An active SSH connection to the remote machine. + remote_enc_file (str): Path to the encrypted file on the remote machine. + password (str): Password used to decrypt the file. + bteq_command_str (str): BTEQ command to execute with the decrypted content. + + Returns: + tuple[int, str, str]: A tuple containing: + - exit_status (int): The exit code of the remote command + - output (str): The stdout output from the command + - err (str): The stderr output from the command + + Note: + The decryption uses OpenSSL with the same parameters as encryption: + - AES-256-CBC + - PBKDF2 key derivation + - Salt + + Example: + >>> exit_code, output, error = decrypt_remote_file_to_string( + ... ssh_client, "remote.enc", "password", "bteq < commands.sql" + ... ) + """ + # Run openssl decrypt command on remote machine + quoted_password = shell_quote_single(password) + + decrypt_cmd = ( + f"openssl enc -d -aes-256-cbc -salt -pbkdf2 -pass pass:{quoted_password} -in {remote_enc_file} | " + + bteq_command_str + ) + stdin, stdout, stderr = ssh_client.exec_command(decrypt_cmd) + # Wait for command to finish + exit_status = stdout.channel.recv_exit_status() + output = stdout.read().decode() + err = stderr.read().decode() + return exit_status, output, err + + +def shell_quote_single(s: str) -> str: + """ + Properly escape and single-quote a string for safe use in shell commands. + + Args: + s (str): The string to be quoted. + + Returns: + str: The safely quoted string. + + Example: + >>> shell_quote_single("don't break") + "'don'\\''t break'" + """ + # Escape single quotes in s, then wrap in single quotes + # In shell, to include a single quote inside single quotes, close, add '\'' and reopen + return "'" + s.replace("'", "'\\''") + "'" + + +def store_credentials(self, host: str, user: str, password: str) -> bool: + """ + Securely store SSH credentials in an encrypted JSON file. + + Args: + host (str): The hostname or IP address of the server. + user (str): The username for authentication. + password (str): The password to store (will be encrypted). + + Returns: + bool: True if credentials were successfully stored, False otherwise. + + Note: + - Credentials are stored in ~/.ssh/bteq_credentials.json + - The directory is created with 700 permissions if it doesn't exist + - The credentials file has 600 permissions (owner read/write only) + - Password is encrypted before storage + - Uses atomic write (writes to temp file then renames) to prevent corruption + + Example: + >>> store_credentials("example.com", "user1", "securepassword") + """ + cred_file = os.path.expanduser("~/.ssh/bteq_credentials.json") + os.makedirs(os.path.dirname(cred_file), mode=0o700, exist_ok=True) + + creds = {} + if os.path.exists(cred_file): + try: + with open(cred_file, "r") as f: + creds = json.load(f) + except Exception as e: + self.log.error(f"Failed to read credentials: {e}") + + cred_key = f"{user}@{host}" + creds[cred_key] = { + "username": user, + "password": self.cred_manager.encrypt(password), + "timestamp": datetime.utcnow().isoformat(), + } + + try: + temp_file = f"{cred_file}.tmp" + with open(temp_file, "w") as f: + json.dump(creds, f, indent=2) + os.chmod(temp_file, 0o600) + os.replace(temp_file, cred_file) + return True + except Exception as e: + self.log.error(f"Failed to store credentials: {e}") + return False + + +def get_stored_credentials(self, host: str, user: str) -> Optional[dict]: + """ + Retrieve stored SSH credentials if they exist. + + Args: + host (str): The hostname or IP address of the server. + user (str): The username for which to retrieve credentials. + + Returns: + Optional[dict]: A dictionary containing the stored credentials if found, + or None if no credentials exist or an error occurred. + The dictionary contains: + - username: The stored username + - password: The encrypted password + - timestamp: When the credentials were stored + + Example: + >>> creds = get_stored_credentials("example.com", "user1") + >>> if creds: + ... password = cred_manager.decrypt(creds["password"]) + """ + cred_file = os.path.expanduser("~/.ssh/bteq_credentials.json") + cred_key = f"{user}@{host}" + + try: + if os.path.exists(cred_file): + with open(cred_file, "r") as f: + return json.load(f).get(cred_key) + except Exception as e: + self.log.warning(f"Failed to read credentials: {e}") + return None diff --git a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py new file mode 100644 index 00000000..b6fff127 --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -0,0 +1,244 @@ +import os + +import pytest +from dagster import job, op, DagsterError +from dagster_teradata import TeradataResource + +td_resource = TeradataResource( + host=os.getenv("TERADATA_HOST"), + user=os.getenv("TERADATA_USER"), + password=os.getenv("TERADATA_PASSWORD"), + database=os.getenv("TERADATA_DATABASE"), +) + + +def test_local_bteq_script_execution(): + @op(required_resource_keys={"teradata"}) + def example_test_local_script(context): + result = context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;" + ) + context.log.info(result) + assert result == 0 + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_local_script() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_local_expected_return_code(): + @op(required_resource_keys={"teradata"}) + def example_test_local_expected_return_code(context): + result = context.resources.teradata.bteq_operator( + sql="delete from abcdefgh;", bteq_quit_rc=8 + ) + context.log.info(result) + assert result == 8 + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_local_expected_return_code() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_remote_expected_return_code(): + @op(required_resource_keys={"teradata"}) + def example_test_local_expected_return_code(context): + result = context.resources.teradata.bteq_operator( + sql="select * from dbc.dbcinfo;", + remote_host="host", + remote_user="user", + remote_password="password", + bteq_quit_rc=0, + ) + context.log.info(result) + assert result == 0 + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_local_expected_return_code() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_remote_file_path(): + @op(required_resource_keys={"teradata"}) + def example_test_local_expected_return_code(context): + result = context.resources.teradata.bteq_operator( + file_path="/tmp/abcd", + remote_host="host", + remote_user="user", + remote_password="password", + bteq_quit_rc=0, + ) + context.log.info(result) + assert result == 0 + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_local_expected_return_code() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_remote_expected_multiple_return_code(): + @op(required_resource_keys={"teradata"}) + def example_test_local_expected_multiple_return_code(context): + result = context.resources.teradata.bteq_operator( + sql="delete from abcdefgh;", + remote_host="host", + remote_user="username", + remote_password="password", + bteq_quit_rc=[0, 8], + ) + context.log.info(result) + assert result is None + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_local_expected_multiple_return_code() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_local_bteq_file_execution(tmp_path): + script_file = tmp_path / "script.bteq" + script_file.write_text("SELECT * FROM dbc.dbcinfo;") + + @op(required_resource_keys={"teradata"}) + def example_test_local_file(context): + result = context.resources.teradata.bteq_operator(file_path=str(script_file)) + context.log.info(result) + assert result == 0 + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_local_file() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_remote_bteq_password_auth(): + @op(required_resource_keys={"teradata"}) + def example_test_remote_password(context): + try: + result = context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;", + remote_host="host", + remote_user="username", + remote_password="password", + ) + context.log.info(result) + except DagsterError as e: + context.log.info(str(e)) + assert "SSH connection failed" in str(e) + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_remote_password() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_remote_bteq_ssh_key_auth(): + @op(required_resource_keys={"teradata"}) + def example_test_remote_ssh_key(context): + result = context.resources.teradata.bteq_operator( + file_path="SELECT * FROM dbc.dbcinfo;", + remote_host="host", + remote_user="username", + ssh_key_path="c:\\users\\\\.ssh\\id_rsa", + ) + context.log.info(result) + assert result is None + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_remote_ssh_key() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_no_bteq_script_provided(): + @op(required_resource_keys={"teradata"}) + def example_test_no_script(context): + with pytest.raises(ValueError): + context.resources.teradata.bteq_operator() + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_no_script() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_conflicting_script_sources(tmp_path): + script_file = tmp_path / "script.bteq" + script_file.write_text("SELECT * FROM dbc.dbcinfo;") + + @op(required_resource_keys={"teradata"}) + def example_test_conflicting_sources(context): + with pytest.raises(ValueError): + context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;", + file_path=str(script_file), + ) + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_conflicting_sources() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_invalid_script_file(): + @op(required_resource_keys={"teradata"}) + def example_test_invalid_file(context): + with pytest.raises(ValueError): + context.resources.teradata.bteq_operator( + file_path="/nonexistent/path/script.bteq" + ) + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_invalid_file() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_remote_missing_credentials(): + @op(required_resource_keys={"teradata"}) + def example_test_missing_creds(context): + with pytest.raises(ValueError): + context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;", + remote_host="remote.teradata.com", + ) + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_missing_creds() + + example_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_conflicting_ssh_auth(): + @op(required_resource_keys={"teradata"}) + def example_test_conflicting_auth(context): + with pytest.raises(ValueError): + context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;", + remote_host="remote.teradata.com", + remote_user="user", + remote_password="password", + ssh_key_path="/path/to/ssh_key", + ) + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_conflicting_auth() + + example_job.execute_in_process(resources={"teradata": td_resource}) diff --git a/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py new file mode 100644 index 00000000..2e411f90 --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py @@ -0,0 +1,292 @@ +import pytest +from dagster import job, op, DagsterError +from dagster_teradata import TeradataResource +from unittest import mock +import tempfile +import os + + +class TestBteq: + def test_local_bteq_script_execution(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test basic local BTEQ script execution.""" + + @op(required_resource_keys={"teradata"}) + def example_test_local_script(context): + with mock.patch( + "dagster_teradata.ttu.bteq.Bteq.execute_bteq_script_at_local" + ) as mock_exec: + mock_exec.return_value = 0 + result = context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;" + ) + assert result == 0 + mock_exec.assert_called_once() + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_local_script() + + example_job.execute_in_process() + + def test_local_file_execution(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test BTEQ execution with local file.""" + + @op(required_resource_keys={"teradata"}) + def example_test_local_file(context): + with tempfile.NamedTemporaryFile(mode="w+", delete=False) as tmp_file: + tmp_file.write("SELECT * FROM dbc.dbcinfo;") + tmp_path = tmp_file.name + + try: + with mock.patch( + "dagster_teradata.ttu.bteq.Bteq._handle_local_bteq_file" + ) as mock_exec: + mock_exec.return_value = 0 + result = context.resources.teradata.bteq_operator( + file_path=tmp_path + ) + assert result == 0 + mock_exec.assert_called_once_with(file_path=tmp_path) + finally: + os.unlink(tmp_path) + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_local_file() + + example_job.execute_in_process() + + def test_invalid_file_path(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test error handling for invalid file path.""" + + @op(required_resource_keys={"teradata"}) + def example_test_invalid_file(context): + with pytest.raises(ValueError, match="invalid or does not exist"): + context.resources.teradata.bteq_operator(file_path="/invalid/path.sql") + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_invalid_file() + + example_job.execute_in_process() + + def test_missing_credentials(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test error handling for missing remote credentials.""" + + @op(required_resource_keys={"teradata"}) + def example_test_missing_creds(context): + with pytest.raises(ValueError): + context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;", remote_host="remote_host" + ) + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_missing_creds() + + example_job.execute_in_process() + + def test_return_code_handling(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test custom return code handling.""" + + @op(required_resource_keys={"teradata"}) + def example_test_return_codes(context): + with mock.patch( + "dagster_teradata.ttu.bteq.Bteq.execute_bteq_script_at_local" + ) as mock_exec: + mock_exec.return_value = 8 + result = context.resources.teradata.bteq_operator( + sql="DELETE FROM non_existent_table;", bteq_quit_rc=[0, 8] + ) + assert result == 8 + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_return_codes() + + example_job.execute_in_process() + + def test_timeout_handling(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test timeout handling.""" + + @op(required_resource_keys={"teradata"}) + def example_test_timeout(context): + with mock.patch( + "dagster_teradata.ttu.bteq.Bteq.execute_bteq_script_at_local" + ) as mock_exec: + mock_exec.side_effect = DagsterError("Timeout") + with pytest.raises(DagsterError, match="Timeout"): + context.resources.teradata.bteq_operator( + sql="SELECT * FROM large_table;", timeout=1 + ) + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_timeout() + + example_job.execute_in_process() + + def test_no_sql_or_file(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test error when neither SQL nor file is provided.""" + + @op(required_resource_keys={"teradata"}) + def example_test_no_input(context): + with pytest.raises( + ValueError, match="requires either the 'sql' or 'file_path' parameter" + ): + context.resources.teradata.bteq_operator() + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_no_input() + + example_job.execute_in_process() + + def test_encoding_validation(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test file encoding validation.""" + + @op(required_resource_keys={"teradata"}) + def example_test_encoding(context): + with tempfile.NamedTemporaryFile(mode="w+", delete=False) as tmp_file: + tmp_file.write("SELECT * FROM dbc.dbcinfo;") + tmp_path = tmp_file.name + + try: + with mock.patch( + "dagster_teradata.ttu.bteq.is_valid_encoding" + ) as mock_encoding: + mock_encoding.side_effect = UnicodeDecodeError( + "utf8", b"", 0, 1, "error" + ) + with pytest.raises(ValueError, match="encoding is different"): + context.resources.teradata.bteq_operator( + file_path=tmp_path, bteq_script_encoding="UTF-8" + ) + finally: + os.unlink(tmp_path) + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_encoding() + + example_job.execute_in_process() + + def test_timeout_with_custom_rc(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test timeout with custom return code handling.""" + + @op(required_resource_keys={"teradata"}) + def example_test_timeout_rc(context): + with mock.patch( + "dagster_teradata.ttu.bteq.Bteq.execute_bteq_script_at_local" + ) as mock_exec: + mock_exec.return_value = 99 + result = context.resources.teradata.bteq_operator( + sql="SELECT * FROM large_table;", + timeout=30, + timeout_rc=99, + bteq_quit_rc=[99], + ) + assert result == 99 + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_timeout_rc() + + example_job.execute_in_process() + + def test_return_code_not_in_quit_rc(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test when return code is not in allowed quit codes.""" + + @op(required_resource_keys={"teradata"}) + def example_test_unexpected_rc(context): + with mock.patch( + "dagster_teradata.ttu.bteq.Bteq.execute_bteq_script_at_local" + ) as mock_exec: + mock_exec.return_value = 42 + result = context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;", bteq_quit_rc=[0, 1] + ) + assert ( + result == 42 + ) # Still returns the code even though it's not in quit_rc + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_unexpected_rc() + + example_job.execute_in_process() + + def test_conflicting_ssh_auth(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test error when both password and key auth are provided.""" + + @op(required_resource_keys={"teradata"}) + def example_test_conflicting_auth(context): + with pytest.raises(ValueError): + context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;", + remote_host="remote_host", + remote_user="remote_user", + remote_password="password", + ssh_key_path="/path/to/key", + ) + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_conflicting_auth() + + example_job.execute_in_process() diff --git a/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py new file mode 100644 index 00000000..96137ddb --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py @@ -0,0 +1,298 @@ +from __future__ import annotations + +import os +import stat +import unittest +from unittest.mock import MagicMock, patch, call + +import pytest + +from dagster import DagsterError + +from dagster_teradata.resources import TeradataResource +from dagster_teradata.ttu.utils.bteq_util import ( + is_valid_encoding, + is_valid_file, + is_valid_remote_bteq_script_file, + prepare_bteq_script_for_local_execution, + prepare_bteq_script_for_remote_execution, + read_file, + transfer_file_sftp, + verify_bteq_installed, + verify_bteq_installed_remote, +) + + +class TestBteqUtils: + @patch("shutil.which") + def test_verify_bteq_installed_success(self, mock_which): + mock_which.return_value = "/usr/bin/bteq" + # Should not raise + verify_bteq_installed() + mock_which.assert_called_with("bteq") + + @patch("shutil.which") + def test_verify_bteq_installed_fail(self, mock_which): + mock_which.return_value = None + with pytest.raises(DagsterError): + verify_bteq_installed() + + def test_prepare_bteq_script_for_remote_execution(self): + teradata_connection_resource = TeradataResource( + host="myhost", user="user", password="pass" + ) + sql = "SELECT * FROM DUAL;" + script = prepare_bteq_script_for_remote_execution( + teradata_connection_resource, sql + ) + assert ".LOGON myhost/user,pass" in script + assert "SELECT * FROM DUAL;" in script + assert ".EXIT" in script + + def test_prepare_bteq_script_for_local_execution(self): + sql = "SELECT 1;" + script = prepare_bteq_script_for_local_execution(sql) + assert "SELECT 1;" in script + assert ".EXIT" in script + + @patch("paramiko.SSHClient.exec_command") + def test_verify_bteq_installed_remote_success(self, mock_exec): + mock_stdin = MagicMock() + mock_stdout = MagicMock() + mock_stderr = MagicMock() + mock_stdout.channel.recv_exit_status.return_value = 0 + mock_stdout.read.return_value = b"/usr/bin/bteq" + mock_stderr.read.return_value = b"" + mock_exec.return_value = (mock_stdin, mock_stdout, mock_stderr) + + ssh_client = MagicMock() + ssh_client.exec_command = mock_exec + + # Should not raise + verify_bteq_installed_remote(ssh_client) + + @patch("paramiko.SSHClient.exec_command") + def test_verify_bteq_installed_remote_fail(self, mock_exec): + mock_stdin = MagicMock() + mock_stdout = MagicMock() + mock_stderr = MagicMock() + mock_stdout.channel.recv_exit_status.return_value = 1 + mock_stdout.read.return_value = b"" + mock_stderr.read.return_value = b"command not found" + mock_exec.return_value = (mock_stdin, mock_stdout, mock_stderr) + + ssh_client = MagicMock() + ssh_client.exec_command = mock_exec + + with pytest.raises(DagsterError): + verify_bteq_installed_remote(ssh_client) + + @patch("paramiko.SSHClient.open_sftp") + def test_transfer_file_sftp(self, mock_open_sftp): + mock_sftp = MagicMock() + mock_open_sftp.return_value = mock_sftp + + ssh_client = MagicMock() + ssh_client.open_sftp = mock_open_sftp + + transfer_file_sftp(ssh_client, "local_file.txt", "remote_file.txt") + + mock_open_sftp.assert_called_once() + mock_sftp.close.assert_called_once() + + def test_is_valid_file(self): + # create temp file + with open("temp_test_file.txt", "w") as f: + f.write("hello") + + assert is_valid_file("temp_test_file.txt") is True + assert is_valid_file("non_existent_file.txt") is False + + os.remove("temp_test_file.txt") + + def test_is_valid_encoding(self): + # Write a file with UTF-8 encoding + with open("temp_utf8_file.txt", "w", encoding="utf-8") as f: + f.write("hello world") + + # Should return True + assert is_valid_encoding("temp_utf8_file.txt", encoding="utf-8") is True + + # Cleanup + os.remove("temp_utf8_file.txt") + + def test_read_file_success(self): + content = "Sample content" + with open("temp_read_file.txt", "w") as f: + f.write(content) + + read_content = read_file("temp_read_file.txt") + assert read_content == content + os.remove("temp_read_file.txt") + + def test_read_file_file_not_found(self): + with pytest.raises(FileNotFoundError): + read_file("non_existent_file.txt") + + @patch("paramiko.SSHClient.open_sftp") + def test_is_valid_remote_bteq_script_file_exists(self, mock_open_sftp): + mock_sftp = MagicMock() + mock_open_sftp.return_value = mock_sftp + + # Mock stat to return a regular file mode + mock_stat = MagicMock() + mock_stat.st_mode = stat.S_IFREG + mock_sftp.stat.return_value = mock_stat + + ssh_client = MagicMock() + ssh_client.open_sftp = mock_open_sftp + + result = is_valid_remote_bteq_script_file(ssh_client, "/remote/path/to/file") + assert result is True + mock_sftp.close.assert_called_once() + + @patch("paramiko.SSHClient.open_sftp") + def test_is_valid_remote_bteq_script_file_not_exists(self, mock_open_sftp): + mock_sftp = MagicMock() + mock_open_sftp.return_value = mock_sftp + + # Raise FileNotFoundError for stat + mock_sftp.stat.side_effect = FileNotFoundError + + ssh_client = MagicMock() + ssh_client.open_sftp = mock_open_sftp + + result = is_valid_remote_bteq_script_file(ssh_client, "/remote/path/to/file") + assert result is False + mock_sftp.close.assert_called_once() + + def test_is_valid_remote_bteq_script_file_none_path(self): + ssh_client = MagicMock() + result = is_valid_remote_bteq_script_file(ssh_client, None) + assert result is False + + @patch("dagster_teradata.ttu.utils.bteq_util.identify_os", return_value="linux") + def test_verify_bteq_installed_remote_linux(self, mock_os): + """Test BTEQ verification on Linux systems.""" + ssh_client = MagicMock() + stdout_mock = MagicMock() + stdout_mock.read.return_value = b"/usr/bin/bteq" + stdout_mock.channel.recv_exit_status.return_value = 0 + ssh_client.exec_command.return_value = (MagicMock(), stdout_mock, MagicMock()) + + verify_bteq_installed_remote(ssh_client) + ssh_client.exec_command.assert_called_once_with("which bteq") + + @patch("dagster_teradata.ttu.utils.bteq_util.identify_os", return_value="windows") + def test_verify_bteq_installed_remote_windows(self, mock_os): + """Test BTEQ verification on Windows systems.""" + ssh_client = MagicMock() + stdout_mock = MagicMock() + stdout_mock.read.return_value = b"C:\\Program Files\\bteq.exe" + stdout_mock.channel.recv_exit_status.return_value = 0 + ssh_client.exec_command.return_value = (MagicMock(), stdout_mock, MagicMock()) + + verify_bteq_installed_remote(ssh_client) + ssh_client.exec_command.assert_called_once_with("where bteq") + + @patch("dagster_teradata.ttu.utils.bteq_util.identify_os", return_value="darwin") + def test_verify_bteq_installed_remote_macos(self, mock_os): + """Test BTEQ verification on macOS systems with zsh.""" + ssh_client = MagicMock() + stdout_mock = MagicMock() + stdout_mock.read.return_value = b"/usr/local/bin/bteq" + stdout_mock.channel.recv_exit_status.return_value = 0 + + ssh_client.exec_command.return_value = (MagicMock(), stdout_mock, MagicMock()) + + verify_bteq_installed_remote(ssh_client) + + ssh_client.exec_command.assert_has_calls( + [call("command -v zsh"), call('zsh -l -c "which bteq"')] + ) + + @patch("dagster_teradata.ttu.utils.bteq_util.identify_os", return_value="darwin") + def test_verify_bteq_installed_remote_macos_no_zsh(self, mock_os): + """Test BTEQ verification on macOS without zsh.""" + ssh_client = MagicMock() + + # First call - no zsh found + stdin_mock1 = MagicMock() + stdout_mock1 = MagicMock() + stdout_mock1.read.return_value = b"" + stdout_mock1.channel.recv_exit_status.return_value = 1 + + # Second call - which bteq + stdin_mock2 = MagicMock() + stdout_mock2 = MagicMock() + stdout_mock2.read.return_value = b"/usr/local/bin/bteq" + stdout_mock2.channel.recv_exit_status.return_value = 0 + + ssh_client.exec_command.side_effect = [ + (stdin_mock1, stdout_mock1, MagicMock()), + (stdin_mock2, stdout_mock2, MagicMock()), + ] + + verify_bteq_installed_remote(ssh_client) + + ssh_client.exec_command.assert_has_calls( + [call("command -v zsh"), call("which bteq")] + ) + + @patch("dagster_teradata.ttu.utils.bteq_util.identify_os", return_value="linux") + def test_verify_bteq_installed_remote_failure(self, mock_os): + """Test BTEQ verification failure when not installed.""" + ssh_client = MagicMock() + stdout_mock = MagicMock() + stderr_mock = MagicMock() + stdout_mock.read.return_value = b"" + stderr_mock.read.return_value = b"command not found" + stdout_mock.channel.recv_exit_status.return_value = 1 + ssh_client.exec_command.return_value = (MagicMock(), stdout_mock, stderr_mock) + + with pytest.raises( + DagsterError, + match="BTEQ is not installed or not available on the remote machine", + ): + verify_bteq_installed_remote(ssh_client) + + ssh_client.exec_command.assert_called_once_with("which bteq") + + @patch("dagster_teradata.ttu.utils.bteq_util.identify_os", return_value="darwin") + def test_verify_bteq_installed_remote_macos_failure(self, mock_os): + """Test BTEQ verification failure on macOS.""" + ssh_client = MagicMock() + + # First call - no zsh found + stdin_mock1 = MagicMock() + stdout_mock1 = MagicMock() + stdout_mock1.read.return_value = b"" + stdout_mock1.channel.recv_exit_status.return_value = 1 + + # Mock the second exec_command call to check for bteq (returns not found) + stdin_mock2 = MagicMock() + stdout_mock2 = MagicMock() + stdout_mock2.read.return_value = b"" + stdout_mock2.channel.recv_exit_status.return_value = 1 + stderr_mock2 = MagicMock() + stderr_mock2.read.return_value = b"command not found" + + # Set the side effect for exec_command to simulate the sequence of remote shell calls + ssh_client.exec_command.side_effect = [ + (stdin_mock1, stdout_mock1, MagicMock()), # First: check for zsh + (stdin_mock2, stdout_mock2, stderr_mock2), # Second: check for bteq + ] + + with pytest.raises( + DagsterError, + match="BTEQ is not installed or not available on the remote machine", + ): + verify_bteq_installed_remote(ssh_client) + + ssh_client.exec_command.assert_has_calls( + [call("command -v zsh"), call("which bteq")] + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/libraries/dagster-teradata/dagster_teradata_tests/test_encryption_utils.py b/libraries/dagster-teradata/dagster_teradata_tests/test_encryption_utils.py new file mode 100644 index 00000000..2e2df9b7 --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata_tests/test_encryption_utils.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import string +import unittest +from unittest.mock import MagicMock, patch + +from dagster_teradata.ttu.utils.encryption_utils import ( + decrypt_remote_file_to_string, + generate_encrypted_file_with_openssl, + generate_random_password, + shell_quote_single, +) + + +class TestEncryptionUtils: + def test_generate_random_password_length(self): + pwd = generate_random_password(16) + assert len(pwd) == 16 + # Check characters are in allowed set + allowed_chars = string.ascii_letters + string.digits + string.punctuation + assert (all(c in allowed_chars for c in pwd)) is True + + @patch("subprocess.run") + def test_generate_encrypted_file_with_openssl_calls_subprocess(self, mock_run): + file_path = "/tmp/plain.txt" + password = "testpass" + out_file = "/tmp/encrypted.enc" + + generate_encrypted_file_with_openssl(file_path, password, out_file) + + mock_run.assert_called_once_with( + [ + "openssl", + "enc", + "-aes-256-cbc", + "-salt", + "-pbkdf2", + "-pass", + f"pass:{password}", + "-in", + file_path, + "-out", + out_file, + ], + check=True, + ) + + def test_shell_quote_single_simple(self): + s = "simple" + quoted = shell_quote_single(s) + assert quoted == "'simple'" + + def test_shell_quote_single_with_single_quote(self): + s = "O'Reilly" + quoted = shell_quote_single(s) + assert quoted == "'O'\\''Reilly'" + + def test_decrypt_remote_file_to_string(self): + password = "mysecret" + remote_enc_file = "/remote/encrypted.enc" + bteq_command_str = "bteq -c UTF-8" + + ssh_client = MagicMock() + mock_stdin = MagicMock() + mock_stdout = MagicMock() + mock_stderr = MagicMock() + + # Setup mock outputs and exit code + mock_stdout.channel.recv_exit_status.return_value = 0 + mock_stdout.read.return_value = b"decrypted output" + mock_stderr.read.return_value = b"" + + ssh_client.exec_command.return_value = (mock_stdin, mock_stdout, mock_stderr) + + exit_status, output, err = decrypt_remote_file_to_string( + ssh_client, remote_enc_file, password, bteq_command_str + ) + + quoted_password = shell_quote_single(password) + expected_cmd = ( + f"openssl enc -d -aes-256-cbc -salt -pbkdf2 -pass pass:{quoted_password} -in {remote_enc_file} | " + + bteq_command_str + ) + + ssh_client.exec_command.assert_called_once_with(expected_cmd) + assert exit_status == 0 + assert output == "decrypted output" + assert err == "" + + +if __name__ == "__main__": + unittest.main() diff --git a/libraries/dagster-teradata/pyproject.toml b/libraries/dagster-teradata/pyproject.toml index 64d36bc0..39283b3e 100644 --- a/libraries/dagster-teradata/pyproject.toml +++ b/libraries/dagster-teradata/pyproject.toml @@ -6,6 +6,7 @@ requires-python = ">=3.8" dependencies = [ "dagster>=1.8.0", "teradatasql", + "paramiko", ] dynamic = ["version"]