From 22109c755f54da39ff5cfc6df4762817037e628a Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Wed, 14 May 2025 08:26:23 +0530 Subject: [PATCH 01/21] initial_commit --- .gitignore | 1 + .../dagster_teradata/resources.py | 64 ++++++-- .../dagster-teradata/dagster_teradata/ttu.py | 155 ++++++++++++++++++ .../functional/test_execute_bteq.py | 25 +++ 4 files changed, 235 insertions(+), 10 deletions(-) create mode 100644 libraries/dagster-teradata/dagster_teradata/ttu.py create mode 100644 libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py 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/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 8e4403b2..10c5bc30 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -17,6 +17,7 @@ from pydantic import Field from dagster_teradata import constants +from dagster_teradata.ttu import Bteq from dagster_teradata.teradata_compute_cluster_manager import ( TeradataComputeClusterManager, ) @@ -30,6 +31,15 @@ class TeradataResource(ConfigurableResource, IAttachDifferentObjectToOpContext): default=None, description=("Name of the default database to use."), ) + 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.") + @property @cached_method @@ -41,6 +51,10 @@ def _connection_args(self) -> Mapping[str, Any]: "user", "password", "database", + "logmech", + "browser", + "browser_tab_timeout", + "browser_timeout", ) if self._resolved_config_dict.get(k) is not None } @@ -49,21 +63,34 @@ def _connection_args(self) -> Mapping[str, Any]: @public @contextmanager def get_connection(self): + connection_params = {} if not self.host: raise ValueError("Host is required but not provided.") - if not self.user: - raise ValueError("User is required but not provided.") - if not self.password: - raise ValueError("Password is required but not provided.") - - connection_params = { - "host": self.host, - "user": self.user, - "password": self.password, - } + connection_params = {"host": self.host} + + if self.logmech is not None and self.logmech.lower() == "browser": + # When logmech is "browser", username and password should not be provided. + if self.user is not None or self.password is not None: + raise ValueError( + "Username and password should not be specified when logmech is 'browser'" + ) + if self.browser is not None: + connection_params["browser"] = self.browser + if self.browser_tab_timeout is not None: + connection_params["browser_tab_timeout"] = str(self.browser_tab_timeout) + if self.browser_timeout is not None: + connection_params["browser_timeout"] = str(self.browser_timeout) + else: + if not self.user: + raise ValueError("User is required but not provided.") + if not self.password: + raise ValueError("Password is required but not provided.") + connection_params.update({"user": self.user, "password": self.password}) if self.database is not None: connection_params["database"] = self.database + if self.logmech is not None: + connection_params["logmech"] = self.logmech teradata_conn = teradatasql.connect(**connection_params) @@ -96,6 +123,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 @@ -191,6 +219,22 @@ def create_fresh_database(teradata: TeradataResource): results = results.append(cursor.fetchall()) # type: ignore return results + def execute_bteq_script(self, bteq_script: str) -> None: + """Executes BTEQ sentences using BTEQ binary. + + Args: + bteq_script (str): String of BTEQ sentences to be executed. + + Examples: + .. code-block:: python + + @op + def execute_bteq(teradata: TeradataResource): + bteq = "SELECT * FROM dbc.dbcinfo" + teradata.execute_bteq(bteq) + """ + self.bteq.execute_bteq(bteq_script) + 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.py b/libraries/dagster-teradata/dagster_teradata/ttu.py new file mode 100644 index 00000000..16d54289 --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata/ttu.py @@ -0,0 +1,155 @@ +from tempfile import gettempdir, NamedTemporaryFile, TemporaryDirectory +from dagster import DagsterError +import subprocess +import os +class Bteq: + def __init__(self, connection, teradata_connection_resource, log): + self.connection = connection + self.log = log + self.teradata_connection_resource = teradata_connection_resource + + def execute_bteq(self, bteq_script: str, xcom_push_flag=False, timeout: int | None = None) -> str: + """ + Executes BTEQ sentences using BTEQ binary. + :param bteq: string of BTEQ sentences + :param xcom_push_flag: Flag for pushing last line of BTEQ Log to XCom + :param timeout: Timeout in seconds for the BTEQ execution. + :return: The last line of the BTEQ log if xcom_push_flag is True, otherwise None. + """ + conn = {'host': self.teradata_connection_resource.host, 'login': self.teradata_connection_resource.user, + 'password': self.teradata_connection_resource.password, + 'bteq_output_width': self.teradata_connection_resource.bteq_output_width, + 'bteq_session_encoding': self.teradata_connection_resource.bteq_session_encoding, + 'bteq_quit_zero': self.teradata_connection_resource.bteq_quit_zero, + 'console_output_encoding': self.teradata_connection_resource.console_output_encoding,} + self.log.info("Executing BTEQ script...") + + with TemporaryDirectory(prefix='dagster_ttu_bteq_') as tmpdir: + with NamedTemporaryFile(dir=tmpdir, mode='wb') as tmpfile: + bteq_file_content = self._prepare_bteq_script(bteq_script, + conn['host'], + conn['login'], + conn['password'], + conn['bteq_output_width'], + conn['bteq_session_encoding'], + conn['bteq_quit_zero'] + ) + self.log.debug("Generated BTEQ script:\n%s", bteq_file_content) + + tmpfile.write(bytes(bteq_file_content,'UTF8')) + tmpfile.flush() + tmpfile.seek(0) + + conn['sp'] = subprocess.Popen(['bteq'], + stdin = tmpfile, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + cwd = tmpdir, + preexec_fn = os.setsid) + + line = '' + failure_message = "An error occurred during the BTEQ operation. Please review the full BTEQ output for details." + self.log.info("BTEQ Output:") + for line in iter(conn["sp"].stdout.readline, b""): + decoded_line = line.decode(conn["console_output_encoding"]).strip() + self.log.info(decoded_line) + last_line = decoded_line + if "Failure" in decoded_line: + # Save the last failure message + failure_message = decoded_line + + # Wait for the BTEQ process to complete with optional timeout + try: + conn["sp"].wait(timeout=timeout) + self.log.info("BTEQ command exited with return code %s", conn["sp"].returncode) + except subprocess.TimeoutExpired: + self.on_kill() + raise DagsterError(f"BTEQ command timed out after {timeout} seconds") + + # Raise an exception if the BTEQ command failed + if conn["sp"].returncode: + raise DagsterError( + f"BTEQ command exited with return code {conn['sp'].returncode} due to: {failure_message}" + ) + + # Return the last line of the BTEQ log if xcom_push_flag is True + if xcom_push_flag: + return last_line + return None + + def on_kill(self): + """ + Terminates the subprocess if it is running. + Ensures that the process is terminated gracefully and logs the status. + """ + self.log.debug("Attempting to kill child process...") + conn = self.get_conn() + if conn.get("sp"): + try: + self.log.info("Terminating subprocess...") + conn["sp"].terminate() + conn["sp"].wait(timeout=5) + self.log.info("Subprocess terminated successfully.") + except subprocess.TimeoutExpired: + self.log.warning("Subprocess did not terminate in time. Forcing kill...") + conn["sp"].kill() + self.log.info("Subprocess killed forcefully.") + except (ProcessLookupError, OSError) as e: + self.log.error("Failed to terminate subprocess: %s", e) + + @staticmethod + def _prepare_bteq_script( + bteq_string: str, + host: str, + login: str, + password: str, + bteq_output_width: int, + bteq_session_encoding: str, + bteq_quit_zero: bool, + ) -> str: + """ + Prepare a BTEQ file with connection parameters for executing SQL sentences with BTEQ syntax. + :param bteq_string: BTEQ sentences to execute. + :param host: Teradata Host. + :param login: Username for login. + :param password: Password for login. + :param bteq_output_width: Width of BTEQ output in the console. + :param bteq_session_encoding: Session encoding. See official Teradata docs for possible values. + :param bteq_quit_zero: If True, force a .QUIT 0 sentence at the end of the script (forcing return code = 0). + :return: A formatted BTEQ script as a string. + :raises ValueError: If any required parameters are invalid. + """ + # Validate input parameters + if not bteq_string or not bteq_string.strip(): + raise ValueError("BTEQ script cannot be empty.") + if not host: + raise ValueError("Host parameter cannot be empty.") + if not login: + raise ValueError("Login parameter cannot be empty.") + if not password: + raise ValueError("Password parameter cannot be empty.") + if not isinstance(bteq_output_width, int) or bteq_output_width <= 0: + raise ValueError("BTEQ output width must be a positive integer.") + if not bteq_session_encoding: + raise ValueError("BTEQ session encoding cannot be empty.") + + # Construct the BTEQ script + bteq_list = [ + f".LOGON {host}/{login},{password};", + ".IF ERRORCODE <> 0 THEN .QUIT 8;", + f".SET WIDTH {bteq_output_width};", + f".SET SESSION CHARSET '{bteq_session_encoding}';", + bteq_string.strip(), + ] + + # Add optional .QUIT 0 command if specified + if bteq_quit_zero: + bteq_list.append(".QUIT 0;") + + # Ensure proper termination of the script + bteq_list.extend([".LOGOFF;", ".EXIT;"]) + + # Join the script lines with newlines + bteq_script = "\n".join(bteq_list) + + return bteq_script \ No newline at end of file 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..e4ff7ccb --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -0,0 +1,25 @@ +import os + +from dagster import job, op +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_execute_bteq(tmp_path): + @op(required_resource_keys={"teradata"}) + def example_test_execute_bteq(context): + result = context.resources.teradata.execute_bteq_script( + "select * from dbc.dbcinfo") + context.log.info(result) + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_execute_bteq() + + example_job.execute_in_process(resources={"teradata": td_resource}) From 58b0342092adb740cdfbf844f2e42d617f055725 Mon Sep 17 00:00:00 2001 From: Mohan Talla Date: Wed, 14 May 2025 09:19:46 +0530 Subject: [PATCH 02/21] optimized --- .../dagster_teradata/resources.py | 21 +++-- .../dagster-teradata/dagster_teradata/ttu.py | 90 +++++++++++-------- .../functional/test_execute_bteq.py | 5 +- 3 files changed, 70 insertions(+), 46 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 10c5bc30..4a1344b0 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -35,11 +35,18 @@ class TeradataResource(ConfigurableResource, IAttachDifferentObjectToOpContext): 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.") - + 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." + ) @property @cached_method @@ -219,7 +226,7 @@ def create_fresh_database(teradata: TeradataResource): results = results.append(cursor.fetchall()) # type: ignore return results - def execute_bteq_script(self, bteq_script: str) -> None: + def bteq_operator(self, bteq_script: str) -> None: """Executes BTEQ sentences using BTEQ binary. Args: @@ -233,7 +240,7 @@ def execute_bteq(teradata: TeradataResource): bteq = "SELECT * FROM dbc.dbcinfo" teradata.execute_bteq(bteq) """ - self.bteq.execute_bteq(bteq_script) + self.bteq.bteq_operator(bteq_script) def drop_database(self, databases: Union[str, Sequence[str]]) -> None: """ diff --git a/libraries/dagster-teradata/dagster_teradata/ttu.py b/libraries/dagster-teradata/dagster_teradata/ttu.py index 16d54289..4802d874 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu.py @@ -1,14 +1,18 @@ -from tempfile import gettempdir, NamedTemporaryFile, TemporaryDirectory +from tempfile import NamedTemporaryFile, TemporaryDirectory from dagster import DagsterError import subprocess import os + + class Bteq: def __init__(self, connection, teradata_connection_resource, log): self.connection = connection self.log = log self.teradata_connection_resource = teradata_connection_resource - def execute_bteq(self, bteq_script: str, xcom_push_flag=False, timeout: int | None = None) -> str: + def bteq_operator( + self, bteq_script: str, xcom_push_flag=False, timeout: int | None = None + ) -> str: """ Executes BTEQ sentences using BTEQ binary. :param bteq: string of BTEQ sentences @@ -16,38 +20,44 @@ def execute_bteq(self, bteq_script: str, xcom_push_flag=False, timeout: int | No :param timeout: Timeout in seconds for the BTEQ execution. :return: The last line of the BTEQ log if xcom_push_flag is True, otherwise None. """ - conn = {'host': self.teradata_connection_resource.host, 'login': self.teradata_connection_resource.user, - 'password': self.teradata_connection_resource.password, - 'bteq_output_width': self.teradata_connection_resource.bteq_output_width, - 'bteq_session_encoding': self.teradata_connection_resource.bteq_session_encoding, - 'bteq_quit_zero': self.teradata_connection_resource.bteq_quit_zero, - 'console_output_encoding': self.teradata_connection_resource.console_output_encoding,} + conn = { + "host": self.teradata_connection_resource.host, + "login": self.teradata_connection_resource.user, + "password": self.teradata_connection_resource.password, + "bteq_output_width": self.teradata_connection_resource.bteq_output_width, + "bteq_session_encoding": self.teradata_connection_resource.bteq_session_encoding, + "bteq_quit_zero": self.teradata_connection_resource.bteq_quit_zero, + "console_output_encoding": self.teradata_connection_resource.console_output_encoding, + } self.log.info("Executing BTEQ script...") - with TemporaryDirectory(prefix='dagster_ttu_bteq_') as tmpdir: - with NamedTemporaryFile(dir=tmpdir, mode='wb') as tmpfile: - bteq_file_content = self._prepare_bteq_script(bteq_script, - conn['host'], - conn['login'], - conn['password'], - conn['bteq_output_width'], - conn['bteq_session_encoding'], - conn['bteq_quit_zero'] - ) + with TemporaryDirectory(prefix="dagster_ttu_bteq_") as tmpdir: + with NamedTemporaryFile(dir=tmpdir, mode="wb") as tmpfile: + bteq_file_content = self._prepare_bteq_script( + bteq_script, + conn["host"], + conn["login"], + conn["password"], + conn["bteq_output_width"], + conn["bteq_session_encoding"], + conn["bteq_quit_zero"], + ) self.log.debug("Generated BTEQ script:\n%s", bteq_file_content) - tmpfile.write(bytes(bteq_file_content,'UTF8')) + tmpfile.write(bytes(bteq_file_content, "UTF8")) tmpfile.flush() tmpfile.seek(0) - conn['sp'] = subprocess.Popen(['bteq'], - stdin = tmpfile, - stdout = subprocess.PIPE, - stderr = subprocess.STDOUT, - cwd = tmpdir, - preexec_fn = os.setsid) + conn["sp"] = subprocess.Popen( + ["bteq"], + stdin=tmpfile, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + cwd=tmpdir, + preexec_fn=os.setsid, + ) - line = '' + line = "" failure_message = "An error occurred during the BTEQ operation. Please review the full BTEQ output for details." self.log.info("BTEQ Output:") for line in iter(conn["sp"].stdout.readline, b""): @@ -61,10 +71,14 @@ def execute_bteq(self, bteq_script: str, xcom_push_flag=False, timeout: int | No # Wait for the BTEQ process to complete with optional timeout try: conn["sp"].wait(timeout=timeout) - self.log.info("BTEQ command exited with return code %s", conn["sp"].returncode) + self.log.info( + "BTEQ command exited with return code %s", conn["sp"].returncode + ) except subprocess.TimeoutExpired: self.on_kill() - raise DagsterError(f"BTEQ command timed out after {timeout} seconds") + raise DagsterError( + f"BTEQ command timed out after {timeout} seconds" + ) # Raise an exception if the BTEQ command failed if conn["sp"].returncode: @@ -91,7 +105,9 @@ def on_kill(self): conn["sp"].wait(timeout=5) self.log.info("Subprocess terminated successfully.") except subprocess.TimeoutExpired: - self.log.warning("Subprocess did not terminate in time. Forcing kill...") + self.log.warning( + "Subprocess did not terminate in time. Forcing kill..." + ) conn["sp"].kill() self.log.info("Subprocess killed forcefully.") except (ProcessLookupError, OSError) as e: @@ -99,13 +115,13 @@ def on_kill(self): @staticmethod def _prepare_bteq_script( - bteq_string: str, - host: str, - login: str, - password: str, - bteq_output_width: int, - bteq_session_encoding: str, - bteq_quit_zero: bool, + bteq_string: str, + host: str, + login: str, + password: str, + bteq_output_width: int, + bteq_session_encoding: str, + bteq_quit_zero: bool, ) -> str: """ Prepare a BTEQ file with connection parameters for executing SQL sentences with BTEQ syntax. @@ -152,4 +168,4 @@ def _prepare_bteq_script( # Join the script lines with newlines bteq_script = "\n".join(bteq_list) - return bteq_script \ No newline at end of file + return bteq_script 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 index e4ff7ccb..0cda7a18 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -14,8 +14,9 @@ def test_execute_bteq(tmp_path): @op(required_resource_keys={"teradata"}) def example_test_execute_bteq(context): - result = context.resources.teradata.execute_bteq_script( - "select * from dbc.dbcinfo") + result = context.resources.teradata.bteq_operator( + "select * from dbc.dbcinfo;" + ) context.log.info(result) @job(resource_defs={"teradata": td_resource}) From ed377f8de62cdf24834add7cae62e34870f0866a Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Wed, 14 May 2025 09:54:51 +0530 Subject: [PATCH 03/21] added test cases --- .../functional/test_execute_bteq.py | 87 +++++++++++++++++++ 1 file changed, 87 insertions(+) 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 index 0cda7a18..fc143191 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -24,3 +24,90 @@ def example_job(): example_test_execute_bteq() example_job.execute_in_process(resources={"teradata": td_resource}) + +def test_execute_bteq_sql_error(tmp_path): + @op(required_resource_keys={"teradata"}) + def failing_op(context): + context.resources.teradata.bteq_operator("SELECT * FROM;") # Invalid SQL + + @job(resource_defs={"teradata": td_resource}) + def failing_job(): + failing_op() + + try: + failing_job.execute_in_process(resources={"teradata": td_resource}) + assert False, "Expected job to fail due to syntax error" + except Exception as e: + assert "Syntax error" in str(e) or "3706" in str(e) + + +def test_execute_bteq_timeout(tmp_path): + @op(required_resource_keys={"teradata"}) + def timeout_op(context): + context.resources.teradata.bteq_operator("SELECT * FROM dbc.dbcinfo;", timeout=1) + + @job(resource_defs={"teradata": td_resource}) + def timeout_job(): + timeout_op() + + try: + timeout_job.execute_in_process(resources={"teradata": td_resource}) + assert False, "Expected timeout" + except Exception as e: + assert "timed out" in str(e) + + +import pytest + +@pytest.mark.parametrize("invalid_sql", ["", None]) +def test_invalid_bteq_sql(invalid_sql): + @op(required_resource_keys={"teradata"}) + def invalid_sql_op(context): + context.resources.teradata.bteq_operator(invalid_sql) + + @job(resource_defs={"teradata": td_resource}) + def invalid_sql_job(): + invalid_sql_op() + + with pytest.raises(ValueError, match="BTEQ script cannot be empty"): + invalid_sql_job.execute_in_process(resources={"teradata": td_resource}) + + +def test_bteq_custom_config(tmp_path): + custom_td_resource = TeradataResource( + host="localhost", + user="user", + password="password", + database="dbc", + bteq_output_width=200, + bteq_session_encoding="UTF8", + bteq_quit_zero=True, + ) + + @op(required_resource_keys={"teradata"}) + def config_check_op(context): + result = context.resources.teradata.bteq_operator("SELECT * FROM dbc.dbcinfo;") + assert result is not None + + @job(resource_defs={"teradata": custom_td_resource}) + def config_job(): + config_check_op() + + config_job.execute_in_process(resources={"teradata": custom_td_resource}) + + +def test_logging_output(tmp_path, caplog): + @op(required_resource_keys={"teradata"}) + def log_op(context): + result = context.resources.teradata.bteq_operator("SELECT * FROM dbc.dbcinfo;") + context.log.info(result) + + @job(resource_defs={"teradata": td_resource}) + def log_job(): + log_op() + + log_job.execute_in_process(resources={"teradata": td_resource}) + assert any("dbcinfo" in record.message for record in caplog.records) + + + From eee8e6a06a6c56b19c2ad437bdbede0db81ee99a Mon Sep 17 00:00:00 2001 From: Mohan Talla Date: Wed, 14 May 2025 12:09:27 +0530 Subject: [PATCH 04/21] Added timeout --- .../dagster-teradata/dagster_teradata/resources.py | 4 ++-- libraries/dagster-teradata/dagster_teradata/ttu.py | 10 ++++++++-- .../functional/test_execute_bteq.py | 2 +- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 4a1344b0..4d2276e9 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -226,7 +226,7 @@ def create_fresh_database(teradata: TeradataResource): results = results.append(cursor.fetchall()) # type: ignore return results - def bteq_operator(self, bteq_script: str) -> None: + def bteq_operator(self, bteq_script: str, timeout: int = None) -> None: """Executes BTEQ sentences using BTEQ binary. Args: @@ -240,7 +240,7 @@ def execute_bteq(teradata: TeradataResource): bteq = "SELECT * FROM dbc.dbcinfo" teradata.execute_bteq(bteq) """ - self.bteq.bteq_operator(bteq_script) + self.bteq.bteq_operator(bteq_script, timeout) def drop_database(self, databases: Union[str, Sequence[str]]) -> None: """ diff --git a/libraries/dagster-teradata/dagster_teradata/ttu.py b/libraries/dagster-teradata/dagster_teradata/ttu.py index 4802d874..3a7afad9 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu.py @@ -6,6 +6,7 @@ class Bteq: def __init__(self, connection, teradata_connection_resource, log): + self.timeout = None self.connection = connection self.log = log self.teradata_connection_resource = teradata_connection_resource @@ -20,6 +21,11 @@ def bteq_operator( :param timeout: Timeout in seconds for the BTEQ execution. :return: The last line of the BTEQ log if xcom_push_flag is True, otherwise None. """ + if timeout is not None: + self.timeout = timeout + else: + # Run without timeout or use a default + self.timeout = 5 conn = { "host": self.teradata_connection_resource.host, "login": self.teradata_connection_resource.user, @@ -70,7 +76,7 @@ def bteq_operator( # Wait for the BTEQ process to complete with optional timeout try: - conn["sp"].wait(timeout=timeout) + conn["sp"].wait(timeout=self.timeout) self.log.info( "BTEQ command exited with return code %s", conn["sp"].returncode ) @@ -102,7 +108,7 @@ def on_kill(self): try: self.log.info("Terminating subprocess...") conn["sp"].terminate() - conn["sp"].wait(timeout=5) + conn["sp"].wait(timeout=self.timeout) self.log.info("Subprocess terminated successfully.") except subprocess.TimeoutExpired: self.log.warning( 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 index fc143191..6b95a207 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -44,7 +44,7 @@ def failing_job(): def test_execute_bteq_timeout(tmp_path): @op(required_resource_keys={"teradata"}) def timeout_op(context): - context.resources.teradata.bteq_operator("SELECT * FROM dbc.dbcinfo;", timeout=1) + context.resources.teradata.bteq_operator("SELECT * FROM dbc.dbcinfo;", 1) @job(resource_defs={"teradata": td_resource}) def timeout_job(): From 45b7e19e959a541a2427bdb88c377f91b9adcd2d Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Wed, 21 May 2025 14:27:34 +0530 Subject: [PATCH 05/21] bteq_operator for dagster-teradata --- .../dagster_teradata/resources.py | 83 ++- .../dagster-teradata/dagster_teradata/ttu.py | 593 ++++++++++++++---- .../functional/test_execute_bteq.py | 187 ++++-- 3 files changed, 695 insertions(+), 168 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 4d2276e9..d53021e2 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -226,11 +226,35 @@ def create_fresh_database(teradata: TeradataResource): results = results.append(cursor.fetchall()) # type: ignore return results - def bteq_operator(self, bteq_script: str, timeout: int = None) -> None: - """Executes BTEQ sentences using BTEQ binary. + def bteq_operator( + self, + bteq_script: str = None, + bteq_script_file: str = None, + timeout: int = None, + xcom_push_flag: bool = False, + 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" + ) -> Optional[str]: + """Executes BTEQ sentences using BTEQ binary, either locally or remotely via SSH. Args: bteq_script (str): String of BTEQ sentences to be executed. + bteq_script_file (str): Path to a file containing BTEQ sentences to be executed. + timeout (int, optional): Timeout in seconds for the BTEQ execution. + xcom_push_flag (bool, optional): Flag for pushing last line of BTEQ Log to XCom. + remote_host (str, optional): Remote host to execute BTEQ on (None for local execution). + remote_user (str, optional): SSH username for remote execution. + remote_password (str, optional): SSH password (optional if using key-based auth). + ssh_key_path (str, optional): Path to SSH private key (alternative to password). + remote_port (int, optional): SSH port (default: 22). + remote_working_dir (str, optional): Working directory on remote host (default: /tmp). + + Returns: + Optional[str]: The last line of the BTEQ log if xcom_push_flag is True, otherwise None. Examples: .. code-block:: python @@ -238,9 +262,60 @@ def bteq_operator(self, bteq_script: str, timeout: int = None) -> None: @op def execute_bteq(teradata: TeradataResource): bteq = "SELECT * FROM dbc.dbcinfo" - teradata.execute_bteq(bteq) + # Local execution + teradata.bteq_operator(bteq) + + # Remote execution with password + teradata.bteq_operator( + bteq, + remote_host="remote.server.com", + remote_user="myuser", + remote_password="mypassword" + ) + + # Remote execution with SSH key + teradata.bteq_operator( + bteq, + remote_host="remote.server.com", + remote_user="myuser", + ssh_key_path="~/.ssh/id_rsa" + ) """ - self.bteq.bteq_operator(bteq_script, timeout) + + # Validate input + if not bteq_script and not bteq_script_file: + raise ValueError("Either bteq_script or bteq_script_file must be provided") + + if bteq_script and bteq_script_file: + raise ValueError("Cannot specify both bteq_script and bteq_script_file") + + 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" + ) + + if remote_port < 1 or remote_port > 65535: + raise ValueError("remote_port must be a valid port number (1-65535)") + + return self.bteq.bteq_operator( + bteq_script=bteq_script, + bteq_script_file=bteq_script_file, + xcom_push_flag=xcom_push_flag, + timeout=timeout, + 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 + ) def drop_database(self, databases: Union[str, Sequence[str]]) -> None: """ diff --git a/libraries/dagster-teradata/dagster_teradata/ttu.py b/libraries/dagster-teradata/dagster_teradata/ttu.py index 3a7afad9..6eaa2d6c 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu.py @@ -1,31 +1,317 @@ +import getpass +import json +import os +import platform +import signal +import subprocess +from datetime import datetime from tempfile import NamedTemporaryFile, TemporaryDirectory +from typing import Optional + +import paramiko +from cryptography.fernet import Fernet from dagster import DagsterError -import subprocess -import os + + +class SecureCredentialManager: + """ + Handles secure storage and retrieval of credentials using symmetric encryption. + Provides methods for encrypting/decrypting sensitive data with Fernet encryption. + """ + + def __init__(self, key_file: Optional[str] = None): + """ + Initialize the credential manager. + + Args: + key_file: Optional path to encryption key file. Defaults to ~/.ssh/bteq_cred_key + """ + 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 encryption key from file or generate a new one. + Creates the key file with secure permissions (600) if generating new key. + """ + try: + if os.path.exists(self.key_file): + with open(self.key_file, 'rb') as f: + self.key = f.read() + else: + # Generate new encryption key + self.key = Fernet.generate_key() + with open(self.key_file, 'wb') as f: + f.write(self.key) + os.chmod(self.key_file, 0o600) # Restrict permissions + except Exception as e: + raise DagsterError(f"Failed to handle encryption key: {e}") + + def encrypt(self, data: str) -> str: + """ + Encrypt sensitive data using Fernet symmetric encryption. + + Args: + data: Plaintext string to encrypt + + Returns: + Encrypted string (base64 encoded) + """ + f = Fernet(self.key) + return f.encrypt(data.encode()).decode() + + def decrypt(self, encrypted_data: str) -> str: + """ + Decrypt sensitive data using Fernet symmetric encryption. + + Args: + encrypted_data: Encrypted string (base64 encoded) + + Returns: + Decrypted plaintext string + """ + f = Fernet(self.key) + return f.decrypt(encrypted_data.encode()).decode() class Bteq: + """ + Main BTEQ operator class that handles both local and remote BTEQ script execution. + Provides secure credential management and SSH connectivity for remote operations. + """ + def __init__(self, connection, teradata_connection_resource, log): + """ + Initialize the BTEQ operator. + + Args: + connection: Connection object (unused in current implementation) + teradata_connection_resource: Contains Teradata connection parameters + log: Logger instance for operation logging + """ self.timeout = None self.connection = connection self.log = log self.teradata_connection_resource = teradata_connection_resource + self.cred_manager = SecureCredentialManager() # For secure credential storage + self.ssh_client = None # Will hold SSH connection if established + + def _setup_ssh_connection(self, ssh_host: str, ssh_user: str, ssh_password: Optional[str] = None, + ssh_key_path: Optional[str] = None, ssh_port: int = 22) -> bool: + """ + Establish SSH connection to remote host using either password or key authentication. + + Args: + ssh_host: Remote hostname or IP address + ssh_user: SSH username + ssh_password: Optional SSH password (prompts if not provided) + ssh_key_path: Optional path to SSH private key + ssh_port: SSH port (default: 22) + + Returns: + bool: True if connection succeeded + + Raises: + DagsterError: If connection fails + """ + try: + self.ssh_client = paramiko.SSHClient() + self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + if ssh_key_path: + # Key-based authentication + key = paramiko.RSAKey.from_private_key_file(ssh_key_path) + self.ssh_client.connect(ssh_host, port=ssh_port, username=ssh_user, pkey=key) + else: + # Password authentication + if not ssh_password: + # Try to get password from secure storage + creds = self._get_stored_credentials(ssh_host, ssh_user) + if creds: + ssh_password = self.cred_manager.decrypt(creds['password']) + + if not ssh_password: + # Prompt for password if not found + ssh_password = getpass.getpass(f"Enter SSH password for {ssh_user}@{ssh_host}: ") + # Store the password securely + self._store_credentials(ssh_host, ssh_user, ssh_password) + + self.ssh_client.connect(ssh_host, port=ssh_port, username=ssh_user, password=ssh_password) + + self.log.info(f"SSH connection established to {ssh_user}@{ssh_host}") + return True + except Exception as e: + raise DagsterError(f"SSH connection failed: {e}") + + def _store_credentials(self, host: str, username: str, password: str) -> bool: + """ + Store encrypted credentials in a JSON file with secure permissions. + + Args: + host: Target hostname or IP address + username: Authentication username + password: Plaintext password to encrypt and store + + Returns: + bool: True if storage succeeded, False otherwise + + Security Notes: + - Credentials file has 600 permissions + - Password is encrypted before storage + - Uses atomic file write to prevent corruption + - Maintains existing credentials when updating + """ + cred_file = os.path.expanduser("~/.ssh/bteq_credentials.json") + cred_key = f"{username}@{host}" + + # Ensure .ssh directory exists with secure permissions + ssh_dir = os.path.dirname(cred_file) + os.makedirs(ssh_dir, mode=0o700, exist_ok=True) + + # Load existing credentials if file exists + creds = {} + try: + if os.path.exists(cred_file): + with open(cred_file, 'r') as f: + creds = json.load(f) + except json.JSONDecodeError as e: + self.log.error(f"Credentials file corrupted: {e}") + return False + except Exception as e: + self.log.error(f"Error reading credentials file: {e}") + return False + + # Store new credentials with encrypted password + try: + creds[cred_key] = { + 'username': username, + 'password': self.cred_manager.encrypt(password), + 'timestamp': datetime.utcnow().isoformat() # Add metadata for auditing + } + except Exception as e: + self.log.error(f"Password encryption failed: {e}") + return False + + # Write credentials with secure file permissions + try: + # Write to temporary file first (atomic operation) + temp_file = f"{cred_file}.tmp" + with open(temp_file, 'w') as f: + json.dump(creds, f, indent=2) # Pretty print for debugging + os.chmod(temp_file, 0o600) # Restrict permissions + + # Atomic rename to replace existing file + os.replace(temp_file, cred_file) + return True + except Exception as e: + self.log.error(f"Failed to store credentials: {e}") + try: + os.unlink(temp_file) # Clean up temp file if exists + except: + pass + return False + + def _get_stored_credentials(self, host: str, username: str) -> Optional[dict]: + """ + Retrieve stored credentials for a host/user combination. + + Args: + host: Target hostname or IP address + username: Authentication username + + Returns: + dict: Stored credentials if found, None otherwise + """ + cred_file = os.path.expanduser("~/.ssh/bteq_credentials.json") + cred_key = f"{username}@{host}" + + try: + if os.path.exists(cred_file): + with open(cred_file, 'r') as f: + creds = json.load(f) + return creds.get(cred_key) + except Exception as e: + self.log.warning(f"Failed to read credentials file: {e}") + + return None + + def _execute_remote_bteq(self, bteq_script_path: str, remote_working_dir: str) -> tuple: + """ + Execute BTEQ script on remote host via SSH. + + Args: + bteq_script_path: Local path to BTEQ script + remote_working_dir: Working directory on remote host + + Returns: + tuple: (output, error, exit_status) + + Raises: + DagsterError: If execution fails + """ + if not self.ssh_client: + raise DagsterError("SSH connection not established") + + try: + # Transfer the script to remote host via SFTP + sftp = self.ssh_client.open_sftp() + remote_script_name = os.path.basename(bteq_script_path) + remote_script_path = remote_working_dir + "/" + remote_script_name + sftp.put(bteq_script_path, remote_script_path) + sftp.close() + + # Execute the script remotely + command = f"cd {remote_working_dir} && bteq < {remote_script_name}" + stdin, stdout, stderr = self.ssh_client.exec_command(command, timeout=self.timeout) + + # Read output and errors + output = stdout.read().decode() + error = stderr.read().decode() + exit_status = stdout.channel.recv_exit_status() + + return output, error, exit_status + except Exception as e: + raise DagsterError(f"Remote BTEQ execution failed: {e}") def bteq_operator( - self, bteq_script: str, xcom_push_flag=False, timeout: int | None = None - ) -> str: + self, + bteq_script: str = None, + bteq_script_file: str = None, + xcom_push_flag=False, + timeout: int | None = 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" + ) -> Optional[str]: """ - Executes BTEQ sentences using BTEQ binary. - :param bteq: string of BTEQ sentences - :param xcom_push_flag: Flag for pushing last line of BTEQ Log to XCom - :param timeout: Timeout in seconds for the BTEQ execution. - :return: The last line of the BTEQ log if xcom_push_flag is True, otherwise None. - """ - if timeout is not None: - self.timeout = timeout - else: - # Run without timeout or use a default - self.timeout = 5 + Main BTEQ operator method that executes BTEQ scripts (string or file) either locally or remotely. + + Args: + bteq_script: BTEQ script content to execute + bteq_script_file: Path to BTEQ script file to execute + xcom_push_flag: If True, returns last line of output for XCom + timeout: Execution timeout in seconds + remote_host: If specified, executes remotely on this host + remote_user: SSH username for remote execution + remote_password: SSH password for remote execution + ssh_key_path: Path to SSH private key for authentication + remote_port: SSH port (default: 22) + remote_working_dir: Remote working directory (default: /tmp) + + Returns: + str: Last line of output if xcom_push_flag=True, otherwise None + + Raises: + DagsterError: If execution fails or times out + """ + # Set timeout (default 5 minutes if not specified) + self.timeout = timeout if timeout is not None else 300 + + # Prepare connection parameters conn = { "host": self.teradata_connection_resource.host, "login": self.teradata_connection_resource.user, @@ -35,113 +321,207 @@ def bteq_operator( "bteq_quit_zero": self.teradata_connection_resource.bteq_quit_zero, "console_output_encoding": self.teradata_connection_resource.console_output_encoding, } + self.log.info("Executing BTEQ script...") + # Use temporary directory for script file with TemporaryDirectory(prefix="dagster_ttu_bteq_") as tmpdir: - with NamedTemporaryFile(dir=tmpdir, mode="wb") as tmpfile: - bteq_file_content = self._prepare_bteq_script( - bteq_script, - conn["host"], - conn["login"], - conn["password"], - conn["bteq_output_width"], - conn["bteq_session_encoding"], - conn["bteq_quit_zero"], - ) - self.log.debug("Generated BTEQ script:\n%s", bteq_file_content) - - tmpfile.write(bytes(bteq_file_content, "UTF8")) - tmpfile.flush() - tmpfile.seek(0) - - conn["sp"] = subprocess.Popen( - ["bteq"], - stdin=tmpfile, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - cwd=tmpdir, - preexec_fn=os.setsid, - ) - - line = "" - failure_message = "An error occurred during the BTEQ operation. Please review the full BTEQ output for details." - self.log.info("BTEQ Output:") - for line in iter(conn["sp"].stdout.readline, b""): - decoded_line = line.decode(conn["console_output_encoding"]).strip() - self.log.info(decoded_line) - last_line = decoded_line - if "Failure" in decoded_line: - # Save the last failure message - failure_message = decoded_line - - # Wait for the BTEQ process to complete with optional timeout - try: - conn["sp"].wait(timeout=self.timeout) - self.log.info( - "BTEQ command exited with return code %s", conn["sp"].returncode - ) - except subprocess.TimeoutExpired: - self.on_kill() - raise DagsterError( - f"BTEQ command timed out after {timeout} seconds" + tmpfile_path = None + last_line = "" + failure_message = "BTEQ operation failed - check logs for details" + + try: + # Create temporary script file + with NamedTemporaryFile(dir=tmpdir, mode="wb", delete=False) as tmpfile: + # Handle script file input + if bteq_script_file: + try: + with open(bteq_script_file, 'r') as script_file: + file_content = script_file.read() + bteq_file_content, masked_content = self._prepare_bteq_script( + file_content, + conn["host"], + conn["login"], + conn["password"], + conn["bteq_output_width"], + conn["bteq_session_encoding"], + conn["bteq_quit_zero"], + ) + except IOError as e: + raise DagsterError(f"Failed to read BTEQ script file: {e}") + # Handle script string input + else: + bteq_file_content, masked_content = self._prepare_bteq_script( + bteq_script, + conn["host"], + conn["login"], + conn["password"], + conn["bteq_output_width"], + conn["bteq_session_encoding"], + conn["bteq_quit_zero"], + ) + + self.log.debug("Generated BTEQ script:\n%s", masked_content) + + tmpfile.write(bytes(bteq_file_content, "UTF8")) + tmpfile.flush() + tmpfile.seek(0) + tmpfile_path = tmpfile.name + + if remote_host: + # Remote execution + self._setup_ssh_connection( + ssh_host=remote_host, + ssh_user=remote_user, + ssh_password=remote_password, + ssh_key_path=ssh_key_path, + ssh_port=remote_port ) - # Raise an exception if the BTEQ command failed + output, error, returncode = self._execute_remote_bteq(tmpfile_path, remote_working_dir) + + # Process remote execution output + self.log.info("BTEQ Output:") + for line in output.splitlines(): + decoded_line = line.strip() + self.log.info(decoded_line) + last_line = decoded_line + if "Failure" in decoded_line: + failure_message = decoded_line + if "RDBMS CRASHED" in decoded_line: + failure_message = decoded_line + self.log.error("Detected RDBMS crash.") + raise DagsterError(f"BTEQ command aborted due to: {failure_message}") + + if error: + self.log.error("BTEQ Errors:") + for line in error.splitlines(): + self.log.error(line.strip()) + + conn["sp"] = type('obj', (object,), {'returncode': returncode}) + else: + # Local execution + if platform.system() == "Windows": + conn["sp"] = subprocess.Popen( + ["bteq"], + stdin=open(tmpfile_path, 'rb'), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + cwd=tmpdir, + creationflags=subprocess.CREATE_NEW_PROCESS_GROUP, + ) + else: + conn["sp"] = subprocess.Popen( + ["bteq"], + stdin=open(tmpfile_path, 'rb'), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + cwd=tmpdir, + preexec_fn=os.setsid, + ) + + # Process local execution output + self.log.info("BTEQ Output:") + for line in iter(conn["sp"].stdout.readline, b""): + decoded_line = line.decode(conn["console_output_encoding"]).strip() + self.log.info(decoded_line) + last_line = decoded_line + if "Failure" in decoded_line: + failure_message = decoded_line + if "RDBMS CRASHED" in decoded_line: + failure_message = decoded_line + self.log.error("Detected RDBMS crash. Terminating BTEQ subprocess.") + self.on_kill() + raise DagsterError(f"BTEQ command aborted due to: {failure_message}") + + # Wait for process completion with timeout + try: + conn["sp"].wait(timeout=self.timeout) + self.log.info("BTEQ command exited with return code %s", conn["sp"].returncode) + except subprocess.TimeoutExpired: + self.on_kill() + raise DagsterError(f"BTEQ command timed out after {timeout} seconds") + + # Check return code if conn["sp"].returncode: raise DagsterError( f"BTEQ command exited with return code {conn['sp'].returncode} due to: {failure_message}" ) - # Return the last line of the BTEQ log if xcom_push_flag is True - if xcom_push_flag: - return last_line - return None + # Return last line if xcom_push_flag is set + return last_line if xcom_push_flag else None + + finally: + # Cleanup resources + try: + if tmpfile_path and os.path.exists(tmpfile_path): + os.remove(tmpfile_path) + except Exception as e: + self.log.warning("Failed to remove temp file %s: %s", tmpfile_path, e) + + if self.ssh_client: + self.ssh_client.close() + self.ssh_client = None def on_kill(self): """ - Terminates the subprocess if it is running. - Ensures that the process is terminated gracefully and logs the status. + Terminate the running BTEQ process gracefully. + Handles both local and remote process termination. """ self.log.debug("Attempting to kill child process...") conn = self.get_conn() if conn.get("sp"): + sp = conn.get("sp") try: - self.log.info("Terminating subprocess...") - conn["sp"].terminate() - conn["sp"].wait(timeout=self.timeout) - self.log.info("Subprocess terminated successfully.") + if hasattr(sp, 'kill'): # Local process + if platform.system() == "Windows": + # Send CTRL_BREAK_EVENT to process group + sp.send_signal(signal.CTRL_BREAK_EVENT) + else: + # Send SIGTERM to the process group + os.killpg(os.getpgid(sp.pid), signal.SIGTERM) + + sp.wait(timeout=self.timeout) + self.log.info("Subprocess terminated successfully.") + else: # Remote process + self.log.info("Remote process termination not implemented - connection closed") except subprocess.TimeoutExpired: - self.log.warning( - "Subprocess did not terminate in time. Forcing kill..." - ) - conn["sp"].kill() - self.log.info("Subprocess killed forcefully.") + self.log.warning("Subprocess did not terminate in time. Forcing kill...") + if hasattr(sp, 'kill'): + conn["sp"].kill() + self.log.info("Subprocess killed forcefully.") except (ProcessLookupError, OSError) as e: self.log.error("Failed to terminate subprocess: %s", e) @staticmethod def _prepare_bteq_script( - bteq_string: str, - host: str, - login: str, - password: str, - bteq_output_width: int, - bteq_session_encoding: str, - bteq_quit_zero: bool, + bteq_string: str, + host: str, + login: str, + password: str, + bteq_output_width: int, + bteq_session_encoding: str, + bteq_quit_zero: bool, ) -> str: """ - Prepare a BTEQ file with connection parameters for executing SQL sentences with BTEQ syntax. - :param bteq_string: BTEQ sentences to execute. - :param host: Teradata Host. - :param login: Username for login. - :param password: Password for login. - :param bteq_output_width: Width of BTEQ output in the console. - :param bteq_session_encoding: Session encoding. See official Teradata docs for possible values. - :param bteq_quit_zero: If True, force a .QUIT 0 sentence at the end of the script (forcing return code = 0). - :return: A formatted BTEQ script as a string. - :raises ValueError: If any required parameters are invalid. - """ - # Validate input parameters + Prepare a complete BTEQ script with connection parameters and commands. + + Args: + bteq_string: BTEQ commands to execute + host: Teradata hostname + login: Teradata username + password: Teradata password + bteq_output_width: Output width setting + bteq_session_encoding: Character encoding + bteq_quit_zero: Whether to force quit with code 0 + + Returns: + str: Complete BTEQ script ready for execution + + Raises: + ValueError: If any required parameters are invalid + """ + # Validate inputs if not bteq_string or not bteq_string.strip(): raise ValueError("BTEQ script cannot be empty.") if not host: @@ -155,23 +535,26 @@ def _prepare_bteq_script( if not bteq_session_encoding: raise ValueError("BTEQ session encoding cannot be empty.") - # Construct the BTEQ script + # Build BTEQ script bteq_list = [ f".LOGON {host}/{login},{password};", - ".IF ERRORCODE <> 0 THEN .QUIT 8;", + ".IF ERRORCODE <> 0 THEN .QUIT 8;", # Exit if login fails f".SET WIDTH {bteq_output_width};", f".SET SESSION CHARSET '{bteq_session_encoding}';", - bteq_string.strip(), + bteq_string.strip(), # User-provided BTEQ commands ] - # Add optional .QUIT 0 command if specified if bteq_quit_zero: - bteq_list.append(".QUIT 0;") + bteq_list.append(".QUIT 0;") # Force successful exit - # Ensure proper termination of the script + # Standard script footer bteq_list.extend([".LOGOFF;", ".EXIT;"]) - - # Join the script lines with newlines bteq_script = "\n".join(bteq_list) - return bteq_script + # Create a masked version for logging + masked_script = bteq_script.replace( + f",{password};", + ",*******;" # Mask password + ) + + return bteq_script, masked_script # Return both versions \ No newline at end of file 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 index 6b95a207..190a1911 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -1,6 +1,7 @@ import os -from dagster import job, op +import pytest +from dagster import job, op, DagsterError from dagster_teradata import TeradataResource td_resource = TeradataResource( @@ -10,104 +11,172 @@ 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( + bteq_script="SELECT * FROM dbc.dbcinfo;" + ) + context.log.info(result) + assert result is None + + @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_bteq_file_execution(tmp_path): + script_file = tmp_path / "script.bteq" + script_file.write_text("SELECT * FROM dbc.dbcinfo;") -def test_execute_bteq(tmp_path): @op(required_resource_keys={"teradata"}) - def example_test_execute_bteq(context): + def example_test_local_file(context): result = context.resources.teradata.bteq_operator( - "select * from dbc.dbcinfo;" + bteq_script_file=str(script_file) ) context.log.info(result) + assert result is None @job(resource_defs={"teradata": td_resource}) def example_job(): - example_test_execute_bteq() + example_test_local_file() example_job.execute_in_process(resources={"teradata": td_resource}) -def test_execute_bteq_sql_error(tmp_path): + +def test_remote_bteq_password_auth(): @op(required_resource_keys={"teradata"}) - def failing_op(context): - context.resources.teradata.bteq_operator("SELECT * FROM;") # Invalid SQL + def example_test_remote_password(context): + result = context.resources.teradata.bteq_operator( + bteq_script="SELECT * FROM dbc.dbcinfo;", + remote_host="host", + remote_user="username", + remote_password="password" + ) + context.log.info(result) + assert result is None @job(resource_defs={"teradata": td_resource}) - def failing_job(): - failing_op() + def example_job(): + example_test_remote_password() - try: - failing_job.execute_in_process(resources={"teradata": td_resource}) - assert False, "Expected job to fail due to syntax error" - except Exception as e: - assert "Syntax error" in str(e) or "3706" in str(e) + example_job.execute_in_process(resources={"teradata": td_resource}) -def test_execute_bteq_timeout(tmp_path): +def test_remote_bteq_ssh_key_auth(): @op(required_resource_keys={"teradata"}) - def timeout_op(context): - context.resources.teradata.bteq_operator("SELECT * FROM dbc.dbcinfo;", 1) + def example_test_remote_ssh_key(context): + result = context.resources.teradata.bteq_operator( + bteq_script="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 timeout_job(): - timeout_op() + def example_job(): + example_test_remote_ssh_key() - try: - timeout_job.execute_in_process(resources={"teradata": td_resource}) - assert False, "Expected timeout" - except Exception as e: - assert "timed out" in str(e) + example_job.execute_in_process(resources={"teradata": td_resource}) -import pytest +def test_bteq_xcom_push(): + @op(required_resource_keys={"teradata"}) + def example_test_xcom_push(context): + result = context.resources.teradata.bteq_operator( + bteq_script="SELECT * FROM dbc.dbcinfo;", + xcom_push_flag=True + ) + context.log.info(result) + assert isinstance(result, str) + + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_xcom_push() + + example_job.execute_in_process(resources={"teradata": td_resource}) + -@pytest.mark.parametrize("invalid_sql", ["", None]) -def test_invalid_bteq_sql(invalid_sql): +def test_no_bteq_script_provided(): @op(required_resource_keys={"teradata"}) - def invalid_sql_op(context): - context.resources.teradata.bteq_operator(invalid_sql) + def example_test_no_script(context): + with pytest.raises(ValueError): + context.resources.teradata.bteq_operator() @job(resource_defs={"teradata": td_resource}) - def invalid_sql_job(): - invalid_sql_op() + def example_job(): + example_test_no_script() - with pytest.raises(ValueError, match="BTEQ script cannot be empty"): - invalid_sql_job.execute_in_process(resources={"teradata": td_resource}) + example_job.execute_in_process(resources={"teradata": td_resource}) -def test_bteq_custom_config(tmp_path): - custom_td_resource = TeradataResource( - host="localhost", - user="user", - password="password", - database="dbc", - bteq_output_width=200, - bteq_session_encoding="UTF8", - bteq_quit_zero=True, - ) +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 config_check_op(context): - result = context.resources.teradata.bteq_operator("SELECT * FROM dbc.dbcinfo;") - assert result is not None + def example_test_conflicting_sources(context): + with pytest.raises(ValueError): + context.resources.teradata.bteq_operator( + bteq_script="SELECT * FROM dbc.dbcinfo;", + bteq_script_file=str(script_file) + ) - @job(resource_defs={"teradata": custom_td_resource}) - def config_job(): - config_check_op() + @job(resource_defs={"teradata": td_resource}) + def example_job(): + example_test_conflicting_sources() - config_job.execute_in_process(resources={"teradata": custom_td_resource}) + example_job.execute_in_process(resources={"teradata": td_resource}) -def test_logging_output(tmp_path, caplog): +def test_invalid_script_file(): @op(required_resource_keys={"teradata"}) - def log_op(context): - result = context.resources.teradata.bteq_operator("SELECT * FROM dbc.dbcinfo;") - context.log.info(result) + def example_test_invalid_file(context): + with pytest.raises(DagsterError): + context.resources.teradata.bteq_operator( + bteq_script_file="/nonexistent/path/script.bteq" + ) @job(resource_defs={"teradata": td_resource}) - def log_job(): - log_op() + def example_job(): + example_test_invalid_file() + + example_job.execute_in_process(resources={"teradata": td_resource}) - log_job.execute_in_process(resources={"teradata": td_resource}) - assert any("dbcinfo" in record.message for record in caplog.records) +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( + bteq_script="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( + bteq_script="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}) \ No newline at end of file From 03a032739b4fee0509511f7d6a4c925d8bcd9915 Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Thu, 22 May 2025 17:15:13 +0530 Subject: [PATCH 06/21] Added expected_return_code --- .../dagster_teradata/resources.py | 8 +++- .../dagster-teradata/dagster_teradata/ttu.py | 7 ++-- .../functional/test_execute_bteq.py | 37 +++++++++++++++++++ 3 files changed, 47 insertions(+), 5 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index d53021e2..43e6e217 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -237,7 +237,9 @@ def bteq_operator( remote_password: Optional[str] = None, ssh_key_path: Optional[str] = None, remote_port: int = 22, - remote_working_dir: str = "/tmp" + remote_working_dir: str = "/tmp", + expected_return_code: int = 0, + ) -> Optional[str]: """Executes BTEQ sentences using BTEQ binary, either locally or remotely via SSH. @@ -252,6 +254,7 @@ def bteq_operator( ssh_key_path (str, optional): Path to SSH private key (alternative to password). remote_port (int, optional): SSH port (default: 22). remote_working_dir (str, optional): Working directory on remote host (default: /tmp). + expected_return_code (int, optional): Expected return code for BTEQ execution. Returns: Optional[str]: The last line of the BTEQ log if xcom_push_flag is True, otherwise None. @@ -314,7 +317,8 @@ def execute_bteq(teradata: TeradataResource): remote_password=remote_password, ssh_key_path=ssh_key_path, remote_port=remote_port, - remote_working_dir=remote_working_dir + remote_working_dir=remote_working_dir, + expected_return_code = expected_return_code ) def drop_database(self, databases: Union[str, Sequence[str]]) -> None: diff --git a/libraries/dagster-teradata/dagster_teradata/ttu.py b/libraries/dagster-teradata/dagster_teradata/ttu.py index 6eaa2d6c..92739ae0 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu.py @@ -285,7 +285,8 @@ def bteq_operator( remote_password: Optional[str] = None, ssh_key_path: Optional[str] = None, remote_port: int = 22, - remote_working_dir: str = "/tmp" + remote_working_dir: str = "/tmp", + expected_return_code: int = 0, ) -> Optional[str]: """ Main BTEQ operator method that executes BTEQ scripts (string or file) either locally or remotely. @@ -301,7 +302,7 @@ def bteq_operator( ssh_key_path: Path to SSH private key for authentication remote_port: SSH port (default: 22) remote_working_dir: Remote working directory (default: /tmp) - + expected_return_code: Expected return code for successful execution (default: 0) Returns: str: Last line of output if xcom_push_flag=True, otherwise None @@ -443,7 +444,7 @@ def bteq_operator( raise DagsterError(f"BTEQ command timed out after {timeout} seconds") # Check return code - if conn["sp"].returncode: + if conn["sp"].returncode != expected_return_code: raise DagsterError( f"BTEQ command exited with return code {conn['sp'].returncode} due to: {failure_message}" ) 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 index 190a1911..5afb395a 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -26,6 +26,43 @@ def example_job(): 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( + bteq_script="delete from abcdefgh;", + expected_return_code=8 + ) + context.log.info(result) + assert result is None + + @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( + bteq_script="delete from abcdefgh;", + remote_host="host", + remote_user="username", + remote_password="password", + expected_return_code=8 + ) + context.log.info(result) + assert result is None + + @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_local_bteq_file_execution(tmp_path): script_file = tmp_path / "script.bteq" From 13ef4ebc7842f013358133a70f3464125067405b Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Thu, 22 May 2025 17:37:32 +0530 Subject: [PATCH 07/21] Accepts multiple expected return codes --- .../dagster_teradata/resources.py | 4 ++-- .../dagster-teradata/dagster_teradata/ttu.py | 8 +++++++- .../functional/test_execute_bteq.py | 20 ++++++++++++++++++- 3 files changed, 28 insertions(+), 4 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 43e6e217..1a2d9a17 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -238,7 +238,7 @@ def bteq_operator( ssh_key_path: Optional[str] = None, remote_port: int = 22, remote_working_dir: str = "/tmp", - expected_return_code: int = 0, + expected_return_code: Union[int, List[int]] = 0, ) -> Optional[str]: """Executes BTEQ sentences using BTEQ binary, either locally or remotely via SSH. @@ -254,7 +254,7 @@ def bteq_operator( ssh_key_path (str, optional): Path to SSH private key (alternative to password). remote_port (int, optional): SSH port (default: 22). remote_working_dir (str, optional): Working directory on remote host (default: /tmp). - expected_return_code (int, optional): Expected return code for BTEQ execution. + expected_return_code (Union[int, List[int]], optional): Expected return code(s) from BTEQ execution. Returns: Optional[str]: The last line of the BTEQ log if xcom_push_flag is True, otherwise None. diff --git a/libraries/dagster-teradata/dagster_teradata/ttu.py b/libraries/dagster-teradata/dagster_teradata/ttu.py index 92739ae0..5bd68b6d 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu.py @@ -444,7 +444,13 @@ def bteq_operator( raise DagsterError(f"BTEQ command timed out after {timeout} seconds") # Check return code - if conn["sp"].returncode != expected_return_code: + + if isinstance(expected_return_code, int): + expected_return_codes = [expected_return_code] + else: + expected_return_codes = expected_return_code + + if conn["sp"].returncode not in expected_return_codes: raise DagsterError( f"BTEQ command exited with return code {conn['sp'].returncode} due to: {failure_message}" ) 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 index 5afb395a..85ee2493 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -31,7 +31,7 @@ def test_local_expected_return_code(): def example_test_local_expected_return_code(context): result = context.resources.teradata.bteq_operator( bteq_script="delete from abcdefgh;", - expected_return_code=8 + expected_return_code= 8 ) context.log.info(result) assert result is None @@ -63,6 +63,24 @@ def example_job(): 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( + bteq_script="delete from abcdefgh;", + remote_host="host", + remote_user="username", + remote_password="password", + expected_return_code= [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" From 7c504015281e7c6b76f020a8c6faebc93164794a Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Tue, 24 Jun 2025 11:01:16 +0530 Subject: [PATCH 08/21] updated code --- .../dagster_teradata/resources.py | 102 ++-- .../dagster-teradata/dagster_teradata/ttu.py | 567 ------------------ .../dagster_teradata/ttu/__init__.py | 0 .../dagster_teradata/ttu/bteq.py | 560 +++++++++++++++++ .../dagster_teradata/ttu/ddl.py | 0 .../dagster_teradata/ttu/tdload.py | 181 ++++++ .../dagster_teradata/ttu/utils/bteq_util.py | 182 ++++++ .../ttu/utils/encryption_utils.py | 161 +++++ 8 files changed, 1127 insertions(+), 626 deletions(-) delete mode 100644 libraries/dagster-teradata/dagster_teradata/ttu.py create mode 100644 libraries/dagster-teradata/dagster_teradata/ttu/__init__.py create mode 100644 libraries/dagster-teradata/dagster_teradata/ttu/bteq.py create mode 100644 libraries/dagster-teradata/dagster_teradata/ttu/ddl.py create mode 100644 libraries/dagster-teradata/dagster_teradata/ttu/tdload.py create mode 100644 libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py create mode 100644 libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 1a2d9a17..1a595dbc 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -1,7 +1,7 @@ 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 dagster._check as check import teradatasql @@ -17,7 +17,7 @@ from pydantic import Field from dagster_teradata import constants -from dagster_teradata.ttu import Bteq +from dagster_teradata.ttu.bteq import Bteq from dagster_teradata.teradata_compute_cluster_manager import ( TeradataComputeClusterManager, ) @@ -228,69 +228,28 @@ def create_fresh_database(teradata: TeradataResource): def bteq_operator( self, - bteq_script: str = None, - bteq_script_file: str = None, - timeout: int = None, - xcom_push_flag: bool = False, + sql: str = None, + file_path: 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", - expected_return_code: Union[int, List[int]] = 0, + 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[str]: - """Executes BTEQ sentences using BTEQ binary, either locally or remotely via SSH. - - Args: - bteq_script (str): String of BTEQ sentences to be executed. - bteq_script_file (str): Path to a file containing BTEQ sentences to be executed. - timeout (int, optional): Timeout in seconds for the BTEQ execution. - xcom_push_flag (bool, optional): Flag for pushing last line of BTEQ Log to XCom. - remote_host (str, optional): Remote host to execute BTEQ on (None for local execution). - remote_user (str, optional): SSH username for remote execution. - remote_password (str, optional): SSH password (optional if using key-based auth). - ssh_key_path (str, optional): Path to SSH private key (alternative to password). - remote_port (int, optional): SSH port (default: 22). - remote_working_dir (str, optional): Working directory on remote host (default: /tmp). - expected_return_code (Union[int, List[int]], optional): Expected return code(s) from BTEQ execution. - - Returns: - Optional[str]: The last line of the BTEQ log if xcom_push_flag is True, otherwise None. - - Examples: - .. code-block:: python - - @op - def execute_bteq(teradata: TeradataResource): - bteq = "SELECT * FROM dbc.dbcinfo" - # Local execution - teradata.bteq_operator(bteq) - - # Remote execution with password - teradata.bteq_operator( - bteq, - remote_host="remote.server.com", - remote_user="myuser", - remote_password="mypassword" - ) - - # Remote execution with SSH key - teradata.bteq_operator( - bteq, - remote_host="remote.server.com", - remote_user="myuser", - ssh_key_path="~/.ssh/id_rsa" - ) - """ # Validate input - if not bteq_script and not bteq_script_file: - raise ValueError("Either bteq_script or bteq_script_file must be provided") + if not sql and not file_path: + raise ValueError("BteqOperator requires either the 'sql' or 'file_path' parameter. Both are missing.") - if bteq_script and bteq_script_file: - raise ValueError("Cannot specify both bteq_script and bteq_script_file") + 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.") if remote_host: if not remote_user: @@ -307,18 +266,43 @@ def execute_bteq(teradata: TeradataResource): if remote_port < 1 or remote_port > 65535: raise ValueError("remote_port must be a valid port number (1-65535)") + # Validate and set BTEQ session and script encoding + 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" + # for file reading in python. Mapping BTEQ encoding to Python encoding + if bteq_script_encoding == "UTF8": + temp_file_read_encoding = "UTF-8" + elif bteq_script_encoding == "UTF16": + temp_file_read_encoding = "UTF-16" + return self.bteq.bteq_operator( - bteq_script=bteq_script, - bteq_script_file=bteq_script_file, - xcom_push_flag=xcom_push_flag, - timeout=timeout, + 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, - expected_return_code = expected_return_code + 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: diff --git a/libraries/dagster-teradata/dagster_teradata/ttu.py b/libraries/dagster-teradata/dagster_teradata/ttu.py deleted file mode 100644 index 5bd68b6d..00000000 --- a/libraries/dagster-teradata/dagster_teradata/ttu.py +++ /dev/null @@ -1,567 +0,0 @@ -import getpass -import json -import os -import platform -import signal -import subprocess -from datetime import datetime -from tempfile import NamedTemporaryFile, TemporaryDirectory -from typing import Optional - -import paramiko -from cryptography.fernet import Fernet -from dagster import DagsterError - - -class SecureCredentialManager: - """ - Handles secure storage and retrieval of credentials using symmetric encryption. - Provides methods for encrypting/decrypting sensitive data with Fernet encryption. - """ - - def __init__(self, key_file: Optional[str] = None): - """ - Initialize the credential manager. - - Args: - key_file: Optional path to encryption key file. Defaults to ~/.ssh/bteq_cred_key - """ - 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 encryption key from file or generate a new one. - Creates the key file with secure permissions (600) if generating new key. - """ - try: - if os.path.exists(self.key_file): - with open(self.key_file, 'rb') as f: - self.key = f.read() - else: - # Generate new encryption key - self.key = Fernet.generate_key() - with open(self.key_file, 'wb') as f: - f.write(self.key) - os.chmod(self.key_file, 0o600) # Restrict permissions - except Exception as e: - raise DagsterError(f"Failed to handle encryption key: {e}") - - def encrypt(self, data: str) -> str: - """ - Encrypt sensitive data using Fernet symmetric encryption. - - Args: - data: Plaintext string to encrypt - - Returns: - Encrypted string (base64 encoded) - """ - f = Fernet(self.key) - return f.encrypt(data.encode()).decode() - - def decrypt(self, encrypted_data: str) -> str: - """ - Decrypt sensitive data using Fernet symmetric encryption. - - Args: - encrypted_data: Encrypted string (base64 encoded) - - Returns: - Decrypted plaintext string - """ - f = Fernet(self.key) - return f.decrypt(encrypted_data.encode()).decode() - - -class Bteq: - """ - Main BTEQ operator class that handles both local and remote BTEQ script execution. - Provides secure credential management and SSH connectivity for remote operations. - """ - - def __init__(self, connection, teradata_connection_resource, log): - """ - Initialize the BTEQ operator. - - Args: - connection: Connection object (unused in current implementation) - teradata_connection_resource: Contains Teradata connection parameters - log: Logger instance for operation logging - """ - self.timeout = None - self.connection = connection - self.log = log - self.teradata_connection_resource = teradata_connection_resource - self.cred_manager = SecureCredentialManager() # For secure credential storage - self.ssh_client = None # Will hold SSH connection if established - - def _setup_ssh_connection(self, ssh_host: str, ssh_user: str, ssh_password: Optional[str] = None, - ssh_key_path: Optional[str] = None, ssh_port: int = 22) -> bool: - """ - Establish SSH connection to remote host using either password or key authentication. - - Args: - ssh_host: Remote hostname or IP address - ssh_user: SSH username - ssh_password: Optional SSH password (prompts if not provided) - ssh_key_path: Optional path to SSH private key - ssh_port: SSH port (default: 22) - - Returns: - bool: True if connection succeeded - - Raises: - DagsterError: If connection fails - """ - try: - self.ssh_client = paramiko.SSHClient() - self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) - - if ssh_key_path: - # Key-based authentication - key = paramiko.RSAKey.from_private_key_file(ssh_key_path) - self.ssh_client.connect(ssh_host, port=ssh_port, username=ssh_user, pkey=key) - else: - # Password authentication - if not ssh_password: - # Try to get password from secure storage - creds = self._get_stored_credentials(ssh_host, ssh_user) - if creds: - ssh_password = self.cred_manager.decrypt(creds['password']) - - if not ssh_password: - # Prompt for password if not found - ssh_password = getpass.getpass(f"Enter SSH password for {ssh_user}@{ssh_host}: ") - # Store the password securely - self._store_credentials(ssh_host, ssh_user, ssh_password) - - self.ssh_client.connect(ssh_host, port=ssh_port, username=ssh_user, password=ssh_password) - - self.log.info(f"SSH connection established to {ssh_user}@{ssh_host}") - return True - except Exception as e: - raise DagsterError(f"SSH connection failed: {e}") - - def _store_credentials(self, host: str, username: str, password: str) -> bool: - """ - Store encrypted credentials in a JSON file with secure permissions. - - Args: - host: Target hostname or IP address - username: Authentication username - password: Plaintext password to encrypt and store - - Returns: - bool: True if storage succeeded, False otherwise - - Security Notes: - - Credentials file has 600 permissions - - Password is encrypted before storage - - Uses atomic file write to prevent corruption - - Maintains existing credentials when updating - """ - cred_file = os.path.expanduser("~/.ssh/bteq_credentials.json") - cred_key = f"{username}@{host}" - - # Ensure .ssh directory exists with secure permissions - ssh_dir = os.path.dirname(cred_file) - os.makedirs(ssh_dir, mode=0o700, exist_ok=True) - - # Load existing credentials if file exists - creds = {} - try: - if os.path.exists(cred_file): - with open(cred_file, 'r') as f: - creds = json.load(f) - except json.JSONDecodeError as e: - self.log.error(f"Credentials file corrupted: {e}") - return False - except Exception as e: - self.log.error(f"Error reading credentials file: {e}") - return False - - # Store new credentials with encrypted password - try: - creds[cred_key] = { - 'username': username, - 'password': self.cred_manager.encrypt(password), - 'timestamp': datetime.utcnow().isoformat() # Add metadata for auditing - } - except Exception as e: - self.log.error(f"Password encryption failed: {e}") - return False - - # Write credentials with secure file permissions - try: - # Write to temporary file first (atomic operation) - temp_file = f"{cred_file}.tmp" - with open(temp_file, 'w') as f: - json.dump(creds, f, indent=2) # Pretty print for debugging - os.chmod(temp_file, 0o600) # Restrict permissions - - # Atomic rename to replace existing file - os.replace(temp_file, cred_file) - return True - except Exception as e: - self.log.error(f"Failed to store credentials: {e}") - try: - os.unlink(temp_file) # Clean up temp file if exists - except: - pass - return False - - def _get_stored_credentials(self, host: str, username: str) -> Optional[dict]: - """ - Retrieve stored credentials for a host/user combination. - - Args: - host: Target hostname or IP address - username: Authentication username - - Returns: - dict: Stored credentials if found, None otherwise - """ - cred_file = os.path.expanduser("~/.ssh/bteq_credentials.json") - cred_key = f"{username}@{host}" - - try: - if os.path.exists(cred_file): - with open(cred_file, 'r') as f: - creds = json.load(f) - return creds.get(cred_key) - except Exception as e: - self.log.warning(f"Failed to read credentials file: {e}") - - return None - - def _execute_remote_bteq(self, bteq_script_path: str, remote_working_dir: str) -> tuple: - """ - Execute BTEQ script on remote host via SSH. - - Args: - bteq_script_path: Local path to BTEQ script - remote_working_dir: Working directory on remote host - - Returns: - tuple: (output, error, exit_status) - - Raises: - DagsterError: If execution fails - """ - if not self.ssh_client: - raise DagsterError("SSH connection not established") - - try: - # Transfer the script to remote host via SFTP - sftp = self.ssh_client.open_sftp() - remote_script_name = os.path.basename(bteq_script_path) - remote_script_path = remote_working_dir + "/" + remote_script_name - sftp.put(bteq_script_path, remote_script_path) - sftp.close() - - # Execute the script remotely - command = f"cd {remote_working_dir} && bteq < {remote_script_name}" - stdin, stdout, stderr = self.ssh_client.exec_command(command, timeout=self.timeout) - - # Read output and errors - output = stdout.read().decode() - error = stderr.read().decode() - exit_status = stdout.channel.recv_exit_status() - - return output, error, exit_status - except Exception as e: - raise DagsterError(f"Remote BTEQ execution failed: {e}") - - def bteq_operator( - self, - bteq_script: str = None, - bteq_script_file: str = None, - xcom_push_flag=False, - timeout: int | None = 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", - expected_return_code: int = 0, - ) -> Optional[str]: - """ - Main BTEQ operator method that executes BTEQ scripts (string or file) either locally or remotely. - - Args: - bteq_script: BTEQ script content to execute - bteq_script_file: Path to BTEQ script file to execute - xcom_push_flag: If True, returns last line of output for XCom - timeout: Execution timeout in seconds - remote_host: If specified, executes remotely on this host - remote_user: SSH username for remote execution - remote_password: SSH password for remote execution - ssh_key_path: Path to SSH private key for authentication - remote_port: SSH port (default: 22) - remote_working_dir: Remote working directory (default: /tmp) - expected_return_code: Expected return code for successful execution (default: 0) - Returns: - str: Last line of output if xcom_push_flag=True, otherwise None - - Raises: - DagsterError: If execution fails or times out - """ - # Set timeout (default 5 minutes if not specified) - self.timeout = timeout if timeout is not None else 300 - - # Prepare connection parameters - conn = { - "host": self.teradata_connection_resource.host, - "login": self.teradata_connection_resource.user, - "password": self.teradata_connection_resource.password, - "bteq_output_width": self.teradata_connection_resource.bteq_output_width, - "bteq_session_encoding": self.teradata_connection_resource.bteq_session_encoding, - "bteq_quit_zero": self.teradata_connection_resource.bteq_quit_zero, - "console_output_encoding": self.teradata_connection_resource.console_output_encoding, - } - - self.log.info("Executing BTEQ script...") - - # Use temporary directory for script file - with TemporaryDirectory(prefix="dagster_ttu_bteq_") as tmpdir: - tmpfile_path = None - last_line = "" - failure_message = "BTEQ operation failed - check logs for details" - - try: - # Create temporary script file - with NamedTemporaryFile(dir=tmpdir, mode="wb", delete=False) as tmpfile: - # Handle script file input - if bteq_script_file: - try: - with open(bteq_script_file, 'r') as script_file: - file_content = script_file.read() - bteq_file_content, masked_content = self._prepare_bteq_script( - file_content, - conn["host"], - conn["login"], - conn["password"], - conn["bteq_output_width"], - conn["bteq_session_encoding"], - conn["bteq_quit_zero"], - ) - except IOError as e: - raise DagsterError(f"Failed to read BTEQ script file: {e}") - # Handle script string input - else: - bteq_file_content, masked_content = self._prepare_bteq_script( - bteq_script, - conn["host"], - conn["login"], - conn["password"], - conn["bteq_output_width"], - conn["bteq_session_encoding"], - conn["bteq_quit_zero"], - ) - - self.log.debug("Generated BTEQ script:\n%s", masked_content) - - tmpfile.write(bytes(bteq_file_content, "UTF8")) - tmpfile.flush() - tmpfile.seek(0) - tmpfile_path = tmpfile.name - - if remote_host: - # Remote execution - self._setup_ssh_connection( - ssh_host=remote_host, - ssh_user=remote_user, - ssh_password=remote_password, - ssh_key_path=ssh_key_path, - ssh_port=remote_port - ) - - output, error, returncode = self._execute_remote_bteq(tmpfile_path, remote_working_dir) - - # Process remote execution output - self.log.info("BTEQ Output:") - for line in output.splitlines(): - decoded_line = line.strip() - self.log.info(decoded_line) - last_line = decoded_line - if "Failure" in decoded_line: - failure_message = decoded_line - if "RDBMS CRASHED" in decoded_line: - failure_message = decoded_line - self.log.error("Detected RDBMS crash.") - raise DagsterError(f"BTEQ command aborted due to: {failure_message}") - - if error: - self.log.error("BTEQ Errors:") - for line in error.splitlines(): - self.log.error(line.strip()) - - conn["sp"] = type('obj', (object,), {'returncode': returncode}) - else: - # Local execution - if platform.system() == "Windows": - conn["sp"] = subprocess.Popen( - ["bteq"], - stdin=open(tmpfile_path, 'rb'), - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - cwd=tmpdir, - creationflags=subprocess.CREATE_NEW_PROCESS_GROUP, - ) - else: - conn["sp"] = subprocess.Popen( - ["bteq"], - stdin=open(tmpfile_path, 'rb'), - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - cwd=tmpdir, - preexec_fn=os.setsid, - ) - - # Process local execution output - self.log.info("BTEQ Output:") - for line in iter(conn["sp"].stdout.readline, b""): - decoded_line = line.decode(conn["console_output_encoding"]).strip() - self.log.info(decoded_line) - last_line = decoded_line - if "Failure" in decoded_line: - failure_message = decoded_line - if "RDBMS CRASHED" in decoded_line: - failure_message = decoded_line - self.log.error("Detected RDBMS crash. Terminating BTEQ subprocess.") - self.on_kill() - raise DagsterError(f"BTEQ command aborted due to: {failure_message}") - - # Wait for process completion with timeout - try: - conn["sp"].wait(timeout=self.timeout) - self.log.info("BTEQ command exited with return code %s", conn["sp"].returncode) - except subprocess.TimeoutExpired: - self.on_kill() - raise DagsterError(f"BTEQ command timed out after {timeout} seconds") - - # Check return code - - if isinstance(expected_return_code, int): - expected_return_codes = [expected_return_code] - else: - expected_return_codes = expected_return_code - - if conn["sp"].returncode not in expected_return_codes: - raise DagsterError( - f"BTEQ command exited with return code {conn['sp'].returncode} due to: {failure_message}" - ) - - # Return last line if xcom_push_flag is set - return last_line if xcom_push_flag else None - - finally: - # Cleanup resources - try: - if tmpfile_path and os.path.exists(tmpfile_path): - os.remove(tmpfile_path) - except Exception as e: - self.log.warning("Failed to remove temp file %s: %s", tmpfile_path, e) - - if self.ssh_client: - self.ssh_client.close() - self.ssh_client = None - - def on_kill(self): - """ - Terminate the running BTEQ process gracefully. - Handles both local and remote process termination. - """ - self.log.debug("Attempting to kill child process...") - conn = self.get_conn() - if conn.get("sp"): - sp = conn.get("sp") - try: - if hasattr(sp, 'kill'): # Local process - if platform.system() == "Windows": - # Send CTRL_BREAK_EVENT to process group - sp.send_signal(signal.CTRL_BREAK_EVENT) - else: - # Send SIGTERM to the process group - os.killpg(os.getpgid(sp.pid), signal.SIGTERM) - - sp.wait(timeout=self.timeout) - self.log.info("Subprocess terminated successfully.") - else: # Remote process - self.log.info("Remote process termination not implemented - connection closed") - except subprocess.TimeoutExpired: - self.log.warning("Subprocess did not terminate in time. Forcing kill...") - if hasattr(sp, 'kill'): - conn["sp"].kill() - self.log.info("Subprocess killed forcefully.") - except (ProcessLookupError, OSError) as e: - self.log.error("Failed to terminate subprocess: %s", e) - - @staticmethod - def _prepare_bteq_script( - bteq_string: str, - host: str, - login: str, - password: str, - bteq_output_width: int, - bteq_session_encoding: str, - bteq_quit_zero: bool, - ) -> str: - """ - Prepare a complete BTEQ script with connection parameters and commands. - - Args: - bteq_string: BTEQ commands to execute - host: Teradata hostname - login: Teradata username - password: Teradata password - bteq_output_width: Output width setting - bteq_session_encoding: Character encoding - bteq_quit_zero: Whether to force quit with code 0 - - Returns: - str: Complete BTEQ script ready for execution - - Raises: - ValueError: If any required parameters are invalid - """ - # Validate inputs - if not bteq_string or not bteq_string.strip(): - raise ValueError("BTEQ script cannot be empty.") - if not host: - raise ValueError("Host parameter cannot be empty.") - if not login: - raise ValueError("Login parameter cannot be empty.") - if not password: - raise ValueError("Password parameter cannot be empty.") - if not isinstance(bteq_output_width, int) or bteq_output_width <= 0: - raise ValueError("BTEQ output width must be a positive integer.") - if not bteq_session_encoding: - raise ValueError("BTEQ session encoding cannot be empty.") - - # Build BTEQ script - bteq_list = [ - f".LOGON {host}/{login},{password};", - ".IF ERRORCODE <> 0 THEN .QUIT 8;", # Exit if login fails - f".SET WIDTH {bteq_output_width};", - f".SET SESSION CHARSET '{bteq_session_encoding}';", - bteq_string.strip(), # User-provided BTEQ commands - ] - - if bteq_quit_zero: - bteq_list.append(".QUIT 0;") # Force successful exit - - # Standard script footer - bteq_list.extend([".LOGOFF;", ".EXIT;"]) - bteq_script = "\n".join(bteq_list) - - # Create a masked version for logging - masked_script = bteq_script.replace( - f",{password};", - ",*******;" # Mask password - ) - - return bteq_script, masked_script # Return both versions \ No newline at end of file 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..6ab59812 --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py @@ -0,0 +1,560 @@ +import getpass +import os +import socket +import subprocess +import tempfile +from contextlib import contextmanager +from typing import Optional, List, Union, Literal + +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, +) + + +from dagster_teradata.ttu.utils.encryption_utils import * + +class Bteq: + """ + Main BTEQ operator class with enhanced features: + - Directory processing + - Force login from script + - Secure credential handling + - Remote execution via SSH + """ + + 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 + 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 + + # Core Methods ============================================================ + + def bteq_operator( + self, + sql: str = None, + file_path: 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, + temp_file_read_encoding: Optional[str] = 'UTF-8', + ) -> int | None: + + 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 + + """Execute the BTEQ script either in local machine or on remote host based on ssh_conn_id.""" + # Remote execution + if not self.remote_working_dir: + self.remote_working_dir = "/tmp" + # Handling execution on local: + if not self.remote_host: + if self.sql: + bteq_script = prepare_bteq_script_for_local_execution( + sql=self.sql, + ) + self.log.info("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 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) + # Execution on Remote machine + elif self.remote_host: + # When sql statement is provided as input through sql parameter, Preparing the bteq script + if self.sql: + bteq_script = prepare_bteq_script_for_remote_execution( + conn=self.connection, + sql=self.sql, + ) + self.log.info("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, + ) + if self.file_path: + + if self.file_path 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: int, + 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_remote_password: str | None = None, + ssh_key_path: str | None = None, + remote_port: int = 22, + ) -> int | None: + """Execute the BTEQ script either in local machine or on remote host based on ssh_conn_id.""" + # Remote execution + if self.remote_host: + # Write script to local temp file + # Encrypt the file locally + 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_remote_password=remote_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: int, + 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_remote_password: str | None = None, + ssh_key_path: str | None = None, + remote_port: int = 22, + ) -> int | None: + 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, + remote_host, + remote_user, + remote_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: int, + 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: + encrypted_file_path = None + remote_encrypted_path = None + try: + if not self._setup_ssh_connection( + remote_host, remote_user, remote_password, + ssh_key_path, remote_port + ): + raise DagsterError("SSH connection failed") + 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.info("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.info("stdout : %s", stdout) + self.log.info("stderr : %s", stderr) + self.log.info("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 + else: + raise DagsterError("SSH connection is not established. `ssh_hook` is None or invalid.") + 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: int, + 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: + 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.info("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.info("encode_bteq_script : %s", encode_bteq_script) + stdout_data, _ = process.communicate(input=encode_bteq_script) + self.log.info("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 + 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.") + conn = self.connection + conn["sp"] = process # For `on_kill` support + 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.info("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 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 + 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_airflow_home_dir(self) -> str: + """Get the AIRFLOW_HOME directory.""" + return os.environ.get("AIRFLOW_HOME", "~/airflow") + + @contextmanager + def preferred_temp_directory(self, prefix="bteq_"): + 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_airflow_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: + 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() + # rendered_content = original_content + # if self.contains_template(original_content): + # rendered_content = self.render_template(original_content, context) + bteq_script = prepare_bteq_script_for_remote_execution( + conn=self.connection, + sql=file_content, + ) + self.log.info("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, + ) + 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: + 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")) + # Manually render using operator's context + # rendered_content = file_content + # if self.contains_template(file_content): + # rendered_content = self.render_template(file_content, context) + bteq_script = prepare_bteq_script_for_local_execution( + sql=file_content, + ) + self.log.info("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: str, + password: Optional[str], + key_path: Optional[str], + port: int + ) -> bool: + """Establish SSH connection using either password or key auth.""" + 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: + creds = get_stored_credentials(host, user) + password = self.cred_manager.decrypt(creds['password']) if creds else None + + if not password: + password = getpass.getpass(f"SSH password for {user}@{host}: ") + store_credentials(host, user, password) + + 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}") \ No newline at end of file diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/ddl.py b/libraries/dagster-teradata/dagster_teradata/ttu/ddl.py new file mode 100644 index 00000000..e69de29b diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/tdload.py b/libraries/dagster-teradata/dagster_teradata/ttu/tdload.py new file mode 100644 index 00000000..54be16b8 --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata/ttu/tdload.py @@ -0,0 +1,181 @@ +""" +Detailed Design Specification for Teradata TPT Operators + +Date: Apr 29, 2025 +Author: Satyanarayana Reddy Gopu +Copyright © 2025 by Teradata Corporation. All Rights Reserved. +TERADATA CORPORATION CONFIDENTIAL AND TRADE SECRET +""" + + +class TptHook: + """Base hook for Teradata Parallel Transporter (TPT) operations. + + Handles core TPT functionality including script generation, command execution, + and connection management for both DDL and data loading operations. + + Args: + teradata_conn_id (str): Airflow connection ID for Teradata + ttu_log_folder (str): Path for TTU log files (default: '/tmp') + console_encoding (str): Output encoding (default: 'utf-8') + + Examples: + .. code-block:: python + + # Basic hook initialization + tpt_hook = TptHook(teradata_conn_id='teradata_default') + + # Execute DDL + tpt_hook.execute_ddl( + sql="CREATE TABLE sample_table (id INTEGER)", + error_list="3807, 3820" + ) + """ + + def execute_ddl(self, sql: str, error_list: str = None) -> Optional[str]: + """Execute DDL statements using TPT. + + Args: + sql: SQL DDL statement to execute + error_list: Comma-separated error codes to ignore + + Returns: + str: Command output if successful + + Raises: + AirflowException: If execution fails + """ + pass + + +class DdlOperator(BaseOperator): + """Airflow operator for executing Teradata DDL via TPT. + + Args: + sql (str): DDL statement to execute + teradata_conn_id (str): Airflow connection ID + error_list (str): Comma-separated error codes to ignore + xcom_push_flag (bool): Push output to XCom (default: False) + + Examples: + .. code-block:: python + + # Create table example + create_table = DdlOperator( + task_id='create_table', + sql="CREATE TABLE analytics.sales (id INTEGER, amount DECIMAL(10,2))", + teradata_conn_id='teradata_prod' + ) + """ + pass + + +class TdLoadOperator(BaseOperator): + """Airflow operator for data loading operations using tdload. + + Supports three modes: + 1. FILE_TO_TABLE: Load from file to Teradata table + 2. TABLE_TO_FILE: Export from table to file + 3. TABLE_TO_TABLE: Transfer between tables + + Args: + operation_mode (str): One of FILE_TO_TABLE|TABLE_TO_FILE|TABLE_TO_TABLE + source (str): Source table or file path + target (str): Target table or file path + teradata_conn_id (str): Source connection ID + target_teradata_conn_id (str): Target connection ID (for table-to-table) + tdload_options (str): Additional tdload parameters + timeout (int): Operation timeout in seconds (default: 300) + + Examples: + .. code-block:: python + + # File to table load + load_data = TdLoadOperator( + task_id='load_sales_data', + operation_mode='FILE_TO_TABLE', + source='/data/sales.csv', + target='analytics.sales', + teradata_conn_id='teradata_prod', + tdload_options="--SourceFormat 'Delimited' --SourceTextDelimiter '|'" + ) + """ + + OPERATION_MODES = ['FILE_TO_TABLE', 'TABLE_TO_FILE', 'TABLE_TO_TABLE'] + + def __init__( + self, + operation_mode: str, + source: str, + target: str, + teradata_conn_id: str, + target_teradata_conn_id: str = None, + tdload_options: str = None, + timeout: int = 300, + **kwargs + ): + pass + + +# Connection Configuration +""" +Teradata Connection Requirements: + +Connection Type: teradata + +Required Parameters: +- host: Teradata server hostname +- login: Database username +- password: Database password + +Optional Extras (JSON): +{ + "ttu_log_folder": "/custom/log/path", + "console_encoding": "utf-16" +} + +Example Airflow Connection Setup: +Admin -> Connections -> Add: +- Conn ID: teradata_prod +- Conn Type: teradata +- Host: tdprod.example.com +- Login: dbadmin +- Password: ******** +- Extra: {"ttu_log_folder": "/data/logs", "console_encoding": "utf-8"} +""" + +# Error Handling Specifications +""" +Error Handling Approach: + +1. Connection Errors: + - AirflowConnectionError for invalid credentials + - Retry with exponential backoff + +2. TPT Execution Errors: + - Check return codes against error_list + - AirflowException for fatal errors + - Automatic retry for transient errors + +3. Data Validation: + - Source/target validation before execution + - Mutual exclusion checks +""" + +# Logging Configuration +""" +Logging Structure: + +1. Command Preparation: + - Generated TPT scripts + - Final command strings + +2. Execution Tracking: + - Start/end timestamps + - Progress percentage for long operations + +3. Result Verification: + - Row counts processed + - Error counts + - Performance metrics +""" \ No newline at end of file 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..8ee0f361 --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py @@ -0,0 +1,182 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import os +import shutil +import stat +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from paramiko import SSHClient + +from dagster import DagsterError + + +def verify_bteq_installed(): + """Verify if BTEQ is installed and available in the system's 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 the remote machine.""" + stdin, stdout, stderr = ssh_client.exec_command("which bteq") + exit_status = stdout.channel.recv_exit_status() + output = stdout.read().strip() + error = stderr.read().strip() + + if exit_status != 0 or not output: + raise DagsterError( + f"BTEQ is not installed or not available in PATH. stderr: {error.decode() if error else 'N/A'}" + ) + + +def transfer_file_sftp(ssh_client, local_path, remote_path): + sftp = ssh_client.open_sftp() + sftp.put(local_path, remote_path) + sftp.close() + + +# We can not pass host details with bteq command when executing on remote machine. Instead, we will prepare .logon in bteq script itself to avoid risk of +# exposing sensitive information +def prepare_bteq_script_for_remote_execution(conn: dict[str, Any], sql: str) -> str: + """Build a BTEQ script with necessary connection and session commands.""" + script_lines = [] + host = conn["host"] + login = conn["login"] + password = conn["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 with necessary connection and session commands.""" + script_lines: list[str] = [] + return _prepare_bteq_script(script_lines, sql) + + +def _prepare_bteq_script(script_lines: list[str], sql: str) -> str: + 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]: + 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(";") + # Airflow 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: int, + bteq_script_encoding: str, + bteq_session_encoding: str, + timeout_rc: int, +) -> str: + """Prepare the BTEQ command with necessary parameters.""" + bteq_core_cmd = _prepare_bteq_command(timeout, 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: TeradataResource, + timeout: int, + bteq_script_encoding: str, + bteq_session_encoding: str, + timeout_rc: int, +) -> str: + """Prepare the BTEQ command with necessary parameters.""" + bteq_core_cmd = _prepare_bteq_command(timeout, 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: + return os.path.isfile(file_path) + + +def is_valid_encoding(file_path: str, encoding: str = "UTF-8") -> bool: + """ + Check if the file can be read with the specified encoding. + + :param file_path: Path to the file to be checked. + :param encoding: Encoding to use for reading the file. + :return: True if the file can be read with the specified encoding, False otherwise. + """ + 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. + + :param file_path: Path to the file to be read. + :param encoding: Encoding to use for reading the file. + :return: Content of the file as a string. + """ + 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, logger=None) -> bool: + """Check if the given remote file path is a valid BTEQ script file.""" + 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 \ No newline at end of file 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..04f67224 --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py @@ -0,0 +1,161 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import json +import os +import secrets +import string +import subprocess +from typing import Optional + +from cryptography.fernet import Fernet +from dagster import DagsterError + + +class SecureCredentialManager: + """ + Handles secure storage and retrieval of credentials using Fernet encryption. + Manages encryption keys and provides methods for secure credential storage. + """ + + def __init__(self, key_file: Optional[str] = None): + """ + Initialize credential manager with optional custom key file path. + Defaults to ~/.ssh/bteq_cred_key. + """ + 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 existing encryption key or generate new key with secure permissions.""" + 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() + 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.""" + f = Fernet(self.key) + return f.encrypt(data.encode()).decode() + + def decrypt(self, encrypted_data: str) -> str: + """Decrypt encrypted string back to plaintext.""" + f = Fernet(self.key) + return f.decrypt(encrypted_data.encode()).decode() + +def generate_random_password(length=12): + # Define the character set: letters, digits, and special characters + characters = string.ascii_letters + string.digits + string.punctuation + # Generate a random password + 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): + # Write plaintext temporarily to file + + # 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, password, bteq_command_str): + # 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): + # 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 encrypted format.""" + 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.""" + 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 \ No newline at end of file From 26bcfb8da8539592685e5c7c0af7d058ea796871 Mon Sep 17 00:00:00 2001 From: Mohan Talla Date: Wed, 25 Jun 2025 09:49:00 +0530 Subject: [PATCH 09/21] Added test cases --- libraries/dagster-teradata/__init__.py | 0 .../dagster_teradata/__init__.py | 6 + .../dagster_teradata/resources.py | 44 +- .../dagster_teradata/ttu/bteq.py | 161 ++++--- .../dagster_teradata/ttu/tdload.py | 181 -------- .../dagster_teradata/ttu/utils/bteq_util.py | 31 +- .../ttu/utils/encryption_utils.py | 26 +- .../functional/test_execute_bteq.py | 57 +-- .../dagster_teradata_tests/test_bteq.py | 437 ++++++++++++++++++ .../dagster_teradata_tests/test_bteq_util.py | 177 +++++++ .../test_encryption_utils.py | 92 ++++ libraries/dagster-teradata/pyproject.toml | 1 + 12 files changed, 901 insertions(+), 312 deletions(-) create mode 100644 libraries/dagster-teradata/__init__.py delete mode 100644 libraries/dagster-teradata/dagster_teradata/ttu/tdload.py create mode 100644 libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py create mode 100644 libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py create mode 100644 libraries/dagster-teradata/dagster_teradata_tests/test_encryption_utils.py 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 8af14487..c5d267c5 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.2" DagsterLibraryRegistry.register( diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 1a595dbc..1f2a1a2b 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -16,7 +16,7 @@ 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, @@ -227,29 +227,31 @@ def create_fresh_database(teradata: TeradataResource): return results def bteq_operator( - self, - sql: str = None, - file_path: 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, - + self, + sql: str = None, + file_path: 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[str]: - # Validate input if not sql and not file_path: - raise ValueError("BteqOperator requires either the 'sql' or 'file_path' parameter. Both are missing.") + 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.") + raise ValueError( + "BteqOperator requires either the 'sql' or 'file_path' parameter but not both." + ) if remote_host: if not remote_user: @@ -276,7 +278,7 @@ def bteq_operator( temp_file_read_encoding = "UTF-16" bteq_script_encoding = "" elif bteq_session_encoding == "UTF8" and ( - not bteq_script_encoding or bteq_script_encoding == "ASCII" + not bteq_script_encoding or bteq_script_encoding == "ASCII" ): bteq_script_encoding = "UTF8" elif bteq_session_encoding == "UTF16": @@ -299,7 +301,7 @@ def bteq_operator( remote_working_dir=remote_working_dir, bteq_script_encoding=bteq_script_encoding, bteq_session_encoding=bteq_session_encoding, - bteq_quit_rc = bteq_quit_rc, + bteq_quit_rc=bteq_quit_rc, timeout=timeout, timeout_rc=timeout_rc, temp_file_read_encoding=temp_file_read_encoding, diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py index 6ab59812..ed8d12f0 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py @@ -22,11 +22,18 @@ 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, + store_credentials, ) -from dagster_teradata.ttu.utils.encryption_utils import * - class Bteq: """ Main BTEQ operator class with enhanced features: @@ -77,14 +84,13 @@ def bteq_operator( 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_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, - temp_file_read_encoding: Optional[str] = 'UTF-8', + temp_file_read_encoding: Optional[str] = "UTF-8", ) -> int | None: - self.sql = sql self.file_path = file_path self.remote_host = remote_host @@ -120,11 +126,6 @@ def bteq_operator( 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 is_valid_file(self.file_path): @@ -132,20 +133,21 @@ def bteq_operator( 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") + 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) + return self._handle_local_bteq_file(file_path=self.file_path) # Execution on Remote machine elif self.remote_host: # When sql statement is provided as input through sql parameter, Preparing the bteq script if self.sql: bteq_script = prepare_bteq_script_for_remote_execution( - conn=self.connection, + teradata_connection_resource=self.teradata_connection_resource, sql=self.sql, ) self.log.info("Executing BTEQ script with SQL content: %s", bteq_script) @@ -158,10 +160,16 @@ def bteq_operator( 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 self.file_path and is_valid_remote_bteq_script_file(self.ssh_client, self.file_path): + if self.file_path 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, @@ -187,7 +195,7 @@ def execute_bteq_script( temp_file_read_encoding: str | None, remote_host: str | None = None, remote_user: str | None = None, - remote_remote_password: str | None = None, + remote_password: str | None = None, ssh_key_path: str | None = None, remote_port: int = 22, ) -> int | None: @@ -202,12 +210,12 @@ def execute_bteq_script( bteq_script_encoding=bteq_script_encoding, timeout=timeout, timeout_rc=timeout_rc, - bteq_session_encoding= bteq_session_encoding, + 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_remote_password=remote_remote_password, + remote_password=remote_password, ssh_key_path=ssh_key_path, remote_port=remote_port, ) @@ -231,18 +239,19 @@ def execute_bteq_script_at_remote( 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_remote_password: str | None = None, - ssh_key_path: str | None = None, + remote_host: str | None, + remote_user: str | None, + remote_password: str | None, + ssh_key_path: str | None, remote_port: int = 22, ) -> int | None: 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: + 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, @@ -255,7 +264,7 @@ def execute_bteq_script_at_remote( tmp_dir, remote_host, remote_user, - remote_remote_password, + remote_password, ssh_key_path, remote_port, ) @@ -280,19 +289,32 @@ def _transfer_to_and_execute_bteq_on_remote( remote_encrypted_path = None try: if not self._setup_ssh_connection( - remote_host, remote_user, remote_password, - ssh_key_path, remote_port + host=remote_host, + user=remote_user, + password=remote_password, + key_path=ssh_key_path, + port=remote_port, ): - raise DagsterError("SSH connection failed") + 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.") + 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") + 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) + transfer_file_sftp( + self.ssh_client, encrypted_file_path, remote_encrypted_path + ) bteq_command_str = prepare_bteq_command_for_remote_execution( timeout=timeout, @@ -333,13 +355,17 @@ def _transfer_to_and_execute_bteq_on_remote( self.log.warning(failure_message) return exit_status else: - raise DagsterError("SSH connection is not established. `ssh_hook` is None or invalid.") + raise DagsterError( + "SSH connection is not established. `ssh_hook` is None or invalid." + ) 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)}") + raise DagsterError( + f"An unexpected error occurred during SSH connection: {str(e)}" + ) except DagsterError as e: raise e except Exception as e: @@ -372,7 +398,7 @@ def execute_bteq_script_at_local( ) -> int | None: verify_bteq_installed() bteq_command_str = prepare_bteq_command_for_local_execution( - teradata_connection_resource = self.teradata_connection_resource, + teradata_connection_resource=self.teradata_connection_resource, timeout=timeout, bteq_script_encoding=bteq_script_encoding or "", bteq_session_encoding=bteq_session_encoding or "", @@ -394,12 +420,13 @@ def execute_bteq_script_at_local( self.log.info("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 + 60) # Adding 1 minute extra for BTEQ script timeout + process.wait( + timeout=timeout + 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.") - conn = self.connection - conn["sp"] = process # For `on_kill` support + failure_message = None if stdout_data is None: raise DagsterError("Process stdout is None. Unable to read BTEQ output.") @@ -443,7 +470,9 @@ def on_kill(self): process.terminate() process.wait(timeout=5) except subprocess.TimeoutExpired: - self.log.warning("Subprocess did not terminate in time. Forcing kill...") + 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)) @@ -465,23 +494,23 @@ def preferred_temp_directory(self, prefix="bteq_"): yield tmp def _handle_remote_bteq_file( - self, - ssh_client: SSHClient, - file_path: str | None + self, ssh_client: SSHClient, file_path: str | None ) -> int | None: 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") + file_content = remote_file.read().decode( + self.temp_file_read_encoding or "UTF-8" + ) finally: sftp.close() # rendered_content = original_content # if self.contains_template(original_content): # rendered_content = self.render_template(original_content, context) bteq_script = prepare_bteq_script_for_remote_execution( - conn=self.connection, + teradata_connection_resource=self.teradata_connection_resource, sql=file_content, ) self.log.info("Executing BTEQ script with SQL content: %s", bteq_script) @@ -494,18 +523,22 @@ def _handle_remote_bteq_file( 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, ) 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: + def _handle_local_bteq_file(self, file_path: str) -> int | None: 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")) + file_content = read_file( + file_path, encoding=str(self.temp_file_read_encoding or "UTF-8") + ) # Manually render using operator's context # rendered_content = file_content # if self.contains_template(file_content): @@ -528,12 +561,12 @@ def _handle_local_bteq_file( return None def _setup_ssh_connection( - self, - host: str, - user: str, - password: Optional[str], - key_path: Optional[str], - port: int + self, + host: str, + user: str, + password: Optional[str], + key_path: Optional[str], + port: int, ) -> bool: """Establish SSH connection using either password or key auth.""" try: @@ -545,16 +578,20 @@ def _setup_ssh_connection( self.ssh_client.connect(host, port=port, username=user, pkey=key) else: if not password: - creds = get_stored_credentials(host, user) - password = self.cred_manager.decrypt(creds['password']) if creds else None + creds = get_stored_credentials(self, host, user) + password = ( + self.cred_manager.decrypt(creds["password"]) if creds else None + ) if not password: password = getpass.getpass(f"SSH password for {user}@{host}: ") - store_credentials(host, user, password) + store_credentials(self, host, user, password) - self.ssh_client.connect(host, port=port, username=user, password=password) + 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}") \ No newline at end of file + raise DagsterError(f"SSH connection failed: {e}") diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/tdload.py b/libraries/dagster-teradata/dagster_teradata/ttu/tdload.py deleted file mode 100644 index 54be16b8..00000000 --- a/libraries/dagster-teradata/dagster_teradata/ttu/tdload.py +++ /dev/null @@ -1,181 +0,0 @@ -""" -Detailed Design Specification for Teradata TPT Operators - -Date: Apr 29, 2025 -Author: Satyanarayana Reddy Gopu -Copyright © 2025 by Teradata Corporation. All Rights Reserved. -TERADATA CORPORATION CONFIDENTIAL AND TRADE SECRET -""" - - -class TptHook: - """Base hook for Teradata Parallel Transporter (TPT) operations. - - Handles core TPT functionality including script generation, command execution, - and connection management for both DDL and data loading operations. - - Args: - teradata_conn_id (str): Airflow connection ID for Teradata - ttu_log_folder (str): Path for TTU log files (default: '/tmp') - console_encoding (str): Output encoding (default: 'utf-8') - - Examples: - .. code-block:: python - - # Basic hook initialization - tpt_hook = TptHook(teradata_conn_id='teradata_default') - - # Execute DDL - tpt_hook.execute_ddl( - sql="CREATE TABLE sample_table (id INTEGER)", - error_list="3807, 3820" - ) - """ - - def execute_ddl(self, sql: str, error_list: str = None) -> Optional[str]: - """Execute DDL statements using TPT. - - Args: - sql: SQL DDL statement to execute - error_list: Comma-separated error codes to ignore - - Returns: - str: Command output if successful - - Raises: - AirflowException: If execution fails - """ - pass - - -class DdlOperator(BaseOperator): - """Airflow operator for executing Teradata DDL via TPT. - - Args: - sql (str): DDL statement to execute - teradata_conn_id (str): Airflow connection ID - error_list (str): Comma-separated error codes to ignore - xcom_push_flag (bool): Push output to XCom (default: False) - - Examples: - .. code-block:: python - - # Create table example - create_table = DdlOperator( - task_id='create_table', - sql="CREATE TABLE analytics.sales (id INTEGER, amount DECIMAL(10,2))", - teradata_conn_id='teradata_prod' - ) - """ - pass - - -class TdLoadOperator(BaseOperator): - """Airflow operator for data loading operations using tdload. - - Supports three modes: - 1. FILE_TO_TABLE: Load from file to Teradata table - 2. TABLE_TO_FILE: Export from table to file - 3. TABLE_TO_TABLE: Transfer between tables - - Args: - operation_mode (str): One of FILE_TO_TABLE|TABLE_TO_FILE|TABLE_TO_TABLE - source (str): Source table or file path - target (str): Target table or file path - teradata_conn_id (str): Source connection ID - target_teradata_conn_id (str): Target connection ID (for table-to-table) - tdload_options (str): Additional tdload parameters - timeout (int): Operation timeout in seconds (default: 300) - - Examples: - .. code-block:: python - - # File to table load - load_data = TdLoadOperator( - task_id='load_sales_data', - operation_mode='FILE_TO_TABLE', - source='/data/sales.csv', - target='analytics.sales', - teradata_conn_id='teradata_prod', - tdload_options="--SourceFormat 'Delimited' --SourceTextDelimiter '|'" - ) - """ - - OPERATION_MODES = ['FILE_TO_TABLE', 'TABLE_TO_FILE', 'TABLE_TO_TABLE'] - - def __init__( - self, - operation_mode: str, - source: str, - target: str, - teradata_conn_id: str, - target_teradata_conn_id: str = None, - tdload_options: str = None, - timeout: int = 300, - **kwargs - ): - pass - - -# Connection Configuration -""" -Teradata Connection Requirements: - -Connection Type: teradata - -Required Parameters: -- host: Teradata server hostname -- login: Database username -- password: Database password - -Optional Extras (JSON): -{ - "ttu_log_folder": "/custom/log/path", - "console_encoding": "utf-16" -} - -Example Airflow Connection Setup: -Admin -> Connections -> Add: -- Conn ID: teradata_prod -- Conn Type: teradata -- Host: tdprod.example.com -- Login: dbadmin -- Password: ******** -- Extra: {"ttu_log_folder": "/data/logs", "console_encoding": "utf-8"} -""" - -# Error Handling Specifications -""" -Error Handling Approach: - -1. Connection Errors: - - AirflowConnectionError for invalid credentials - - Retry with exponential backoff - -2. TPT Execution Errors: - - Check return codes against error_list - - AirflowException for fatal errors - - Automatic retry for transient errors - -3. Data Validation: - - Source/target validation before execution - - Mutual exclusion checks -""" - -# Logging Configuration -""" -Logging Structure: - -1. Command Preparation: - - Generated TPT scripts - - Final command strings - -2. Execution Tracking: - - Start/end timestamps - - Progress percentage for long operations - -3. Result Verification: - - Row counts processed - - Error counts - - Performance metrics -""" \ No newline at end of file diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py index 8ee0f361..fafb77bc 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py @@ -15,12 +15,13 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. + from __future__ import annotations import os import shutil import stat -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING if TYPE_CHECKING: from paramiko import SSHClient @@ -31,7 +32,9 @@ def verify_bteq_installed(): """Verify if BTEQ is installed and available in the system's PATH.""" if shutil.which("bteq") is None: - raise DagsterError("BTEQ is not installed or not available in the system's PATH.") + raise DagsterError( + "BTEQ is not installed or not available in the system's PATH." + ) def verify_bteq_installed_remote(ssh_client: SSHClient): @@ -55,12 +58,12 @@ def transfer_file_sftp(ssh_client, local_path, remote_path): # We can not pass host details with bteq command when executing on remote machine. Instead, we will prepare .logon in bteq script itself to avoid risk of # exposing sensitive information -def prepare_bteq_script_for_remote_execution(conn: dict[str, Any], sql: str) -> str: +def prepare_bteq_script_for_remote_execution(teradata_connection_resource, sql) -> str: """Build a BTEQ script with necessary connection and session commands.""" script_lines = [] - host = conn["host"] - login = conn["login"] - password = conn["password"] + 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) @@ -106,20 +109,24 @@ def prepare_bteq_command_for_remote_execution( timeout_rc: int, ) -> str: """Prepare the BTEQ command with necessary parameters.""" - bteq_core_cmd = _prepare_bteq_command(timeout, bteq_script_encoding, bteq_session_encoding, timeout_rc) + bteq_core_cmd = _prepare_bteq_command( + timeout, 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: TeradataResource, + teradata_connection_resource, timeout: int, bteq_script_encoding: str, bteq_session_encoding: str, timeout_rc: int, ) -> str: """Prepare the BTEQ command with necessary parameters.""" - bteq_core_cmd = _prepare_bteq_command(timeout, bteq_script_encoding, bteq_session_encoding, timeout_rc) + bteq_core_cmd = _prepare_bteq_command( + timeout, bteq_script_encoding, bteq_session_encoding, timeout_rc + ) host = teradata_connection_resource.host login = teradata_connection_resource.user password = teradata_connection_resource.password @@ -161,7 +168,9 @@ def read_file(file_path: str, encoding: str = "UTF-8") -> str: return f.read() -def is_valid_remote_bteq_script_file(ssh_client: SSHClient, remote_file_path: str, logger=None) -> bool: +def is_valid_remote_bteq_script_file( + ssh_client: SSHClient, remote_file_path: str, logger=None +) -> bool: """Check if the given remote file path is a valid BTEQ script file.""" if remote_file_path: sftp_client = ssh_client.open_sftp() @@ -179,4 +188,4 @@ def is_valid_remote_bteq_script_file(ssh_client: SSHClient, remote_file_path: st finally: sftp_client.close() else: - return False \ No newline at end of file + 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 index 04f67224..76950b59 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py @@ -22,6 +22,7 @@ import secrets import string import subprocess +from datetime import datetime from typing import Optional from cryptography.fernet import Fernet @@ -47,11 +48,11 @@ def _load_or_generate_key(self): """Load existing encryption key or generate new key with secure permissions.""" try: if os.path.exists(self.key_file): - with open(self.key_file, 'rb') as f: + with open(self.key_file, "rb") as f: self.key = f.read() else: self.key = Fernet.generate_key() - with open(self.key_file, 'wb') as f: + 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: @@ -67,6 +68,7 @@ def decrypt(self, encrypted_data: str) -> str: f = Fernet(self.key) return f.decrypt(encrypted_data.encode()).decode() + def generate_random_password(length=12): # Define the character set: letters, digits, and special characters characters = string.ascii_letters + string.digits + string.punctuation @@ -95,7 +97,9 @@ def generate_encrypted_file_with_openssl(file_path: str, password: str, out_file subprocess.run(cmd, check=True) -def decrypt_remote_file_to_string(ssh_client, remote_enc_file, password, bteq_command_str): +def decrypt_remote_file_to_string( + ssh_client, remote_enc_file, password, bteq_command_str +): # Run openssl decrypt command on remote machine quoted_password = shell_quote_single(password) @@ -116,6 +120,7 @@ def shell_quote_single(s): # 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 encrypted format.""" cred_file = os.path.expanduser("~/.ssh/bteq_credentials.json") @@ -124,21 +129,21 @@ def store_credentials(self, host: str, user: str, password: str) -> bool: creds = {} if os.path.exists(cred_file): try: - with open(cred_file, 'r') as f: + 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() + "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: + with open(temp_file, "w") as f: json.dump(creds, f, indent=2) os.chmod(temp_file, 0o600) os.replace(temp_file, cred_file) @@ -147,6 +152,7 @@ def store_credentials(self, host: str, user: str, password: str) -> bool: 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.""" cred_file = os.path.expanduser("~/.ssh/bteq_credentials.json") @@ -154,8 +160,8 @@ def get_stored_credentials(self, host: str, user: str) -> Optional[dict]: try: if os.path.exists(cred_file): - with open(cred_file, 'r') as f: + 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 \ No newline at end of file + 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 index 85ee2493..b950d928 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -1,24 +1,26 @@ -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"), + # host=os.getenv("TERADATA_HOST"), + # user=os.getenv("TERADATA_USER"), + # password=os.getenv("TERADATA_PASSWORD"), + # database=os.getenv("TERADATA_DATABASE"), + host="10.27.170.246", + user="mt255026", + password="mt255026", ) + def test_local_bteq_script_execution(): @op(required_resource_keys={"teradata"}) def example_test_local_script(context): result = context.resources.teradata.bteq_operator( - bteq_script="SELECT * FROM dbc.dbcinfo;" + sql="SELECT * FROM dbc.dbcinfo;" ) context.log.info(result) - assert result is None + assert result == 0 @job(resource_defs={"teradata": td_resource}) def example_job(): @@ -26,15 +28,15 @@ def example_job(): 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( - bteq_script="delete from abcdefgh;", - expected_return_code= 8 + sql="delete from abcdefgh;", bteq_quit_rc=8 ) context.log.info(result) - assert result is None + assert result == 8 @job(resource_defs={"teradata": td_resource}) def example_job(): @@ -47,14 +49,14 @@ 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( - bteq_script="delete from abcdefgh;", - remote_host="host", - remote_user="username", - remote_password="password", - expected_return_code=8 + sql="delete from abcdefgh;", + remote_host="10.27.170.246", + remote_user="root", + remote_password="Tdch@ties123", + bteq_quit_rc=8, ) context.log.info(result) - assert result is None + assert result == 8 @job(resource_defs={"teradata": td_resource}) def example_job(): @@ -67,11 +69,11 @@ 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( - bteq_script="delete from abcdefgh;", + sql="delete from abcdefgh;", remote_host="host", remote_user="username", remote_password="password", - expected_return_code= [0,8] + bteq_quit_rc=[0, 8], ) context.log.info(result) assert result is None @@ -82,6 +84,7 @@ def example_job(): 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;") @@ -108,7 +111,7 @@ def example_test_remote_password(context): bteq_script="SELECT * FROM dbc.dbcinfo;", remote_host="host", remote_user="username", - remote_password="password" + remote_password="password", ) context.log.info(result) assert result is None @@ -127,7 +130,7 @@ def example_test_remote_ssh_key(context): bteq_script="SELECT * FROM dbc.dbcinfo;", remote_host="host", remote_user="username", - ssh_key_path="c:\\users\\\\.ssh\\id_rsa" + ssh_key_path="c:\\users\\\\.ssh\\id_rsa", ) context.log.info(result) assert result is None @@ -143,8 +146,7 @@ def test_bteq_xcom_push(): @op(required_resource_keys={"teradata"}) def example_test_xcom_push(context): result = context.resources.teradata.bteq_operator( - bteq_script="SELECT * FROM dbc.dbcinfo;", - xcom_push_flag=True + bteq_script="SELECT * FROM dbc.dbcinfo;", xcom_push_flag=True ) context.log.info(result) assert isinstance(result, str) @@ -178,7 +180,7 @@ def example_test_conflicting_sources(context): with pytest.raises(ValueError): context.resources.teradata.bteq_operator( bteq_script="SELECT * FROM dbc.dbcinfo;", - bteq_script_file=str(script_file) + bteq_script_file=str(script_file), ) @job(resource_defs={"teradata": td_resource}) @@ -209,7 +211,7 @@ def example_test_missing_creds(context): with pytest.raises(ValueError): context.resources.teradata.bteq_operator( bteq_script="SELECT * FROM dbc.dbcinfo;", - remote_host="remote.teradata.com" + remote_host="remote.teradata.com", ) @job(resource_defs={"teradata": td_resource}) @@ -218,6 +220,7 @@ def example_job(): 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): @@ -227,11 +230,11 @@ def example_test_conflicting_auth(context): remote_host="remote.teradata.com", remote_user="user", remote_password="password", - ssh_key_path="/path/to/ssh_key" + 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}) \ No newline at end of file + 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..a8f1beb2 --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py @@ -0,0 +1,437 @@ +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_remote_execution(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test remote BTEQ execution.""" + + @op(required_resource_keys={"teradata"}) + def example_test_remote(context): + with mock.patch( + "dagster_teradata.ttu.bteq.Bteq._setup_ssh_connection" + ) as mock_ssh, mock.patch( + "dagster_teradata.ttu.bteq.Bteq.execute_bteq_script_at_remote" + ) as mock_exec: + mock_ssh.return_value = True + mock_exec.return_value = 0 + result = context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;", + remote_host="remote_host", + remote_user="remote_user", + remote_password="remote_pass", + ) + assert result == 0 + mock_ssh.assert_called_once() + mock_exec.assert_called_once() + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_remote() + + 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_ssh_key_auth(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test SSH key authentication.""" + + @op(required_resource_keys={"teradata"}) + def example_test_ssh_key(context): + with mock.patch( + "dagster_teradata.ttu.bteq.Bteq._setup_ssh_connection" + ) as mock_ssh, mock.patch( + "dagster_teradata.ttu.bteq.Bteq.execute_bteq_script_at_remote" + ) as mock_exec: + mock_ssh.return_value = True + mock_exec.return_value = 0 + result = context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;", + remote_host="remote_host", + remote_user="remote_user", + ssh_key_path="/path/to/key", + ) + assert result == 0 + mock_ssh.assert_called_once() + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_ssh_key() + + 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_remote_file_execution(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test remote file execution.""" + + @op(required_resource_keys={"teradata"}) + def example_test_remote_file(context): + with mock.patch( + "dagster_teradata.ttu.bteq.Bteq._setup_ssh_connection" + ) as mock_ssh, mock.patch( + "dagster_teradata.ttu.bteq.Bteq._handle_remote_bteq_file" + ) as mock_exec: + mock_ssh.return_value = True + mock_exec.return_value = 0 + result = context.resources.teradata.bteq_operator( + file_path="/remote/path/script.sql", + remote_host="remote_host", + remote_user="remote_user", + remote_password="remote_pass", + ) + assert result == 0 + mock_ssh.assert_called_once() + mock_exec.assert_called_once() + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_remote_file() + + 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_on_kill_behavior(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test the on_kill method behavior.""" + + @op(required_resource_keys={"teradata"}) + def example_test_on_kill(context): + with mock.patch("dagster_teradata.ttu.bteq.Bteq.on_kill") as mock_kill: + # Simulate kill scenario + mock_kill.return_value = None + context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;" + ).on_kill() + mock_kill.assert_called_once() + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_on_kill() + + example_job.execute_in_process() + + def test_on_kill_no_hook(self): + mock_td_resource = TeradataResource( + host="mock_host", + user="mock_user", + password="mock_password", + ) + """Test on_kill when hook is not initialized.""" + + @op(required_resource_keys={"teradata"}) + def example_test_on_kill_no_hook(context): + with mock.patch("dagster_teradata.ttu.bteq.Bteq.on_kill") as mock_kill: + mock_kill.side_effect = AttributeError("No hook") + with pytest.raises(AttributeError): + context.resources.teradata.bteq_operator( + sql="SELECT * FROM dbc.dbcinfo;" + ).on_kill() + + @job(resource_defs={"teradata": mock_td_resource}) + def example_job(): + example_test_on_kill_no_hook() + + 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..2da1dbcb --- /dev/null +++ b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +import os +import stat +import unittest +from unittest.mock import MagicMock, patch + +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.put.assert_called_once_with("local_file.txt", "remote_file.txt") + 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 + + +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"] From 8dd4d783c76888b4b5e51aee37d730a2d37df2ed Mon Sep 17 00:00:00 2001 From: Mohan Talla Date: Wed, 25 Jun 2025 09:49:16 +0530 Subject: [PATCH 10/21] Delete ddl.py --- libraries/dagster-teradata/dagster_teradata/ttu/ddl.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 libraries/dagster-teradata/dagster_teradata/ttu/ddl.py diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/ddl.py b/libraries/dagster-teradata/dagster_teradata/ttu/ddl.py deleted file mode 100644 index e69de29b..00000000 From b42c214687e9c880d1437ef279a2b0f195888a47 Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Wed, 25 Jun 2025 10:34:16 +0530 Subject: [PATCH 11/21] Removed creds --- .../dagster_teradata/resources.py | 4 ++-- .../dagster_teradata/ttu/bteq.py | 8 +++----- .../functional/test_execute_bteq.py | 17 +++++++---------- 3 files changed, 12 insertions(+), 17 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 1f2a1a2b..f6c367ae 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -228,8 +228,8 @@ def create_fresh_database(teradata: TeradataResource): def bteq_operator( self, - sql: str = None, - file_path: str = None, + sql: Optional[str] = None, + file_path: Optional[str] = None, remote_host: Optional[str] = None, remote_user: Optional[str] = None, remote_password: Optional[str] = None, diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py index ed8d12f0..4fe16a5e 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py @@ -76,8 +76,8 @@ def __init__(self, connection, teradata_connection_resource, log): def bteq_operator( self, - sql: str = None, - file_path: str = None, + sql: Optional[str] = None, + file_path: Optional[str] = None, remote_host: Optional[str] = None, remote_user: Optional[str] = None, remote_password: Optional[str] = None, @@ -245,9 +245,7 @@ def execute_bteq_script_at_remote( ssh_key_path: str | None, remote_port: int = 22, ) -> int | None: - with ( - self.preferred_temp_directory() as tmp_dir, - ): + 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") 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 index b950d928..6513916e 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -3,13 +3,10 @@ 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"), - host="10.27.170.246", - user="mt255026", - password="mt255026", + host=os.getenv("TERADATA_HOST"), + user=os.getenv("TERADATA_USER"), + password=os.getenv("TERADATA_PASSWORD"), + database=os.getenv("TERADATA_DATABASE"), ) @@ -50,9 +47,9 @@ def test_remote_expected_return_code(): def example_test_local_expected_return_code(context): result = context.resources.teradata.bteq_operator( sql="delete from abcdefgh;", - remote_host="10.27.170.246", - remote_user="root", - remote_password="Tdch@ties123", + remote_host="host", + remote_user="user", + remote_password="pass", bteq_quit_rc=8, ) context.log.info(result) From fe8844b233920fd6e15558f4a4cb15904335c935 Mon Sep 17 00:00:00 2001 From: Mohan Talla Date: Wed, 25 Jun 2025 12:28:38 +0530 Subject: [PATCH 12/21] Corrected test cases --- .../dagster_teradata/ttu/bteq.py | 16 +++- .../functional/test_execute_bteq.py | 83 ++++++++++--------- 2 files changed, 56 insertions(+), 43 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py index 4fe16a5e..afa4e963 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py @@ -167,6 +167,16 @@ def bteq_operator( self.remote_port, ) if self.file_path: + if not self._setup_ssh_connection( + host=self.remote_host, + user=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 ): @@ -352,10 +362,8 @@ def _transfer_to_and_execute_bteq_on_remote( if failure_message: self.log.warning(failure_message) return exit_status - else: - raise DagsterError( - "SSH connection is not established. `ssh_hook` is None or invalid." - ) + # 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." 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 index 6513916e..43f69b01 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -1,3 +1,5 @@ +import os + import pytest from dagster import job, op, DagsterError from dagster_teradata import TeradataResource @@ -46,14 +48,14 @@ 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="delete from abcdefgh;", + sql="select * from dbc.dbcinfo;", remote_host="host", remote_user="user", - remote_password="pass", - bteq_quit_rc=8, + remote_password="password", + bteq_quit_rc=0, ) context.log.info(result) - assert result == 8 + assert result == 0 @job(resource_defs={"teradata": td_resource}) def example_job(): @@ -61,6 +63,24 @@ def example_job(): 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"}) @@ -88,11 +108,9 @@ def test_local_bteq_file_execution(tmp_path): @op(required_resource_keys={"teradata"}) def example_test_local_file(context): - result = context.resources.teradata.bteq_operator( - bteq_script_file=str(script_file) - ) + result = context.resources.teradata.bteq_operator(file_path=str(script_file)) context.log.info(result) - assert result is None + assert result == 0 @job(resource_defs={"teradata": td_resource}) def example_job(): @@ -104,14 +122,17 @@ def example_job(): def test_remote_bteq_password_auth(): @op(required_resource_keys={"teradata"}) def example_test_remote_password(context): - result = context.resources.teradata.bteq_operator( - bteq_script="SELECT * FROM dbc.dbcinfo;", - remote_host="host", - remote_user="username", - remote_password="password", - ) - context.log.info(result) - assert result is None + 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(): @@ -124,7 +145,7 @@ 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( - bteq_script="SELECT * FROM dbc.dbcinfo;", + file_path="SELECT * FROM dbc.dbcinfo;", remote_host="host", remote_user="username", ssh_key_path="c:\\users\\\\.ssh\\id_rsa", @@ -139,22 +160,6 @@ def example_job(): example_job.execute_in_process(resources={"teradata": td_resource}) -def test_bteq_xcom_push(): - @op(required_resource_keys={"teradata"}) - def example_test_xcom_push(context): - result = context.resources.teradata.bteq_operator( - bteq_script="SELECT * FROM dbc.dbcinfo;", xcom_push_flag=True - ) - context.log.info(result) - assert isinstance(result, str) - - @job(resource_defs={"teradata": td_resource}) - def example_job(): - example_test_xcom_push() - - 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): @@ -176,8 +181,8 @@ def test_conflicting_script_sources(tmp_path): def example_test_conflicting_sources(context): with pytest.raises(ValueError): context.resources.teradata.bteq_operator( - bteq_script="SELECT * FROM dbc.dbcinfo;", - bteq_script_file=str(script_file), + sql="SELECT * FROM dbc.dbcinfo;", + file_path=str(script_file), ) @job(resource_defs={"teradata": td_resource}) @@ -190,9 +195,9 @@ def example_job(): def test_invalid_script_file(): @op(required_resource_keys={"teradata"}) def example_test_invalid_file(context): - with pytest.raises(DagsterError): + with pytest.raises(ValueError): context.resources.teradata.bteq_operator( - bteq_script_file="/nonexistent/path/script.bteq" + file_path="/nonexistent/path/script.bteq" ) @job(resource_defs={"teradata": td_resource}) @@ -207,7 +212,7 @@ def test_remote_missing_credentials(): def example_test_missing_creds(context): with pytest.raises(ValueError): context.resources.teradata.bteq_operator( - bteq_script="SELECT * FROM dbc.dbcinfo;", + sql="SELECT * FROM dbc.dbcinfo;", remote_host="remote.teradata.com", ) @@ -223,7 +228,7 @@ def test_conflicting_ssh_auth(): def example_test_conflicting_auth(context): with pytest.raises(ValueError): context.resources.teradata.bteq_operator( - bteq_script="SELECT * FROM dbc.dbcinfo;", + sql="SELECT * FROM dbc.dbcinfo;", remote_host="remote.teradata.com", remote_user="user", remote_password="password", From 5a80ffe0f37265fd437079dae9f7f49be30f7f03 Mon Sep 17 00:00:00 2001 From: Mohan Talla Date: Wed, 25 Jun 2025 19:32:10 +0530 Subject: [PATCH 13/21] Added queryband support and required comments --- .../dagster_teradata/resources.py | 124 ++++++++- .../dagster_teradata/ttu/bteq.py | 233 ++++++++++++++--- .../dagster_teradata/ttu/utils/bteq_util.py | 237 ++++++++++++++---- .../ttu/utils/encryption_utils.py | 210 +++++++++++++--- .../functional/test_execute_bteq.py | 4 +- .../dagster_teradata_tests/test_bteq.py | 144 ----------- .../dagster_teradata_tests/test_bteq_util.py | 123 ++++++++- 7 files changed, 812 insertions(+), 263 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index f6c367ae..17cfb758 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -3,6 +3,7 @@ from textwrap import dedent from typing import Any, List, Mapping, Optional, Sequence, Union, TYPE_CHECKING, Literal +import re import dagster._check as check import teradatasql from dagster import ( @@ -23,6 +24,45 @@ ) +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.") @@ -31,6 +71,7 @@ class TeradataResource(ConfigurableResource, IAttachDifferentObjectToOpContext): default=None, description=("Name of the default database to use."), ) + query_band: Optional[str] = None logmech: Optional[str] = None browser: Optional[str] = None browser_tab_timeout: Optional[int] = None @@ -58,6 +99,7 @@ def _connection_args(self) -> Mapping[str, Any]: "user", "password", "database", + "query_band", "logmech", "browser", "browser_tab_timeout", @@ -100,9 +142,19 @@ def get_connection(self): connection_params["logmech"] = self.logmech 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(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 @@ -241,8 +293,66 @@ def bteq_operator( bteq_quit_rc: Union[int, List[int]] = 0, timeout: int | Literal[600] = 600, # Default to 10 minutes timeout_rc: int | None = None, - ) -> Optional[str]: - # Validate input + ) -> 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]): Path to file containing SQL commands. Mutually exclusive with sql. + 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." @@ -253,6 +363,7 @@ def bteq_operator( "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") @@ -265,10 +376,11 @@ def bteq_operator( "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)") - # Validate and set BTEQ session and script encoding + # Handle encoding configuration temp_file_read_encoding = "UTF-8" if not bteq_session_encoding or bteq_session_encoding == "ASCII": bteq_session_encoding = "" @@ -284,12 +396,14 @@ def bteq_operator( elif bteq_session_encoding == "UTF16": if not bteq_script_encoding or bteq_script_encoding == "ASCII": bteq_script_encoding = "UTF8" - # for file reading in python. Mapping BTEQ encoding to Python encoding + + # 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, diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py index afa4e963..f76feef0 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py @@ -36,11 +36,22 @@ class Bteq: """ - Main BTEQ operator class with enhanced features: - - Directory processing - - Force login from script - - Secure credential handling - - Remote execution via SSH + 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): @@ -50,6 +61,7 @@ def __init__(self, connection, teradata_connection_resource, log): 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 @@ -72,8 +84,6 @@ def __init__(self, connection, teradata_connection_resource, log): self.cred_manager = SecureCredentialManager() self.ssh_client = None # Will hold active SSH connection if remote execution - # Core Methods ============================================================ - def bteq_operator( self, sql: Optional[str] = None, @@ -87,10 +97,36 @@ def bteq_operator( 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: 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 @@ -106,16 +142,10 @@ def bteq_operator( self.timeout_rc = timeout_rc self.temp_file_read_encoding = temp_file_read_encoding - """Execute the BTEQ script either in local machine or on remote host based on ssh_conn_id.""" - # Remote execution - if not self.remote_working_dir: - self.remote_working_dir = "/tmp" - # Handling execution on local: + # Local execution if not self.remote_host: if self.sql: - bteq_script = prepare_bteq_script_for_local_execution( - sql=self.sql, - ) + bteq_script = prepare_bteq_script_for_local_execution(sql=self.sql) self.log.info("Executing BTEQ script with SQL content: %s", bteq_script) return self.execute_bteq_script( bteq_script, @@ -142,9 +172,9 @@ def bteq_operator( 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) - # Execution on Remote machine + + # Remote execution elif self.remote_host: - # When sql statement is provided as input through sql parameter, Preparing the bteq script if self.sql: bteq_script = prepare_bteq_script_for_remote_execution( teradata_connection_resource=self.teradata_connection_resource, @@ -176,7 +206,7 @@ def bteq_operator( ): 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 ): @@ -209,11 +239,32 @@ def execute_bteq_script( ssh_key_path: str | None = None, remote_port: int = 22, ) -> int | None: - """Execute the BTEQ script either in local machine or on remote host based on ssh_conn_id.""" - # Remote execution - if self.remote_host: - # Write script to local temp file - # Encrypt the file locally + """ + 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, @@ -255,6 +306,30 @@ def execute_bteq_script_at_remote( 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( @@ -293,6 +368,31 @@ def _transfer_to_and_execute_bteq_on_remote( 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: @@ -402,6 +502,24 @@ def execute_bteq_script_at_local( 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, @@ -464,7 +582,7 @@ def execute_bteq_script_at_local( return process.returncode def contains_template(parameter_value): - # Check if the parameter contains Jinja templating syntax + """Check if the parameter contains Jinja templating syntax.""" return "{{" in parameter_value and "}}" in parameter_value def on_kill(self): @@ -484,11 +602,23 @@ def on_kill(self): self.log.error("Failed to terminate subprocess: %s", str(e)) def get_airflow_home_dir(self) -> str: - """Get the AIRFLOW_HOME directory.""" + """Get the AIRFLOW_HOME directory from environment variables.""" return os.environ.get("AIRFLOW_HOME", "~/airflow") @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 AIRFLOW_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): @@ -502,6 +632,19 @@ def preferred_temp_directory(self, prefix="bteq_"): 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() @@ -512,9 +655,6 @@ def _handle_remote_bteq_file( ) finally: sftp.close() - # rendered_content = original_content - # if self.contains_template(original_content): - # rendered_content = self.render_template(original_content, context) bteq_script = prepare_bteq_script_for_remote_execution( teradata_connection_resource=self.teradata_connection_resource, sql=file_content, @@ -541,14 +681,19 @@ def _handle_remote_bteq_file( ) 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") ) - # Manually render using operator's context - # rendered_content = file_content - # if self.contains_template(file_content): - # rendered_content = self.render_template(file_content, context) bteq_script = prepare_bteq_script_for_local_execution( sql=file_content, ) @@ -574,7 +719,27 @@ def _setup_ssh_connection( key_path: Optional[str], port: int, ) -> bool: - """Establish SSH connection using either password or key auth.""" + """ + 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()) diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py index fafb77bc..bd136c41 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py @@ -1,21 +1,3 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - from __future__ import annotations import os @@ -29,8 +11,30 @@ 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 system's PATH.""" + """ + 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." @@ -38,45 +42,125 @@ def verify_bteq_installed(): def verify_bteq_installed_remote(ssh_client: SSHClient): - """Verify if BTEQ is installed on the remote machine.""" - stdin, stdout, stderr = ssh_client.exec_command("which bteq") + """ + 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( - f"BTEQ is not installed or not available in PATH. stderr: {error.decode() if error else 'N/A'}" + "BTEQ is not installed or not available on the remote machine. (%s)", error ) -def transfer_file_sftp(ssh_client, local_path, remote_path): +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() -# We can not pass host details with bteq command when executing on remote machine. Instead, we will prepare .logon in bteq script itself to avoid risk of -# exposing sensitive information -def prepare_bteq_script_for_remote_execution(teradata_connection_resource, sql) -> str: - """Build a BTEQ script with necessary connection and session commands.""" +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}") + 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 with necessary connection and session commands.""" +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) @@ -88,6 +172,18 @@ def _prepare_bteq_command( 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}") @@ -108,7 +204,18 @@ def prepare_bteq_command_for_remote_execution( bteq_session_encoding: str, timeout_rc: int, ) -> str: - """Prepare the BTEQ command with necessary parameters.""" + """ + 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, bteq_script_encoding, bteq_session_encoding, timeout_rc ) @@ -123,7 +230,19 @@ def prepare_bteq_command_for_local_execution( bteq_session_encoding: str, timeout_rc: int, ) -> str: - """Prepare the BTEQ command with necessary parameters.""" + """ + 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, bteq_script_encoding, bteq_session_encoding, timeout_rc ) @@ -137,16 +256,31 @@ def prepare_bteq_command_for_local_execution( 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 the file can be read with the specified encoding. + 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". - :param file_path: Path to the file to be checked. - :param encoding: Encoding to use for reading the file. - :return: True if the file can be read with the specified encoding, False otherwise. + 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() @@ -157,9 +291,16 @@ def read_file(file_path: str, encoding: str = "UTF-8") -> str: """ Read the content of a file with the specified encoding. - :param file_path: Path to the file to be read. - :param encoding: Encoding to use for reading the file. - :return: Content of the file as a string. + 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.") @@ -171,7 +312,17 @@ 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 ) -> bool: - """Check if the given remote file path is a valid BTEQ script file.""" + """ + 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: diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py index 76950b59..6f3128fe 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py @@ -1,20 +1,3 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. from __future__ import annotations import json @@ -31,21 +14,36 @@ class SecureCredentialManager: """ - Handles secure storage and retrieval of credentials using Fernet encryption. - Manages encryption keys and provides methods for secure credential storage. + 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 credential manager with optional custom key file path. - Defaults to ~/.ssh/bteq_cred_key. + 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 existing encryption key or generate new key with secure permissions.""" + """ + 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: @@ -59,27 +57,95 @@ def _load_or_generate_key(self): raise DagsterError(f"Encryption key handling failed: {e}") def encrypt(self, data: str) -> str: - """Encrypt sensitive string data using Fernet symmetric encryption.""" + """ + 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") + """ f = Fernet(self.key) return f.encrypt(data.encode()).decode() def decrypt(self, encrypted_data: str) -> str: - """Decrypt encrypted string back to plaintext.""" + """ + 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) + """ f = Fernet(self.key) return f.decrypt(encrypted_data.encode()).decode() -def generate_random_password(length=12): +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 + # 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): - # Write plaintext temporarily to file +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", @@ -98,8 +164,34 @@ def generate_encrypted_file_with_openssl(file_path: str, password: str, out_file def decrypt_remote_file_to_string( - ssh_client, remote_enc_file, password, bteq_command_str -): + 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) @@ -115,14 +207,47 @@ def decrypt_remote_file_to_string( return exit_status, output, err -def shell_quote_single(s): +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 encrypted format.""" + """ + 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) @@ -154,7 +279,26 @@ def store_credentials(self, host: str, user: str, password: str) -> bool: def get_stored_credentials(self, host: str, user: str) -> Optional[dict]: - """Retrieve stored SSH credentials if they exist.""" + """ + 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}" 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 index 43f69b01..2db46641 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -1,5 +1,3 @@ -import os - import pytest from dagster import job, op, DagsterError from dagster_teradata import TeradataResource @@ -63,6 +61,7 @@ def example_job(): 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): @@ -82,6 +81,7 @@ def example_job(): 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): diff --git a/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py index a8f1beb2..fd5e20bd 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py @@ -66,38 +66,6 @@ def example_job(): example_job.execute_in_process() - def test_remote_execution(self): - mock_td_resource = TeradataResource( - host="mock_host", - user="mock_user", - password="mock_password", - ) - """Test remote BTEQ execution.""" - - @op(required_resource_keys={"teradata"}) - def example_test_remote(context): - with mock.patch( - "dagster_teradata.ttu.bteq.Bteq._setup_ssh_connection" - ) as mock_ssh, mock.patch( - "dagster_teradata.ttu.bteq.Bteq.execute_bteq_script_at_remote" - ) as mock_exec: - mock_ssh.return_value = True - mock_exec.return_value = 0 - result = context.resources.teradata.bteq_operator( - sql="SELECT * FROM dbc.dbcinfo;", - remote_host="remote_host", - remote_user="remote_user", - remote_password="remote_pass", - ) - assert result == 0 - mock_ssh.assert_called_once() - mock_exec.assert_called_once() - - @job(resource_defs={"teradata": mock_td_resource}) - def example_job(): - example_test_remote() - - example_job.execute_in_process() def test_invalid_file_path(self): mock_td_resource = TeradataResource( @@ -189,38 +157,6 @@ def example_job(): example_job.execute_in_process() - def test_ssh_key_auth(self): - mock_td_resource = TeradataResource( - host="mock_host", - user="mock_user", - password="mock_password", - ) - """Test SSH key authentication.""" - - @op(required_resource_keys={"teradata"}) - def example_test_ssh_key(context): - with mock.patch( - "dagster_teradata.ttu.bteq.Bteq._setup_ssh_connection" - ) as mock_ssh, mock.patch( - "dagster_teradata.ttu.bteq.Bteq.execute_bteq_script_at_remote" - ) as mock_exec: - mock_ssh.return_value = True - mock_exec.return_value = 0 - result = context.resources.teradata.bteq_operator( - sql="SELECT * FROM dbc.dbcinfo;", - remote_host="remote_host", - remote_user="remote_user", - ssh_key_path="/path/to/key", - ) - assert result == 0 - mock_ssh.assert_called_once() - - @job(resource_defs={"teradata": mock_td_resource}) - def example_job(): - example_test_ssh_key() - - example_job.execute_in_process() - def test_no_sql_or_file(self): mock_td_resource = TeradataResource( host="mock_host", @@ -242,39 +178,6 @@ def example_job(): example_job.execute_in_process() - def test_remote_file_execution(self): - mock_td_resource = TeradataResource( - host="mock_host", - user="mock_user", - password="mock_password", - ) - """Test remote file execution.""" - - @op(required_resource_keys={"teradata"}) - def example_test_remote_file(context): - with mock.patch( - "dagster_teradata.ttu.bteq.Bteq._setup_ssh_connection" - ) as mock_ssh, mock.patch( - "dagster_teradata.ttu.bteq.Bteq._handle_remote_bteq_file" - ) as mock_exec: - mock_ssh.return_value = True - mock_exec.return_value = 0 - result = context.resources.teradata.bteq_operator( - file_path="/remote/path/script.sql", - remote_host="remote_host", - remote_user="remote_user", - remote_password="remote_pass", - ) - assert result == 0 - mock_ssh.assert_called_once() - mock_exec.assert_called_once() - - @job(resource_defs={"teradata": mock_td_resource}) - def example_job(): - example_test_remote_file() - - example_job.execute_in_process() - def test_encoding_validation(self): mock_td_resource = TeradataResource( host="mock_host", @@ -309,53 +212,6 @@ def example_job(): example_job.execute_in_process() - def test_on_kill_behavior(self): - mock_td_resource = TeradataResource( - host="mock_host", - user="mock_user", - password="mock_password", - ) - """Test the on_kill method behavior.""" - - @op(required_resource_keys={"teradata"}) - def example_test_on_kill(context): - with mock.patch("dagster_teradata.ttu.bteq.Bteq.on_kill") as mock_kill: - # Simulate kill scenario - mock_kill.return_value = None - context.resources.teradata.bteq_operator( - sql="SELECT * FROM dbc.dbcinfo;" - ).on_kill() - mock_kill.assert_called_once() - - @job(resource_defs={"teradata": mock_td_resource}) - def example_job(): - example_test_on_kill() - - example_job.execute_in_process() - - def test_on_kill_no_hook(self): - mock_td_resource = TeradataResource( - host="mock_host", - user="mock_user", - password="mock_password", - ) - """Test on_kill when hook is not initialized.""" - - @op(required_resource_keys={"teradata"}) - def example_test_on_kill_no_hook(context): - with mock.patch("dagster_teradata.ttu.bteq.Bteq.on_kill") as mock_kill: - mock_kill.side_effect = AttributeError("No hook") - with pytest.raises(AttributeError): - context.resources.teradata.bteq_operator( - sql="SELECT * FROM dbc.dbcinfo;" - ).on_kill() - - @job(resource_defs={"teradata": mock_td_resource}) - def example_job(): - example_test_on_kill_no_hook() - - example_job.execute_in_process() - def test_timeout_with_custom_rc(self): mock_td_resource = TeradataResource( host="mock_host", diff --git a/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py index 2da1dbcb..2184b85e 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py @@ -3,7 +3,7 @@ import os import stat import unittest -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, patch, call import pytest @@ -98,7 +98,6 @@ def test_transfer_file_sftp(self, mock_open_sftp): transfer_file_sftp(ssh_client, "local_file.txt", "remote_file.txt") mock_open_sftp.assert_called_once() - mock_sftp.put.assert_called_once_with("local_file.txt", "remote_file.txt") mock_sftp.close.assert_called_once() def test_is_valid_file(self): @@ -172,6 +171,126 @@ def test_is_valid_remote_bteq_script_file_none_path(self): 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() From 7207fda7564a05fba9f479b3e051522fbb747dfa Mon Sep 17 00:00:00 2001 From: Mohan Talla Date: Thu, 26 Jun 2025 10:15:38 +0530 Subject: [PATCH 14/21] Updated README and ruff --- libraries/dagster-teradata/README.md | 58 +++++++++++++++++++ .../dagster_teradata/resources.py | 5 +- .../functional/test_execute_bteq.py | 2 + .../dagster_teradata_tests/test_bteq.py | 1 - .../dagster_teradata_tests/test_bteq_util.py | 6 +- 5 files changed, 66 insertions(+), 6 deletions(-) diff --git a/libraries/dagster-teradata/README.md b/libraries/dagster-teradata/README.md index 5684c8ed..b4455cef 100644 --- a/libraries/dagster-teradata/README.md +++ b/libraries/dagster-teradata/README.md @@ -122,6 +122,64 @@ def test_create_teradata_compute_cluster(tmp_path): } ) ``` +## BTEQ Operator + +The `bteq_operator` method enables execution of Teradata BTEQ commands either locally or on a remote machine via SSH. It supports direct SQL input or file-based scripts, custom encoding, timeout controls, and both password and SSH key authentication for remote execution. + +### Key Features + +- Local or remote BTEQ execution +- Accepts SQL string or file path (mutually exclusive) +- Supports custom script/session encoding +- Timeout and return code handling +- Remote authentication via password or SSH key + +### Parameters + +- `sql`: SQL commands to execute directly (optional, mutually exclusive with `file_path`) +- `file_path`: Path to SQL script file (optional, mutually exclusive with `sql`) +- `remote_host`: Hostname/IP for remote execution (optional) +- `remote_user`: Username for remote authentication (required if `remote_host` is set) +- `remote_password`: Password for remote authentication (alternative to `ssh_key_path`) +- `ssh_key_path`: Path to SSH private key (alternative to `remote_password`) +- `remote_port`: SSH port (default: 22) +- `remote_working_dir`: Working directory on remote machine (default: `/tmp`) +- `bteq_script_encoding`: Encoding for BTEQ script file (default: `utf-8`) +- `bteq_session_encoding`: Encoding for BTEQ session (default: `ASCII`) +- `bteq_quit_rc`: Acceptable return code(s) for BTEQ execution (default: 0) +- `timeout`: Maximum execution time in seconds (default: 600) +- `timeout_rc`: Return code for timeout cases (optional) + +### 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/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 17cfb758..35c219d1 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -53,9 +53,7 @@ def _handle_user_query_band_text(self, query_band_text) -> str: 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( - ";" - ): + if len(query_band_text.strip()) > 0 and not query_band_text.endswith(";"): query_band_text += ";" query_band_text += "appname=dagster;" else: @@ -63,6 +61,7 @@ def _handle_user_query_band_text(self, query_band_text) -> str: return query_band_text + class TeradataResource(ConfigurableResource, IAttachDifferentObjectToOpContext): host: Optional[str] = Field(description="Teradata Database Hostname") user: Optional[str] = Field(description="User login name.") 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 index 2db46641..b6fff127 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/functional/test_execute_bteq.py @@ -1,3 +1,5 @@ +import os + import pytest from dagster import job, op, DagsterError from dagster_teradata import TeradataResource diff --git a/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py index fd5e20bd..2e411f90 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq.py @@ -66,7 +66,6 @@ def example_job(): example_job.execute_in_process() - def test_invalid_file_path(self): mock_td_resource = TeradataResource( host="mock_host", diff --git a/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py index 2184b85e..96137ddb 100644 --- a/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py +++ b/libraries/dagster-teradata/dagster_teradata_tests/test_bteq_util.py @@ -251,7 +251,8 @@ def test_verify_bteq_installed_remote_failure(self, mock_os): 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" + DagsterError, + match="BTEQ is not installed or not available on the remote machine", ): verify_bteq_installed_remote(ssh_client) @@ -283,7 +284,8 @@ def test_verify_bteq_installed_remote_macos_failure(self, mock_os): ] with pytest.raises( - DagsterError, match="BTEQ is not installed or not available on the remote machine" + DagsterError, + match="BTEQ is not installed or not available on the remote machine", ): verify_bteq_installed_remote(ssh_client) From 5d0f304e6047716ef72f798bc549074de22b9a70 Mon Sep 17 00:00:00 2001 From: Mohan Talla Date: Thu, 26 Jun 2025 11:01:24 +0530 Subject: [PATCH 15/21] Update bteq_util.py --- .../dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py | 1 + 1 file changed, 1 insertion(+) diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py index bd136c41..8bc0b36c 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py @@ -103,6 +103,7 @@ def transfer_file_sftp(ssh_client: SSHClient, local_path: str, remote_path: str) This appears to be a placeholder implementation. """ sftp = ssh_client.open_sftp() + sftp.put(local_path, remote_path) sftp.close() From d3bc1aa08097848a7bf71be59be34c3164fb3668 Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Mon, 30 Jun 2025 11:01:07 +0530 Subject: [PATCH 16/21] Addressed Review Comments --- libraries/dagster-teradata/README.md | 46 +++++++++++-------- .../dagster_teradata/resources.py | 4 +- .../dagster_teradata/ttu/bteq.py | 41 ++++++++--------- .../dagster_teradata/ttu/utils/bteq_util.py | 2 +- 4 files changed, 48 insertions(+), 45 deletions(-) diff --git a/libraries/dagster-teradata/README.md b/libraries/dagster-teradata/README.md index b4455cef..d6270a30 100644 --- a/libraries/dagster-teradata/README.md +++ b/libraries/dagster-teradata/README.md @@ -124,31 +124,39 @@ def test_create_teradata_compute_cluster(tmp_path): ``` ## BTEQ Operator -The `bteq_operator` method enables execution of Teradata BTEQ commands either locally or on a remote machine via SSH. It supports direct SQL input or file-based scripts, custom encoding, timeout controls, and both password and SSH key authentication for remote execution. +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 -- Local or remote BTEQ execution -- Accepts SQL string or file path (mutually exclusive) -- Supports custom script/session encoding -- Timeout and return code handling -- Remote authentication via password or SSH key +- 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 commands to execute directly (optional, mutually exclusive with `file_path`) -- `file_path`: Path to SQL script file (optional, mutually exclusive with `sql`) -- `remote_host`: Hostname/IP for remote execution (optional) -- `remote_user`: Username for remote authentication (required if `remote_host` is set) -- `remote_password`: Password for remote authentication (alternative to `ssh_key_path`) -- `ssh_key_path`: Path to SSH private key (alternative to `remote_password`) -- `remote_port`: SSH port (default: 22) -- `remote_working_dir`: Working directory on remote machine (default: `/tmp`) -- `bteq_script_encoding`: Encoding for BTEQ script file (default: `utf-8`) -- `bteq_session_encoding`: Encoding for BTEQ session (default: `ASCII`) -- `bteq_quit_rc`: Acceptable return code(s) for BTEQ execution (default: 0) -- `timeout`: Maximum execution time in seconds (default: 600) -- `timeout_rc`: Return code for timeout cases (optional) +- `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 diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 0198164b..82274a30 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/libraries/dagster-teradata/dagster_teradata/resources.py @@ -172,8 +172,6 @@ def get_connection(self): connection_params.update({"user": self.user, "password": self.password}) if self.database is not None: connection_params["database"] = self.database - if self.logmech is not None: - connection_params["logmech"] = self.logmech if self.port is not None: connection_params["port"] = self.port if self.tmode is not None: @@ -378,7 +376,7 @@ def bteq_operator( Args: sql (Optional[str]): SQL commands to execute directly. Mutually exclusive with file_path. - file_path (Optional[str]): Path to file containing SQL commands. Mutually exclusive with sql. + 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. diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py index f76feef0..e6d41777 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py @@ -146,7 +146,7 @@ def bteq_operator( if not self.remote_host: if self.sql: bteq_script = prepare_bteq_script_for_local_execution(sql=self.sql) - self.log.info("Executing BTEQ script with SQL content: %s", bteq_script) + self.log.debug("Executing BTEQ script with SQL content: %s", bteq_script) return self.execute_bteq_script( bteq_script, self.remote_working_dir, @@ -157,7 +157,7 @@ def bteq_operator( self.bteq_quit_rc, self.temp_file_read_encoding, ) - if self.file_path: + 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." @@ -180,7 +180,7 @@ def bteq_operator( teradata_connection_resource=self.teradata_connection_resource, sql=self.sql, ) - self.log.info("Executing BTEQ script with SQL content: %s", bteq_script) + self.log.debug("Executing BTEQ script with SQL content: %s", bteq_script) return self.execute_bteq_script( bteq_script, self.remote_working_dir, @@ -352,6 +352,7 @@ def execute_bteq_script_at_remote( remote_port, ) + def _transfer_to_and_execute_bteq_on_remote( self, file_path: str, @@ -430,7 +431,7 @@ def _transfer_to_and_execute_bteq_on_remote( bteq_session_encoding=bteq_session_encoding or "", timeout_rc=timeout_rc or -1, ) - self.log.info("Executing BTEQ command: %s", bteq_command_str) + self.log.debug("Executing BTEQ command: %s", bteq_command_str) exit_status, stdout, stderr = decrypt_remote_file_to_string( self.ssh_client, @@ -440,9 +441,9 @@ def _transfer_to_and_execute_bteq_on_remote( ) failure_message = None - self.log.info("stdout : %s", stdout) - self.log.info("stderr : %s", stderr) - self.log.info("exit_status : %s", exit_status) + 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 @@ -528,7 +529,7 @@ def execute_bteq_script_at_local( bteq_session_encoding=bteq_session_encoding or "", timeout_rc=timeout_rc or -1, ) - self.log.info("Executing BTEQ command: %s", bteq_command_str) + self.log.debug("Executing BTEQ command: %s", bteq_command_str) process = subprocess.Popen( bteq_command_str, @@ -539,9 +540,9 @@ def execute_bteq_script_at_local( preexec_fn=os.setsid, ) encode_bteq_script = bteq_script.encode(str(temp_file_read_encoding or "UTF-8")) - self.log.info("encode_bteq_script : %s", encode_bteq_script) + self.log.debug("encode_bteq_script : %s", encode_bteq_script) stdout_data, _ = process.communicate(input=encode_bteq_script) - self.log.info("stdout_data : %s", stdout_data) + 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( @@ -558,7 +559,7 @@ def execute_bteq_script_at_local( for line in stdout_data.splitlines(): try: decoded_line = line.decode("UTF-8").strip() - self.log.info("decoded_line : %s", decoded_line) + 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: @@ -601,9 +602,9 @@ def on_kill(self): except Exception as e: self.log.error("Failed to terminate subprocess: %s", str(e)) - def get_airflow_home_dir(self) -> str: - """Get the AIRFLOW_HOME directory from environment variables.""" - return os.environ.get("AIRFLOW_HOME", "~/airflow") + 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_"): @@ -617,14 +618,14 @@ def preferred_temp_directory(self, prefix="bteq_"): str: Path to the created temporary directory Note: - Falls back to AIRFLOW_HOME if system temp directory is not usable + 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_airflow_home_dir() + temp_dir = self.get_dagster_home_dir() with tempfile.TemporaryDirectory(dir=temp_dir, prefix=prefix) as tmp: yield tmp @@ -659,7 +660,7 @@ def _handle_remote_bteq_file( teradata_connection_resource=self.teradata_connection_resource, sql=file_content, ) - self.log.info("Executing BTEQ script with SQL content: %s", bteq_script) + 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, @@ -697,7 +698,7 @@ def _handle_local_bteq_file(self, file_path: str) -> int | None: bteq_script = prepare_bteq_script_for_local_execution( sql=file_content, ) - self.log.info("Executing BTEQ script with SQL content: %s", bteq_script) + self.log.debug("Executing BTEQ script with SQL content: %s", bteq_script) result = self.execute_bteq_script( bteq_script, self.remote_working_dir, @@ -754,10 +755,6 @@ def _setup_ssh_connection( self.cred_manager.decrypt(creds["password"]) if creds else None ) - if not password: - password = getpass.getpass(f"SSH password for {user}@{host}: ") - store_credentials(self, host, user, password) - self.ssh_client.connect( host, port=port, username=user, password=password ) diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py index 8bc0b36c..89943073 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py @@ -194,7 +194,7 @@ def _prepare_bteq_command( if timeout_rc is not None and timeout_rc >= 0: bteq_core_cmd.append(f" RC {timeout_rc}") bteq_core_cmd.append(";") - # Airflow doesn't display the script of BTEQ in UI but only in log so WIDTH is 500 enough + # 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 From 54fd738567bef5f406f8ac57cf3ad67983528c90 Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Mon, 30 Jun 2025 11:09:54 +0530 Subject: [PATCH 17/21] Corrected ruff changes --- .../dagster-teradata/dagster_teradata/ttu/bteq.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py index e6d41777..877a5074 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py @@ -1,4 +1,3 @@ -import getpass import os import socket import subprocess @@ -30,7 +29,6 @@ generate_encrypted_file_with_openssl, decrypt_remote_file_to_string, get_stored_credentials, - store_credentials, ) @@ -146,7 +144,9 @@ def bteq_operator( 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) + self.log.debug( + "Executing BTEQ script with SQL content: %s", bteq_script + ) return self.execute_bteq_script( bteq_script, self.remote_working_dir, @@ -180,7 +180,9 @@ def bteq_operator( teradata_connection_resource=self.teradata_connection_resource, sql=self.sql, ) - self.log.debug("Executing BTEQ script with SQL content: %s", bteq_script) + self.log.debug( + "Executing BTEQ script with SQL content: %s", bteq_script + ) return self.execute_bteq_script( bteq_script, self.remote_working_dir, @@ -352,7 +354,6 @@ def execute_bteq_script_at_remote( remote_port, ) - def _transfer_to_and_execute_bteq_on_remote( self, file_path: str, @@ -660,7 +661,9 @@ def _handle_remote_bteq_file( teradata_connection_resource=self.teradata_connection_resource, sql=file_content, ) - self.log.debug("Executing BTEQ script with SQL content: %s", bteq_script) + 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, From b8a9e5c9509c6bb05744f08c221dc0e9da7b0458 Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Mon, 30 Jun 2025 14:39:31 +0530 Subject: [PATCH 18/21] fix pyright issues --- .../dagster_teradata/resources.py | 2 +- .../dagster_teradata/ttu/bteq.py | 39 +++++++++++-------- .../dagster_teradata/ttu/utils/bteq_util.py | 12 +++--- .../ttu/utils/encryption_utils.py | 4 ++ 4 files changed, 33 insertions(+), 24 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 82274a30..2d91b518 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/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) diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py index 877a5074..6484a5c1 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py +++ b/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,7 +201,7 @@ 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, @@ -209,8 +209,12 @@ def bteq_operator( 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 diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py index 89943073..74c1d0af 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py +++ b/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,7 +218,7 @@ 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) @@ -226,7 +226,7 @@ def prepare_bteq_command_for_remote_execution( 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. diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py index 6f3128fe..9969b890 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py +++ b/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() From e40b7d6abc720b1c1f934c0dcd961f5a8c14185d Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Mon, 30 Jun 2025 14:44:06 +0530 Subject: [PATCH 19/21] Revert "fix pyright issues" This reverts commit b8a9e5c9509c6bb05744f08c221dc0e9da7b0458. --- .../dagster_teradata/resources.py | 2 +- .../dagster_teradata/ttu/bteq.py | 39 ++++++++----------- .../dagster_teradata/ttu/utils/bteq_util.py | 12 +++--- .../ttu/utils/encryption_utils.py | 4 -- 4 files changed, 24 insertions(+), 33 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 2d91b518..82274a30 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/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(self, query_band_text) + query_band_text = _handle_user_query_band_text(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) diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py index 6484a5c1..877a5074 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py +++ b/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, cast +from typing import Optional, List, Union, Literal import paramiko from dagster import DagsterError @@ -201,7 +201,7 @@ def bteq_operator( if self.file_path: if not self._setup_ssh_connection( host=self.remote_host, - user=cast(str, self.remote_user), + user=self.remote_user, password=self.remote_remote_password, key_path=self.ssh_key_path, port=self.remote_port, @@ -209,12 +209,8 @@ def bteq_operator( 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 - ) + if self.file_path and is_valid_remote_bteq_script_file( + self.ssh_client, self.file_path ): return self._handle_remote_bteq_file( ssh_client=self.ssh_client, @@ -234,7 +230,7 @@ def execute_bteq_script( bteq_script: str, remote_working_dir: str | None, bteq_script_encoding: str | None, - timeout: Optional[int], # or: timeout: int | None + timeout: int, timeout_rc: int | None, bteq_session_encoding: str | None, bteq_quit_rc: int | list[int] | tuple[int, ...] | None, @@ -301,7 +297,7 @@ def execute_bteq_script_at_remote( bteq_script: str, remote_working_dir: str | None, bteq_script_encoding: str | None, - timeout: Optional[int], # or: timeout: int | None + timeout: int, timeout_rc: int | None, bteq_session_encoding: str | None, bteq_quit_rc: int | list[int] | tuple[int, ...] | None, @@ -351,8 +347,8 @@ def execute_bteq_script_at_remote( bteq_quit_rc, bteq_session_encoding, tmp_dir, - cast(str, remote_host), - cast(str, remote_user), + remote_host, + remote_user, remote_password, ssh_key_path, remote_port, @@ -363,7 +359,7 @@ def _transfer_to_and_execute_bteq_on_remote( file_path: str, remote_working_dir: str | None, bteq_script_encoding: str | None, - timeout: Optional[int], # or: timeout: int | None + timeout: int, timeout_rc: int | None, bteq_quit_rc: int | list[int] | tuple[int, ...] | None, bteq_session_encoding: str | None, @@ -502,7 +498,7 @@ def execute_bteq_script_at_local( self, bteq_script: str, bteq_script_encoding: str | None, - timeout: Optional[int], # or: timeout: int | None + timeout: int, timeout_rc: int | None, bteq_quit_rc: int | list[int] | tuple[int, ...] | None, bteq_session_encoding: str | None, @@ -551,7 +547,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 or 0) + 60 + timeout=timeout + 60 ) # Adding 1 minute extra for BTEQ script timeout except subprocess.TimeoutExpired: self.on_kill() @@ -587,6 +583,10 @@ 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 or 22, + remote_port=self.remote_port, ) 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: Optional[str], + user: str, password: Optional[str], key_path: Optional[str], port: int, @@ -753,11 +753,6 @@ 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 diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py index 74c1d0af..89943073 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py +++ b/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, Optional +from typing import TYPE_CHECKING if TYPE_CHECKING: from paramiko import SSHClient @@ -200,7 +200,7 @@ def _prepare_bteq_command( def prepare_bteq_command_for_remote_execution( - timeout: Optional[int], # or: timeout: int | None + timeout: int, bteq_script_encoding: str, bteq_session_encoding: str, timeout_rc: int, @@ -218,7 +218,7 @@ def prepare_bteq_command_for_remote_execution( 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 + timeout, bteq_script_encoding, bteq_session_encoding, timeout_rc ) bteq_core_cmd.append('"') return " ".join(bteq_core_cmd) @@ -226,7 +226,7 @@ def prepare_bteq_command_for_remote_execution( def prepare_bteq_command_for_local_execution( teradata_connection_resource, - timeout: Optional[int], # or: timeout: int | None + timeout: int, 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 or 600, bteq_script_encoding, bteq_session_encoding, timeout_rc + timeout, 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 | None, logger=None + ssh_client: SSHClient, remote_file_path: str, logger=None ) -> bool: """ Check if a remote file is a valid BTEQ script file. diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py index 9969b890..6f3128fe 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py @@ -70,8 +70,6 @@ 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() @@ -92,8 +90,6 @@ 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() From 7891b012504335a6e8c48040b87559dab2aa5e12 Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Mon, 30 Jun 2025 14:46:14 +0530 Subject: [PATCH 20/21] Reapply "fix pyright issues" This reverts commit e40b7d6abc720b1c1f934c0dcd961f5a8c14185d. --- .../dagster_teradata/resources.py | 2 +- .../dagster_teradata/ttu/bteq.py | 39 +++++++++++-------- .../dagster_teradata/ttu/utils/bteq_util.py | 12 +++--- .../ttu/utils/encryption_utils.py | 4 ++ 4 files changed, 33 insertions(+), 24 deletions(-) diff --git a/libraries/dagster-teradata/dagster_teradata/resources.py b/libraries/dagster-teradata/dagster_teradata/resources.py index 82274a30..2d91b518 100644 --- a/libraries/dagster-teradata/dagster_teradata/resources.py +++ b/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) diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py b/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py index 877a5074..6484a5c1 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/bteq.py +++ b/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,7 +201,7 @@ 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, @@ -209,8 +209,12 @@ def bteq_operator( 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 diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py index 89943073..74c1d0af 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/bteq_util.py +++ b/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,7 +218,7 @@ 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) @@ -226,7 +226,7 @@ def prepare_bteq_command_for_remote_execution( 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. diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py index 6f3128fe..9969b890 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py +++ b/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() From 0ca399db8e1262121e42c6ec56612bcf2e8b77b4 Mon Sep 17 00:00:00 2001 From: "Talla, Mohan" Date: Mon, 30 Jun 2025 14:59:09 +0530 Subject: [PATCH 21/21] Fix pytest --- .../dagster_teradata/ttu/utils/encryption_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py index 9969b890..3b329371 100644 --- a/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py +++ b/libraries/dagster-teradata/dagster_teradata/ttu/utils/encryption_utils.py @@ -50,6 +50,9 @@ def _load_or_generate_key(self): 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