@@ -29,10 +29,10 @@ def test_basic_shard_key(self) -> None:
29
29
30
30
def test_multi_column_shard_key (self ) -> None :
31
31
"""Test shard key with multiple columns."""
32
- shard_key = ShardKey ([ 'user_id' , 'category_id' ] )
32
+ shard_key = ShardKey ('user_id' , 'category_id' )
33
33
assert shard_key .columns == ('user_id' , 'category_id' )
34
34
assert shard_key .metadata_only is False
35
- assert repr (shard_key ) == "ShardKey([ 'user_id', 'category_id'] )"
35
+ assert repr (shard_key ) == "ShardKey('user_id', 'category_id')"
36
36
37
37
def test_empty_shard_key (self ) -> None :
38
38
"""Test empty shard key for keyless sharding."""
@@ -50,30 +50,22 @@ def test_shard_key_metadata_only_single_column(self) -> None:
50
50
51
51
def test_shard_key_metadata_only_multi_column (self ) -> None :
52
52
"""Test SHARD KEY ONLY with multiple columns."""
53
- shard_key = ShardKey ([ 'user_id' , 'category_id' ] , metadata_only = True )
53
+ shard_key = ShardKey ('user_id' , 'category_id' , metadata_only = True )
54
54
assert shard_key .columns == ('user_id' , 'category_id' )
55
55
assert shard_key .metadata_only is True
56
- expected_repr = "ShardKey([ 'user_id', 'category_id'] , metadata_only=True)"
56
+ expected_repr = "ShardKey('user_id', 'category_id', metadata_only=True)"
57
57
assert repr (shard_key ) == expected_repr
58
58
59
59
def test_shard_key_metadata_only_empty (self ) -> None :
60
60
"""Test SHARD KEY ONLY with no columns (should fallback to empty)."""
61
61
shard_key = ShardKey (metadata_only = True )
62
62
assert shard_key .columns == ()
63
63
assert shard_key .metadata_only is True
64
- assert repr (shard_key ) == 'ShardKey(None, metadata_only=True)'
64
+ assert repr (shard_key ) == 'ShardKey(metadata_only=True)'
65
65
66
66
67
67
class TestShardKeyCompiler :
68
- """Test ShardKey DDL compilation."""
69
-
70
- def setup_method (self ) -> None :
71
- """Set up mock engine for DDL compilation."""
72
- def dump (sql : Any , * multiparams : Any , ** params : Any ) -> None :
73
- self .compiled_sql = str (sql .compile (dialect = self .mock_engine .dialect ))
74
-
75
- self .mock_engine = create_mock_engine ('singlestoredb://' , dump )
76
- self .compiled_sql = ''
68
+ """Test ShardKey SQL compilation."""
77
69
78
70
def test_compile_basic_shard_key (self ) -> None :
79
71
"""Test compilation of basic shard key."""
@@ -87,7 +79,7 @@ def test_compile_multi_column_shard_key(self) -> None:
87
79
"""Test compilation of multi-column shard key."""
88
80
from sqlalchemy_singlestoredb .ddlelement import compile_shard_key
89
81
90
- shard_key = ShardKey ([ 'user_id' , 'category_id' ] )
82
+ shard_key = ShardKey ('user_id' , 'category_id' )
91
83
result = compile_shard_key (shard_key , None )
92
84
assert result == 'SHARD KEY (user_id, category_id)'
93
85
@@ -111,7 +103,7 @@ def test_compile_shard_key_metadata_only_multi_column(self) -> None:
111
103
"""Test compilation of SHARD KEY ONLY with multiple columns."""
112
104
from sqlalchemy_singlestoredb .ddlelement import compile_shard_key
113
105
114
- shard_key = ShardKey ([ 'user_id' , 'category_id' ] , metadata_only = True )
106
+ shard_key = ShardKey ('user_id' , 'category_id' , metadata_only = True )
115
107
result = compile_shard_key (shard_key , None )
116
108
assert result == 'SHARD KEY (user_id, category_id) METADATA_ONLY'
117
109
@@ -138,7 +130,7 @@ def test_compile_shard_key_with_special_column_names(self) -> None:
138
130
assert result == 'SHARD KEY (`user id`)'
139
131
140
132
# Test multiple columns with special characters
141
- shard_key = ShardKey ([ 'user-id' , 'tenant id' , 'normal_column' ] )
133
+ shard_key = ShardKey ('user-id' , 'tenant id' , 'normal_column' )
142
134
result = compile_shard_key (shard_key , None )
143
135
assert result == 'SHARD KEY (`user-id`, `tenant id`, normal_column)'
144
136
@@ -247,7 +239,7 @@ class MyTable(Base): # type: ignore
247
239
data = Column (String (50 ))
248
240
249
241
__table_args__ = {
250
- 'info' : {'singlestoredb_shard_key' : ShardKey ([ 'user_id' , 'category_id' ] )},
242
+ 'info' : {'singlestoredb_shard_key' : ShardKey ('user_id' , 'category_id' )},
251
243
}
252
244
253
245
Base .metadata .create_all (self .mock_engine , checkfirst = False )
@@ -640,7 +632,7 @@ def test_table_constructor_multi_column_shard_key(self) -> None:
640
632
Column ('user_id' , Integer , primary_key = True ),
641
633
Column ('category_id' , Integer , primary_key = True ),
642
634
Column ('amount' , Integer ),
643
- ShardKey ([ 'user_id' , 'category_id' ] ),
635
+ ShardKey ('user_id' , 'category_id' ),
644
636
)
645
637
646
638
# Verify info is set correctly
@@ -661,7 +653,7 @@ def test_table_constructor_with_both_keys(self) -> None:
661
653
Column ('order_id' , Integer , primary_key = True ),
662
654
Column ('created_at' , String (50 )),
663
655
ShardKey ('user_id' ),
664
- SortKey ([ 'created_at' ] ),
656
+ SortKey ('created_at' ),
665
657
)
666
658
667
659
# Verify info is set correctly
@@ -737,6 +729,6 @@ def test_table_constructor_multiple_sort_keys_error(self) -> None:
737
729
Column ('id' , Integer , primary_key = True ),
738
730
Column ('created_at' , String (50 )),
739
731
Column ('updated_at' , String (50 )),
740
- SortKey ([ 'created_at' ] ),
741
- SortKey ([ 'updated_at' ] ), # This should cause an error
732
+ SortKey ('created_at' ),
733
+ SortKey ('updated_at' ), # This should cause an error
742
734
)
0 commit comments