From c4c60707b67c82905557b88e28aea4c0908dbc64 Mon Sep 17 00:00:00 2001
From: Vitaly Markov <wild.desu@gmail.com>
Date: Wed, 1 Nov 2023 10:16:39 +0000
Subject: [PATCH] Introduce custom application and --query-tag CLI option

---
 CHANGELOG.md             |  9 ++++++++-
 snowddl/app/base.py      | 15 +++++++++++++++
 snowddl/config.py        |  2 +-
 snowddl/resolver/view.py |  4 +++-
 snowddl/version.py       |  2 +-
 5 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index baa227c..be82d73 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,8 +1,15 @@
 # Changelog
 
+## [0.21.0] - 2023-11-01
+
+- Introduced custom value for application option (`SnowDDL <version>`) while opening Snowflake connection. Now it should be possible to find sessions created by SnowDDL using `SESSIONS` system view.
+- Added `--query-tag` CLI option to set custom `QUERY_TAG` session parameter.
+- Fixed pydantic deprecation warning related to `__fields__`.
+- Added explicit `.close()` call for Snowflake connection after execution of CLI commands. It should help to terminate SnowDDL sessions earlier, regardless of `CLIENT_SESSION_KEEP_ALIVE` parameter.
+
 ## [0.20.1] - 2023-10-20
 
-- Add additional debug logs for `VIEW` resolver in attempt to diagnose rare unnecessary re-creation problem.
+- Added additional debug logs for `VIEW` resolver in attempt to diagnose rare unnecessary re-creation problem.
 
 ## [0.20.0] - 2023-10-13
 
diff --git a/snowddl/app/base.py b/snowddl/app/base.py
index 13ce490..1f953b6 100644
--- a/snowddl/app/base.py
+++ b/snowddl/app/base.py
@@ -18,6 +18,9 @@
 
 
 class BaseApp:
+    application_name = "SnowDDL"
+    application_version = __version__
+
     parser_sequence = default_parser_sequence
     resolver_sequence = default_resolver_sequence
 
@@ -100,6 +103,11 @@ def init_arguments_parser(self):
         parser.add_argument(
             "--max-workers", help="Maximum number of workers to resolve objects in parallel", default=None, type=int
         )
+        parser.add_argument(
+            "--query-tag",
+            help="Add QUERY_TAG to all queries produced by SnowDDL",
+            default=environ.get("SNOWFLAKE_QUERY_TAG"),
+        )
 
         # Logging
         parser.add_argument(
@@ -363,6 +371,7 @@ def get_connection(self):
             "user": self.args["u"],
             "role": self.args["r"],
             "warehouse": self.args["w"],
+            "application": f"{self.application_name} {self.application_version}",
         }
 
         if self.args.get("k"):
@@ -381,6 +390,11 @@ def get_connection(self):
         else:
             options["password"] = self.args["p"]
 
+        if self.args.get("query_tag"):
+            options["session_parameters"] = {
+                "QUERY_TAG": self.args.get("query_tag"),
+            }
+
         return connect(**options)
 
     def execute(self):
@@ -407,6 +421,7 @@ def execute(self):
 
                     error_count += len(resolver.errors)
 
+            self.engine.connection.close()
             self.output_engine_stats()
 
             if self.args.get("show_sql"):
diff --git a/snowddl/config.py b/snowddl/config.py
index 40e1e53..6376814 100644
--- a/snowddl/config.py
+++ b/snowddl/config.py
@@ -31,7 +31,7 @@ def get_blueprints_by_type_and_pattern(self, cls: Type[T_Blueprint], pattern: st
         type_blueprints = self.blueprints.get(cls, {})
 
         # Add env prefix to pattern IF blueprint supports it
-        if "full_name" in cls.__fields__ and issubclass(cls.__fields__["full_name"].annotation, AbstractIdentWithPrefix):
+        if "full_name" in cls.model_fields and issubclass(cls.model_fields["full_name"].annotation, AbstractIdentWithPrefix):
             pattern = f"{self.env_prefix}{pattern}"
 
         # Use Unix-style wildcards if any special characters detected in pattern
diff --git a/snowddl/resolver/view.py b/snowddl/resolver/view.py
index 0d90399..901df30 100644
--- a/snowddl/resolver/view.py
+++ b/snowddl/resolver/view.py
@@ -67,7 +67,9 @@ def compare_object(self, bp: ViewBlueprint, row: dict):
                     },
                 )
             except SnowDDLExecuteError as e:
-                self.engine.logger.debug(f"View [{bp.full_name}] caused describe error [{e.snow_exc.errno}]: {e.snow_exc.raw_msg}")
+                self.engine.logger.debug(
+                    f"View [{bp.full_name}] caused describe error [{e.snow_exc.errno}]: {e.snow_exc.raw_msg}"
+                )
             else:
                 # Comments on views are broken and must be applied separately
                 if bp.comment != row["comment"]:
diff --git a/snowddl/version.py b/snowddl/version.py
index ac82404..6a726d8 100644
--- a/snowddl/version.py
+++ b/snowddl/version.py
@@ -1 +1 @@
-__version__ = "0.20.1"
+__version__ = "0.21.0"