Merge pull request #10 from remoteoss/fix-column-name-renaming

fix: add option to not replace chars
remoteoss · Aug 20, 2024 · a558c8b · a558c8b
2 parents 6c7e28a + a9a9157
commit a558c8b
Showing 1 changed file with 15 additions and 7 deletions.
diff --git a/snowflake_utils/models.py b/snowflake_utils/models.py
@@ -56,11 +56,17 @@ class TableStructure(BaseModel):
     columns: dict = [str, Column]
 
     @property
-    def parsed_columns(self):
-        return ", ".join(
-            f'"{str.upper(k).strip().replace("-","_")}" {v.data_type}'
-            for k, v in self.columns.items()
-        )
+    def parsed_columns(self, replace_chars: bool = False) -> str:
+        if replace_chars:
+            return ", ".join(
+                f'"{str.upper(k).strip().replace("-","_")}" {v.data_type}'
+                for k, v in self.columns.items()
+            )
+        else:
+            return ", ".join(
+                f'"{str.upper(k).strip()}" {v.data_type}'
+                for k, v in self.columns.items()
+            )
 
     def parse_from_json(self):
         raise NotImplementedError("Not implemented yet")
@@ -406,8 +412,10 @@ def current_tags(self) -> dict[str, dict[str, str]]:
         tags = defaultdict(dict)
         with connect() as connection:
             cursor = connection.cursor()
-            cursor.execute(f"""select lower(column_name) as column_name, lower(tag_name) as tag_name, tag_value
-                from table(information_schema.tag_references_all_columns('{self.fqn}', 'table'))""")
+            cursor.execute(
+                f"""select lower(column_name) as column_name, lower(tag_name) as tag_name, tag_value
+                from table(information_schema.tag_references_all_columns('{self.fqn}', 'table'))"""
+            )
             for column_name, tag_name, tag_value in cursor.fetchall():
                 tags[column_name][tag_name] = tag_value
         return tags