dataops-testgen 2.2.0__py3-none-any.whl

testgen/common/database/flavor/mssql_flavor_service.py

@@ -0,0 +1,34 @@
+from urllib.parse import quote_plus
+
+from testgen.common.database.flavor.flavor_service import FlavorService
+
+
+class MssqlFlavorService(FlavorService):
+    def get_connection_string_head(self, strPW):
+        username = self.username
+        password = quote_plus(strPW)
+
+        strConnect = f"mssql+pyodbc://{username}:{password}@"
+
+        return strConnect
+
+    def get_connection_string_from_fields(self, strPW, is_password_overwritten: bool = False):    # NOQA ARG002
+        password = quote_plus(strPW)
+
+        strConnect = (
+            f"mssql+pyodbc://{self.username}:{password}@{self.host}:{self.port}/{self.dbname}?driver=ODBC+Driver+18+for+SQL+Server"
+        )
+
+        if "synapse" in self.host:
+            strConnect += "&autocommit=True"
+
+        return strConnect
+
+    def get_pre_connection_queries(self):  # ARG002
+        return [
+            "SET ANSI_DEFAULTS ON;",
+            "SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED;",
+        ]
+
+    def get_concat_operator(self):
+        return "+"

testgen/common/database/flavor/postgresql_flavor_service.py

@@ -0,0 +1,5 @@
+from testgen.common.database.flavor.redshift_flavor_service import RedshiftFlavorService
+
+
+class PostgresqlFlavorService(RedshiftFlavorService):
+    pass

testgen/common/database/flavor/redshift_flavor_service.py

@@ -0,0 +1,22 @@
+from urllib.parse import quote_plus
+
+from testgen.common.database.flavor.flavor_service import FlavorService
+
+
+class RedshiftFlavorService(FlavorService):
+    def get_connection_string_head(self, strPW):
+        strConnect = f"{self.flavor}://{self.username}:{quote_plus(strPW)}@"
+        return strConnect
+
+    def get_connection_string_from_fields(self, strPW, is_password_overwritten: bool = False):  # NOQA ARG002
+        # STANDARD FORMAT:  strConnect = 'flavor://username:password@host:port/database'
+        strConnect = f"{self.flavor}://{self.username}:{quote_plus(strPW)}@{self.host}:{self.port}/{self.dbname}"
+        return strConnect
+
+    def get_pre_connection_queries(self):
+        return [
+            "SET SEARCH_PATH = '" + self.dbschema + "'",
+        ]
+
+    def get_connect_args(self, is_password_overwritten: bool = False):  # NOQA ARG002
+        return {}

testgen/common/database/flavor/snowflake_flavor_service.py

@@ -0,0 +1,69 @@
+from urllib.parse import quote_plus
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+
+from testgen.common.database.flavor.flavor_service import FlavorService
+
+
+class SnowflakeFlavorService(FlavorService):
+
+    def get_connect_args(self, is_password_overwritten: bool = False):
+        connect_args = super().get_connect_args(is_password_overwritten)
+
+        if self.connect_by_key and not is_password_overwritten:
+            # https://docs.snowflake.com/en/developer-guide/python-connector/sqlalchemy#key-pair-authentication-support
+            private_key_passphrase = self.private_key_passphrase.encode() if self.private_key_passphrase else None
+            private_key = serialization.load_pem_private_key(
+                self.private_key.encode(),
+                password=private_key_passphrase,
+                backend=default_backend(),
+            )
+
+            private_key_bytes = private_key.private_bytes(
+                encoding=serialization.Encoding.DER,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption(),
+            )
+
+            connect_args.update({"private_key": private_key_bytes})
+        return connect_args
+
+    def get_connection_string_head(self, strPW):
+        if self.connect_by_key and not strPW:
+            strConnect = f"snowflake://{self.username}@"
+        else:
+            strConnect = f"snowflake://{self.username}:{quote_plus(strPW)}@"
+        return strConnect
+
+    def get_connection_string_from_fields(self, strPW, is_password_overwritten: bool = False):
+        # SNOWFLAKE FORMAT:  strConnect = 'flavor://username:password@host/database'
+        #   optionally + '/[schema]' + '?warehouse=xxx'
+        #   NOTE:  Snowflake host should NOT include ".snowflakecomputing.com"
+
+        def get_raw_host_name(host):
+            endings = [
+                ".snowflakecomputing.com",
+            ]
+            for ending in endings:
+                if host.endswith(ending):
+                    i = host.index(ending)
+                    return host[0:i]
+            return host
+
+        raw_host = get_raw_host_name(self.host)
+        host = raw_host
+        if self.port != "443":
+            host += ":" + self.port
+
+        if self.connect_by_key and not is_password_overwritten:
+            strConnect = f"snowflake://{self.username}@{host}/{self.dbname}/{self.dbschema}"
+        else:
+            strConnect = f"snowflake://{self.username}:{quote_plus(strPW)}@{host}/{self.dbname}/{self.dbschema}"
+        return strConnect
+
+    def get_pre_connection_queries(self):  # ARG002
+        return [
+            "ALTER SESSION SET MULTI_STATEMENT_COUNT = 0;",
+            "ALTER SESSION SET WEEK_START = 7;",
+        ]

