MindsDB 25.5.3.0__py3-none-any.whl → 25.5.4.0__py3-none-any.whl

mindsdb/integrations/handlers/snowflake_handler/snowflake_handler.py

@@ -1,14 +1,16 @@
 import psutil
 import pandas
 from pandas import DataFrame
+from pandas.api import types as pd_types
 from snowflake.sqlalchemy import snowdialect
 from snowflake import connector
 from snowflake.connector.errors import NotSupportedError
+from snowflake.connector.cursor import SnowflakeCursor, ResultMetadata
 
-from mindsdb.utilities import log
 from mindsdb_sql_parser.ast.base import ASTNode
 from mindsdb_sql_parser.ast import Select, Identifier
 
+from mindsdb.utilities import log
 from mindsdb.integrations.libs.base import DatabaseHandler
 from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender
 from mindsdb.integrations.libs.response import (
@@ -61,10 +63,105 @@ def _map_type(internal_type_name: str) -> MYSQL_DATA_TYPE:
         if internal_type_name in db_types_list:
             return mysql_data_type
 
-    logger.warning(f"Snowflake handler type mapping: unknown type: {internal_type_name}, use VARCHAR as fallback.")
+    logger.debug(f"Snowflake handler type mapping: unknown type: {internal_type_name}, use VARCHAR as fallback.")
     return MYSQL_DATA_TYPE.VARCHAR
 
 
+def _make_table_response(result: DataFrame, cursor: SnowflakeCursor) -> Response:
+    """Build response from result and cursor.
+    NOTE: Snowflake return only 'general' type in description, so look on result's
+          DF types and use types from description only if DF type is 'object'
+
+    Args:
+        result (DataFrame): result of the query.
+        cursor (SnowflakeCursor): cursor object.
+
+    Returns:
+        Response: response object.
+    """
+    description: list[ResultMetadata] = cursor.description
+    mysql_types: list[MYSQL_DATA_TYPE] = []
+    for column in description:
+        column_dtype = result[column.name].dtype
+        description_column_type = connector.constants.FIELD_ID_TO_NAME.get(column.type_code)
+        if pd_types.is_integer_dtype(column_dtype):
+            column_dtype_name = column_dtype.name
+            if column_dtype_name in ('int8', 'Int8'):
+                mysql_types.append(MYSQL_DATA_TYPE.TINYINT)
+            elif column_dtype in ('int16', 'Int16'):
+                mysql_types.append(MYSQL_DATA_TYPE.SMALLINT)
+            elif column_dtype in ('int32', 'Int32'):
+                mysql_types.append(MYSQL_DATA_TYPE.MEDIUMINT)
+            elif column_dtype in ('int64', 'Int64'):
+                mysql_types.append(MYSQL_DATA_TYPE.BIGINT)
+            else:
+                mysql_types.append(MYSQL_DATA_TYPE.INT)
+            continue
+        if pd_types.is_float_dtype(column_dtype):
+            column_dtype_name = column_dtype.name
+            if column_dtype_name in ('float16', 'Float16'):  # Float16 does not exists so far
+                mysql_types.append(MYSQL_DATA_TYPE.FLOAT)
+            elif column_dtype_name in ('float32', 'Float32'):
+                mysql_types.append(MYSQL_DATA_TYPE.FLOAT)
+            elif column_dtype_name in ('float64', 'Float64'):
+                mysql_types.append(MYSQL_DATA_TYPE.DOUBLE)
+            else:
+                mysql_types.append(MYSQL_DATA_TYPE.FLOAT)
+            continue
+        if pd_types.is_bool_dtype(column_dtype):
+            mysql_types.append(MYSQL_DATA_TYPE.BOOLEAN)
+            continue
+        if pd_types.is_datetime64_any_dtype(column_dtype):
+            mysql_types.append(MYSQL_DATA_TYPE.DATETIME)
+            series = result[column.name]
+            # snowflake use pytz.timezone
+            if series.dt.tz is not None and getattr(series.dt.tz, 'zone', 'UTC') != 'UTC':
+                series = series.dt.tz_convert('UTC')
+                result[column.name] = series.dt.tz_localize(None)
+            continue
+
+        if pd_types.is_object_dtype(column_dtype):
+            if description_column_type == 'TEXT':
+                # we can also check column.internal_size, if == 16777216 then it is TEXT, else VARCHAR(internal_size)
+                mysql_types.append(MYSQL_DATA_TYPE.TEXT)
+                continue
+            elif description_column_type == 'BINARY':
+                # if column.internal_size == 8388608 then BINARY, else VARBINARY(internal_size)
+                mysql_types.append(MYSQL_DATA_TYPE.BINARY)
+                continue
+            elif description_column_type == 'DATE':
+                mysql_types.append(MYSQL_DATA_TYPE.DATE)
+                continue
+            elif description_column_type == 'TIME':
+                mysql_types.append(MYSQL_DATA_TYPE.TIME)
+                continue
+
+        if description_column_type == 'FIXED':
+            if column.scale == 0:
+                mysql_types.append(MYSQL_DATA_TYPE.INT)
+            else:
+                # It is NUMBER, DECIMAL or NUMERIC with scale > 0
+                mysql_types.append(MYSQL_DATA_TYPE.FLOAT)
+            continue
+        elif description_column_type == 'REAL':
+            mysql_types.append(MYSQL_DATA_TYPE.FLOAT)
+            continue
+
+        mysql_types.append(MYSQL_DATA_TYPE.TEXT)
+
+    df = DataFrame(
+        result,
+        columns=[column.name for column in description],
+    )
+
+    return Response(
+        RESPONSE_TYPE.TABLE,
+        data_frame=df,
+        affected_rows=None,
+        mysql_types=mysql_types
+    )
+
+
 class SnowflakeHandler(DatabaseHandler):
     """
     This handler handles connection and execution of the Snowflake statements.
@@ -214,10 +311,7 @@ class SnowflakeHandler(DatabaseHandler):
                                 raise MemoryError('Not enought memory')
                         # endregion
                     if len(batches) > 0:
-                        response = Response(
-                            RESPONSE_TYPE.TABLE,
-                            pandas.concat(batches, ignore_index=True)
-                        )
+                        response = _make_table_response(result=pandas.concat(batches, ignore_index=True), cursor=cur)
                     else:
                         response = Response(
                             RESPONSE_TYPE.TABLE,

mindsdb/integrations/handlers/solr_handler/connection_args.py

@@ -36,11 +36,11 @@ connection_args = OrderedDict(
 )
 
 connection_args_example = OrderedDict(
-    username = "demo_user",
-    password =  "demo_password",
-    host = "127.0.0.1",
-    port = 8981,
-    server_path = "solr",
-    collection = "gettingstarted",
-    use_ssl = False,
+    username="demo_user",
+    password="demo_password",
+    host="127.0.0.1",
+    port=8981,
+    server_path="solr",
+    collection="gettingstarted",
+    use_ssl=False,
 )

mindsdb/integrations/handlers/solr_handler/solr_handler.py

@@ -17,6 +17,7 @@ from mindsdb.integrations.libs.response import (
 
 logger = log.getLogger(__name__)
 
+
 class SolrHandler(DatabaseHandler):
     """
     This handler handles connection and execution of the Solr SQL statements.
@@ -131,7 +132,7 @@ class SolrHandler(DatabaseHandler):
                     RESPONSE_TYPE.TABLE,
                     pd.DataFrame(
                         result,
-                        columns = columns
+                        columns=columns
                     )
                 )
             else:

mindsdb/integrations/handlers/solr_handler/tests/test_solr_handler.py

@@ -34,5 +34,6 @@ class SolrHandlerTest(unittest.TestCase):
         result = self.handler.query(query)
         assert result['type'] is RESPONSE_TYPE.TABLE
 
+
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()

mindsdb/integrations/handlers/sqlany_handler/sqlany_handler.py

@@ -18,6 +18,7 @@ from mindsdb.integrations.libs.response import (
 
 logger = log.getLogger(__name__)
 
+
 class SQLAnyHandler(DatabaseHandler):
     """
     This handler handles connection and execution of the SAP SQL Anywhere statements.
@@ -61,7 +62,7 @@ class SQLAnyHandler(DatabaseHandler):
             self.encryption = "SIMPLE"
         else:
             self.encryption = "NONE"
-        
+
         connection = sqlanydb.connect(
             host=self.host,
             userid=self.userid,
@@ -160,7 +161,7 @@ class SQLAnyHandler(DatabaseHandler):
         List all tables in SAP SQL Anywhere in the current schema
         """
 
-        return self.native_query(f"""
+        return self.native_query("""
             SELECT USER_NAME(ob.UID) AS SCHEMA_NAME
             , st.TABLE_NAME
             , st.TABLE_TYPE

mindsdb/integrations/handlers/sqlite_handler/sqlite_handler.py

@@ -20,6 +20,7 @@ from mindsdb.integrations.libs.response import (
 
 logger = log.getLogger(__name__)
 
+
 class SQLiteHandler(DatabaseHandler):
     """
     This handler handles connection and execution of the SQLite statements.

mindsdb/integrations/handlers/sqreamdb_handler/connection_args.py

@@ -20,7 +20,7 @@ connection_args = OrderedDict(
     port={
         'type': ARG_TYPE.INT,
         'description': 'Specify port to connect SQreamDB server'
-    }, 
+    },
     database={
         'type': ARG_TYPE.STR,
         'description': 'Specify database name  to connect SQreamDB server'

mindsdb/integrations/handlers/sqreamdb_handler/sqreamdb_handler.py

@@ -14,15 +14,14 @@ from mindsdb.integrations.libs.response import (
 import pandas as pd
 import pysqream as db
 
-from pysqream_sqlalchemy.dialect import SqreamDialect 
+from pysqream_sqlalchemy.dialect import SqreamDialect
 
 logger = log.getLogger(__name__)
 
 
 class SQreamDBHandler(DatabaseHandler):
 
-
-    name= 'sqreamdb'
+    name = 'sqreamdb'
 
     def __init__(self, name: str, connection_data: Optional[dict], **kwargs):
         """ Initialize the handler
@@ -32,34 +31,30 @@ class SQreamDBHandler(DatabaseHandler):
             **kwargs: arbitrary keyword arguments.
         """
         super().__init__(name)
-        
-        
+
         self.connection_data = connection_data
 
         self.connection = None
         self.is_connected = False
-          
- 
+
     def connect(self):
         """
         Handles the connection to a YugabyteSQL database insance.
         """
         if self.is_connected is True:
             return self.connection
-        
-        args={
+
+        args = {
             "database": self.connection_data.get('database'),
             "host": self.connection_data.get('host'),
             "port": self.connection_data.get('port'),
             "username": self.connection_data.get('user'),
             "password": self.connection_data.get('password'),
-            "clustered":self.connection_data.get('clustered',False),
-            "use_ssl":self.connection_data.get('use_ssl',False), 
-            "service":self.connection_data.get('service','sqream')
+            "clustered": self.connection_data.get('clustered', False),
+            "use_ssl": self.connection_data.get('use_ssl', False),
+            "service": self.connection_data.get('service', 'sqream')
         }
 
-
-
         connection = db.connect(**args)
 
         self.is_connected = True
@@ -103,9 +98,9 @@ class SQreamDBHandler(DatabaseHandler):
         with conn.cursor() as cur:
             try:
                 cur.execute(query)
-                   
-                if cur.rowcount >0 and query.upper().startswith('SELECT') :
-                    result = cur.fetchall() 
+
+                if cur.rowcount > 0 and query.upper().startswith('SELECT'):
+                    result = cur.fetchall()
                     response = Response(
                         RESPONSE_TYPE.TABLE,
                         data_frame=pd.DataFrame(
@@ -139,7 +134,7 @@ class SQreamDBHandler(DatabaseHandler):
 
     def get_tables(self) -> Response:
         """
-        List all tables in SQreamDB stored in 'sqream_catalog' 
+        List all tables in SQreamDB stored in 'sqream_catalog'
         """
 
         query = "SELECT table_name FROM sqream_catalog.tables"
@@ -147,8 +142,8 @@ class SQreamDBHandler(DatabaseHandler):
         return self.query(query)
 
     def get_columns(self, table_name):
-        query = f"""SELECT column_name, type_name 
-        FROM sqream_catalog.columns 
+        query = f"""SELECT column_name, type_name
+        FROM sqream_catalog.columns
         WHERE table_name = '{table_name}';
         """
         return self.query(query)

mindsdb/integrations/handlers/sqreamdb_handler/tests/test_sqreamdb_handler.py

@@ -31,15 +31,15 @@ class SQreamDBHandlerTest(unittest.TestCase):
     def test_3_get_tables(self):
         tables = self.handler.get_tables()
         assert tables.type is not RESPONSE_TYPE.ERROR
- 
+
     def test_4_select_query(self):
         query = "SELECT * FROM AUTHORS"
         result = self.handler.native_query(query)
         assert result.type is RESPONSE_TYPE.TABLE
 
     def test_5_check_connection(self):
-         self.handler.check_connection()
+        self.handler.check_connection()
+
 
-        
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()

mindsdb/integrations/handlers/stabilityai_handler/__init__.py

@@ -16,4 +16,4 @@ permanent = False
 
 __all__ = [
     'Handler', 'version', 'name', 'type', 'title', 'description', 'import_error', 'icon_path'
-]
+]

mindsdb/integrations/handlers/starrocks_handler/starrocks_handler.py

@@ -15,7 +15,6 @@ class StarRocksHandler(MysqlHandler):
         super().__init__(name, **kwargs)
 
 
-
 connection_args = OrderedDict(
     user={
         'type': ARG_TYPE.STR,

mindsdb/integrations/handlers/starrocks_handler/tests/test_starrocks_handler.py

@@ -2,6 +2,7 @@ import unittest
 from mindsdb.integrations.handlers.starrocks_handler.starrocks_handler import StarRocksHandler
 from mindsdb.integrations.libs.response import RESPONSE_TYPE
 
+
 class StarRocksHandlerTest(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
@@ -21,36 +22,33 @@ class StarRocksHandlerTest(unittest.TestCase):
 
     def test_1_connect(self):
         assert self.handler.connect()
-    
+
     def test_2_create_table(self):
         query = "CREATE Table IF NOT EXISTS Lover(name varchar(101));"
         result = self.handler.query(query)
-        assert result.type is not RESPONSE_TYPE.ERROR 
+        assert result.type is not RESPONSE_TYPE.ERROR
 
     def test_3_insert(self):
         query = "INSERT INTO LOVER VALUES('Shiv Shakti');"
         result = self.handler.query(query)
-        assert result.type is not RESPONSE_TYPE.ERROR 
+        assert result.type is not RESPONSE_TYPE.ERROR
 
     def test_4_native_query_select(self):
         query = "SELECT * FROM LOVER;"
         result = self.handler.query(query)
-        assert result.type is RESPONSE_TYPE.TABLE 
+        assert result.type is RESPONSE_TYPE.TABLE
 
     def test_5_get_tables(self):
         tables = self.handler.get_tables()
-        assert tables.type is  RESPONSE_TYPE.TABLE
+        assert tables.type is RESPONSE_TYPE.TABLE
 
     def test_6_get_columns(self):
         columns = self.handler.get_columns('LOVER')
-        
+
         query = "DROP Table IF  EXISTS Lover;"
-        result = self.handler.query(query)
+        self.handler.query(query)
         assert columns.type is not RESPONSE_TYPE.ERROR
 
 
 if __name__ == '__main__':
     unittest.main()
-
-
-    

mindsdb/integrations/handlers/statsforecast_handler/statsforecast_handler.py

@@ -99,7 +99,7 @@ class StatsForecastHandler(BaseMLEngine):
         time_settings = args["timeseries_settings"]
         using_args = args["using"]
         assert time_settings["is_timeseries"], "Specify time series settings in your query"
-        ###### store model args and time series settings in the model folder
+        # store model args and time series settings in the model folder
         model_args = {}
         model_args.update(using_args)
         model_args["target"] = target
@@ -130,7 +130,7 @@ class StatsForecastHandler(BaseMLEngine):
         sf = StatsForecast([model], freq=model_args["frequency"], df=training_df)
         fitted_models = sf.fit().fitted_
 
-        ###### persist changes to handler folder
+        # persist changes to handler folder
         self.model_storage.json_set("model_args", model_args)
         self.model_storage.file_set("training_df", dill.dumps(training_df))
         self.model_storage.file_set("fitted_models", dill.dumps(fitted_models))

mindsdb/integrations/handlers/strava_handler/strava_handler.py

@@ -7,9 +7,6 @@ from mindsdb.integrations.libs.response import (
 from mindsdb.utilities import log
 from mindsdb_sql_parser import parse_sql
 
-import requests
-import pandas as pd
-import json
 from collections import OrderedDict
 from mindsdb.integrations.libs.const import HANDLER_CONNECTION_ARG_TYPE as ARG_TYPE
 
@@ -17,6 +14,7 @@ from stravalib.client import Client
 
 logger = log.getLogger(__name__)
 
+
 class StravaHandler(APIHandler):
     """Strava handler implementation"""
 
@@ -39,7 +37,7 @@ class StravaHandler(APIHandler):
 
         strava_all_clubs_data = StravaAllClubsTable(self)
         self._register_table("all_clubs", strava_all_clubs_data)
-        
+
         strava_club_activites_data = StravaClubActivitesTable(self)
         self._register_table("club_activities", strava_club_activites_data)
 
@@ -54,7 +52,6 @@ class StravaHandler(APIHandler):
             return self.connection
 
         client = Client()
-        url = client.authorization_url(client_id=self.connection_data['strava_client_id'], redirect_uri='http://127.0.0.1:5000/authorization')
         client.access_token = self.connection_data['strava_access_token']
         self.connection = client
 
@@ -68,7 +65,6 @@ class StravaHandler(APIHandler):
             Status confirmation
         """
         response = StatusResponse(False)
-        need_to_close = self.is_connected is False
 
         try:
             self.connect()
@@ -108,6 +104,6 @@ connection_args = OrderedDict(
 )
 
 connection_args_example = OrderedDict(
-    strava_client_id ='<your-strava-client_id>',
-    strava_access_token ='<your-strava-api-token>'
+    strava_client_id='<your-strava-client_id>',
+    strava_access_token='<your-strava-api-token>'
 )

mindsdb/integrations/handlers/strava_handler/strava_tables.py

@@ -8,12 +8,10 @@ from mindsdb.utilities import log
 
 from mindsdb_sql_parser import ast
 
-import requests
-import pandas as pd
-import json
 
 logger = log.getLogger(__name__)
 
+
 class StravaAllClubsTable(APITable):
     """Strava List all Clubs Table implementation"""
 
@@ -32,7 +30,6 @@ class StravaAllClubsTable(APITable):
         ValueError
             If the query contains an unsupported condition
         """
-        conditions = extract_comparison_conditions(query.where)
 
         order_by_conditions = {}
 
@@ -42,7 +39,7 @@ class StravaAllClubsTable(APITable):
 
             for an_order in query.order_by:
                 if an_order.field.parts[0] != "id":
-                    next    
+                    next
                 if an_order.field.parts[1] in self.get_columns():
                     order_by_conditions["columns"].append(an_order.field.parts[1])
 
@@ -66,7 +63,6 @@ class StravaAllClubsTable(APITable):
             else:
                 raise ValueError(f"Unknown query target {type(target)}")
 
-
         if len(strava_clubs_df) == 0:
             strava_clubs_df = pd.DataFrame([], columns=selected_columns)
         else:
@@ -93,24 +89,24 @@ class StravaAllClubsTable(APITable):
             List of columns
         """
         return [
-        'id',
-        'name',
-        'sport_type',
-        'city',
-        'state',
-        'country',
-        'member_count',
+            'id',
+            'name',
+            'sport_type',
+            'city',
+            'state',
+            'country',
+            'member_count',
         ]
 
     def call_strava_allclubs_api(self):
         """Pulls all the records from the given and returns it select()
-    
+
         Returns
         -------
         pd.DataFrame of all the records of the "List Athlete Clubs" API end point
         """
 
-        clubs = self.handler.connect().get_athlete_clubs() 
+        clubs = self.handler.connect().get_athlete_clubs()
 
         club_cols = self.get_columns()
         data = []
@@ -153,7 +149,7 @@ class StravaClubActivitesTable(APITable):
 
             for an_order in query.order_by:
                 if an_order.field.parts[0] != "id":
-                    next    
+                    next
                 if an_order.field.parts[1] in self.get_columns():
                     order_by_conditions["columns"].append(an_order.field.parts[1])
 
@@ -165,7 +161,7 @@ class StravaClubActivitesTable(APITable):
                     raise ValueError(
                         f"Order by unknown column {an_order.field.parts[1]}"
                     )
-        
+
         for a_where in conditions:
             if a_where[1] == "strava_club_id":
                 if a_where[0] != "=":
@@ -174,7 +170,6 @@ class StravaClubActivitesTable(APITable):
             else:
                 raise ValueError(f"Unsupported where argument {a_where[1]}")
 
-
         strava_club_activities_df = self.call_strava_clubactivities_api(a_where[2])
 
         selected_columns = []
@@ -187,7 +182,6 @@ class StravaClubActivitesTable(APITable):
             else:
                 raise ValueError(f"Unknown query target {type(target)}")
 
-
         if len(strava_club_activities_df) == 0:
             strava_club_activities_df = pd.DataFrame([], columns=selected_columns)
         else:
@@ -214,18 +208,18 @@ class StravaClubActivitesTable(APITable):
             List of columns
         """
         return [
-        'name',
-        'distance',
-        'moving_time',	
-        'elapsed_time',
-        'total_elevation_gain',
-        'sport_type',
-        'athlete.firstname',
+            'name',
+            'distance',
+            'moving_time',
+            'elapsed_time',
+            'total_elevation_gain',
+            'sport_type',
+            'athlete.firstname',
         ]
 
-    def call_strava_clubactivities_api(self,club_id):
+    def call_strava_clubactivities_api(self, club_id):
         """Pulls all the records from the given and returns it select()
-    
+
         Returns
         -------
         pd.DataFrame of all the records of the "getClubActivitiesById" API end point
@@ -243,5 +237,3 @@ class StravaClubActivitesTable(APITable):
         all_strava_club_activities_df = pd.DataFrame(data, columns=club_cols)
 
         return all_strava_club_activities_df
-
-

mindsdb/integrations/handlers/stripe_handler/stripe_handler.py

@@ -10,6 +10,7 @@ from mindsdb_sql_parser import parse_sql
 
 logger = log.getLogger(__name__)
 
+
 class StripeHandler(APIHandler):
     """
     The Stripe handler implementation.
@@ -41,7 +42,7 @@ class StripeHandler(APIHandler):
 
         payment_intents_data = PaymentIntentsTable(self)
         self._register_table("payment_intents", payment_intents_data)
-        
+
         payouts_data = PayoutsTable(self)
         self._register_table("payouts", payouts_data)
 
@@ -80,7 +81,7 @@ class StripeHandler(APIHandler):
             stripe.Account.retrieve()
             response.success = True
         except Exception as e:
-            logger.error(f'Error connecting to Stripe!')
+            logger.error('Error connecting to Stripe!')
             response.error_message = str(e)
 
         self.is_connected = response.success