Add sanitizier

Author: BenConstable9

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/databricks_sql.py

@@ -56,6 +56,19 @@ def invalid_identifiers(self) -> list[str]:
             "SHOW DATABASES",
         ]
 
+    def sanitize_identifier(self, identifier: str) -> str:
+        """Sanitize the identifier to ensure it is valid.
+
+        Args:
+        ----
+            identifier (str): The identifier to sanitize.
+
+        Returns:
+        -------
+            str: The sanitized identifier.
+        """
+        return f"`{identifier}`"
+
     async def query_execution(
         self,
         sql_query: Annotated[

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/postgres_sql.py

@@ -48,6 +48,19 @@ def invalid_identifiers(self) -> list[str]:
             "PGP_PUB_DECRYPT()",  # (from pgcrypto extension) Asymmetric decryption function
         ]
 
+    def sanitize_identifier(self, identifier: str) -> str:
+        """Sanitize the identifier to ensure it is valid.
+
+        Args:
+        ----
+            identifier (str): The identifier to sanitize.
+
+        Returns:
+        -------
+            str: The sanitized identifier.
+        """
+        return f'"{identifier}"'
+
     async def query_execution(
         self,
         sql_query: Annotated[str, "The SQL query to run against the database."],

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/snowflake_sql.py

@@ -72,6 +72,19 @@ def invalid_identifiers(self) -> list[str]:
             "QUERY_MEMORY_USAGE",
         ]
 
+    def sanitize_identifier(self, identifier: str) -> str:
+        """Sanitize the identifier to ensure it is valid.
+
+        Args:
+        ----
+            identifier (str): The identifier to sanitize.
+
+        Returns:
+        -------
+            str: The sanitized identifier.
+        """
+        return f'"{identifier}"'
+
     async def query_execution(
         self,
         sql_query: Annotated[

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/sql.py

@@ -11,6 +11,7 @@
 from jinja2 import Template
 import json
 from text_2_sql_core.utils.database import DatabaseEngineSpecificFields
+import re
 
 
 class SqlConnector(ABC):
@@ -40,19 +41,16 @@ def __init__(self):
     @abstractmethod
     def engine_specific_rules(self) -> str:
         """Get the engine specific rules."""
-        pass
 
     @property
     @abstractmethod
     def invalid_identifiers(self) -> list[str]:
         """Get the invalid identifiers upon which a sql query is rejected."""
-        pass
 
     @property
     @abstractmethod
     def engine_specific_fields(self) -> list[str]:
         """Get the engine specific fields."""
-        pass
 
     @property
     def excluded_engine_specific_fields(self):
@@ -85,6 +83,19 @@ async def query_execution(
             list[dict]: The results of the SQL query.
         """
 
+    @abstractmethod
+    def sanitize_identifier(self, identifier: str) -> str:
+        """Sanitize the identifier to ensure it is valid.
+
+        Args:
+        ----
+            identifier (str): The identifier to sanitize.
+
+        Returns:
+        -------
+            str: The sanitized identifier.
+        """
+
     async def get_column_values(
         self,
         text: Annotated[
@@ -204,6 +215,26 @@ async def query_execution_with_limit(
                 default=str,
             )
 
+    def clean_query(self, sql_query: str) -> str:
+        """Clean the SQL query to ensure it is valid.
+
+        Args:
+        ----
+            sql_query (str): The SQL query to clean.
+
+        Returns:
+        -------
+            str: The cleaned SQL query.
+        """
+        single_line_query = sql_query.strip().replace("\n", " ")
+        cleaned_query = re.sub(
+            r'(?<!["\[\w])\b([a-zA-Z_][a-zA-Z0-9_-]*)\b(?!["\]])',
+            lambda m: self.sanitize_identifier(m.group(1)),
+            single_line_query,
+        )
+
+        return cleaned_query
+
     async def query_validation(
         self,
         sql_query: Annotated[
@@ -213,7 +244,7 @@ async def query_validation(
     ) -> Union[bool | list[dict]]:
         """Validate the SQL query."""
         try:
-            cleaned_query = sql_query.strip().replace("\n", " ")
+            cleaned_query = self.clean_query(sql_query)
             logging.info("Validating SQL Query: %s", cleaned_query)
             parsed_queries = sqlglot.parse(
                 cleaned_query,
@@ -249,14 +280,12 @@ def handle_node(node):
                     detected_invalid_identifiers.append(identifier)
 
             if len(detected_invalid_identifiers) > 0:
-                logging.error(
-                    "SQL Query contains invalid identifiers: %s",
-                    detected_invalid_identifiers,
-                )
-                return (
+                error_message = (
                     "SQL Query contains invalid identifiers: %s"
                     % detected_invalid_identifiers
                 )
+                logging.error(error_message)
+                return False, None, error_message
 
         except sqlglot.errors.ParseError as e:
             logging.error("SQL Query is invalid: %s", e.errors)

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/sqlite_sql.py

@@ -44,6 +44,19 @@ def engine_specific_fields(self) -> list[str]:
         """Get the engine specific fields."""
         return []  # SQLite doesn't use warehouses, catalogs, or separate databases
 
+    def sanitize_identifier(self, identifier: str) -> str:
+        """Sanitize the identifier to ensure it is valid.
+
+        Args:
+        ----
+            identifier (str): The identifier to sanitize.
+
+        Returns:
+        -------
+            str: The sanitized identifier.
+        """
+        return f'"{identifier}"'
+
     async def query_execution(
         self,
         sql_query: Annotated[

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/tsql_sql.py

@@ -65,6 +65,19 @@ def invalid_identifiers(self) -> list[str]:
             "VERSION",
         ]
 
+    def sanitize_identifier(self, identifier: str) -> str:
+        """Sanitize the identifier to ensure it is valid.
+
+        Args:
+        ----
+            identifier (str): The identifier to sanitize.
+
+        Returns:
+        -------
+            str: The sanitized identifier.
+        """
+        return f"[{identifier}]"
+
     async def query_execution(
         self,
         sql_query: Annotated[