dataops-testgen 2.2.0__py3-none-any.whl

testgen/ui/services/authentication_service.py

@@ -0,0 +1,119 @@
+# ruff: noqa: S105
+
+import datetime
+import logging
+import typing
+
+import extra_streamlit_components as stx
+import jwt
+import streamlit as st
+
+from testgen.common.encrypt import encrypt_ui_password
+from testgen.ui.queries import authentication_queries
+from testgen.ui.session import session
+
+RoleType = typing.Literal["admin", "edit", "read"]
+
+JWT_HASHING_KEY = "dk_signature_key"
+AUTH_TOKEN_COOKIE_NAME = "dk_cookie_name"
+AUTH_TOKEN_EXPIRATION_DAYS = 5
+
+LOG = logging.getLogger("testgen")
+
+
+def load_user_session() -> None:
+    cookies = stx.CookieManager(key="testgen.cookies.get")
+    token = cookies.get(AUTH_TOKEN_COOKIE_NAME)
+    if token is not None:
+        try:
+            token = jwt.decode(token, JWT_HASHING_KEY, algorithms=["HS256"])
+            if token["exp_date"] > datetime.datetime.utcnow().timestamp():
+                start_user_session(token["name"], token["username"])
+        except Exception:
+            LOG.debug("Invalid auth token found on cookies", exc_info=True, stack_info=True)
+
+
+def start_user_session(name: str, username: str) -> None:
+    session.name = name
+    session.username = username
+    session.auth_role = get_role_for_user(get_auth_data(), username)
+    session.authentication_status = True
+    if not session.current_page or session.current_page == "login":
+        session.current_page = "overview"
+        session.current_page_args = {}
+    session.logging_out = False
+
+
+def end_user_session() -> None:
+    session.auth_role = None
+    session.authentication_status = None
+    session.current_page = "login"
+    session.current_page_args = {}
+    session.logging_out = True
+
+    del session.name
+    del session.username
+
+
+def add_user(user):
+    encrypted_password = encrypt_ui_password(user["password"])
+    schema = st.session_state["dbschema"]
+    authentication_queries.add_user(schema, user, encrypted_password)
+
+
+def delete_users(user_ids):
+    schema = st.session_state["dbschema"]
+    return authentication_queries.delete_users(schema, user_ids)
+
+
+def edit_user(user):
+    encrypted_password = encrypt_ui_password(user["password"])
+    schema = st.session_state["dbschema"]
+    authentication_queries.edit_user(schema, user, encrypted_password)
+
+
+def get_auth_data():
+    auth_data = authentication_queries.get_users(session.dbschema)
+
+    usernames = {}
+    preauthorized_list = []
+
+    for item in auth_data.itertuples():
+        usernames[item.username] = {
+            "email": item.email,
+            "name": item.name,
+            "password": item.password,
+            "role": item.role,
+        }
+        if item.preauthorized:
+            preauthorized_list.append(item.email)
+
+    return {
+        "credentials": {"usernames": usernames},
+        "cookie": {"expiry_days": AUTH_TOKEN_EXPIRATION_DAYS, "key": JWT_HASHING_KEY, "name": AUTH_TOKEN_COOKIE_NAME},
+        "preauthorized": {"emails": preauthorized_list},
+    }
+
+
+def get_users():
+    return authentication_queries.get_users(session.dbschema)
+
+
+def get_role_for_user(auth_data, username):
+    return auth_data["credentials"]["usernames"][username]["role"]
+
+
+def current_user_has_admin_role():
+    return session.auth_role == "admin"
+
+
+def current_user_has_edit_role():
+    return session.auth_role in ("edit", "admin")
+
+
+def current_user_has_read_role():
+    return not session.auth_role or session.auth_role == "read"
+
+
+def current_user_has_role(role: RoleType) -> bool:
+    return session.auth_role == role

testgen/ui/services/connection_service.py

