Merge pull request #49 from DataKitchen/release/4.26.1

Release/4.26.1

Author: datakitchen-devops

pyproject.toml

@@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dataops-testgen"
-version = "4.22.2"
+version = "4.26.1"
 description = "DataKitchen's Data Quality DataOps TestGen"
 authors = [
     { "name" = "DataKitchen, Inc.", "email" = "[email protected]" },
@@ -58,6 +58,7 @@ dependencies = [
     "pydantic==1.10.13",
     "streamlit-pydantic==0.6.0",
     "cron-converter==1.2.1",
+    "cron-descriptor==2.0.5",
 
     # Pinned to match the manually compiled libs or for security
     "pyarrow==18.1.0",

testgen/commands/run_execute_tests.py

@@ -125,9 +125,12 @@ def run_execution_steps_in_background(project_code, test_suite):
     if settings.IS_DEBUG:
         LOG.info(msg + ". Running in debug mode (new thread instead of new process).")
         empty_cache()
+        username = None
+        if session.auth:
+            username = session.auth.user_display
         background_thread = threading.Thread(
             target=run_execution_steps,
-            args=(project_code, test_suite, session.auth.user_display),
+            args=(project_code, test_suite, username),
         )
         background_thread.start()
     else:

testgen/commands/run_launch_db_config.py

@@ -72,10 +72,10 @@ def _get_params_mapping() -> dict:
 
 
 @with_database_session
-def run_launch_db_config(delete_db: bool) -> None:
+def run_launch_db_config(delete_db: bool, drop_users_and_roles: bool = True) -> None:
     params_mapping = _get_params_mapping()
 
-    create_database(get_tg_db(), params_mapping, drop_existing=delete_db, drop_users_and_roles=True)
+    create_database(get_tg_db(), params_mapping, drop_existing=delete_db, drop_users_and_roles=drop_users_and_roles)
 
     queries = get_queries_for_command("dbsetup", params_mapping)
 
@@ -91,4 +91,8 @@ def run_launch_db_config(delete_db: bool) -> None:
             project_code=settings.PROJECT_KEY,
             table_groups_name=settings.DEFAULT_TABLE_GROUPS_NAME,
         )
-    ).save()
+    ).save()
+
+
+def get_app_db_params_mapping() -> dict:
+    return _get_params_mapping()

testgen/commands/run_profiling_bridge.py

@@ -10,6 +10,8 @@
 import testgen.common.process_service as process_service
 from testgen import settings
 from testgen.commands.queries.profiling_query import CProfilingSQL