testgen/common/database/flavor/trino_flavor_service.py

@@ -0,0 +1,21 @@
+from urllib.parse import quote_plus
+
+from testgen.common.database.flavor.flavor_service import FlavorService
+
+
+class TrinoFlavorService(FlavorService):
+    def get_connection_string_head(self, strPW):
+        strConnect = f"{self.flavor}://{self.username}:{quote_plus(strPW)}@"
+        return strConnect
+
+    def get_connection_string_from_fields(self, strPW, is_password_overwritten: bool = False):  # NOQA ARG002
+        # STANDARD FORMAT:  strConnect = 'flavor://username:password@host:port/catalog'
+        return f"{self.flavor}://{self.username}:{quote_plus(strPW)}@{self.host}:{self.port}/{self.catalog}"
+
+    def get_pre_connection_queries(self):
+        return [
+            "USE " + self.catalog + "." + self.dbschema,
+        ]
+
+    def get_connect_args(self, is_password_overwritten: bool = False):  # NOQA ARG002
+        return {}

testgen/common/date_service.py

@@ -0,0 +1,68 @@
+from datetime import datetime, timedelta
+
+import pandas as pd
+
+
+def get_today_as_string():
+    return datetime.utcnow().strftime("%Y-%m-%d")
+
+
+def get_now_as_string():
+    return datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
+
+
+def get_now_as_string_with_offset(minutes_offset):
+    ret = datetime.utcnow()
+    if minutes_offset > 0:
+        ret = ret + timedelta(minutes=minutes_offset)
+    return ret.strftime("%Y-%m-%d %H:%M:%S")
+
+
+def get_now_as_iso_timestamp():
+    return as_iso_timestamp(datetime.utcnow())
+
+
+def as_iso_timestamp(date: datetime) -> str | None:
+    if date is None:
+        return None
+    return date.strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+def accommodate_dataframe_to_timezone(df, streamlit_session, time_columns=None):
+    if time_columns is None:
+        time_columns = []
+        for column_name in df.columns:
+            if df[column_name].dtype == "datetime64[ns]":
+                time_columns.append(column_name)
+
+    if time_columns and "browser_timezone" in streamlit_session:
+        timezone = streamlit_session["browser_timezone"]
+        for time_column in time_columns:
+            df[time_column] = pd.to_datetime(df[time_column], errors="coerce")
+            df[time_column] = df[time_column].dt.tz_localize("UTC")
+            df[time_column] = df[time_column].dt.tz_convert(timezone)
+            df[time_column] = df[time_column].dt.strftime("%Y-%m-%d %H:%M:%S")
+
+
+def create_timezoned_column_in_dataframe(streamlit_session, df, new_column_name, existing_column_name):
+    if new_column_name and existing_column_name and "browser_timezone" in streamlit_session:
+        timezone = streamlit_session["browser_timezone"]
+        df[new_column_name] = (
+            df[existing_column_name].dt.tz_localize("UTC").dt.tz_convert(timezone).dt.strftime("%Y-%m-%d %H:%M:%S")
+        )
+
+
+def get_timezoned_timestamp(streamlit_session, value):
+    ret = None
+    if value and "browser_timezone" in streamlit_session:
+        data = {"value": [value]}
+        df = pd.DataFrame(data)
+        timezone = streamlit_session["browser_timezone"]
+        df["value"] = df["value"].dt.tz_localize("UTC").dt.tz_convert(timezone).dt.strftime("%Y-%m-%d %H:%M:%S")
+        ret = df.iloc[0, 0]
+    return ret
+
+
+def get_timezoned_now(streamlit_session):
+    value = datetime.utcnow()
+    return get_timezoned_timestamp(streamlit_session, value)

