Add MariaDB integration with native VECTOR support

Implement MariaDB database connector and vector store with native VECTOR(384) data type support for efficient similarity search. Includes improvements to SQL extraction and DDL retrieval.

Author: imash

mindsql/_helper/helper.py

@@ -62,13 +62,21 @@ def log_and_return(extracted_sql: str) -> str:
         log.info(LOG_AND_RETURN_CONSTANT.format(llm_response, extracted_sql))
         return extracted_sql
 
+    if "SQLQuery:" in llm_response:
+        sql_part = llm_response.split("SQLQuery:", 1)[1].strip()
+        if "\n\n" in sql_part:
+            sql_part = sql_part.split("\n\n")[0].strip()
+        return log_and_return(sql_part)
+    
     sql_match = re.search(r"```(sql)?\n(.+?)```", llm_response, re.DOTALL)
     if sql_match:
         return log_and_return(sql_match.group(2).replace("`", ""))
+    
     elif has_select_and_semicolon(llm_response):
         start_sql = llm_response.find("SELECT")
         end_sql = llm_response.find(";")
         return log_and_return(llm_response[start_sql:end_sql + 1].replace("`", ""))
+    
     return llm_response
 
 

mindsql/_utils/constants.py

@@ -18,6 +18,9 @@
 MYSQL_SHOW_DATABASE_QUERY = "SHOW DATABASES;"
 MYSQL_DB_TABLES_INFO_SCHEMA_QUERY = "SELECT table_name FROM information_schema.tables WHERE table_schema = '{}';"
 MYSQL_SHOW_CREATE_TABLE_QUERY = "SHOW CREATE TABLE `{}`;"
+MARIADB_SHOW_DATABASE_QUERY = "SHOW DATABASES;"
+MARIADB_DB_TABLES_INFO_SCHEMA_QUERY = "SELECT table_name FROM information_schema.tables WHERE table_schema = '{}';"
+MARIADB_SHOW_CREATE_TABLE_QUERY = "SHOW CREATE TABLE `{}`;"
 POSTGRESQL_SHOW_DATABASE_QUERY = "SELECT datname as DATABASE_NAME FROM pg_database WHERE datistemplate = false;"
 POSTGRESQL_DB_TABLES_INFO_SCHEMA_QUERY = "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' AND table_catalog = '{db}';"
 ERROR_DOWNLOADING_SQLITE_DB_CONSTANT = "Error downloading sqlite db: {}"

mindsql/_utils/prompts.py

@@ -1,6 +1,10 @@
-DEFAULT_PROMPT: str = """As a {dialect_name} expert, your task is to generate SQL queries based on user questions. Ensure that your {dialect_name} queries are syntactically correct and tailored to the user's inquiry. Retrieve at most 10 results using the LIMIT clause and order them for relevance. Avoid querying for all columns from a table. Select only the necessary columns wrapped in backticks (`). Use CURDATE() to handle 'today' queries and employ the LIKE clause for precise matches in {dialect_name}. Carefully consider column names and their respective tables to avoid querying non-existent columns. Stop after delivering the SQLQuery, avoiding follow-up questions.
+DEFAULT_PROMPT: str = """As a {dialect_name} expert, your task is to generate accurate SQL queries based on user questions and the provided table schemas.
 