@@ -0,0 +1,220 @@
+import streamlit as st
+
+import testgen.ui.queries.connection_queries as connection_queries
+import testgen.ui.services.table_group_service as table_group_service
+from testgen.commands.run_profiling_bridge import InitializeProfilingSQL
+from testgen.commands.run_setup_profiling_tools import run_setup_profiling_tools
+from testgen.common.database.database_service import (
+    AssignConnectParms,
+    RetrieveDBResultsToList,
+    empty_cache,
+    get_db_type,
+    get_flavor_service,
+)
+from testgen.common.encrypt import DecryptText, EncryptText
+
+
+def get_by_id(connection_id, hide_passwords: bool = True):
+    connections_df = connection_queries.get_by_id(connection_id)
+    decrypt_connections(connections_df, hide_passwords)
+    connection = connections_df.to_dict(orient="records")[0]
+    return connection
+
+
+def get_connections(project_code, hide_passwords: bool = False):
+    connections = connection_queries.get_connections(project_code)
+    decrypt_connections(connections, hide_passwords)
+    return connections
+
+
+def decrypt_connections(connections, hide_passwords: bool = False):
+    for index, connection in connections.iterrows():
+        if hide_passwords:
+            password = "***"  # noqa S105
+            private_key = "***"  # S105
+            private_key_passphrase = "***"  # noqa S105
+        else:
+            password = DecryptText(connection["project_pw_encrypted"]) if connection["project_pw_encrypted"] else None
+            private_key = DecryptText(connection["private_key"]) if connection["private_key"] else None
+            private_key_passphrase = DecryptText(connection["private_key_passphrase"]) if connection["private_key_passphrase"] else ""
+        connections.at[index, "password"] = password
+        connections.at[index, "private_key"] = private_key
+        connections.at[index, "private_key_passphrase"] = private_key_passphrase
+
+
+def encrypt_credentials(connection):
+    encrypted_password = EncryptText(connection["password"]) if connection["password"] else None
+    encrypted_private_key = EncryptText(connection["private_key"]) if connection["private_key"] else None
+    encrypted_private_key_passphrase = EncryptText(connection["private_key_passphrase"]) if connection["private_key_passphrase"] else None
+    return encrypted_password, encrypted_private_key, encrypted_private_key_passphrase
+
+
+def edit_connection(connection):
+    empty_cache()
+    schema = st.session_state["dbschema"]
+    connection = pre_save_connection_process(connection)
+    encrypted_password, encrypted_private_key, encrypted_private_key_passphrase = encrypt_credentials(connection)
+    connection_queries.edit_connection(schema, connection, encrypted_password, encrypted_private_key, encrypted_private_key_passphrase)
+
+
+def add_connection(connection):
+    empty_cache()
+    schema = st.session_state["dbschema"]
+    connection = pre_save_connection_process(connection)
+    encrypted_password, encrypted_private_key, encrypted_private_key_passphrase = encrypt_credentials(connection)
+    connection_queries.add_connection(schema, connection, encrypted_password, encrypted_private_key, encrypted_private_key_passphrase)
+
+
+def pre_save_connection_process(connection):
+    if connection["connect_by_url"]:
+        url = connection["url"]
+        if url:
+            url_sections = url.split("/")
+            if len(url_sections) > 0:
+                host_port = url_sections[0]
+                host_port_sections = host_port.split(":")
+                if len(host_port_sections) > 0:
+                    connection["project_host"] = host_port_sections[0]
+                    connection["project_port"] = "".join(host_port_sections[1:])
+                else:
+                    connection["project_host"] = host_port
+                    connection["project_port"] = ""
+            if len(url_sections) > 1:
+                connection["project_db"] = url_sections[1]
+    return connection
+
+
+def delete_connections(connection_ids):
+    empty_cache()
+    schema = st.session_state["dbschema"]
+    return connection_queries.delete_connections(schema, connection_ids)
+
+
+def cascade_delete(connection_ids, dry_run=False):
+    schema = st.session_state["dbschema"]
+    can_be_deleted = True
+    table_group_names = get_table_group_names_by_connection(connection_ids)
+    connection_has_dependencies = table_group_names is not None and len(table_group_names) > 0
+    if connection_has_dependencies:
+        can_be_deleted = False
+    if not dry_run:
+        if connection_has_dependencies:
+            table_group_service.cascade_delete(table_group_names)
+        connection_queries.delete_connections(schema, connection_ids)
+    return can_be_deleted
+
+
+def are_connections_in_use(connection_ids):
+    table_group_names = get_table_group_names_by_connection(connection_ids)
+    table_groups_in_use = table_group_service.are_table_groups_in_use(table_group_names)
+    return table_groups_in_use
+
+
+def get_table_group_names_by_connection(connection_ids):
+    if not connection_ids:
+        return []
+    schema = st.session_state["dbschema"]
+    table_group_names = connection_queries.get_table_group_names_by_connection(schema, connection_ids)
+    return table_group_names.to_dict()["table_groups_name"].values()
+
+
+def init_profiling_sql(project_code, connection, table_group_schema=None):
+    # get connection data
+    empty_cache()
+    connection_id = str(connection["connection_id"]) if connection["connection_id"] else None
+    sql_flavor = connection["sql_flavor"]
+    url = connection["url"]
+    connect_by_url = connection["connect_by_url"]
+    connect_by_key = connection["connect_by_key"]
+    private_key = connection["private_key"]
+    private_key_passphrase = connection["private_key_passphrase"]
+    project_host = connection["project_host"]
+    project_port = connection["project_port"]
+    project_db = connection["project_db"]
+    project_user = connection["project_user"]
+    project_qc_schema = connection["project_qc_schema"]
+    password = connection["password"]
+
+    # prepare the profiling query
+    clsProfiling = InitializeProfilingSQL(project_code, sql_flavor)
+
+    AssignConnectParms(
+        project_code,
+        connection_id,
+        project_host,
+        project_port,
+        project_db,
+        table_group_schema if table_group_schema else project_qc_schema,
+        project_user,
+        sql_flavor,
+        url,
+        connect_by_url,
+        connect_by_key,
+        private_key,
+        private_key_passphrase,
+        connectname="PROJECT",
+        password=password,
+    )
+
+    return clsProfiling
+
+
+def test_qc_connection(project_code, connection, init_profiling=True):
+    qc_results = {}
+
+    if init_profiling:
+        init_profiling_sql(project_code, connection)
+
+    project_qc_schema = connection["project_qc_schema"]
+    query_isnum_true = f"select {project_qc_schema}.fndk_isnum('32')"
+    query_isnum_true_result_raw = RetrieveDBResultsToList("PROJECT", query_isnum_true)
+    isnum_true_result = query_isnum_true_result_raw[0][0][0] == 1
+    qc_results["isnum_true_result"] = isnum_true_result
+
+    query_isnum_false = f"select {project_qc_schema}.fndk_isnum('HELLO')"
+    query_isnum_false_result_raw = RetrieveDBResultsToList("PROJECT", query_isnum_false)
+    isnum_false_result = query_isnum_false_result_raw[0][0][0] == 0
+    qc_results["isnum_false_result"] = isnum_false_result
+
+    query_isdate_true = f"select {project_qc_schema}.fndk_isdate('2013-05-18')"
+    query_isdate_true_result_raw = RetrieveDBResultsToList("PROJECT", query_isdate_true)
+    isdate_true_result = query_isdate_true_result_raw[0][0][0] == 1
+    qc_results["isdate_true_result"] = isdate_true_result
+
+    query_isdate_false = f"select {project_qc_schema}.fndk_isdate('HELLO')"
+    query_isdate_false_result_raw = RetrieveDBResultsToList("PROJECT", query_isdate_false)
+    isdate_false_result = query_isdate_false_result_raw[0][0][0] == 0
+    qc_results["isdate_false_result"] = isdate_false_result
+
+    return qc_results
+
+
+def create_qc_schema(connection_id, create_qc_schema, db_user, db_password, skip_granting_privileges, admin_private_key_passphrase=None, admin_private_key=None, user_role=None):
+    dry_run = False
+    empty_cache()
+    run_setup_profiling_tools(connection_id, dry_run, create_qc_schema, db_user, db_password, skip_granting_privileges, admin_private_key_passphrase, admin_private_key, user_role)
+
+
+def form_overwritten_connection_url(connection):
+    flavor = connection["sql_flavor"]
+
+    connection_credentials = {
+        "flavor": flavor,
+        "user": "<user>",
+        "host": connection["project_host"],
+        "port": connection["project_port"],
+        "dbname": connection["project_db"],
+        "url": None,
+        "connect_by_url": None,
+        "connect_by_key": connection["connect_by_key"],
+        "private_key": None,
+        "private_key_passphrase": "",
+        "dbschema": "",
+    }
+
+    db_type = get_db_type(flavor)
+    flavor_service = get_flavor_service(db_type)
+    flavor_service.init(connection_credentials)
+    connection_string = flavor_service.get_connection_string("<password>")
+
+    return connection_string