testgen/common/display_service.py

@@ -0,0 +1,85 @@
+import csv
+import logging
+import os
+
+import click
+import yaml
+from prettytable import PrettyTable
+
+LOG = logging.getLogger("testgen")
+
+
+def print_table(rows: list[dict], column_names: list[str]):
+    table = PrettyTable(column_names)
+    table.max_width = 80
+    table.align = "l"
+
+    for row in rows:
+        table.add_row(row)
+    click.echo(table)
+
+
+def to_csv(file_name: str, rows: list[dict], column_names: list[str]):
+    _, file_out_path = get_in_out_paths()
+    full_path = os.path.join(file_out_path, file_name)
+    with open(full_path, "w", newline="") as file:
+        writer = csv.writer(file)
+        writer.writerow(column_names)
+        for row in rows:
+            writer.writerow(row)
+    echo(f"Output written to: ~/testgen/file-out/{file_name}")
+
+
+def get_in_out_paths():
+    # create the paths if not exist
+    home = os.path.expanduser("~")
+    file_in_path = os.path.join(home, "testgen", "file-in")
+    file_out_path = os.path.join(home, "testgen", "file-out")
+    os.makedirs(file_in_path, exist_ok=True)
+    os.makedirs(file_out_path, exist_ok=True)
+    return file_in_path, file_out_path
+
+
+def write_to_file(full_path_and_name: str, file_content: str):
+    with open(full_path_and_name, "w") as file:
+        file.write(file_content)
+
+
+def to_yaml(file_name: str, yaml_dict: dict, display: bool):
+    yaml_content = yaml.dump(yaml_dict, sort_keys=False)
+    yaml_content.replace("None", "null")
+
+    _, file_out_path = get_in_out_paths()
+    full_path = os.path.join(file_out_path, file_name)
+    with open(full_path, "w", newline="") as file:
+        file.write(yaml_content)
+
+    if display:
+        echo(yaml_content + "\n")
+
+    echo(f"Output written to: ~/testgen/file-out/{file_name}")
+
+
+def echo(message: str):
+    click.echo(message)
+
+
+def from_yaml(file_name: str, display: bool):
+    echo(f"Attempting to read from : ~/testgen/file-in/{file_name}")
+    file_in_path, _ = get_in_out_paths()
+    full_path = os.path.join(file_in_path, file_name)
+    with open(full_path, newline="") as file:
+        yaml_content = yaml.safe_load(file)
+
+    if display:
+        data = yaml.dump(yaml_content, sort_keys=False)
+        echo(data)
+
+    return yaml_content
+
+
+def check_config_file_presence(file_name: str) -> None:
+    file_in_path, _ = get_in_out_paths()
+    full_path = os.path.join(file_in_path, file_name)
+    if not os.path.exists(full_path):
+        echo(click.style(f"Warning: File ~/testgen/file-in/{file_name} is not present.", fg="yellow"))

testgen/common/docker_service.py