-Follow this format:
+Make sure you look at the DDL statements to understand what columns are available in each table. When someone asks to "show" or "display" a specific table, just select from that table directly. Only query columns that actually exist - this is really important because querying non-existent columns will cause errors. Pick the columns that make sense for answering what the user is asking about.
+
+Always wrap your table and column names with backticks (`) since that's the {dialect_name} way of handling identifiers. Try to add LIMIT 10 to your queries so we don't return too much data at once. If it makes sense, use ORDER BY to sort the results in a logical way. For queries about "today", you can use the CURDATE() function. Also make sure any filter values match exactly what's in the schema.
+
+Your response should follow this format:
 Question: User's question here
 SQLQuery: Your SQL query without preamble
 
@@ -54,6 +58,5 @@
     - Ensure that the code is well-commented for readability and syntactically correct.
     """
 
-SQL_EXCEPTION_RESPONSE = """Apologies for the inconvenience! 🙏 It seems the database is currently experiencing a bit 
-of a hiccup and isn't cooperating as we'd like. 🤖"""
+SQL_EXCEPTION_RESPONSE = """Apologies for the inconvenience! It seems the database is currently experiencing a bit of a hiccup and isn't cooperating as we'd like."""
 

mindsql/core/mindsql_core.py

@@ -49,7 +49,7 @@ def create_database_query(self, question: str, connection, tables: list, **kwarg
         question_sql_list = self.vectorstore.retrieve_relevant_question_sql(question, **kwargs)
         prompt = self.build_sql_prompt(question=question, connection=connection, question_sql_list=question_sql_list,
                                        tables=tables, **kwargs)
-        log.info(prompt)
+        # log.info(prompt)
         llm_response = self.llm.invoke(prompt, **kwargs)
         return _helper.helper.extract_sql(llm_response)
 
@@ -176,13 +176,25 @@ def __get_ddl_statements(self, connection: any, tables: list[str], question: str
         Returns:
             list[str]: The list of DDL statements.
         """
+        vector_ddls = []
+        try:
+            vector_ddls = self.vectorstore.retrieve_relevant_ddl(question, **kwargs)
+        except Exception as e:
+            log.info(f"Vector store retrieval failed: {e}")
+        
+        if vector_ddls and len(vector_ddls) > 0:
+            return vector_ddls
+        
         if tables and connection:
             ddl_statements = []
             for table_name in tables:
-                ddl_statements.append(self.database.get_ddl(connection=connection, table_name=table_name))
-        else:
-            ddl_statements = self.vectorstore.retrieve_relevant_ddl(question, **kwargs)
-        return ddl_statements
+                try:
+                    ddl_statements.append(self.database.get_ddl(connection=connection, table_name=table_name))
+                except Exception as e:
+                    log.info(f"Failed to get DDL for table {table_name}: {e}")
+            return ddl_statements
+        
+        return []
 
     def ask_db(self, connection, question: Union[str, None] = None, table_names: list = None, visualize: bool = False,
                **kwargs) -> dict:

mindsql/databases/__init__.py

@@ -1,4 +1,5 @@
 from .idatabase import IDatabase
+from .mariadb import MariaDB
 from .mysql import MySql
 from .postgres import Postgres
 from .sqlite import Sqlite

mindsql/databases/mariadb.py

@@ -0,0 +1,99 @@
+from typing import List
+from urllib.parse import urlparse
+
+import mariadb
+import pandas as pd
+
+from .._utils import logger
+from .._utils.constants import SUCCESSFULLY_CONNECTED_TO_DB_CONSTANT, ERROR_CONNECTING_TO_DB_CONSTANT, \
+    INVALID_DB_CONNECTION_OBJECT, ERROR_WHILE_RUNNING_QUERY, MARIADB_DB_TABLES_INFO_SCHEMA_QUERY, \
+    MARIADB_SHOW_DATABASE_QUERY, MARIADB_SHOW_CREATE_TABLE_QUERY, CONNECTION_ESTABLISH_ERROR_CONSTANT
+from . import IDatabase
+
+log = logger.init_loggers("MariaDB")
+
+
+class MariaDB(IDatabase):
+    def create_connection(self, url: str, **kwargs) -> any:
+        url = urlparse(url)
+        try:
+            connection_params = {
+                'host': url.hostname,
+                'port': url.port or int(kwargs.get('port', 3306)),
+                'user': url.username,
+                'password': url.password,
+                'database': url.path.lstrip('/') if url.path else None,
+                'autocommit': True,
+            }
+            
+            connection_params = {k: v for k, v in connection_params.items() if v is not None}
+            connection_params.update({k: v for k, v in kwargs.items() if k not in ['port']})
+            
+            conn = mariadb.connect(**connection_params)
+            log.info(SUCCESSFULLY_CONNECTED_TO_DB_CONSTANT.format("MariaDB"))
+            return conn
+
+        except mariadb.Error as e:
+            log.info(ERROR_CONNECTING_TO_DB_CONSTANT.format("MariaDB", str(e)))
+            return None
+
+    def validate_connection(self, connection: any) -> None:
+        if connection is None:
+            raise ValueError(CONNECTION_ESTABLISH_ERROR_CONSTANT)
+        if not hasattr(connection, 'cursor'):
+            raise ValueError(INVALID_DB_CONNECTION_OBJECT.format("MariaDB"))
+
+    def execute_sql(self, connection, sql: str) -> pd.DataFrame:
+        try:
+            self.validate_connection(connection)
+            cursor = connection.cursor()
+            cursor.execute(sql)
+            
+            if sql.strip().upper().startswith(('CREATE', 'INSERT', 'UPDATE', 'DELETE', 'DROP', 'ALTER')):
+                connection.commit()
+                cursor.close()
+                return pd.DataFrame()
+            
+            results = cursor.fetchall()
+            if cursor.description:
+                column_names = [i[0] for i in cursor.description]
+                df = pd.DataFrame(results, columns=column_names)
+            else:
+                df = pd.DataFrame()
+            cursor.close()
+            return df
+        except mariadb.Error as e:
+            log.info(ERROR_WHILE_RUNNING_QUERY.format(e))
+            return pd.DataFrame()
+
+    def get_databases(self, connection) -> List[str]:
+        try:
+            self.validate_connection(connection)
+            df_databases = self.execute_sql(connection=connection, sql=MARIADB_SHOW_DATABASE_QUERY)
+        except Exception as e:
+            log.info(e)
+            return []
+        return df_databases["Database"].unique().tolist()
+
+    def get_table_names(self, connection, database: str) -> pd.DataFrame:
+        self.validate_connection(connection)
+        df_tables = self.execute_sql(connection, MARIADB_DB_TABLES_INFO_SCHEMA_QUERY.format(database))
+        return df_tables
+
+    def get_all_ddls(self, connection, database: str) -> pd.DataFrame:
+        self.validate_connection(connection)
+        df_tables = self.get_table_names(connection, database)
+        df_ddl = pd.DataFrame(columns=['Table', 'DDL'])
+        for index, row in df_tables.iterrows():
+            table_name = row.get('TABLE_NAME') or row.get('table_name')
+            if table_name:
+                ddl_df = self.get_ddl(connection, table_name)
+                df_ddl = df_ddl._append({'Table': table_name, 'DDL': ddl_df}, ignore_index=True)
+        return df_ddl
+
+    def get_ddl(self, connection: any, table_name: str, **kwargs) -> str:
+        ddl_df = self.execute_sql(connection, MARIADB_SHOW_CREATE_TABLE_QUERY.format(table_name))
+        return ddl_df["Create Table"].iloc[0]
+
+    def get_dialect(self) -> str:
+        return 'mysql'

mindsql/vectorstores/__init__.py

@@ -2,3 +2,4 @@
 from .chromadb import ChromaDB
 from .faiss_db import Faiss
 from .qdrant import Qdrant
+from .mariadb_vector import MariaDBVectorStore