Skip to content

Commit 96b8b60

Browse files
kesmit13claude
kesmit13
and
claude
committed
feat: Enhance reflection system with comprehensive DESC sort key and table type support
- Add comprehensive declarative reflection tests covering ShardKey, SortKey, and VectorKey - Update reflection parser to properly handle ASC/DESC directions in CREATE TABLE statements - Enhance base.py with get_table_options method for converting parsed features to dialect options - Fix parser handling of new (column_name, direction) tuple format for key specifications - Update test expectations in test_shard_key.py and test_sort_key.py for new tuple format - Add support for table type reflection (RowStore, ColumnStore) with modifiers - Improve CREATE TABLE parsing with metadata_only, named keys, and direction parsing - Ensure full round-trip compatibility between table creation and reflection 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
1 parent f94dfab commit 96b8b60

File tree

5 files changed

+1596
-37
lines changed

5 files changed

+1596
-37
lines changed

‎sqlalchemy_singlestoredb/base.py

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -745,6 +745,137 @@ def is_disconnect(
745745

746746
return any(indicator in error_msg for indicator in disconnect_indicators)
747747

748+
def get_table_options(
749+
self, connection: Any, table_name: str, schema: Optional[str] = None, **kw: Any,
750+
) -> Dict[str, Any]:
751+
"""Reflect table options including SingleStore-specific features."""
752+
options = super().get_table_options(
753+
connection, table_name, schema=schema, **kw,
754+
)
755+
756+
# Parse the CREATE TABLE statement to extract SingleStore features
757+
parsed_state = self._parsed_state_or_create(
758+
connection, table_name, schema, **kw,
759+
)
760+
761+
# Convert parsed SingleStore features back to dialect options
762+
if hasattr(parsed_state, 'singlestore_features'):
763+
from sqlalchemy_singlestoredb.ddlelement import (
764+
ShardKey, SortKey, VectorKey, RowStore, ColumnStore,
765+
)
766+
767+
for feature_type, spec in parsed_state.singlestore_features.items():
768+
if feature_type == 'shard_key':
769+
# Convert parsed spec back to ShardKey object
770+
# Handle multiple formats:
771+
# 1. New SingleStore format: [(col_name, direction), ...]
772+
# 2. MySQL fallback format: [(col_name, direction, extra), ...]
773+
# 3. Legacy format: [col_name, ...]
774+
columns = spec['columns']
775+
column_specs = []
776+
777+
if columns and isinstance(columns[0], tuple):
778+
# Check if this is our new format or MySQL format
779+
first_tuple = columns[0]
780+
if (
781+
len(first_tuple) == 2 and isinstance(first_tuple[1], str) and
782+
first_tuple[1] in ('ASC', 'DESC')
783+
):
784+
# New SingleStore format: [(col_name, direction), ...]
785+
column_specs = columns
786+
else:
787+
# MySQL parser format: [(col_name, direction, extra), ...]
788+
# Extract just the column names
789+
column_specs = [(col[0], 'ASC') for col in columns if col[0]]
790+
else:
791+
# Legacy SingleStore parser format: [col_name, ...]
792+
column_specs = [(col, 'ASC') for col in columns]
793+
794+
# Check for metadata_only flag
795+
metadata_only = spec.get('metadata_only', False)
796+
shard_key = ShardKey(*column_specs, metadata_only=metadata_only)
797+
options['singlestoredb_shard_key'] = shard_key
798+
799+
elif feature_type == 'sort_key':
800+
# Convert parsed spec back to SortKey object
801+
# Handle multiple formats (same logic as shard_key)
802+
columns = spec['columns']
803+
column_specs = []
804+
805+
if columns and isinstance(columns[0], tuple):
806+
# Check if this is our new format or MySQL format
807+
first_tuple = columns[0]
808+
if (
809+
len(first_tuple) == 2 and isinstance(first_tuple[1], str) and
810+
first_tuple[1] in ('ASC', 'DESC')
811+
):
812+
# New SingleStore format: [(col_name, direction), ...]
813+
column_specs = columns
814+
else:
815+
# MySQL parser format: [(col_name, direction, extra), ...]
816+
# Extract just the column names
817+
column_specs = [(col[0], 'ASC') for col in columns if col[0]]
818+
else:
819+
# Legacy SingleStore parser format: [col_name, ...]
820+
column_specs = [(col, 'ASC') for col in columns]
821+
822+
sort_key = SortKey(*column_specs)
823+
options['singlestoredb_sort_key'] = sort_key
824+
825+
elif feature_type == 'vector_key':
826+
# Convert parsed spec back to VectorKey object
827+
# Handle both SingleStore format (list of strings) and MySQL
828+
# fallback format (list of tuples)
829+
columns = spec['columns']
830+
if columns and isinstance(columns[0], tuple):
831+
# MySQL parser format: [(col_name, direction, extra), ...]
832+
# Extract just the column names
833+
column_names = [col[0] for col in columns if col[0]]
834+
else:
835+
# SingleStore parser format: [col_name, ...]
836+
column_names = columns
837+
838+
vector_key = VectorKey(
839+
*column_names,
840+
name=spec.get('name'),
841+
index_options=spec.get('index_options'),
842+
)
843+
# For vector keys, store as list if multiple exist
844+
existing = options.get('singlestoredb_vector_key')
845+
if existing:
846+
if isinstance(existing, list):
847+
existing.append(vector_key)
848+
else:
849+
options['singlestoredb_vector_key'] = [existing, vector_key]
850+
else:
851+
options['singlestoredb_vector_key'] = vector_key
852+
853+
elif feature_type == 'table_type':
854+
# Convert parsed table type spec back to TableType object
855+
is_rowstore = spec.get('is_rowstore', False)
856+
is_temporary = spec.get('is_temporary', False)
857+
is_global_temporary = spec.get('is_global_temporary', False)
858+
is_reference = spec.get('is_reference', False)
859+
860+
if is_rowstore:
861+
# Create RowStore with appropriate modifiers
862+
table_type = RowStore(
863+
temporary=is_temporary,
864+
global_temporary=is_global_temporary,
865+
reference=is_reference,
866+
)
867+
else:
868+
# Default to ColumnStore (handles CREATE TABLE without ROWSTORE)
869+
# Note: ColumnStore doesn't support global_temporary
870+
table_type = ColumnStore(
871+
temporary=is_temporary,
872+
reference=is_reference,
873+
)
874+
875+
options['singlestoredb_table_type'] = table_type
876+
877+
return options
878+
748879
def on_connect(self) -> Optional[Callable[[Any], None]]:
749880
"""Return a callable that will be executed on new connections."""
750881
def connect(dbapi_connection: Any) -> None:

‎sqlalchemy_singlestoredb/reflection.py

Lines changed: 72 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -119,24 +119,41 @@ def _parse_constraints(self, line: str) -> Tuple[str, Dict[str, Any]]:
119119
spec = m.groupdict()
120120
key_type = spec['type'].lower() # 'shard' or 'sort'
121121

122-
# Parse columns - handle both empty () and actual column lists
122+
# Parse columns - handle both empty () and actual column lists with ASC/DESC
123123
columns_str = spec.get('columns', '').strip()
124+
columns = []
124125
if columns_str:
125-
# Parse column names - simple split for now,
126-
# may need enhancement for quoted identifiers
127-
columns = [
128-
col.strip().strip('`') for col in columns_str.split(',')
129-
if col.strip()
130-
]
131-
else:
132-
columns = []
126+
# Parse column names with ASC/DESC directions
127+
# Example: "user_id DESC, category_id" ->
128+
# [("user_id", "DESC"), ("category_id", "ASC")]
129+
import re
130+
131+
# Split by comma and process each column specification
132+
for col_spec in columns_str.split(','):
133+
col_spec = col_spec.strip()
134+
if not col_spec:
135+
continue
136+
137+
# Check for explicit DESC or ASC at the end
138+
direction_match = re.match(
139+
r'^(.+?)\s+(DESC|ASC)\s*$', col_spec, re.IGNORECASE,
140+
)
141+
if direction_match:
142+
col_name = direction_match.group(1).strip().strip('`')
143+
direction = direction_match.group(2).upper()
144+
columns.append((col_name, direction))
145+
else:
146+
# No explicit direction, default to ASC
147+
col_name = col_spec.strip().strip('`')
148+
columns.append((col_name, 'ASC'))
133149

134150
# Create spec dictionary with parsed information
135151
parsed_spec = {
136152
'type': spec['type'], # 'SHARD' or 'SORT'
137-
'name': None,
153+
'name': spec.get('name'), # Optional key name
138154
'columns': columns,
139155
'only': spec.get('only') is not None, # True if ONLY was specified
156+
'metadata_only': spec.get('metadata_only') is not None,
140157
}
141158

142159
return f'{key_type}_key', parsed_spec
@@ -157,6 +174,43 @@ def _parse_constraints(self, line: str) -> Tuple[str, Dict[str, Any]]:
157174

158175
return type_, spec
159176

177+
def _parse_table_name(self, line: str, state: Any) -> None:
178+
"""Parse CREATE TABLE line to extract table name and table type information."""
179+
# Call parent to handle standard parsing
180+
super()._parse_table_name(line, state)
181+
182+
# Parse table type information from CREATE statement
183+
# Patterns we need to detect:
184+
# CREATE ROWSTORE TABLE -> RowStore()
185+
# CREATE ROWSTORE TEMPORARY TABLE -> RowStore(temporary=True)
186+
# CREATE ROWSTORE GLOBAL TEMPORARY TABLE -> RowStore(global_temporary=True)
187+
# CREATE ROWSTORE REFERENCE TABLE -> RowStore(reference=True)
188+
# CREATE TEMPORARY TABLE -> ColumnStore(temporary=True)
189+
# CREATE REFERENCE TABLE -> ColumnStore(reference=True)
190+
# CREATE TABLE -> ColumnStore() (default)
191+
192+
line_clean = line.strip()
193+
194+
# Parse table type specification
195+
is_rowstore = 'ROWSTORE' in line_clean
196+
is_global_temporary = 'GLOBAL TEMPORARY' in line_clean
197+
# Exclude GLOBAL TEMPORARY from regular TEMPORARY
198+
is_temporary = 'TEMPORARY' in line_clean and not is_global_temporary
199+
is_reference = 'REFERENCE' in line_clean
200+
201+
# Store table type info for later conversion to dialect options
202+
if not hasattr(state, 'singlestore_features'):
203+
state.singlestore_features = {}
204+
205+
table_type_spec = {
206+
'is_rowstore': is_rowstore,
207+
'is_temporary': is_temporary,
208+
'is_global_temporary': is_global_temporary,
209+
'is_reference': is_reference,
210+
}
211+
212+
state.singlestore_features['table_type'] = table_type_spec
213+
160214
def parse(self, show_create: str, charset: str) -> ReflectedState:
161215
state = ReflectedState()
162216
state.charset = charset
@@ -186,8 +240,10 @@ def parse(self, show_create: str, charset: str) -> ReflectedState:
186240
elif type_ == 'ck_constraint':
187241
state.ck_constraints.append(spec)
188242
elif type_ in ('shard_key', 'sort_key', 'vector_key'):
189-
# Don't add SHARD KEY, SORT KEY, and VECTOR KEY to the keys list
190-
pass
243+
# Store SingleStore features for later conversion to dialect options
244+
if not hasattr(state, 'singlestore_features'):
245+
state.singlestore_features = {}
246+
state.singlestore_features[type_] = spec
191247
elif type_ == 'singlestore_table_option':
192248
# Silently ignore SingleStore table options to avoid warnings
193249
# as they're not regular indexes
@@ -237,13 +293,15 @@ def _prep_regexes(self) -> None:
237293

238294
# SingleStore specific SHARD KEY and SORT KEY patterns
239295
# Handles: SHARD KEY (columns), SHARD KEY ONLY (columns),
240-
# SHARD KEY (), SORT KEY (columns), etc.
296+
# SHARD KEY (), SORT KEY (columns), SHARD KEY name (columns) METADATA_ONLY, etc.
241297
self._re_singlestore_key = _re_compile(
242298
r' (?:, *)?' # Handle optional leading comma
243299
r'(?P<type>SHARD|SORT)' # Key type
244300
r' +KEY' # KEY immediately after type
245301
r'(?: +(?P<only>ONLY))?' # Optional ONLY modifier for SHARD KEY
246-
r' +\((?P<columns>.*?)\)' # Column list (can be empty)
302+
r'(?:[ `]*(?P<name>[^`\s()]+)[ `]*)?' # Optional key name (may be quoted)
303+
r' *\((?P<columns>.*?)\)' # Column list (can be empty)
304+
r'(?: +(?P<metadata_only>METADATA_ONLY))?' # Optional METADATA_ONLY suffix
247305
r' *,?$', # Handle trailing comma and spaces
248306
)
249307

0 commit comments

Comments
 (0)