@@ -0,0 +1,76 @@
+import logging
+
+import requests
+
+from testgen import settings
+from testgen.common import get_tg_db, get_tg_host, get_tg_password, get_tg_schema, get_tg_username
+
+LOG = logging.getLogger("testgen")
+
+
+
+def check_for_new_docker_release() -> str:
+    if not settings.CHECK_FOR_LATEST_VERSION:
+        return "unknown"
+
+    try:
+        tags = get_docker_tags()
+
+        if len(tags) == 0:
+            LOG.debug("docker_service: No tags to parse, skipping check.")
+            return "unknown"
+
+        ordered_tags = sorted(tags, key=lambda item: item[1], reverse=True)
+        latest_tag = ordered_tags[0][0]
+
+        if latest_tag != settings.VERSION:
+            LOG.warning(
+                f"A new TestGen upgrade is available. Please update to version {latest_tag} for new features and improvements."
+            )
+
+        return latest_tag  # noqa: TRY300
+    except Exception:
+        LOG.warning("Unable to check for latest release", exc_info=True, stack_info=True)
+
+
+def get_docker_tags(url: str = "https://hub.docker.com/v2/repositories/datakitchen/dataops-testgen/tags/"):
+    params = {"page_size": 25, "page": 1, "ordering": "last_updated"}
+    response = requests.get(url, params=params, timeout=3)
+
+    tags_to_return = []
+    if not response.status_code == 200:
+        LOG.warning(f"docker_service: Failed to fetch docker tags. Status code: {response.status_code}")
+        return tags_to_return
+
+    tags_data = response.json()
+    results = tags_data.get("results", [])
+    for result in results:
+        tag_name = result["name"]
+        last_pushed = result["tag_last_pushed"]
+        if tag_name.count(".") >= 2 and "experimental" not in tag_name:
+            tags_to_return.append((tag_name, last_pushed))
+
+    return tags_to_return
+
+
+def check_basic_configuration():
+    ret = True
+    message = ""
+
+    configs = [
+        ("host", get_tg_host),
+        ("username", get_tg_username),
+        ("password", get_tg_password),
+        ("schema", get_tg_schema),
+        ("db", get_tg_db),
+    ]
+
+    for config in configs:
+        if not config[1]():
+            ret = False
+            message += f"\n{config[0]} configuration is missing."
+
+    if message:
+        message = "The system is not properly configured. Please check. Details: \n" + message
+
+    return ret, message

testgen/common/encrypt.py

@@ -0,0 +1,55 @@
+import base64
+
+import streamlit_authenticator as stauth
+from Crypto.Cipher import AES
+from Crypto.Protocol.KDF import PBKDF2
+from Crypto.Random import get_random_bytes
+
+from testgen import settings
+
+
+def EncryptText(strText):
+    block_size = 16
+
+    def pad(s):
+        return s + (block_size - len(s) % block_size) * chr(block_size - len(s) % block_size)
+
+    # Generate a random salt
+    salt = settings.APP_ENCRYPTION_SALT.encode("ascii")
+    strPassword = settings.APP_ENCRYPTION_SECRET.encode("ascii")
+
+    # Derive the key using PBKDF2
+    kdf = PBKDF2(strPassword, salt, 64, 1000)
+    private_key = kdf[:32]
+
+    # Initialize the cipher
+    strText = pad(strText)
+    strText = bytes(strText, "utf-8")
+    iv = get_random_bytes(AES.block_size)
+    cipher = AES.new(private_key, AES.MODE_CBC, iv)
+
+    # Perform encryption
+    encrypted_text = base64.b64encode(iv + cipher.encrypt(strText))
+    return encrypted_text.decode("UTF-8")
+
+
+def DecryptText(baEncrypted):
+    def unpad(s):
+        return s[: -ord(s[len(s) - 1 :])]
+
+    # Calc Private Key from Password
+    salt = settings.APP_ENCRYPTION_SALT.encode("ascii")
+    strPassword = settings.APP_ENCRYPTION_SECRET.encode("ascii")
+    kdf = PBKDF2(strPassword, salt, 64, 1000)
+    private_key = kdf[:32]
+
+    baEncrypted = base64.b64decode(baEncrypted)
+    iv = baEncrypted[:16]
+    cipher = AES.new(private_key, AES.MODE_CBC, iv)
+
+    return bytes.decode(unpad(cipher.decrypt(baEncrypted[16:])))
+
+
+def encrypt_ui_password(plain_password):
+    hashed_passwords = stauth.Hasher([plain_password]).generate()
+    return hashed_passwords.pop()