+from testgen.commands.run_execute_tests import run_execution_steps_in_background
+from testgen.commands.run_generate_tests import run_test_gen_queries
 from testgen.commands.run_refresh_score_cards_results import run_refresh_score_cards_results
 from testgen.common import (
     date_service,
@@ -25,6 +27,7 @@
 from testgen.common.mixpanel_service import MixpanelService
 from testgen.common.models import with_database_session
 from testgen.common.models.connection import Connection
+from testgen.common.models.test_suite import TestSuite
 from testgen.ui.session import session
 
 LOG = logging.getLogger("testgen")
@@ -211,7 +214,7 @@ def run_profiling_in_background(table_group_id):
         empty_cache()
         background_thread = threading.Thread(
             target=run_profiling_queries,
-            args=(table_group_id, session.auth.user_display),
+            args=(table_group_id, session.auth.user_display if session.auth else None),
         )
         background_thread.start()
     else:
@@ -238,6 +241,9 @@ def run_profiling_queries(table_group_id: str, username: str | None = None, spin
     profiling_run_id = str(uuid.uuid4())
 
     params = get_profiling_params(table_group_id)
+    needs_monitor_tests_generated = (
+        bool(params["monitor_test_suite_id"]) and not params["last_complete_profile_run_id"]
+    )
 
     LOG.info("CurrentStep: Initializing Query Generator")
     clsProfiling = CProfilingSQL(params["project_code"], connection.sql_flavor, minutes_offset=minutes_offset)
@@ -471,7 +477,24 @@ def run_profiling_queries(table_group_id: str, username: str | None = None, spin
             scoring_duration=(datetime.now(UTC) - end_time).total_seconds(),
         )
 
+        if needs_monitor_tests_generated:
+            _generate_monitor_tests(params["project_code"], table_group_id, params["monitor_test_suite_id"])
+
     return f"""
         Profiling completed {"with errors. Check log for details." if has_errors else "successfully."}
         Run ID: {profiling_run_id}
     """
+
+
+@with_database_session
+def _generate_monitor_tests(project_code: str, table_group_id: str, test_suite_id: str) -> None:
+    try:
+        monitor_test_suite = TestSuite.get(test_suite_id)
+        if not monitor_test_suite:
+            LOG.info("Skipping test generation on missing monitor test suite")
+        else:
+            LOG.info("Generating monitor tests")
+            run_test_gen_queries(table_group_id, monitor_test_suite.test_suite, "Monitor")
+            run_execution_steps_in_background(project_code, monitor_test_suite.test_suite)
+    except Exception:
+        LOG.exception("Error generating monitor tests")

testgen/commands/run_quick_start.py

@@ -3,7 +3,7 @@
 import click
 
 from testgen import settings
-from testgen.commands.run_launch_db_config import run_launch_db_config
+from testgen.commands.run_launch_db_config import get_app_db_params_mapping, run_launch_db_config
 from testgen.common.credentials import get_tg_schema
 from testgen.common.database.database_service import (
     create_database,
@@ -117,6 +117,14 @@ def run_quick_start(delete_target_db: bool) -> None:
     delete_db = True
     run_launch_db_config(delete_db)
 
+    click.echo("Seeding the application db")
+    app_db_params = get_app_db_params_mapping()
+    execute_db_queries(
+        [
+            (replace_params(read_template_sql_file("initial_data_seeding.sql", "quick_start"), app_db_params), app_db_params),
+        ],
+    )
+
     # Schema and Populate target db
     click.echo(f"Populating target db : {target_db_name}")
     execute_db_queries(

testgen/commands/run_refresh_score_cards_results.py

@@ -72,7 +72,7 @@ def run_refresh_score_cards_results(
                         history_entry.add_as_cutoff()
             definition.save()
             LOG.info(
-                "CurrentStep: Done rereshing scorecard %s in project %s",
+                "CurrentStep: Done refreshing scorecard %s in project %s",
                 definition.name,
                 definition.project_code,
             )

testgen/common/database/flavor/flavor_service.py

@@ -1,5 +1,6 @@
 from abc import abstractmethod
-from typing import Literal, TypedDict
+from typing import Any, Literal, TypedDict
+from urllib.parse import parse_qs, urlparse
 
 from testgen.common.encrypt import DecryptText
 
@@ -37,19 +38,21 @@ class FlavorService:
     private_key_passphrase = None
     http_path = None
     catalog = None
+    warehouse = None
 
     def init(self, connection_params: ConnectionParams):
-        self.url = connection_params.get("url", None)
+        self.url = connection_params.get("url") or ""
         self.connect_by_url = connection_params.get("connect_by_url", False)
-        self.username = connection_params.get("project_user")
-        self.host = connection_params.get("project_host")
-        self.port = connection_params.get("project_port")
-        self.dbname = connection_params.get("project_db")
+        self.username = connection_params.get("project_user") or ""
+        self.host = connection_params.get("project_host") or ""
+        self.port = connection_params.get("project_port") or ""
+        self.dbname = connection_params.get("project_db") or ""
         self.flavor = connection_params.get("sql_flavor")
         self.dbschema = connection_params.get("table_group_schema", None)
         self.connect_by_key = connection_params.get("connect_by_key", False)
-        self.http_path = connection_params.get("http_path", None)
-        self.catalog = connection_params.get("catalog", None)
+        self.http_path = connection_params.get("http_path") or ""
+        self.catalog = connection_params.get("catalog") or ""
+        self.warehouse = connection_params.get("warehouse") or ""
 
         password = connection_params.get("project_pw_encrypted", None)
         if isinstance(password, memoryview) or isinstance(password, bytes):
@@ -90,3 +93,45 @@ def get_connection_string_from_fields(self) -> str:
     @abstractmethod
     def get_connection_string_head(self) -> str:
         raise NotImplementedError("Subclasses must implement this method")
+
+    def get_parts_from_connection_string(self) -> dict[str, Any]:
+        if self.connect_by_url:
+            if not self.url:
+                return {}
+
+            parsed_url = urlparse(self.get_connection_string())
+            credentials, location = (
+                parsed_url.netloc if "@" in parsed_url.netloc else f"@{parsed_url.netloc}"
+            ).split("@")
+            username, password = (
+                credentials if ":" in credentials else f"{credentials}:"
+            ).split(":")
+            host, port = (
+                location if ":" in location else f"{location}:"
+            ).split(":")
+
+            database = (path_patrs[0] if (path_patrs := parsed_url.path.strip("/").split("/")) else "")
+
+            extras = {
+                param_name: param_values[0]
+                for param_name, param_values in parse_qs(parsed_url.query or "").items()
+            }
+
+            return {
+                "username": username,
+                "password": password,
+                "host": host,
+                "port": port,
+                "dbname": database,
+                **extras,
+            }
+
+        return {
+            "username": self.username,
+            "password": self.password,
+            "host": self.host,
+            "port": self.port,
+            "dbname": self.dbname,
+            "http_path": self.http_path,
+            "catalog": self.catalog,
+        }

testgen/common/database/flavor/snowflake_flavor_service.py

@@ -2,6 +2,7 @@
 
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
+from snowflake.sqlalchemy import URL
 
 from testgen.common.database.flavor.flavor_service import FlavorService
 
@@ -38,25 +39,27 @@ def get_connection_string_from_fields(self):
         #   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
+        account, _ = self.host.split(".", maxsplit=1) if "." in self.host else ("", "")
+        host = self.host
+        if ".snowflakecomputing.com" not in host:
+            host = f"{host}.snowflakecomputing.com"
 
-        raw_host = get_raw_host_name(self.host)
-        host = raw_host
-        if self.port != "443":
-            host += ":" + self.port
+        extra_params = {}
+        if self.warehouse:
+            extra_params["warehouse"] = self.warehouse
 
-        if self.connect_by_key:
-            return f"snowflake://{self.username}@{host}/{self.dbname}/{self.dbschema}"
-        else:
-            return f"snowflake://{self.username}:{quote_plus(self.password)}@{host}/{self.dbname}/{self.dbschema}"
+        connection_url = URL(
+            host=host,
+            port=int(self.port if str(self.port).isdigit() else 443),
+            account=account,
+            user=self.username,
+            password="" if self.connect_by_key else self.password,
+            database=self.dbname,
+            schema=self.dbschema or "",
+            **extra_params,
+        )
+
+        return connection_url
 
     def get_pre_connection_queries(self):
         return [

testgen/common/get_pipeline_parms.py

@@ -21,6 +21,8 @@ class ProfilingParams(BaseParams):
     profile_sample_min_count: int
     profile_do_pair_rules: str
     profile_pair_rule_pct: int
+    monitor_test_suite_id: str | None
+    last_complete_profile_run_id: str | None
 
 
 class TestGenerationParams(BaseParams):

testgen/common/mixpanel_service.py

@@ -57,7 +57,7 @@ def send_event(self, event_name, include_usage=False, **properties):
         properties.setdefault("instance_id", self.instance_id)
         properties.setdefault("edition", settings.DOCKER_HUB_REPOSITORY)
         properties.setdefault("version", settings.VERSION)
-        properties.setdefault("username", session.auth.user_display)
+        properties.setdefault("username", session.auth.user_display if session.auth else None)
         properties.setdefault("distinct_id", self.get_distinct_id(properties["username"]))
         if include_usage:
             properties.update(self.get_usage())