testgen/ui/services/database_service.py

@@ -0,0 +1,282 @@
+from urllib.parse import quote_plus
+
+import pandas as pd
+from sqlalchemy import create_engine, text
+
+from testgen.common.credentials import (
+    get_tg_db,
+    get_tg_host,
+    get_tg_password,
+    get_tg_port,
+    get_tg_schema,
+    get_tg_username,
+)
+from testgen.common.database.database_service import get_flavor_service
+from testgen.common.encrypt import DecryptText
+
+"""
+ Shared database access and utility functions
+"""
+
+
+def get_schema():
+    return get_tg_schema()
+
+
+def _start_engine():
+    # TestGen database
+    dbhost = get_tg_host()
+    dbport = get_tg_port()
+    dbname = get_tg_db()
+    # User Information
+    dbuser = get_tg_username()
+    dbpw = get_tg_password()
+
+    conn_str = "postgresql://" + dbuser + ":" + quote_plus(dbpw) + "@" + dbhost + ":" + dbport + "/" + dbname
+    return create_engine(conn_str)
+
+
+def _make_connection():
+    engine = _start_engine()
+    return engine
+
+
+def make_header_db_friendly(str_header):
+    return str_header.replace(" ", "_").lower()
+
+
+def make_value_db_friendly(value):
+    if value is None or pd.isna(value):
+        newval = "NULL"
+    else:
+        newval = str(value) if isinstance(value, int | float) else f"'{value}'"
+    return newval
+
+
+def retrieve_data(str_sql):
+    tg_engine = _start_engine()
+    # Retrieve data from Postgres
+    return pd.read_sql_query(str_sql, tg_engine)
+
+
+def retrieve_data_list(str_sql):
+    tg_engine = _start_engine()
+    # Retrieve data from Postgres
+    with tg_engine.connect() as con:
+        return con.execute(text(str_sql)).fetchall()
+
+
+def retrieve_single_result(str_sql):
+    tg_engine = _start_engine()
+    with tg_engine.connect() as con:
+        lstResult = con.execute(text(str_sql)).fetchone()
+        if lstResult:
+            return lstResult[0]
+
+
+def execute_sql(str_sql):
+    if str_sql > "":
+        tg_engine = _start_engine()
+        tg_engine.execute(text(str_sql))
+
+
+def execute_sql_raw(str_sql):
+    # For special cases where SQLAlchemy can't handle query syntax
+    if str_sql > "":
+        tg_engine = _start_engine()
+        con = tg_engine.raw_connection()
+        with con.cursor() as cur:
+            cur.execute(str_sql)
+        con.commit()
+
+
+def _get_df_edits(df_original: pd.DataFrame, df_edited: pd.DataFrame, lst_id_columns: list) -> tuple:
+    # Rows in df_edited that exist in df_original but have had any column changed
+    #  based on composite ID columns
+
+    # Merge the two dataframes based on the composite ID columns
+    merged_df = df_edited.merge(df_original, on=lst_id_columns, how="outer", indicator=True, suffixes=("", "_original"))
+    # Filter the merged dataframe to only keep rows that are changed
+    # Step 1: Filter rows that exist in both dataframes
+    both_rows = merged_df[merged_df["_merge"] == "both"]
+
+    # Step 2: Identify changed rows
+    def has_changes(row):
+        for col in df_original.columns:
+            # Skip the ID columns
+            if col in lst_id_columns:
+                continue
+            if row[col] != row[col + "_original"]:
+                return True
+        return False
+
+    changed_rows_mask = both_rows.apply(has_changes, axis=1)
+
+    # Step 3: Combine the filters
+    changed_rows = both_rows[changed_rows_mask]
+
+    # All rows in df_edited that are newly created and don't exist in df_original
+    new_rows = merged_df[merged_df["_merge"] == "left_only"].drop(
+        columns=["_merge"] + [col + "_original" for col in df_original.columns if col not in lst_id_columns]
+    )
+
+    # All rows in df_original that have been deleted from df_edited
+    deleted_rows = merged_df[merged_df["_merge"] == "right_only"][df_original.columns]
+
+    return changed_rows, new_rows, deleted_rows
+
+
+def _gen_df_update_sql(
+    changed_rows: pd.DataFrame, table_name: str, lst_id_columns: list, no_update_columns: list
+) -> list:
+    # Generate a list of SQL UPDATE statements based on the changed rows.
+
+    # Extract the original column names by removing the "_original" suffix
+    original_columns = [col.replace("_original", "") for col in changed_rows.columns if col.endswith("_original")]
+    # Drop columns we aren't updating from list
+    update_columns = [col for col in original_columns if col not in no_update_columns]
+
+    # Generate SQL UPDATE statements
+    sql_statements = []
+    for _, row in changed_rows.iterrows():
+        set_statements = []
+        for col in update_columns:
+            # If the value is different from the original value
+            if row[col] != row[col + "_original"]:
+                value = make_value_db_friendly(row[col])
+                set_statements.append(f"{col} = {value}")
+
+        # Handle composite keys for the WHERE clause
+        where_statements = []
+        for col in lst_id_columns:
+            value = make_value_db_friendly(row[col])
+            # value = f"'{row[col]}'" if isinstance(row[col], str) else row[col]
+            where_statements.append(f"{col} = {value}")
+
+        update_statement = f"UPDATE {get_schema()}.{table_name} SET {', '.join(set_statements)} WHERE {' AND '.join(where_statements)};"
+        sql_statements.append(update_statement)
+
+    return sql_statements
+
+
+def _gen_df_delete_sql(deleted_rows: pd.DataFrame, table_name: str, lst_id_columns: list) -> list:
+    # Generate a list of SQL DELETE statements based on the deleted rows.
+
+    # Generate SQL DELETE statements
+    sql_statements = []
+    for _, row in deleted_rows.iterrows():
+        # Handle composite keys for the WHERE clause
+        where_statements = []
+        for col in lst_id_columns:
+            value = make_value_db_friendly(row[col])
+            # value = f"'{row[col]}'" if isinstance(row[col], str) else row[col]
+            where_statements.append(f"{col} = {value}")
+
+        delete_statement = f"DELETE FROM {get_schema()}.{table_name} WHERE {' AND '.join(where_statements)};"
+        sql_statements.append(delete_statement)
+
+    return sql_statements
+
+
+def _gen_insert_sql(
+    new_rows: pd.DataFrame,
+    table_name: str,
+    lst_id_columns: list,
+    no_update_columns: list,
+    dct_hard_default_columns: dict,
+) -> str:
+    # Generate a SQL INSERT statement for the new rows, ensuring strings are properly quoted.
+
+    # Remove the id column as it will be generated by the server
+    if lst_id_columns:
+        new_rows = new_rows.drop(columns=lst_id_columns)
+    if no_update_columns:
+        # Remove columns we aren't updating
+        new_rows = new_rows.drop(columns=no_update_columns)
+    if dct_hard_default_columns:
+        # Add and default all columns
+        new_rows = new_rows.assign(**dct_hard_default_columns)
+
+    # Generate column names and values for the INSERT statement
+    columns = ", ".join(new_rows.columns)
+
+    # Ensure strings are quoted
+    values = []
+    for _, row in new_rows.iterrows():
+        row_values = []
+        for val in row:
+            row_values.append(make_value_db_friendly(val))
+            # if isinstance(val, str):
+            #     row_values.append(f"'{val}'")
+            # else:
+            #     row_values.append(str(val))
+        values.append(f"({', '.join(row_values)})")
+
+    if values:
+        values_str = ", ".join(values)
+        # Construct the SQL INSERT statement
+        sql_statement = f"INSERT INTO {get_schema()}.{table_name} ({columns}) VALUES {values_str};"
+        return sql_statement
+
+
+def apply_df_edits(df_original, df_edited, str_table, lst_id_columns, no_update_columns, dct_hard_default_columns):
+    booStatus = False
+    df_changed, df_new, df_deleted = _get_df_edits(df_original, df_edited, lst_id_columns)
+
+    # Generate SQL UPDATE statements
+    lst_update_SQL = _gen_df_update_sql(df_changed, str_table, lst_id_columns, no_update_columns)
+    if lst_update_SQL:
+        for str_sql in lst_update_SQL:
+            execute_sql(str_sql)
+            booStatus = True
+    # Generate SQL DELETE statements
+    lst_delete_SQL = _gen_df_delete_sql(df_deleted, str_table, lst_id_columns)
+    if lst_delete_SQL:
+        for str_sql in lst_delete_SQL:
+            execute_sql(str_sql)
+            booStatus = True
+    # Generate SQL INSERT statements
+    str_insert_sql = _gen_insert_sql(df_new, str_table, lst_id_columns, no_update_columns, dct_hard_default_columns)
+    if str_insert_sql:
+        execute_sql(str_insert_sql)
+        booStatus = True
+
+    return booStatus
+
+
+def _start_target_db_engine(flavor, host, port, db_name, user, password, url, connect_by_url, connect_by_key, private_key, private_key_passphrase):
+    connection_params = {
+        "flavor": flavor if flavor != "redshift" else "postgresql",
+        "user": user,
+        "host": host,
+        "port": port,
+        "dbname": db_name,
+        "url": url,
+        "connect_by_url": connect_by_url,
+        "connect_by_key": connect_by_key,
+        "private_key": private_key,
+        "private_key_passphrase": private_key_passphrase,
+        "dbschema": None,
+    }
+    flavor_service = get_flavor_service(flavor)
+    flavor_service.init(connection_params)
+    connection_string = flavor_service.get_connection_string(password)
+    connect_args = {"connect_timeout": 3600}
+    connect_args.update(flavor_service.get_connect_args())
+    return create_engine(connection_string, connect_args=connect_args)
+
+
+def retrieve_target_db_data(flavor, host, port, db_name, user, password, url, connect_by_url, connect_by_key, private_key, private_key_passphrase, sql_query, decrypt=False):
+    if decrypt:
+        password = DecryptText(password)
+    db_engine = _start_target_db_engine(flavor, host, port, db_name, user, password, url, connect_by_url, connect_by_key, private_key, private_key_passphrase)
+    with db_engine.connect() as connection:
+        query_result = connection.execute(text(sql_query))
+        return query_result.fetchall()
+
+
+def retrieve_target_db_df(flavor, host, port, db_name, user, password, sql_query, url, connect_by_url, connect_by_key, private_key, private_key_passphrase):
+    if password:
+        password = DecryptText(password)
+    db_engine = _start_target_db_engine(flavor, host, port, db_name, user, password, url, connect_by_url, connect_by_key, private_key, private_key_passphrase)
+    return pd.read_sql_query(text(sql_query), db_engine)