testgen/common/get_pipeline_parms.py

@@ -0,0 +1,57 @@
+from testgen.common.database.database_service import RetrieveDBResultsToDictList
+from testgen.common.read_file import read_template_sql_file
+
+
+def RetrieveProfilingParms(strTableGroupsID):
+    strSQL = read_template_sql_file("parms_profiling.sql", "parms")
+    # Replace Parameters
+    strSQL = strSQL.replace("{TABLE_GROUPS_ID}", strTableGroupsID)
+
+    # Execute Query
+    lstParms = RetrieveDBResultsToDictList("DKTG", strSQL)
+
+    if lstParms is None:
+        raise ValueError("Project Connection Parameters not found")
+    elif (
+        lstParms[0]["project_code"] == ""
+        or lstParms[0]["connection_id"] == ""
+        or lstParms[0]["sql_flavor"] == ""
+        or lstParms[0]["project_user"] == ""
+        or lstParms[0]["profile_use_sampling"] == ""
+        or lstParms[0]["profile_sample_percent"] == ""
+        or lstParms[0]["profile_sample_min_count"] == ""
+        or lstParms[0]["project_qc_schema"] == ""
+        or lstParms[0]["table_group_schema"] == ""
+    ):
+        raise ValueError("Project Connection parameters not correctly set")
+    else:
+        return lstParms[0]
+
+
+def RetrieveTestGenParms(strTableGroupsID, strTestSuite):
+    strSQL = read_template_sql_file("parms_test_gen.sql", "parms")
+    # Replace Parameters
+    strSQL = strSQL.replace("{TABLE_GROUPS_ID}", strTableGroupsID)
+    strSQL = strSQL.replace("{TEST_SUITE}", strTestSuite)
+
+    # Execute Query
+    lstParms = RetrieveDBResultsToDictList("DKTG", strSQL)
+    if len(lstParms) == 0:
+        raise ValueError("SQL retrieved 0 records")
+    return lstParms[0]
+
+
+def RetrieveTestExecParms(strProjectCode, strTestSuite):
+    strSQL = read_template_sql_file("parms_test_execution.sql", "parms")
+    # Replace Parameters
+    strSQL = strSQL.replace("{PROJECT_CODE}", strProjectCode)
+    strSQL = strSQL.replace("{TEST_SUITE}", strTestSuite)
+
+    # Execute Query
+    lstParms = RetrieveDBResultsToDictList("DKTG", strSQL)
+    if len(lstParms) == 0:
+        raise ValueError("Test Execution parameters could not be retrieved")
+    elif len(lstParms) > 1:
+        raise ValueError("Test Execution parameters returned too many records")
+
+    return lstParms[0]

testgen/common/logs.py

@@ -0,0 +1,79 @@
+__all__ = ["configure_logging"]
+
+import io
+import logging
+import logging.handlers
+import os
+import sys
+import threading
+
+from concurrent_log_handler import ConcurrentTimedRotatingFileHandler
+
+from testgen import settings
+
+
+def configure_logging(
+    level: int = logging.DEBUG,
+    log_format: str = "[PID: %(process)s] %(asctime)s - %(levelname)s - %(message)s",
+) -> None:
+    """
+    Configures the testgen logger.
+    """
+    logger = logging.getLogger("testgen")
+    logger.setLevel(level)
+
+    formatter = logging.Formatter(log_format)
+
+    console_out_handler = logging.StreamHandler(stream=sys.stdout)
+    if settings.IS_DEBUG:
+        console_out_handler.setLevel(level)
+    else:
+        console_out_handler.setLevel(logging.WARNING)
+    console_out_handler.setFormatter(formatter)
+
+    console_err_handler = logging.StreamHandler(stream=sys.stderr)
+    console_err_handler.setLevel(logging.WARNING)
+    console_err_handler.setFormatter(formatter)
+
+    logger.addHandler(console_out_handler)
+    logger.addHandler(console_err_handler)
+
+    if settings.LOG_TO_FILE:
+        os.makedirs(settings.LOG_FILE_PATH, exist_ok=True)
+
+        file_handler = ConcurrentTimedRotatingFileHandler(
+            get_log_full_path(),
+            when="D",
+            interval=1,
+            backupCount=int(settings.LOG_FILE_MAX_QTY),
+        )
+        file_handler.setLevel(level)
+        file_handler.setFormatter(formatter)
+
+        logger.addHandler(file_handler)
+
+
+def get_log_full_path() -> str:
+    return os.path.join(settings.LOG_FILE_PATH, "app.log")
+
+
+class LogPipe(threading.Thread, io.TextIOBase):
+    def __init__(self, logger: logging.Logger, log_level: int) -> None:
+        threading.Thread.__init__(self)
+
+        self.daemon = False
+        self.logger = logger
+        self.level = log_level
+        self.readDescriptor, self.writeDescriptor = os.pipe()
+        self.start()
+
+    def run(self) -> None:
+        with os.fdopen(self.readDescriptor) as reader:
+            for line in iter(reader.readline, ""):
+                self.logger.log(self.level, line.strip("\n"))
+
+    def fileno(self) -> int:
+        return self.writeDescriptor
+
+    def close(self) -> None:
+        os.close(self.writeDescriptor)

testgen/common/process_service.py

@@ -0,0 +1,62 @@
+import logging
+import os
+
+import psutil
+
+from testgen import settings
+
+LOG = logging.getLogger("testgen")
+
+
+def get_current_process_id():
+    return os.getpid()
+
+
+def kill_profile_run(process_id):
+    keywords = ["run-profile"]
+    status, message = kill_process(process_id, keywords)
+    return status, message
+
+
+def kill_test_run(process_id):
+    keywords = ["run-tests"]
+    status, message = kill_process(process_id, keywords)
+    return status, message
+
+
+def kill_process(process_id, keywords=None):
+    if settings.IS_DEBUG:
+        msg = "Cannot kill processes in debug mode (threads are used instead of new process)"
+        LOG.warn(msg)
+        return False, msg
+    try:
+        process = psutil.Process(process_id)
+        if process.name().lower() != "testgen":
+            message = f"The process was not killed because the process_id {process_id} is not a testgen process. Details: {process.name()}"
+            LOG.error(f"kill_process: {message}")
+            return False, message
+
+        if keywords:
+            for keyword in keywords:
+                if keyword.lower() not in process.cmdline():
+                    message = f"The process was not killed because the keyword {keyword} was not found. Details: {process.cmdline()}"
+                    LOG.error(f"kill_process: {message}")
+                    return False, message
+
+        process.terminate()
+        process.wait(timeout=10)
+        message = f"Process {process_id} has been terminated."
+    except psutil.NoSuchProcess:
+        message = f"No such process with PID {process_id}."
+        LOG.exception(f"kill_process: {message}")
+        return False, message
+    except psutil.AccessDenied:
+        message = f"Access denied when trying to terminate process {process_id}."
+        LOG.exception(f"kill_process: {message}")
+        return False, message
+    except psutil.TimeoutExpired:
+        message = f"Process {process_id} did not terminate within the timeout period."
+        LOG.exception(f"kill_process: {message}")
+        return False, message
+    LOG.info(f"kill_process: Success. {message}")
+    return True, message