From 2982c3408b03a7d1cdfe4d50bf47485d4e36341c Mon Sep 17 00:00:00 2001 From: mpf82 Date: Wed, 17 Jun 2026 01:04:21 +0200 Subject: [PATCH 1/2] refactor: Improve type hints, comments, simplify code Changelog: refactor --- sqlalchemy_cauldron/__init__.py | 16 +- sqlalchemy_cauldron/cli.py | 42 +- sqlalchemy_cauldron/config.py | 5 +- .../generator/core_generator.py | 976 +++++++----------- sqlalchemy_cauldron/generator/helpers.py | 83 +- .../generator/orm_generator.py | 316 +++--- .../generator/orm_test_generator.py | 58 +- .../generator/schema_extractor.py | 95 +- .../generator/sqlalchemy_runtime.py | 24 +- .../generator/type_inference.py | 192 ++-- sqlalchemy_cauldron/main.py | 7 +- sqlalchemy_cauldron/parser/normalizer.py | 21 +- 12 files changed, 774 insertions(+), 1061 deletions(-) diff --git a/sqlalchemy_cauldron/__init__.py b/sqlalchemy_cauldron/__init__.py index 9e8de0d..b451a70 100644 --- a/sqlalchemy_cauldron/__init__.py +++ b/sqlalchemy_cauldron/__init__.py @@ -62,16 +62,16 @@ __author__ = "Michael Faust" __all__ = [ + "Dialect", + "Generator", + # Configuration + "GeneratorConfig", + # Data structures + "ParseResult", + "SchemaExtractResult", "__version__", # Main functions "brew", - "normalize_and_parse", "extract_schema_and_tables", - # Data structures - "ParseResult", - "SchemaExtractResult", - # Configuration - "GeneratorConfig", - "Generator", - "Dialect", + "normalize_and_parse", ] diff --git a/sqlalchemy_cauldron/cli.py b/sqlalchemy_cauldron/cli.py index 691c1ee..7eaf0df 100644 --- a/sqlalchemy_cauldron/cli.py +++ b/sqlalchemy_cauldron/cli.py @@ -33,10 +33,11 @@ class SchemaOutput: ) -def version_callback(value: bool): +def version_callback(value: bool) -> None: + """Callback to display version information and exit.""" if value: - print(f"SQLAlchemy-Cauldron {package_version}") - raise typer.Exit() + typer.echo(f"sqlalchemy-cauldron {package_version}") + raise typer.Exit def get_sql(sql: str | None, file: str | None) -> str: @@ -55,7 +56,7 @@ def get_sql(sql: str | None, file: str | None) -> str: typer.Exit: If neither sql nor file provided """ if file: - return Path(file).read_text() + return Path(file).read_text(encoding="utf-8") if sql: return sql typer.echo("Error: Provide SQL as argument or use --file", err=True) @@ -78,7 +79,7 @@ def output_code( success_msg: Success message to display when writing to file. """ if output: - Path(output).write_text(code) + Path(output).write_text(code, encoding="utf-8") typer.echo(f"[OK] {success_msg} written to {output}") else: typer.echo(code) @@ -90,7 +91,8 @@ def main( bool | None, typer.Option("--version", "-V", callback=version_callback, help="Show the version and exit"), ] = None, -): ... +) -> None: + """Main entry point for the CLI.""" @app.command() @@ -98,7 +100,8 @@ def orm( sql: Annotated[str | None, typer.Argument(help="SQL query string")] = None, file: Annotated[str | None, typer.Option("--file", "-f", help="Read SQL from file")] = None, dialect: Annotated[ - str, typer.Option("--dialect", "-d", help=f"SQL dialect (default: {Dialect.POSTGRES})") + str, + typer.Option("--dialect", "-d", help=f"SQL dialect (default: {Dialect.POSTGRES})"), ] = Dialect.POSTGRES.value, optimize: Annotated[bool, typer.Option("--optimize", help="Optimize SQL before generation")] = False, resolve_aliases: Annotated[bool, typer.Option("--resolve-aliases", help="Resolve table aliases")] = False, @@ -135,7 +138,7 @@ def orm( output_code(code, output, "ORM code") except Exception as e: typer.echo(f"Error: {e}", err=True) - raise typer.Exit(code=1) + raise typer.Exit(code=1) from None @app.command() @@ -143,11 +146,11 @@ def core( sql: Annotated[str | None, typer.Argument(help="SQL query string")] = None, file: Annotated[str | None, typer.Option("--file", "-f", help="Read SQL from file")] = None, dialect: Annotated[ - str, typer.Option("--dialect", "-d", help=f"SQL dialect (default: {Dialect.POSTGRES})") + str, + typer.Option("--dialect", "-d", help=f"SQL dialect (default: {Dialect.POSTGRES})"), ] = Dialect.POSTGRES.value, optimize: Annotated[bool, typer.Option("--optimize", help="Optimize SQL before generation")] = False, resolve_aliases: Annotated[bool, typer.Option("--resolve-aliases", help="Resolve table aliases")] = False, - include_test: Annotated[bool, typer.Option("--include-test", help="Include self-contained test code")] = False, output: Annotated[str | None, typer.Option("--output", "-o", help="Write to file instead of stdout")] = None, ) -> None: """Generate SQLAlchemy CORE code from SQL. @@ -178,7 +181,7 @@ def core( output_code(code, output, "CORE code") except Exception as e: typer.echo(f"Error: {e}", err=True) - raise typer.Exit(code=1) + raise typer.Exit(code=1) from None @app.command() @@ -186,7 +189,8 @@ def normalize( sql: Annotated[str | None, typer.Argument(help="SQL query string")] = None, file: Annotated[str | None, typer.Option("--file", "-f", help="Read SQL from file")] = None, dialect: Annotated[ - str, typer.Option("--dialect", "-d", help=f"SQL dialect (default: {Dialect.POSTGRES})") + str, + typer.Option("--dialect", "-d", help=f"SQL dialect (default: {Dialect.POSTGRES})"), ] = Dialect.POSTGRES.value, optimize: Annotated[bool, typer.Option("--optimize", help="Optimize SQL before generation")] = False, resolve_aliases: Annotated[bool, typer.Option("--resolve-aliases", help="Resolve table aliases")] = False, @@ -211,7 +215,7 @@ def normalize( output_code(sql_out, output, "Normalized SQL") except Exception as e: typer.echo(f"Error: {e}", err=True) - raise typer.Exit(code=1) + raise typer.Exit(code=1) from None @app.command() @@ -219,7 +223,8 @@ def schema( sql: Annotated[str | None, typer.Argument(help="SQL query string")] = None, file: Annotated[str | None, typer.Option("--file", "-f", help="Read SQL from file")] = None, dialect: Annotated[ - str, typer.Option("--dialect", "-d", help=f"SQL dialect (default: {Dialect.POSTGRES})") + str, + typer.Option("--dialect", "-d", help=f"SQL dialect (default: {Dialect.POSTGRES})"), ] = Dialect.POSTGRES.value, output: Annotated[str | None, typer.Option("--output", "-o", help="Write to file instead of stdout")] = None, ) -> None: @@ -252,7 +257,7 @@ def schema( output_code(code, output, "Schema") except Exception as e: typer.echo(f"Error: {e}", err=True) - raise typer.Exit(code=1) + raise typer.Exit(code=1) from None @app.command(name="json") @@ -260,7 +265,8 @@ def schema_json( sql: Annotated[str | None, typer.Argument(help="SQL query string")] = None, file: Annotated[str | None, typer.Option("--file", "-f", help="Read SQL from file")] = None, dialect: Annotated[ - str, typer.Option("--dialect", "-d", help=f"SQL dialect (default: {Dialect.POSTGRES})") + str, + typer.Option("--dialect", "-d", help=f"SQL dialect (default: {Dialect.POSTGRES})"), ] = Dialect.POSTGRES.value, output: Annotated[str | None, typer.Option("--output", "-o", help="Write to file instead of stdout")] = None, ) -> None: @@ -287,14 +293,14 @@ def schema_json( "columns": [col.name for col in (table_info.columns or [])], } for table_info in schema_result.table_infos - ] + ], ) json_output = json.dumps(schema_data.__dict__, indent=2) output_code(json_output, output, "Schema JSON") except Exception as e: typer.echo(f"Error: {e}", err=True) - raise typer.Exit(code=1) + raise typer.Exit(code=1) from None if __name__ == "__main__": diff --git a/sqlalchemy_cauldron/config.py b/sqlalchemy_cauldron/config.py index 1b66b57..165d161 100644 --- a/sqlalchemy_cauldron/config.py +++ b/sqlalchemy_cauldron/config.py @@ -2,7 +2,10 @@ from dataclasses import dataclass from enum import StrEnum -from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path FALLBACK_PK: str = "recid" diff --git a/sqlalchemy_cauldron/generator/core_generator.py b/sqlalchemy_cauldron/generator/core_generator.py index 0b3d2f0..063425b 100644 --- a/sqlalchemy_cauldron/generator/core_generator.py +++ b/sqlalchemy_cauldron/generator/core_generator.py @@ -3,7 +3,8 @@ import json import keyword -from typing import Callable, TypeAlias +from collections.abc import Callable +from typing import Literal, TypeAlias import libcst as cst @@ -34,7 +35,7 @@ DispatchCreateDictType: TypeAlias = dict[str, Callable[[exp.Expression], list[cst.SimpleStatementLine]]] TablesType: TypeAlias = dict[str, Table | AliasedTable] -# Comparison operator mapping from sqlglot expression types to cst nodes +# Comparison operator mapping from SQLGlot expression types to cst nodes CMP_OPS: CmpOpsType = { exp.EQ: cst.Equal(), exp.NEQ: cst.NotEqual(), @@ -44,7 +45,7 @@ exp.LTE: cst.LessThanEqual(), } -# Logical operator mapping from sqlglot expression types to cst nodes +# Logical operator mapping from SQLGlot expression types to cst nodes ARITHMETIC_OPS: ArithmeticOpsType = { exp.Add: cst.Add, exp.Sub: cst.Subtract, @@ -52,7 +53,7 @@ exp.Div: cst.Divide, } -# Bitwise operator mapping from sqlglot expression types to cst nodes +# Bitwise operator mapping from SQLGlot expression types to cst nodes BITWISE_OPS: BitwiseOpsType = { exp.BitwiseAnd: cst.BitAnd, exp.BitwiseOr: cst.BitOr, @@ -61,7 +62,7 @@ exp.BitwiseRightShift: cst.RightShift, } -# Mapping from sqlglot DType to SQLAlchemy type names +# Mapping from SQLGlot DType to SQLAlchemy type names # Used for CAST expressions DTYPE_TO_SQLALCHEMY: DtypeMappingType = { DType.INT: "Integer", @@ -108,7 +109,7 @@ def cst_attribute(value: cst.BaseExpression | str, attr: str) -> cst.Attribute: - """Helper to create a cst.Attribute node. + """Create a cst.Attribute node. Args: value: The base expression (e.g., a Name or another Attribute) @@ -123,10 +124,11 @@ def cst_attribute(value: cst.BaseExpression | str, attr: str) -> cst.Attribute: def cst_call_literal_column(value: str) -> cst.Call: - """Helper to create a cst.Call node for literal_column with a string argument. + """Create a cst.Call node for literal_column with a string argument. Args: value: The string value to pass to literal_column + Returns: A cst.Call node representing literal_column("value") """ @@ -134,10 +136,10 @@ def cst_call_literal_column(value: str) -> cst.Call: def cst_call_func(func: str | cst.BaseExpression, args: ArgsType) -> cst.Call: - """Helper to create a cst.Call node for a function call. + """Create a cst.Call node for a function call. Args: - func_name: The function name as a string or a cst.BaseExpression + func: The function name as a string or a cst.BaseExpression args: List of arguments as cst.BaseExpression or cst.Arg Returns: @@ -149,6 +151,18 @@ def cst_call_func(func: str | cst.BaseExpression, args: ArgsType) -> cst.Call: return cst.Call(func=func, args=args) +def cst_bool(value: bool) -> cst.Name: + """Convert a boolean value to a cst.Name node representing True or False. + + Args: + value: The boolean value to convert + + Returns: + A cst.Name node representing the boolean value + """ + return cst.Name("True" if value else "False") + + def build_import_from(module: str, import_names: list[str]) -> cst.SimpleStatementLine: """Build an ImportFrom statement for a module and list of items. @@ -174,8 +188,8 @@ def build_import_from(module: str, import_names: list[str]) -> cst.SimpleStateme cst.ImportFrom( module=module_ref, names=[cst.ImportAlias(cst.Name(name)) for name in import_names], - ) - ] + ), + ], ) @@ -187,14 +201,21 @@ def metadata_definitions() -> list[cst.SimpleStatementLine]: cst.Assign( targets=[cst.AssignTarget(cst.Name("metadata"))], value=cst_call_func("MetaData", []), - ) - ] - ) + ), + ], + ), ] def column_definitions(table: Table) -> list[cst.Arg]: - """Build column definition arguments for a Table.""" + """Build column definition arguments for a Table. + + Args: + table: SQLAlchemy Table object to generate columns for + + Returns: + List of cst.Arg nodes representing column definitions for the Table + """ return [ cst.Arg( cst_call_func( @@ -203,14 +224,14 @@ def column_definitions(table: Table) -> list[cst.Arg]: cst.SimpleString(f'"{col.name}"'), cst.Name(type(col.type).__name__), ], - ) + ), ) for col in table.columns ] def get_inner_expr(node: exp.Expression) -> exp.Expression: - """Extract inner expression from a node, handling sqlglot's varied structure. + """Extract inner expression from a node, handling SQLGlot's varied structure. Tries to get the inner expression from either 'this' or 'expression' attributes, whichever is available. This handles nodes like Not, Exists, Paren, and other @@ -231,6 +252,7 @@ def single_line_sql(sql: str, dialect: Dialect) -> str: Args: sql: The original SQL string, potentially with newlines and extra whitespace dialect: The SQL dialect to use for normalization + Returns: A single-line version of the SQL string with normalized whitespace """ @@ -238,13 +260,12 @@ def single_line_sql(sql: str, dialect: Dialect) -> str: return "" sqlglot_dialect = SQLGLOT_DIALECT_MAPPING.get(dialect.value.lower() if dialect else None, None) sql = transpile_and_parse(sql, sqlglot_dialect).sql(dialect=sqlglot_dialect) - sql = " ".join(line.strip() for line in sql.split("\n") if line.strip()) - return sql + return " ".join(line.strip() for line in sql.split("\n") if line.strip()) class CodeGenerator: - """ - Generates Python source code (LibCST) by walking a sqlglot expression tree. + """Generates Python source code (LibCST) by walking a SQLGlot expression tree. + Supports SELECT, INSERT, UPDATE, DELETE. """ @@ -255,6 +276,16 @@ def generate_module( original_sql: str = "", dialect: Dialect = Dialect.POSTGRES, ) -> str: + """Generate a Python module as a string from a SQLGlot expression and table info. + + Args: + expr: The root SQLGlot expression to generate code for + tables: Dictionary of table objects with column information + original_sql: The original SQL string (used for comments) + dialect: SQL dialect for any dialect-specific code generation + Returns: + A string containing the generated Python source code + """ self.expr = expr self.tables = tables self.table_nodes = self.get_table_nodes(expr) @@ -263,6 +294,7 @@ def generate_module( self.dialect = dialect self._alias_map = self._build_alias_map(expr) self._cte_aliases: set[str] = set() # Track CTE aliases for column emission + self._recursive_cte_aliases: set[str] = set() # Track RECURSIVE CTE aliases self._placeholder_counter = 0 # Counter for generating unique placeholder names body: list = [] @@ -275,10 +307,8 @@ def generate_module( return cst.Module(body=body).code def get_table_nodes(self, expr: exp.Expression) -> list[exp.Table]: - """Extract table nodes from the expression, handling cases where find_all may not be available.""" - if hasattr(expr, "find_all"): - return list(expr.find_all(exp.Table)) - return [] + """Extract table nodes from the expression.""" + return list(expr.find_all(exp.Table)) # ------------------------------------------------------------------ # DISPATCH @@ -310,10 +340,9 @@ def _dispatch_query(self, expr: exp.Expression) -> list[cst.SimpleStatementLine] "INDEX": self._query_create_index, "SEQUENCE": self._query_create_sequence, } - if kind in create_dispatch: - return create_dispatch.get(kind, self._query_create_table)(expr) + return create_dispatch.get(kind, self._query_create_table)(expr) - elif isinstance(expr, exp.Drop): + if isinstance(expr, exp.Drop): kind = expr.args.get("kind") drop_dispatch: DispatchCreateDictType = { "TABLE": self._query_drop_table, @@ -321,21 +350,20 @@ def _dispatch_query(self, expr: exp.Expression) -> list[cst.SimpleStatementLine] "INDEX": self._query_drop_index, "SEQUENCE": self._query_drop_sequence, } - if kind in drop_dispatch: - return drop_dispatch.get(kind, self._query_drop_table)(expr) + return drop_dispatch.get(kind, self._query_drop_table)(expr) - elif isinstance(expr, exp.Alter): + if isinstance(expr, exp.Alter): kind = expr.args.get("kind") alter_dispatch: DispatchCreateDictType = { "TABLE": self._query_alter_table, "VIEW": self._query_alter_view, "INDEX": self._query_alter_index, - "SEQUENCE": self._query_alter_sequence, + # > "SEQUENCE": self._query_alter_sequence, # not supported by SQLGlot } - if kind in alter_dispatch: - return alter_dispatch.get(kind, self._query_alter_table)(expr) + return alter_dispatch.get(kind, self._query_alter_table)(expr) - raise NotImplementedError(f"Unsupported statement type: {expr_type.__name__}") + msg = f"Unsupported statement type: {expr_type.__name__}" + raise NotImplementedError(msg) # ------------------------------------------------------------------ # NAME SAFETY @@ -384,6 +412,12 @@ def _table_name(self, table_name: str) -> str: """Return the safe Python variable name for a given SQL table name. Always converts to a safe name to handle reserved keywords and special characters. + + Args: + table_name: The original SQL table name + + Returns: + A safe Python variable name to use in the generated code """ return safe_table_name(table_name) @@ -391,11 +425,17 @@ def _cst_table_name(self, table_name: str) -> cst.Name: """Return a cst.Name node for a table variable. Convenience helper that combines _table_name with cst.Name creation. + + Args: + table_name: The original SQL table name + + Returns: + A cst.Name node representing the safe Python variable name """ return cst.Name(self._table_name(table_name)) def _is_cte_alias(self, name: str) -> bool: - """Check if a name is a CTE alias""" + """Check if a name is a CTE alias.""" return name in self._cte_aliases def _generate_alias_assignments(self) -> list[cst.SimpleStatementLine]: @@ -403,14 +443,20 @@ def _generate_alias_assignments(self) -> list[cst.SimpleStatementLine]: Uses _table_name() to get the base reference and delegates to _build_aliased_call() for subclass-specific call generation (CORE: table.alias(), ORM: aliased()). + + Returns: + List of cst.SimpleStatementLine nodes for alias assignments + + Returns: + List of cst.SimpleStatementLine nodes for alias assignments """ assignments = [] - for table_ref, table in self.tables.items(): + for table in self.tables.values(): if isinstance(table, AliasedTable): # table_ref is the alias (e.g., 'e', 'm') # table.name is the underlying table name (e.g., 'employees') - alias = table.table_alias + alias = table.alias base_table_name = table.name # Use _table_name to get the appropriate reference (table name for CORE, class name for ORM) @@ -422,8 +468,8 @@ def _generate_alias_assignments(self) -> list[cst.SimpleStatementLine]: cst.Assign( targets=[cst.AssignTarget(cst.Name(alias))], value=self._build_aliased_call(base_ref, alias), - ) - ] + ), + ], ) assignments.append(assignment) @@ -457,9 +503,7 @@ def _imports(self) -> list[cst.SimpleStatementLine]: imports = [build_import_from("sqlalchemy", CORE_IMPORTS)] # Add dialect import if available - dialect_import = self._build_dialect_import() - if dialect_import: - imports.append(dialect_import) + imports.extend(self._build_dialect_import(self.dialect or "")) return imports @@ -489,9 +533,9 @@ def _table_definitions(self) -> list[cst.SimpleStatementLine | cst.ClassDef]: *cols, ], ), - ) - ] - ) + ), + ], + ), ) return body @@ -520,15 +564,14 @@ def _compile_and_print(self) -> list[cst.SimpleStatementLine]: compile_args = [] # Add dialect argument if available and not None - if self.dialect and self.dialect.value in SQLALCHEMY_DIALECT_MAPPING: - dialect_name = SQLALCHEMY_DIALECT_MAPPING[self.dialect.value] - if dialect_name is not None: - compile_args.append( - cst.Arg( - keyword=cst.Name("dialect"), - value=cst_call_func(cst_attribute(dialect_name, "dialect"), args=[]), - ) - ) + dialect_name = SQLALCHEMY_DIALECT_MAPPING[self.dialect.value] + if dialect_name is not None: + compile_args.append( + cst.Arg( + keyword=cst.Name("dialect"), + value=cst_call_func(cst_attribute(dialect_name, "dialect"), args=[]), + ), + ) # Add compile_kwargs argument compile_args.append( @@ -539,10 +582,10 @@ def _compile_and_print(self) -> list[cst.SimpleStatementLine]: cst.DictElement( key=cst.SimpleString('"literal_binds"'), value=cst.Name("True"), - ) - ] + ), + ], ), - ) + ), ) return [ @@ -554,8 +597,8 @@ def _compile_and_print(self) -> list[cst.SimpleStatementLine]: func=cst_attribute("query", "compile"), args=compile_args, ), - ) - ] + ), + ], ), cst.SimpleStatementLine([cst.Expr(value=cst_call_func("print", args=[cst.Name("sql")]))]), ] @@ -578,44 +621,6 @@ def _chain_method( """ return cst_call_func(cst_attribute(target, method), args=args) - def _reassign(self, method: str, args: list[cst.BaseExpression]) -> cst.SimpleStatementLine: - """Reassign query variable by chaining a method call. - - Args: - method: The method name to call on query - args: The arguments to pass as positional arguments - - Returns: - A SimpleStatementLine reassigning query - """ - return cst.SimpleStatementLine( - [ - cst.Assign( - targets=[cst.AssignTarget(cst.Name("query"))], - value=self._chain_method(cst.Name("query"), method, args), - ) - ] - ) - - def _reassign_kwargs(self, method: str, kwargs: list[cst.Arg]) -> cst.SimpleStatementLine: - """Reassign query variable by chaining a method call with keyword arguments. - - Args: - method: The method name to call on query - kwargs: The keyword arguments as cst.Arg objects with keywords - - Returns: - A SimpleStatementLine reassigning query - """ - return cst.SimpleStatementLine( - [ - cst.Assign( - targets=[cst.AssignTarget(cst.Name("query"))], - value=self._chain_method(cst.Name("query"), method, kwargs), - ) - ] - ) - def _build_assignment( self, target_name: str, @@ -666,30 +671,12 @@ def _finalize_query_statement(self, query_expr: cst.BaseExpression, body: list) # Add the assignment with Select type hint body.append(self._build_assignment("query", query_expr, self._leading_comment(), type_hint="Select")) - - # Add compile and print body.extend(self._compile_and_print()) - def _get_node_arg(self, node: exp.Expression, key: str, default=None) -> exp.Expression | None: - """Safely extract an argument from a node, handling both dict and attribute access patterns. - - Args: - node: The node to extract from - key: The argument key/attribute name - default: Default value if not found - - Returns: - The argument value or default - """ - args = getattr(node, "args", None) - if isinstance(args, dict): - return args.get(key, default) - return getattr(node, key, default) - def _extract_expr_value(self, expr_wrapper: exp.Expression) -> exp.Expression: """Extract expression value from a wrapper object or return direct value. - Some sqlglot expression types wrap their values (e.g., Limit, Offset have + Some SQLGlot expression types wrap their values (e.g., Limit, Offset have an 'expression' attribute), while others are direct values. This helper normalizes the extraction. @@ -703,7 +690,7 @@ def _extract_expr_value(self, expr_wrapper: exp.Expression) -> exp.Expression: return expr_wrapper.expression return expr_wrapper - def _build_dialect_import(self) -> cst.SimpleStatementLine | None: + def _build_dialect_import(self, dialect: str) -> list[cst.SimpleStatementLine]: """Build dialect-specific import statement if needed. Checks DIALECT_MAPPING and returns an import for the dialect module if available. @@ -711,11 +698,11 @@ def _build_dialect_import(self) -> cst.SimpleStatementLine | None: Returns: An import statement, or None if no dialect import is needed """ - if self.dialect and self.dialect.value in SQLALCHEMY_DIALECT_MAPPING: - dialect_name = SQLALCHEMY_DIALECT_MAPPING[self.dialect.value] + if dialect and dialect in SQLALCHEMY_DIALECT_MAPPING: + dialect_name = SQLALCHEMY_DIALECT_MAPPING[dialect] if dialect_name is not None: - return build_import_from("sqlalchemy.dialects", [dialect_name]) - return None + return [build_import_from("sqlalchemy.dialects", [dialect_name])] + return [] def _apply_limit_offset_clauses(self, query_expr: cst.BaseExpression, expr: exp.Expression) -> cst.BaseExpression: """Apply LIMIT, FETCH, and OFFSET clauses to a query expression. @@ -787,7 +774,7 @@ def _process_clause_conditions(self, expr: exp.Expression, clause_name: str) -> Returns: None if no clause, or list of condition expressions if present. """ - clause = self._get_node_arg(expr, clause_name) + clause = expr.args.get(clause_name) if clause: conditions = self._flatten_and(clause.this if hasattr(clause, "this") else clause) return [self._emit_expr(c) for c in conditions] @@ -802,7 +789,14 @@ def _process_having_clause(self, expr: exp.Expression) -> list[cst.BaseExpressio return self._process_clause_conditions(expr, "having") def _wrap_in_parens(self, expr: cst.BaseExpression) -> cst.BaseExpression: - """Wrap an expression in parentheses.""" + """Wrap an expression in parentheses. + + Args: + expr: The expression to wrap + + Returns: + The expression wrapped in parentheses + """ return expr.with_changes( lpar=[cst.LeftParen()], rpar=[cst.RightParen()], @@ -815,21 +809,16 @@ def _build_and_emit_ctes(self, expr: exp.Expression, body: list) -> None: expr: The SQL expression containing potential WITH clause body: The body list to append CTE statements to """ - with_node = self._get_node_arg(expr, "with_") + with_node = expr.args.get("with_") if with_node and getattr(with_node, "expressions", None): recursive = getattr(with_node, "recursive", False) - for cte in with_node.expressions: alias_name = self._extract_cte_alias_name(cte) # Track this CTE alias for column emission self._cte_aliases.add(alias_name) - - cte_select = cte.this - try: - cte_expr = self._build_select_expr(cte_select) - except Exception: - continue - + if recursive: + self._recursive_cte_aliases.add(alias_name) + cte_expr = self._build_select_expr(cte.this) self._emit_cte_definition(alias_name, cte_expr, body, recursive=recursive) def _apply_where_clause(self, query_expr: cst.BaseExpression, expr: exp.Expression) -> cst.BaseExpression: @@ -976,62 +965,40 @@ def _apply_select_clauses( self, query_expr: cst.BaseExpression, expr: exp.Select, - include_from_join: bool = True, ) -> cst.BaseExpression: - """Apply standard SELECT clauses in SQLAlchemy 2.0+ best practice order. + """Apply standard SELECT clauses in SQLAlchemy 2.0+. Order: DISTINCT -> FROM/JOIN -> WHERE -> GROUP BY -> HAVING -> ORDER BY -> LIMIT/OFFSET - This follows SQLAlchemy 2.0 recommendations where select_from() is called early. Args: query_expr: The base query expression to build upon expr: The SELECT expression containing clause information - include_from_join: Whether to include FROM/JOIN clause processing Returns: The modified query expression with clauses applied """ - # DISTINCT (first method after select()) + # DISTINCT if expr.args.get("distinct"): query_expr = self._chain_method(query_expr, "distinct", []) - # FROM / JOIN (early, before WHERE/GROUP BY for best practices) - if include_from_join: - from_tables, join_nodes = self._extract_from_and_joins(expr) - - if from_tables: - # Check if first table has an alias (for self-joins) - first_table_alias = getattr(from_tables[0], "alias", None) - if first_table_alias: - nested_join_expr = cst.Name(first_table_alias) - else: - nested_join_expr = self._cst_table_name(from_tables[0].name) - - if join_nodes: - for j in join_nodes: - jt_name, args_list, join_method = self._process_join_node(j) - nested_join_expr = cst_call_func( - func=cst_attribute(nested_join_expr, join_method), - args=args_list, - ) - else: - for t in from_tables[1:]: - jt_name = t.name - t_alias = getattr(t, "alias", None) - if t_alias: - nested_join_expr = cst_call_func( - func=cst_attribute(nested_join_expr, "join"), - args=[cst.Name(t_alias)], - ) - else: - nested_join_expr = cst_call_func( - func=cst_attribute(nested_join_expr, "join"), - args=[self._cst_table_name(jt_name)], - ) + # FROM / JOIN + from_tables, join_nodes = self._extract_from_and_joins(expr) + if from_tables: + from_table: exp.Table = from_tables[0] + # Check if from table has an alias (for self-joins) + nested_join_expr = cst.Name(from_table.alias) if from_table.alias else self._cst_table_name(from_table.name) + + if join_nodes: + for j in join_nodes: + args_list = self._process_join_node(j) + nested_join_expr = cst_call_func( + func=cst_attribute(nested_join_expr, "join"), + args=args_list, + ) - query_expr = self._chain_method(query_expr, "select_from", [nested_join_expr]) + query_expr = self._chain_method(query_expr, "select_from", [nested_join_expr]) - # WHERE, GROUP BY, HAVING, ORDER BY using shared helpers (after FROM/JOIN) + # WHERE, GROUP BY, HAVING, ORDER BY using shared helpers query_expr = self._apply_where_clause(query_expr, expr) query_expr = self._apply_group_by_clause(query_expr, expr) query_expr = self._apply_having_clause(query_expr, expr) @@ -1040,12 +1007,12 @@ def _apply_select_clauses( query_expr = self._apply_order_by_clause(query_expr, expr) # Apply LIMIT and OFFSET using shared helper - query_expr = self._apply_limit_offset_clauses(query_expr, expr) - - return query_expr + return self._apply_limit_offset_clauses(query_expr, expr) def _flatten_set_operation( - self, select_expr: exp.Select | exp.SetOperation, operation_type: type + self, + select_expr: exp.Select | exp.SetOperation, + operation_type: type, ) -> list[exp.Select]: """Flatten left-associative chains of set operations into a list. @@ -1099,7 +1066,8 @@ def _build_set_operation_expr(self, select_expr: exp.SetOperation) -> cst.BaseEx # Similar logic could be applied to except_all, but not checking for now function_name = "except_" else: - raise NotImplementedError(f"Unsupported set operation: {type(select_expr).__name__}") + msg = f"Unsupported set operation: {type(select_expr).__name__}" + raise NotImplementedError(msg) # Flatten the left-associative chain operands = self._flatten_set_operation(select_expr, operation_type) @@ -1111,40 +1079,25 @@ def _build_set_operation_expr(self, select_expr: exp.SetOperation) -> cst.BaseEx result_expr = cst_call_func(function_name, args=operand_exprs) # Apply ORDER BY clause if present on the set operation - result_expr = self._apply_order_by_clause(result_expr, select_expr) - - return result_expr + return self._apply_order_by_clause(result_expr, select_expr) - def _get_union_method(self, union_expr: exp.Union) -> str: + def _get_union_method(self, expr: exp.Union) -> Literal["union", "union_all"]: """Determine if a UNION should be UNION or UNION ALL. Args: - union_expr: A Union expression + expr: A Union expression Returns: Either "union" or "union_all" """ - all_flag = False - args = getattr(union_expr, "args", None) - - # sqlglot may represent the ALL modifier in different ways: - # - args['all'] == True - # - union_expr.is_all == True - # - args['distinct'] == False (meaning ALL) - if isinstance(args, dict): - if args.get("all"): - all_flag = True - elif "distinct" in args and args.get("distinct") is False: - all_flag = True - elif getattr(union_expr, "is_all", False): - all_flag = True - + all_flag = expr.args.get("distinct") is False return "union_all" if all_flag else "union" def _build_select_expr(self, select_expr: exp.Select) -> cst.BaseExpression: - """ + """Build a select expression, handling both regular SELECTs and set operations. + Build a LibCST expression that represents a SQLAlchemy `select(...)` call - for the provided sqlglot `select_expr`. This returns a `cst.BaseExpression` + for the provided SQLGlot `select_expr`. This returns a `cst.BaseExpression` (not a statement) that can be further chained with `.cte()` or passed into `select_from(...)`. """ @@ -1159,9 +1112,7 @@ def _build_select_expr(self, select_expr: exp.Select) -> cst.BaseExpression: select_call = cst_call_func("select", args=final_select_args) # Apply all SELECT clauses using shared helper - select_call = self._apply_select_clauses(select_call, select_expr, include_from_join=True) - - return select_call + return self._apply_select_clauses(select_call, select_expr) # ------------------------------------------------------------------ # SELECT @@ -1178,7 +1129,7 @@ def _query_select(self, expr: exp.Select) -> list[cst.SimpleStatementLine]: query_expr = cst_call_func("select", args=final_select_args) # Apply all SELECT clauses using shared helper - query_expr = self._apply_select_clauses(query_expr, expr, include_from_join=True) + query_expr = self._apply_select_clauses(query_expr, expr) # Finalize using shared helper self._finalize_query_statement(query_expr, body) @@ -1193,6 +1144,12 @@ def _query_union(self, expr: exp.Union) -> list[cst.SimpleStatementLine]: A UNION query combines results of multiple SELECT statements. SQLAlchemy uses .union() for UNION and .union_all() for UNION ALL. + + Args: + expr: The Union expression to generate code for + + Returns: + List of cst.SimpleStatementLine representing the generated code for the UNION query """ return self._query_set_operation("UNION", expr) @@ -1238,34 +1195,31 @@ def _build_values_clause(self, values_node: exp.Expression, col_names: list[str] # Use runtime table column definitions if available col_names = [c.name for c in self.tables[table_name].columns] - rows = [] - for row in values_node.expressions: - if isinstance(row, exp.Tuple): - rows.append(row.expressions) + rows = [row.expressions for row in values_node.expressions if isinstance(row, exp.Tuple)] if rows: if len(rows) == 1 and col_names: # Single-row: emit keyword args like `id = 1, name = "Alice"`. - for col_name, val_node in zip(col_names, rows[0]): + for col_name, val_node in zip(col_names, rows[0], strict=False): value_args.append( cst.Arg( keyword=cst.Name(col_name), value=self._emit_expr(val_node), - ) + ), ) else: # Multi-row: emit a list of dicts: .values([{"col": val}, ...]) list_elements: list[cst.Element] = [] for r in rows: - dict_elems: list[cst.DictElement] = [] - for col_name, val_node in zip(col_names, r) if col_names else []: - dict_elems.append( + dict_elements: list[cst.DictElement] = [] + for col_name, val_node in zip(col_names, r, strict=False) if col_names else []: + dict_elements.append( cst.DictElement( key=cst.SimpleString(f'"{col_name}"'), value=self._emit_expr(val_node), - ) + ), ) - list_elements.append(cst.Element(value=cst.Dict(elements=dict_elems))) + list_elements.append(cst.Element(value=cst.Dict(elements=dict_elements))) value_args.append(cst.Arg(value=cst.List(elements=list_elements))) return value_args @@ -1295,18 +1249,13 @@ def _build_from_select_clause( # emit `.from_select([cols], )` using the built select # expression for the SELECT side. select_node = None - if isinstance(values_node, exp.Select) or isinstance(values_node, exp.Union): + if isinstance(values_node, (exp.Select, exp.Union)): select_node = values_node - elif isinstance(values_node, exp.Subquery): - select_node = values_node.this if select_node is None: return False - try: - select_expr = self._build_select_expr(select_node) - except Exception: - select_expr = None + select_expr = self._build_select_expr(select_node) # Build column list for from_select (empty list if no column names) cols_list = cst.List(elements=[cst.Element(value=cst.SimpleString(f'"{c}"')) for c in col_names]) @@ -1319,18 +1268,7 @@ def _build_from_select_clause( args=[cst.Arg(cols_list), cst.Arg(select_expr)] if select_expr is not None else [cst.Arg(cols_list)], ) - body.append( - cst.SimpleStatementLine( - [ - cst.Assign( - targets=[cst.AssignTarget(cst.Name("query"))], - value=insert_call, - ) - ], - leading_lines=self._leading_comment(), - ) - ) - + body.append(self._build_assignment("query", insert_call, self._leading_comment(), type_hint="Select")) body.extend(self._compile_and_print()) return True @@ -1347,21 +1285,7 @@ def _query_insert(self, expr: exp.Insert) -> list[cst.SimpleStatementLine]: # WITH / CTE handling for INSERT statements: emit CTEs as variables # so they can be referenced by the SELECT that supplies the rows. - with_node = expr.args.get("with_") - if with_node and getattr(with_node, "expressions", None): - recursive = getattr(with_node, "recursive", False) - - for cte in with_node.expressions: - alias_name = self._extract_cte_alias_name(cte) - - cte_select = cte.this - try: - cte_expr = self._build_select_expr(cte_select) - except Exception: - # If we cannot build the inner select, skip emitting the CTE - continue - - self._emit_cte_definition(alias_name, cte_expr, body, recursive=recursive) + self._build_and_emit_ctes(expr, body) # Get the values/select expression to determine which INSERT type this is values_node = expr.args.get("expression") @@ -1372,24 +1296,14 @@ def _query_insert(self, expr: exp.Insert) -> list[cst.SimpleStatementLine]: # Otherwise, handle regular INSERT...VALUES value_args = self._build_values_clause(values_node, col_names, table_name) - body.append( - cst.SimpleStatementLine( - [ - cst.Assign( - targets=[cst.AssignTarget(cst.Name("query"))], - value=cst_call_func( - func=cst_attribute( - value=cst_call_func("insert", args=[self._cst_table_name(table_name)]), - attr="values", - ), - args=value_args, - ), - ) - ], - leading_lines=self._leading_comment(), - ) + value_call = cst_call_func( + func=cst_attribute( + value=cst_call_func("insert", args=[self._cst_table_name(table_name)]), + attr="values", + ), + args=value_args, ) - + body.append(self._build_assignment("query", value_call, self._leading_comment(), type_hint="Insert")) body.extend(self._compile_and_print()) return body @@ -1413,23 +1327,19 @@ def _query_update(self, expr: exp.Update) -> list[cst.SimpleStatementLine]: query_expr = self._chain_method(query_expr, "where", where_args) # .values(col=val, ...) from SET assignments - # sqlglot stores SET clauses in args["set"] as a list of EQ nodes - assignments = expr.args.get("set") or expr.args.get("expressions") or [] - if isinstance(assignments, exp.Expression): - assignments = [assignments] + expressions: list[exp.EQ] = expr.args.get("expressions") or [] set_kwargs: list[cst.Arg] = [] - for assignment in assignments: - if isinstance(assignment, exp.EQ): - col = assignment.left - val = assignment.right - col_name = self._get_node_name(col) - set_kwargs.append( - cst.Arg( - keyword=cst.Name(col_name), - value=self._emit_expr(val), - ) - ) + for assignment in expressions: + col = assignment.left + val = assignment.right + col_name = self._get_node_name(col) + set_kwargs.append( + cst.Arg( + keyword=cst.Name(col_name), + value=self._emit_expr(val), + ), + ) query_expr = self._chain_method(query_expr, "values", set_kwargs) @@ -1450,16 +1360,12 @@ def _extract_dml_table_name(self, expr: exp.Expression, is_delete: bool = False) if is_delete: # expr.this may be From(Table) or Table directly from_clause = expr.args.get("this") - if isinstance(from_clause, exp.From): - table_node = from_clause.this - else: - table_node = from_clause + table_node = from_clause.this if isinstance(from_clause, exp.From) else from_clause return table_node.name - else: - # UPDATE: expr.this is the table directly - return expr.this.name + # UPDATE: expr.this is the table directly + return expr.this.name - def _extract_cte_alias_name(self, cte: exp.Expression) -> str: + def _extract_cte_alias_name(self, cte: exp.CTE) -> str: """Extract the alias name from a CTE node. Args: @@ -1468,17 +1374,9 @@ def _extract_cte_alias_name(self, cte: exp.Expression) -> str: Returns: The CTE alias name as a string """ - alias_node = cte.args.get("alias") if isinstance(cte.args, dict) else cte.args.get("alias") - if not alias_node: - return str(cte) - - # Try to extract from TableAlias structure - inner = getattr(alias_node, "this", None) - if inner and hasattr(inner, "this"): - return inner.this - if hasattr(alias_node, "this") and hasattr(alias_node.this, "name"): - return alias_node.this.name - return str(alias_node) + alias_node: exp.Alias = cte.args.get("alias") + inner: exp.TableAlias = alias_node.this + return inner.this def _emit_cte_definition( self, @@ -1501,16 +1399,7 @@ def _emit_cte_definition( args.append(cst.Arg(keyword=cst.Name("recursive"), value=cst.Name("True"))) cte_value = cst_call_func(cst_attribute(cte_expr, "cte"), args=args) - body.append( - cst.SimpleStatementLine( - [ - cst.Assign( - targets=[cst.AssignTarget(cst.Name(alias_name))], - value=cte_value, - ) - ] - ) - ) + body.append(self._build_assignment(alias_name, cte_value)) self._alias_map[alias_name] = alias_name def _query_set_operation(self, method_name: str, expr: exp.Expression) -> list[cst.SimpleStatementLine]: @@ -1526,8 +1415,9 @@ def _query_set_operation(self, method_name: str, expr: exp.Expression) -> list[c body = [] try: result_expr = self._build_select_expr(expr) - except NotImplementedError as e: - raise NotImplementedError(f"Cannot generate {method_name.upper()} query: {e}") + except (NotImplementedError, AttributeError) as e: + msg = f"Cannot generate {method_name.upper()} query: {e}" + raise NotImplementedError(msg) from None # Finalize the query with result expression self._finalize_query_statement(result_expr, body) @@ -1546,15 +1436,19 @@ def _query_ddl_stub(self, comment: str) -> list[cst.SimpleStatementLine]: cst.SimpleStatementLine( [cst.Pass()], leading_lines=[cst.EmptyLine(indent=True, comment=cst.Comment(f"# {comment}"))], - ) + ), ] body.extend(self._compile_and_print()) return body def _is_select_star_query(self, expr: exp.Select) -> bool: - """Check if current query is a SELECT * (bare or qualified table.*)""" - if not hasattr(expr, "expressions"): - return False + """Check if current query is a SELECT * (bare or qualified table.*). + + Args: + expr: The SELECT expression to check + Returns: + True if this is a SELECT * query, False otherwise + """ for e in expr.expressions: if isinstance(e, exp.Star): return True @@ -1585,53 +1479,13 @@ def _extract_join_condition(self, join_node: exp.Join) -> cst.BaseExpression | N join_node: The JOIN node to extract condition from Returns: - Condition as cst.BaseExpression or None if not found + Condition as cst.BaseExpression if present, otherwise None """ - if not join_node: - return None - - # Try to find EQ nodes for join condition - eq_nodes = list(join_node.find_all(exp.EQ)) if join_node else [] - if eq_nodes: - if len(eq_nodes) == 1: - return self._emit_expr(eq_nodes[0]) - else: - return cst_call_func( - func="and_", - args=[cst.Arg(self._emit_expr(n)) for n in eq_nodes], - ) - - # Try the ON clause - on_node = self._get_node_arg(join_node, "on") + on_node = join_node.args.get("on") if on_node: - target = on_node.this if hasattr(on_node, "this") else on_node - try: - return self._emit_expr(target) - except NotImplementedError: - pass - + return self._emit_expr(on_node) return None - def _get_join_method(self, join_node: exp.Join) -> str: - """Determine the join method (join or outerjoin) from a JOIN node. - - Args: - join_node: The JOIN node to analyze - - Returns: - 'join' or 'outerjoin' based on join side - """ - join_method = "join" - side = self._get_node_arg(join_node, "side") - if side: - try: - side_s = str(side).upper() - if "LEFT" in side_s: - join_method = "outerjoin" - except Exception: - pass - return join_method - def _resolve_join_table_ref(self, jt_name: str | None, jt_alias: str | None) -> str: """Resolve the table reference to use in a JOIN clause. @@ -1650,7 +1504,7 @@ def _resolve_join_table_ref(self, jt_name: str | None, jt_alias: str | None) -> return jt_alias return self._table_name(jt_name) if jt_name else "table" - def _process_join_node(self, join_node: exp.Join) -> tuple[str, list[cst.Arg], str]: + def _process_join_node(self, join_node: exp.Join) -> list[cst.Arg]: """Process a JOIN node and extract table name, args list, and join method. For self-joins, uses the alias if available (preserves distinction between repeated tables). @@ -1660,7 +1514,7 @@ def _process_join_node(self, join_node: exp.Join) -> tuple[str, list[cst.Arg], s join_node: The JOIN node to process Returns: - Tuple of (table_name, args_list, join_method) + List of cst.Arg for the join """ jt = join_node.this jt_name = jt.name @@ -1669,13 +1523,19 @@ def _process_join_node(self, join_node: exp.Join) -> tuple[str, list[cst.Arg], s cond_expr = self._extract_join_condition(join_node) table_ref = self._resolve_join_table_ref(jt_name, jt_alias) + # Possible args, in order + # > right: _FromClauseArgument, + # > onclause: Optional[_ColumnExpressionArgument[bool]] = None, + # > isouter: bool = False, + # > full: bool = False, args_list = [cst.Arg(cst.Name(table_ref))] - if cond_expr is not None: + if cond_expr: args_list.append(cst.Arg(cond_expr)) - - join_method = self._get_join_method(join_node) - - return (jt_name, args_list, join_method) + if join_node.kind.upper() == "OUTER" or join_node.side.upper() in ("LEFT", "RIGHT"): + args_list.append(cst.Arg(cst.Name("True"), keyword=cst.Name("isouter"))) + if join_node.side.upper() == "FULL": + args_list.append(cst.Arg(cst.Name("True"), keyword=cst.Name("full"))) + return args_list def _process_select_star_expr(self, star_expr: exp.Expression) -> list[cst.BaseExpression]: """Process SELECT * or SELECT table.* expressions. @@ -1700,23 +1560,23 @@ def _process_select_star_expr(self, star_expr: exp.Expression) -> list[cst.BaseE if table_part: # Qualified star: resolve table name with alias mapping - table_name = getattr(table_part, "name", None) or getattr(table_part, "this", None) or str(table_part) - if isinstance(table_name, exp.Expression): - table_name = str(table_name) + table_name = table_part.name if isinstance(table_part, exp.Expression) else str(table_part) # Resolve alias if present - if isinstance(table_name, str) and table_name in self._alias_map: + if table_name in self._alias_map: table_name = self._alias_map[table_name] result.append(cst_call_literal_column(f"{table_name}.*")) - else: - # Bare *: check if we have tables to determine whether to emit it - if self.table_nodes: - result.append(cst_call_literal_column("*")) + # Bare *: check if we have tables to determine whether to emit it + elif self.table_nodes: + result.append(cst_call_literal_column("*")) return result - def _extract_from_and_joins(self, expr: exp.Expression) -> tuple[list[exp.Table], list[exp.Join]]: + def _extract_from_and_joins(self, expr: exp.Select) -> tuple[list[exp.Table], list[exp.Join]]: """Extract FROM clause tables and JOIN nodes from a SELECT expression. + Since SQLGlot normalizes even cartesian products into explicit JOINs, + we expect to return only a single table from the FROM clause and all joins as JOIN nodes. + Args: expr: The SELECT expression to extract FROM/JOIN from @@ -1725,7 +1585,7 @@ def _extract_from_and_joins(self, expr: exp.Expression) -> tuple[list[exp.Table] - from_tables: List of Table nodes from FROM clause - join_nodes: List of Join nodes from JOINS clause or FROM clause """ - from_clause = expr.args.get("from_") or expr.args.get("from") + from_clause = expr.args.get("from_") raw_joins = expr.args.get("joins") join_nodes: list[exp.Join] = [] @@ -1753,16 +1613,7 @@ def _columns_reference_tables(self, select_args: list[cst.BaseExpression]) -> bo Returns: True if select_args contain direct table attribute accesses (like Model.column) """ - for arg in select_args: - # Check if this is a direct attribute access (Model.column or Model.column.label(...)) - if isinstance(arg, cst.Attribute): - # E.g., Users.age or Users.id - return True - elif isinstance(arg, cst.Call): - # Could be something like Users.column.label("name") or func.count().label(...) - # For now, be conservative: if it's a function call, assume it doesn't reference tables - pass - return False + return any(isinstance(arg, cst.Attribute) for arg in select_args) def _needs_explicit_select_from( self, @@ -1784,8 +1635,7 @@ def _needs_explicit_select_from( qualified_table = self._get_qualified_star_table(expr) if qualified_table and qualified_table in self.tables: return (True, qualified_table, False) - else: - return (True, next(iter(self.tables.keys())), False) + return (True, next(iter(self.tables.keys())), False) # Always need select_from for CTEs from_clause = expr.args.get("from_") or expr.args.get("from") @@ -1804,19 +1654,15 @@ def _needs_explicit_select_from( if len(from_tables) > 1 or join_nodes: # Use the alias if present, otherwise use the table name first_table = from_tables[0] - first_table_ref = getattr(first_table, "alias", None) or first_table.name + first_table_ref = first_table.alias or first_table.name return (True, first_table_ref, False) # Single table without joins: only need select_from if columns don't already reference the table - if len(from_tables) == 1: - # If select columns reference the table directly (e.g., Users.age), FROM is inferred - if self._columns_reference_tables(select_args): - return (False, None, False) - else: - # Columns are all functions/literals, need explicit FROM - return (True, from_tables[0].name, False) - - return (False, None, False) + # If select columns reference the table directly (e.g., Users.age), FROM is inferred + if self._columns_reference_tables(select_args): + return (False, None, False) + # Columns are all functions/literals, need explicit FROM + return (True, from_tables[0].name, False) def _get_qualified_star_table(self, expr: exp.Expression) -> str | None: """Extract the table name from a qualified star expression. @@ -1824,12 +1670,12 @@ def _get_qualified_star_table(self, expr: exp.Expression) -> str | None: For `SELECT e.*`, returns 'e'. For `SELECT *`, returns None. + Args: + expr: The SELECT expression to check for qualified star + Returns: The table name/alias if the star is qualified, else None. """ - if not hasattr(expr, "expressions"): - return None - for expr_item in expr.expressions: if isinstance(expr_item, exp.Column) and getattr(expr_item, "name", None) == "*": # This is a star - check if it has a table qualifier @@ -1865,34 +1711,16 @@ def _query_delete(self, expr: exp.Delete) -> list[cst.SimpleStatementLine]: def _query_truncate(self, expr: exp.TruncateTable) -> list[cst.SimpleStatementLine]: # expr.expressions contains the table(s) to truncate tables = expr.args.get("expressions", []) - if not tables: - raise NotImplementedError("TRUNCATE without table") table_node = tables[0] - if isinstance(table_node, exp.Table): - table_name = table_node.name - else: - table_name = table_node.name if hasattr(table_node, "name") else str(table_node) + table_name = table_node.name body = [] # For TRUNCATE, we use delete() which achieves similar effect - # query = delete(table) - body.append( - cst.SimpleStatementLine( - [ - cst.Assign( - targets=[cst.AssignTarget(cst.Name("query"))], - value=cst_call_func( - func="delete", - args=[self._cst_table_name(table_name)], - ), - ) - ], - leading_lines=self._leading_comment(), - ) - ) - + # > query = delete(table) + value_call = cst_call_func("delete", args=[self._cst_table_name(table_name)]) + body.append(self._build_assignment("query", value_call, self._leading_comment(), type_hint="Delete")) body.extend(self._compile_and_print()) return body @@ -1908,21 +1736,21 @@ def _extract_ddl_names(self, expr: exp.Expression, is_drop: bool = False) -> str """ if is_drop: # DROP can have multiple items - table_node = self._get_node_arg(expr, "this") - expressions = self._get_node_arg(expr, "expressions", []) + drop_expr: exp.Drop = expr + table_node = drop_expr.args.get("this") names = [] if table_node: names.append(self._get_node_name(table_node)) - if expressions: - names.extend([self._get_node_name(item) for item in expressions]) return ", ".join(names) if names else "unknown" - else: - # CREATE/ALTER have single item - node = expr.this if hasattr(expr, "this") else self._get_node_arg(expr, "this") - return self._get_node_name(node) if node else "unknown" + # CREATE/ALTER have single item + node = expr.args.get("this") + return self._get_node_name(node) if node else "unknown" def _query_ddl( - self, expr: exp.Expression, operation: str, object_type: str = "TABLE" + self, + expr: exp.Expression, + operation: str, + object_type: str = "TABLE", ) -> list[cst.SimpleStatementLine]: """Generic DDL query handler for CREATE, DROP, ALTER operations. @@ -1978,9 +1806,6 @@ def _query_create_sequence(self, expr: exp.Create) -> list[cst.SimpleStatementLi def _query_drop_sequence(self, expr: exp.Drop) -> list[cst.SimpleStatementLine]: return self._query_ddl(expr, "DROP", "SEQUENCE") - def _query_alter_sequence(self, expr: exp.Alter) -> list[cst.SimpleStatementLine]: - return self._query_ddl(expr, "ALTER", "SEQUENCE") - # ------------------------------------------------------------------ # EXPRESSION EMITTERS # ------------------------------------------------------------------ @@ -1991,51 +1816,52 @@ def _get_dispatch_dict(self) -> DispatchDictType: Lazily initialized per instance to support inheritance and method overrides. Bound method references ensure subclass overrides work correctly. """ - if not hasattr(self, "_dispatch_cache"): - self._dispatch_cache: DispatchDictType = { - exp.Alias: self._emit_alias, - exp.Column: self._emit_column, - exp.Count: self._emit_count, - exp.Exists: self._emit_exists, - exp.Select: self._emit_scalar_subquery, - exp.Not: self._emit_not, - exp.Neg: self._emit_neg, - exp.Paren: self._emit_paren, - exp.And: self._emit_and, - exp.Or: self._emit_or, - exp.Like: self._emit_like, - exp.ILike: self._emit_ilike, - exp.Between: self._emit_between, - exp.Case: self._emit_case, - exp.Func: self._emit_func, - exp.Add: self._emit_arithmetic, - exp.Sub: self._emit_arithmetic, - exp.Mul: self._emit_arithmetic, - exp.Div: self._emit_arithmetic, - exp.Mod: self._emit_modulo, - exp.BitwiseAnd: self._emit_bitwise_binary, - exp.BitwiseOr: self._emit_bitwise_binary, - exp.BitwiseXor: self._emit_bitwise_binary, - exp.BitwiseLeftShift: self._emit_bitwise_binary, - exp.BitwiseRightShift: self._emit_bitwise_binary, - exp.BitwiseNot: self._emit_bitwise_not, - exp.Distinct: self._emit_distinct, - exp.Window: self._emit_window, - exp.Ordered: self._emit_ordered, - exp.DPipe: self._emit_concat, - exp.Is: self._emit_is, - exp.In: self._emit_in, - exp.EQ: self._emit_comparison, - exp.NEQ: self._emit_comparison, - exp.GT: self._emit_comparison, - exp.LT: self._emit_comparison, - exp.GTE: self._emit_comparison, - exp.LTE: self._emit_comparison, - exp.Literal: self._emit_literal, - exp.Boolean: self.emit_boolean, - exp.Null: self.emit_none, - exp.Placeholder: self._emit_placeholder, - } + if hasattr(self, "_dispatch_cache"): + return self._dispatch_cache + self._dispatch_cache: DispatchDictType = { + exp.Alias: self._emit_alias, + exp.Column: self._emit_column, + exp.Count: self._emit_count, + exp.Exists: self._emit_exists, + exp.Select: self._emit_scalar_subquery, + exp.Not: self._emit_not, + exp.Neg: self._emit_neg, + exp.Paren: self._emit_paren, + exp.And: self._emit_and, + exp.Or: self._emit_or, + exp.Like: self._emit_like, + exp.ILike: self._emit_ilike, + exp.Between: self._emit_between, + exp.Case: self._emit_case, + exp.Func: self._emit_func, + exp.Add: self._emit_arithmetic, + exp.Sub: self._emit_arithmetic, + exp.Mul: self._emit_arithmetic, + exp.Div: self._emit_arithmetic, + exp.Mod: self._emit_modulo, + exp.BitwiseAnd: self._emit_bitwise_binary, + exp.BitwiseOr: self._emit_bitwise_binary, + exp.BitwiseXor: self._emit_bitwise_binary, + exp.BitwiseLeftShift: self._emit_bitwise_binary, + exp.BitwiseRightShift: self._emit_bitwise_binary, + exp.BitwiseNot: self._emit_bitwise_not, + exp.Distinct: self._emit_distinct, + exp.Window: self._emit_window, + exp.Ordered: self._emit_ordered, + exp.DPipe: self._emit_concat, + exp.Is: self._emit_is, + exp.In: self._emit_in, + exp.EQ: self._emit_comparison, + exp.NEQ: self._emit_comparison, + exp.GT: self._emit_comparison, + exp.LT: self._emit_comparison, + exp.GTE: self._emit_comparison, + exp.LTE: self._emit_comparison, + exp.Literal: self._emit_literal, + exp.Boolean: self._emit_boolean, + exp.Null: self._emit_none, + exp.Placeholder: self._emit_placeholder, + } return self._dispatch_cache def _emit_expr(self, node: exp.Expression) -> cst.BaseExpression: @@ -2047,7 +1873,7 @@ def _emit_expr(self, node: exp.Expression) -> cst.BaseExpression: Uses MRO (Method Resolution Order) to handle subclass types like Sum (subclass of Func). Args: - node: A sqlglot expression node + node: A SQLGlot expression node Returns: A LibCST expression @@ -2070,7 +1896,8 @@ def _emit_expr(self, node: exp.Expression) -> cst.BaseExpression: if handler: return handler(node) - raise NotImplementedError(f"Cannot emit node type {type_node.__name__}: {node}") + msg = f"Cannot emit node type {type_node.__name__}: {node}" + raise NotImplementedError(msg) def _resolve_table_qualified_column(self, node: exp.Column) -> tuple[str | None, str]: """Resolve a column reference to table name and column name. @@ -2122,16 +1949,6 @@ def _emit_column(self, node: exp.Column) -> cst.BaseExpression: Subclasses can override _emit_qualified_column_access to customize how qualified columns are emitted for their dialect. """ - # Handle qualified star like `users.*` represented as a Column with name='*' - if getattr(node, "name", None) == "*": - if node.table: - real_table = self._alias_map.get(node.table, node.table) - return self._cst_table_name(real_table) - # Fallback: if bare '*' somehow arrives here, expand to first FROM table - if self.table_nodes: - return self._cst_table_name(self.table_nodes[0].name) - raise NotImplementedError("Cannot emit bare '*' without FROM tables") - table, column_name = self._resolve_table_qualified_column(node) # If column is unqualified and we have exactly one table in the runtime @@ -2168,10 +1985,7 @@ def _emit_column_access(self, table_or_var: str, column_name: str) -> cst.BaseEx Returns: A CST expression for accessing the column """ - if isinstance(table_or_var, str): - base_var = cst.Name(table_or_var) - else: - base_var = table_or_var + base_var = cst.Name(table_or_var) if isinstance(table_or_var, str) else table_or_var # Get the .c part c_access = cst_attribute(base_var, "c") @@ -2189,20 +2003,27 @@ def _emit_column_access(self, table_or_var: str, column_name: str) -> cst.BaseEx # Use attribute notation: table.c.column_name return cst_attribute(c_access, column_name) - def emit_boolean(self, node: exp.Boolean) -> cst.Name: - return cst.Name("True" if node.this else "False") + def _emit_boolean(self, node: exp.Boolean) -> cst.Call: + """Emit a boolean. + + Since cst.Name("True") or cst.Name("False") would work in comparison, + but would fail in a select list (SELECT True), we emit booleans as false() + or true() function calls to ensure they can be used in all contexts without syntax issues. + """ + return cst_call_func("true" if node.this else "false", args=[]) - def emit_none(self, node: exp.Null) -> cst.Name: + def _emit_none(self, _node: exp.Null) -> cst.Name: + """Emit a NULL literal as cst.Name("None").""" return cst.Name("None") - def _emit_placeholder(self, node: exp.Placeholder) -> cst.BaseExpression: + def _emit_placeholder(self, _node: exp.Placeholder) -> cst.BaseExpression: """Emit a SQL placeholder (?) as SQLAlchemy bindparam(). Generates bindparam('pN') for each placeholder, with N incrementing. This allows users to provide parameter values at execution time. Args: - node: A Placeholder node from sqlglot + node: A Placeholder node from SQLGlot Returns: A Call node representing bindparam('pN') @@ -2227,9 +2048,9 @@ def _emit_alias(self, node: exp.Alias) -> cst.BaseExpression: # Wrap binary and unary expressions in parentheses to ensure correct precedence # Binary: (a * b).label("total") not a * b.label("total") # Unary: (-1).label("alias") not -1.label("alias") - if isinstance(underlying_node, exp.Binary) and isinstance(expr, cst.BinaryOperation): - expr = expr.with_changes(lpar=[cst.LeftParen()], rpar=[cst.RightParen()]) - elif isinstance(underlying_node, (exp.Neg, exp.BitwiseNot)) and isinstance(expr, cst.UnaryOperation): + if (isinstance(underlying_node, exp.Binary) and isinstance(expr, cst.BinaryOperation)) or ( + isinstance(underlying_node, (exp.Neg, exp.BitwiseNot)) and isinstance(expr, cst.UnaryOperation) + ): expr = expr.with_changes(lpar=[cst.LeftParen()], rpar=[cst.RightParen()]) return cst_call_func(cst_attribute(expr, "label"), args=[cst.SimpleString(f'"{node.alias}"')]) @@ -2258,12 +2079,13 @@ def _emit_func(self, node: exp.Func) -> cst.BaseExpression: func_name = type(node).__name__.lower() - # special handling for date/datetime functions that sqlglot parses into TsOrDsToDate nodes with a 'unit' attribute + # special handling for date/datetime functions that SQLGlot parses + # into TsOrDsToDate nodes with a 'unit' attribute if isinstance(node, exp.TsOrDsToDate): func_name = "date" # Generic function argument collection - # Some sqlglot function nodes store the first arg in `this`, + # Some SQLGlot function nodes store the first arg in `this`, # additional args in `expressions`, and some binary function nodes # use `expression` for the second arg. arg_nodes: list = [] @@ -2290,15 +2112,14 @@ def _emit_func(self, node: exp.Func) -> cst.BaseExpression: def _emit_substring(self, node: exp.Substring) -> cst.BaseExpression: """Emit a SUBSTRING expression as func.substring(expr, start, length). - sqlglot parses SUBSTRING from FROM/FOR syntax into dedicated attributes. + SQLGlot parses SUBSTRING from FROM/FOR syntax into dedicated attributes. We extract these and pass them to func.substring(). SQL: SUBSTRING(str FROM start FOR length) SQLAlchemy: func.substring(str, start, length) """ arg_nodes: list = [] - if getattr(node, "this", None) is not None: - arg_nodes.append(node.this) + arg_nodes.append(node.this) if node.args.get("start") is not None: arg_nodes.append(node.args["start"]) if node.args.get("length") is not None: @@ -2310,14 +2131,13 @@ def _emit_substring(self, node: exp.Substring) -> cst.BaseExpression: def _emit_round(self, node: exp.Round) -> cst.BaseExpression: """Emit a ROUND expression as func.round(expr[, decimals]). - sqlglot stores the value in `this` and decimals in args['decimals']. + SQLGlot stores the value in `this` and decimals in args['decimals']. SQL: ROUND(price, 2) SQLAlchemy: func.round(price, 2) """ arg_nodes: list = [] - if getattr(node, "this", None) is not None: - arg_nodes.append(node.this) + arg_nodes.append(node.this) if node.args.get("decimals") is not None: arg_nodes.append(node.args["decimals"]) @@ -2327,7 +2147,7 @@ def _emit_round(self, node: exp.Round) -> cst.BaseExpression: def _emit_trim(self, node: exp.Trim) -> cst.BaseExpression: """Emit a TRIM expression as func.trim(str[, trim_char]). - sqlglot parses TRIM into a Trim node with 'position' (side) and 'this' (string) and optional 'trim_char'. + SQLGlot parses TRIM into a Trim node with 'position' (side) and 'this' (string) and optional 'trim_char'. SQL: TRIM(BOTH 'x' FROM str) SQLAlchemy: func.trim(str, 'x') @@ -2374,10 +2194,8 @@ def _emit_cast(self, node: exp.Cast) -> cst.BaseExpression: # Get the target type target_type = node.to - if not isinstance(target_type, exp.DataType): - raise ValueError(f"Expected DataType in Cast, got {type(target_type)}") - # Map the sqlglot DType to SQLAlchemy type + # Map the SQLGlot DType to SQLAlchemy type dtype_enum = target_type.this sqlalchemy_type_name = DTYPE_TO_SQLALCHEMY.get(dtype_enum, "String") @@ -2404,7 +2222,7 @@ def _emit_comparison(self, node: exp.Binary) -> cst.Comparison: cst.ComparisonTarget( operator=CMP_OPS[type(node)], comparator=self._emit_expr(node.right), - ) + ), ], ) @@ -2436,7 +2254,6 @@ def _is_pure_literal_expr(self, node: exp.Expression) -> bool: if isinstance( node, ( - exp.Binary, exp.Add, exp.Sub, exp.Mul, @@ -2464,36 +2281,21 @@ def _is_pure_literal_expr(self, node: exp.Expression) -> bool: def _emit_is(self, node: exp.Is) -> cst.BaseExpression: # Emit SQL `IS` and `IS NOT` expressions via SQLAlchemy's # `column.is_(None)` / `column.isnot(None)` helpers when comparing - # against NULL, otherwise fall back to equality-style emission. - left, right = self._get_binary_operands(node) + # against NULL. This ensures proper SQL generation of `IS NULL` and `IS NOT NULL`. + left, _right = self._get_binary_operands(node) left_expr = self._emit_expr(left) - # Detect IS [NOT] NULL - is_null = isinstance(right, exp.Null) - not_flag = self._get_node_arg(node, "not", False) or getattr(node, "is_not", False) - - if is_null: - method = "isnot" if not_flag else "is_" - return cst_call_func( - func=cst_attribute(left_expr, method), - args=[cst.Name("None")], - ) - - # Fallback: emit as a normal comparison (e.g., `col == value`) - if not_flag: - return cst.Comparison( - left=left_expr, - comparisons=[cst.ComparisonTarget(operator=cst.NotEqual(), comparator=self._emit_expr(right))], - ) - return cst.Comparison( - left=left_expr, - comparisons=[cst.ComparisonTarget(operator=cst.Equal(), comparator=self._emit_expr(right))], + not_flag = node.args.get("not") or False + method = "isnot" if not_flag else "is_" + return cst_call_func( + func=cst_attribute(left_expr, method), + args=[cst.Name("None")], ) - def _get_binary_operands(self, node: exp.Expression) -> tuple[exp.Expression, exp.Expression]: + def _get_binary_operands(self, node: exp.Binary) -> tuple[exp.Expression, exp.Expression]: """Extract left and right operands from a binary expression node. - Handles different sqlglot structures where operands may be stored + Handles different SQLGlot structures where operands may be stored in different attributes (left/right vs this/expression). Args: @@ -2502,17 +2304,12 @@ def _get_binary_operands(self, node: exp.Expression) -> tuple[exp.Expression, ex Returns: Tuple of (left_operand, right_operand) """ - left = getattr(node, "left", None) or getattr(node, "this", None) - right = getattr(node, "right", None) or getattr(node, "expression", None) - if left is None or right is None: - raise NotImplementedError(f"Cannot extract operands from {type(node).__name__}") - return left, right + return node.left, node.right def _emit_binary_operation( self, - node: exp.Expression, + node: exp.Binary, operator_cls: type | ArithmeticOpsType, - op_type: str = "arithmetic", ) -> cst.BaseExpression: """Emit a generic binary operation with optional literal wrapping. @@ -2542,59 +2339,50 @@ def _emit_binary_operation( op_cls = cls break if op_cls is None: - raise NotImplementedError(f"Unsupported {op_type} operator: {type(node).__name__}") + msg = f"Unsupported operator: {type(node).__name__}" + raise NotImplementedError(msg) else: op_cls = operator_cls return cst.BinaryOperation(left=left_expr, operator=op_cls(), right=right_expr) def _emit_arithmetic(self, node: exp.Expression) -> cst.BaseExpression: - """Emit binary arithmetic: +, -, *, /""" - return self._emit_binary_operation(node, ARITHMETIC_OPS, "arithmetic") + """Emit binary arithmetic: +, -, *, /.""" + return self._emit_binary_operation(node, ARITHMETIC_OPS) def _emit_modulo(self, node: exp.Mod) -> cst.BaseExpression: """Emit modulo operator (%) as a binary operation.""" - return self._emit_binary_operation(node, cst.Modulo, "modulo") + return self._emit_binary_operation(node, cst.Modulo) def _emit_bitwise_binary(self, node: exp.Expression) -> cst.BaseExpression: - """Emit bitwise binary operators: &, |, ^, <<, >>""" - return self._emit_binary_operation(node, BITWISE_OPS, "bitwise") + """Emit bitwise binary operators: &, |, ^, <<, >>.""" + return self._emit_binary_operation(node, BITWISE_OPS) def _emit_neg(self, node: exp.Neg) -> cst.BaseExpression: """Emit negation operator (-) for negative numbers and expressions.""" operand = get_inner_expr(node) - if operand is None: - raise NotImplementedError("Neg node without operand") return cst.UnaryOperation( operator=cst.Minus(), expression=self._emit_expr(operand), ) def _emit_bitwise_not(self, node: exp.BitwiseNot) -> cst.BaseExpression: - """Emit bitwise NOT operator (~)""" + """Emit bitwise NOT operator (~).""" operand = get_inner_expr(node) - if operand is None: - raise NotImplementedError("BitwiseNot node without operand") return cst.UnaryOperation( operator=cst.BitInvert(), expression=self._emit_expr(operand), ) def _emit_distinct(self, node: exp.Distinct) -> cst.BaseExpression: - """Emit DISTINCT modifier (e.g., in COUNT(DISTINCT col))""" + """Emit DISTINCT modifier (e.g., in COUNT(DISTINCT col)).""" # The Distinct node wraps the actual expression expression = getattr(node, "expressions", None) or getattr(node, "this", None) if not expression: - raise NotImplementedError("Distinct node without expression") - - # Handle case where expressions is a list - if isinstance(expression, list) and len(expression) == 1: - expr = expression[0] - elif isinstance(expression, list): - # Multiple expressions - join them - expr = expression - else: - expr = expression + msg = "Distinct node without expression" + raise NotImplementedError(msg) + + expr = expression[0] if isinstance(expression, list) else expression # Emit the inner expression(s) inner = self._emit_expr(expr) if not isinstance(expr, list) else self._emit_expr(expr[0]) @@ -2603,15 +2391,12 @@ def _emit_distinct(self, node: exp.Distinct) -> cst.BaseExpression: return cst_call_func(cst_attribute(value=inner, attr="distinct"), args=[]) def _emit_window(self, node: exp.Window) -> cst.BaseExpression: - """Emit window function with OVER clause (e.g., ROW_NUMBER() OVER (PARTITION BY ... ORDER BY ...))""" + """Emit window function with OVER clause (e.g., ROW_NUMBER() OVER (PARTITION BY ... ORDER BY ...)).""" # Window node structure: # - this: the function (e.g., ROW_NUMBER()) # - args['partition_by']: list of columns to partition by # - args['order']: Order node with ORDER BY expressions - - func_node = node.args.get("this", None) - if not func_node: - raise NotImplementedError("Window node without function") + func_node = node.this # Emit the function call func_expr = self._emit_expr(func_node) @@ -2633,7 +2418,7 @@ def _emit_window(self, node: exp.Window) -> cst.BaseExpression: value=partition_args[0] if len(partition_args) == 1 else cst.List([cst.Element(value=p) for p in partition_args]), - ) + ), ) # Handle ORDER BY @@ -2652,25 +2437,25 @@ def _emit_window(self, node: exp.Window) -> cst.BaseExpression: value=order_args[0] if len(order_args) == 1 else cst.List([cst.Element(value=o) for o in order_args]), - ) + ), ) # Apply .over() method to the function if over_args: return cst_call_func(cst_attribute(value=func_expr, attr="over"), args=over_args) - else: - # If no partition or order, just return the function - return func_expr + # If no partition or order, just return the function + return func_expr def _emit_ordered(self, node: exp.Ordered) -> cst.BaseExpression: - """Emit ordered expression (used in ORDER BY clauses)""" + """Emit ordered expression (used in ORDER BY clauses).""" # Ordered node has: # - this: the column expression # - desc: True if DESC, False if ASC expr = node.this if not expr: - raise NotImplementedError("Ordered node without expression") + msg = "Ordered node without expression" + raise NotImplementedError(msg) # Handle positional references (literal integers like ORDER BY 1, 2) # These need to be wrapped with literal_column() to generate valid Python @@ -2692,10 +2477,9 @@ def _emit_concat(self, node: exp.DPipe) -> cst.BaseExpression: # Collect all concatenated parts (handles chained || operations) parts_final: list[cst.BaseExpression] = [] - def collect_parts(n): + def collect_parts(n: exp.Expression) -> None: if isinstance(n, exp.DPipe): - left = getattr(n, "left", None) or getattr(n, "this", None) - right = getattr(n, "right", None) or getattr(n, "expression", None) + left, right = n.this, n.expression collect_parts(left) collect_parts(right) else: @@ -2742,7 +2526,7 @@ def _emit_and(self, node: exp.And) -> cst.BaseExpression: Converts SQL AND to SQLAlchemy and_() with left and right operands. Args: - node: sqlglot AND expression node. + node: SQLGlot AND expression node. Returns: libcst Call expression for and_(left, right). @@ -2761,7 +2545,7 @@ def _emit_or(self, node: exp.Or) -> cst.BaseExpression: Converts SQL OR to SQLAlchemy or_() with left and right operands. Args: - node: sqlglot OR expression node. + node: SQLGlot OR expression node. Returns: libcst Call expression for or_(left, right). @@ -2780,7 +2564,7 @@ def _emit_like(self, node: exp.Like) -> cst.BaseExpression: Converts SQL LIKE to column.like(pattern) method call. Args: - node: sqlglot LIKE expression node. + node: SQLGlot LIKE expression node. Returns: libcst Call expression for column.like(pattern). @@ -2796,7 +2580,7 @@ def _emit_ilike(self, node: exp.ILike) -> cst.BaseExpression: Converts SQL ILIKE to column.ilike(pattern) method call. Args: - node: sqlglot ILIKE expression node. + node: SQLGlot ILIKE expression node. Returns: libcst Call expression for column.ilike(pattern). @@ -2812,7 +2596,7 @@ def _emit_between(self, node: exp.Between) -> cst.BaseExpression: Converts SQL BETWEEN to column.between(low, high) method call. Args: - node: sqlglot BETWEEN expression node with 'low' and 'high' args. + node: SQLGlot BETWEEN expression node with 'low' and 'high' args. Returns: libcst Call expression for column.between(low, high). @@ -2831,7 +2615,7 @@ def _emit_in(self, node: exp.In) -> cst.BaseExpression: Converts SQL IN to column.in_([values]) method call. Args: - node: sqlglot IN expression node. + node: SQLGlot IN expression node. Returns: libcst Call expression for column.in_([values]) or subquery. @@ -2845,7 +2629,7 @@ def _emit_in_operator(self, node: exp.In, operator: str) -> cst.BaseExpression: for NOT IN and other variants. Args: - node: sqlglot IN expression node. + node: SQLGlot IN expression node. operator: SQLAlchemy method name (e.g., 'in_', 'notin_'). Returns: @@ -2883,11 +2667,11 @@ def _emit_in_operator(self, node: exp.In, operator: str) -> cst.BaseExpression: def _get_in_rhs(self, node: exp.In) -> exp.Expression | list[exp.Expression] | None: """Extract the right-hand side of an IN clause (list or subquery). - Handles various sqlglot attribute names for the right operand. + Handles various SQLGlot attribute names for the right operand. Looks for: expression, right, query, expressions. Args: - node: sqlglot IN expression node. + node: SQLGlot IN expression node. Returns: Expression (subquery), list of expressions, or None. @@ -2934,7 +2718,7 @@ def _emit_case(self, node: exp.Case) -> cst.BaseExpression: cst.ComparisonTarget( operator=cst.Equal(), comparator=when_value, - ) + ), ], ) else: diff --git a/sqlalchemy_cauldron/generator/helpers.py b/sqlalchemy_cauldron/generator/helpers.py index 2e73ffd..177df66 100644 --- a/sqlalchemy_cauldron/generator/helpers.py +++ b/sqlalchemy_cauldron/generator/helpers.py @@ -3,7 +3,7 @@ import builtins import keyword -# Mapping from sqlglot dialect names to SQLAlchemy dialect module names +# Mapping from SQLGlot dialect names to SQLAlchemy dialect module names # Used for importing dialect-specific modules (e.g., sqlalchemy.dialects.postgresql) # None means no specific dialect (use SQLAlchemy default) DIALECT_MAPPING: dict[str | None, str | None] = { @@ -17,44 +17,49 @@ # Shared SQLAlchemy imports used by both CORE and ORM generators COMMON_IMPORTS: list[str] = [ - "Column", - "Integer", - "SmallInteger", + "and_", "BigInteger", - "String", - "NVARCHAR", - "NCHAR", - "Float", - "Double", "Boolean", "BOOLEAN", - "DateTime", - "Numeric", - "Text", + "case", + "cast", + "CHAR", + "Column", "Date", - "Time", + "DateTime", "DATETIME", - "CHAR", - "Select", - "cast", + "delete", + "Delete", + "Double", + "except_", + "except_all", + "exists", + "false", + "Float", "func", - "select", "insert", - "update", - "delete", - "and_", - "or_", - "not_", - "literal", + "Insert", + "Integer", + "intersect_all", + "intersect", "literal_column", - "case", - "exists", - "union", + "literal", + "NCHAR", + "not_", + "Numeric", + "NVARCHAR", + "or_", + "select", + "Select", + "SmallInteger", + "String", + "Text", + "Time", + "true", "union_all", - "intersect", - "intersect_all", - "except_", - "except_all", + "union", + "update", + "Update", ] # CORE-specific imports (CORE generator only) @@ -66,8 +71,8 @@ # ORM-specific imports (ORM generator only) ORM_IMPORTS: list[str] = [ - "desc", "asc", + "desc", *COMMON_IMPORTS, ] @@ -85,7 +90,7 @@ # SQLAlchemy imports from all generators *CORE_IMPORTS, *ORM_IMPORTS, - } + }, ) BUILTIN_NAMES = frozenset(dir(builtins)) @@ -118,11 +123,11 @@ "convert", "create", "cross", - "current", "current_date", "current_time", "current_timestamp", "current_user", + "current", "cursor", "database", "date", @@ -226,8 +231,8 @@ "schema", "select", "sensitive", - "session", "session_user", + "session", "set", "signal", "similar", @@ -241,8 +246,8 @@ "start", "static", "symmetric", - "system", "system_user", + "system", "table", "then", "time", @@ -276,7 +281,7 @@ "without", "year", "zone", - } + }, ) @@ -320,7 +325,7 @@ def safe_table_name(name: str) -> str: or safe in BUILTIN_NAMES or safe.lower() in SQL_KEYWORDS ): - safe = safe + "_t" + safe += "_t" return safe @@ -343,7 +348,7 @@ def safe_column_name(name: str) -> str: """ safe = name while keyword.iskeyword(safe) or keyword.issoftkeyword(safe) or safe in BUILTIN_NAMES: - safe = safe + "_col" + safe += "_col" return safe @@ -368,5 +373,5 @@ def safe_column_name_for_orm(name: str) -> str: safe = name # Only escape hard keywords, not builtins while keyword.iskeyword(safe) or keyword.issoftkeyword(safe): - safe = safe + "_col" + safe += "_col" return safe diff --git a/sqlalchemy_cauldron/generator/orm_generator.py b/sqlalchemy_cauldron/generator/orm_generator.py index ed19363..e710a35 100644 --- a/sqlalchemy_cauldron/generator/orm_generator.py +++ b/sqlalchemy_cauldron/generator/orm_generator.py @@ -1,11 +1,12 @@ from __future__ import annotations +from typing import TYPE_CHECKING, Any + import libcst as cst -from sqlalchemy import Column, Table -from sqlglot import exp +from sqlalchemy import Column, Integer, Table -from sqlalchemy_cauldron.config import Dialect +from sqlalchemy_cauldron.config import FALLBACK_PK, Dialect from sqlalchemy_cauldron.generator.core_generator import ( CodeGenerator, TablesType, @@ -23,14 +24,25 @@ from sqlalchemy_cauldron.generator.type_inference import ( ColumnType, TypeInference, - get_python_type_hint, - get_python_type_hint_for_sqlalchemy_type, ) +if TYPE_CHECKING: + from sqlglot import exp + + +def get_table_columns_safe(table: Table, pk_column: str) -> list[Column[Any]]: + """Get the list of column names for a table, ensuring the primary key is included.""" + columns = list(table.columns) + # If table has no columns, add a primary key placeholder to satisfy SQLAlchemy + if not columns: + if not pk_column: + pk_column = FALLBACK_PK + columns = [Column(pk_column, Integer, primary_key=True)] + return columns + class ORMCodeGenerator(CodeGenerator): - """ - Generates ORM-style Python code using SQLAlchemy declarative models (SQLAlchemy 2.0+). + """Generates ORM-style Python code using SQLAlchemy declarative models (SQLAlchemy 2.0+). Inherits table schema extraction from CodeGenerator (CORE). Generates model class definitions inheriting from DeclarativeBase. @@ -45,8 +57,18 @@ def generate_module( dialect: Dialect = Dialect.POSTGRES, include_self_contained_test: bool = False, ) -> str: - self._recursive_cte_aliases: set[str] = set() # Track RECURSIVE CTE aliases + """Generate a complete Python module with ORM models and query code. + + Args: + expr: The parsed SQL expression to generate code for + tables: The extracted table schemas and metadata + original_sql: The original SQL string (for reference in comments) + dialect: The SQL dialect to target (affects imports and type mapping) + include_self_contained_test: Whether to append a self-contained test function at the end + Returns: + A string containing the complete generated Python module code + """ # Generate the main module code using the parent method, which will call our overridden methods # for imports, table definitions, and query generation. code = super().generate_module( @@ -68,8 +90,7 @@ def generate_module( # ------------------------------------------------------------------ def _imports(self) -> list[cst.SimpleStatementLine]: - """Import DeclarativeBase and necessary SQLAlchemy ORM components (SQLAlchemy 2.0+)""" - + """Import DeclarativeBase and necessary SQLAlchemy ORM components (SQLAlchemy 2.0+).""" # Use shared helper to build imports imports = [ build_import_from("sqlalchemy", ORM_IMPORTS), @@ -77,9 +98,7 @@ def _imports(self) -> list[cst.SimpleStatementLine]: ] # Add dialect import if available - dialect_import = self._build_dialect_import() - if dialect_import: - imports.append(dialect_import) + imports.extend(self._build_dialect_import(self.dialect or "")) # Check if any DateTime columns are used - if so, import datetime module if self._has_datetime_columns(): @@ -92,7 +111,7 @@ def _metadata_definitions(self) -> list[cst.SimpleStatementLine]: return [] def _has_datetime_columns(self) -> bool: - """Check if any columns use DateTime type""" + """Check if any columns use DateTime type.""" for table in self.tables.values(): for col in table.columns: inferred_type = TypeInference.infer(col.name) @@ -113,109 +132,85 @@ class Base(DeclarativeBase): """ body: list = [] - body.append(cst.EmptyLine()) - body.append( - cst.ClassDef( - name=cst.Name("Base"), - bases=[cst.Arg(cst.Name("DeclarativeBase"))], - body=cst.IndentedBlock(body=[cst.SimpleStatementLine([cst.Pass()])]), - ) + body.extend( + ( + cst.EmptyLine(), + cst.ClassDef( + name=cst.Name("Base"), + bases=[cst.Arg(cst.Name("DeclarativeBase"))], + body=cst.IndentedBlock(body=[cst.SimpleStatementLine([cst.Pass()])]), + ), + cst.EmptyLine(), + ), ) - body.append(cst.EmptyLine()) # Generate a model class for each table - for table_name, table in sorted(self.tables.items()): + for _table_name, table in sorted(self.tables.items()): # Skip AliasedTable instances - they're handled by _generate_alias_assignments() if isinstance(table, AliasedTable): continue - # Don't generate models for internal tables - if table_name.startswith("_"): - continue - - class_def = self._generate_model_class(table_name, table) - body.append(class_def) - body.append(cst.EmptyLine()) + class_def = self._generate_model_class(table) + body.extend((class_def, cst.EmptyLine())) return body - def _generate_model_class(self, table_name: str, table: Table) -> cst.ClassDef: - """Generate a single model class definition""" - model_name = safe_model_name(table_name) + def _generate_model_class(self, table: Table) -> cst.ClassDef: + """Generate a single model class definition.""" + model_name = safe_model_name(table.name) # Generate class body class_body_stmts: list = [] - # __tablename__ = "table_name" + # > __tablename__ = "table_name" class_body_stmts.append( - cst.SimpleStatementLine( - [ - cst.Assign( - targets=[cst.AssignTarget(cst.Name("__tablename__"))], - value=cst.SimpleString(f'"{table_name}"'), - ) - ] - ) + self._build_assignment( + target_name="__tablename__", + value_expr=cst.SimpleString(f'"{table.name}"'), + ), ) - # Generate Column definitions with type hints - pk_column = TypeInference.infer_primary_key_from_table(table_name, table) - - # If table has no columns, add a primary key placeholder to satisfy SQLAlchemy - if not table.columns and pk_column: - col_def = self._generate_column_definition(col_name=pk_column, sql_type="Integer", is_primary_key=True) - class_body_stmts.append(col_def) - else: - # Generate columns as usual - for col in table.columns: - col_def = self._generate_column_definition(col=col, is_primary_key=col.name == pk_column) - class_body_stmts.append(col_def) + class_body_stmts.extend(self._generate_columns_definitions(table)) # Create ClassDef with proper body - class_def = cst.ClassDef( + return cst.ClassDef( name=cst.Name(model_name), bases=[cst.Arg(cst.Name("Base"))], body=cst.IndentedBlock(body=class_body_stmts), ) - return class_def + def _generate_columns_definitions(self, table: Table) -> list[cst.SimpleStatementLine]: + pk_column = TypeInference.infer_primary_key_from_table(table) + columns: list[Column[Any]] = get_table_columns_safe(table, pk_column=pk_column) + body: list[cst.SimpleStatementLine] = [] + + # Generate columns + for col in columns: + col_def = self._generate_column_definition(col=col, is_primary_key=col.name == pk_column) + body.append(col_def) + + return body def _generate_column_definition( self, - col: Column | None = None, - col_name: str | None = None, - sql_type: str | None = None, + col: Column, is_primary_key: bool = False, ) -> cst.SimpleStatementLine: - """Generate a single Column definition with type hint. + """Generate a single Column. Can be called with either: - A Column object (from table.columns) - - A column name and SQL type string (for placeholder columns) Args: col: A Column object (SQLAlchemy Table column) - col_name: The column name (when col is None) - sql_type: The SQL type name as string (when col is None) is_primary_key: Whether this is a primary key column Returns: A SimpleStatementLine with the column definition """ - from sqlalchemy_cauldron.generator.type_inference import TYPE_MAP, TypeInference - - # Determine column name - if col is not None: - column_name = col.name - # Infer column type from column name - inferred_type = TypeInference.infer(column_name) - sa_type = TYPE_MAP[inferred_type] - type_name = sa_type.__name__ - else: - if col_name is None or sql_type is None: - raise ValueError("Must provide either col or both col_name and sql_type") - column_name = col_name - type_name = sql_type + # Determine column name and type + column_name = col.name + type_name = col.type.__class__.__name__ # Build Column() call with arguments # Get safe Python attribute name for this column @@ -235,35 +230,18 @@ def _generate_column_definition( cst.Arg( keyword=cst.Name("primary_key"), value=cst.Name("True"), - ) + ), ) col_call = cst_call_func(func="Column", args=col_args) - # Get Python type hint - if col is not None: - inferred_type = TypeInference.infer(column_name) - type_hint = get_python_type_hint(inferred_type) - else: - # Map SQL type name to Python type hint using the shared helper - type_hint = get_python_type_hint_for_sqlalchemy_type(type_name) - - # Generate: safe_attr_name: python_type = Column(...) - # Use safe_column_name_for_orm to escape Python keywords in attribute names - return self._build_assignment(target_name=safe_attr_name, value_expr=col_call, type_hint=type_hint) - - # ------------------------------------------------------------------ - # ORM-specific query generation (override CORE methods) - # ------------------------------------------------------------------ + # > safe_attr_name = Column(...) + return self._build_assignment(target_name=safe_attr_name, value_expr=col_call, type_hint=None) def _model_class_name(self, table_name: str) -> str: - """Convert table name to ORM model class name""" + """Convert table name to ORM model class name.""" return safe_model_name(table_name) - # ------------------------------------------------------------------ - # ALIAS TRACKING (simplified with AliasedTable) - # ------------------------------------------------------------------ - def _table_name(self, table_name: str) -> str: """Override parent to return model class name, or the alias if it's an AliasedTable. @@ -275,7 +253,7 @@ def _table_name(self, table_name: str) -> str: table = self.tables[table_name] # If it's an AliasedTable, return its alias (e.g., 'e', 'm') if isinstance(table, AliasedTable): - return table.table_alias + return table.alias # Otherwise it's a regular Table, so return the model class name return self._model_class_name(table_name) # Fallback: convert to class name @@ -323,12 +301,11 @@ def _resolve_join_table_ref(self, jt_name: str | None, jt_alias: str | None) -> if jt_name and jt_name in self._recursive_cte_aliases: # RECURSIVE CTE: use class name to avoid forward reference return self._table_name(jt_name) - elif jt_name and self._is_cte_alias(jt_name): + if jt_name and self._is_cte_alias(jt_name): # Non-recursive CTE: use variable name return jt_name - else: - # For regular tables/aliases, use parent's logic - return super()._resolve_join_table_ref(jt_name, jt_alias) + # For regular tables/aliases, use parent's logic + return super()._resolve_join_table_ref(jt_name, jt_alias) def _emit_qualified_column_access(self, table_or_var: str, column_name: str) -> cst.BaseExpression: """Override parent to emit ORM column access (ModelClass.column) instead of CORE syntax (table.c.column). @@ -358,8 +335,7 @@ def _emit_qualified_column_access(self, table_or_var: str, column_name: str) -> def _apply_select_clauses( self, query_expr: cst.BaseExpression, - expr: exp.Expression, - include_from_join: bool = True, + expr: exp.Select, ) -> cst.BaseExpression: """Override parent to handle ORM-specific FROM/JOIN logic. @@ -370,104 +346,75 @@ def _apply_select_clauses( if expr.args.get("distinct"): query_expr = self._chain_method(query_expr, "distinct", []) + # FROM / JOIN + from_tables, _join_nodes = self._extract_from_and_joins(expr) + # For ORM, don't call .join() on class names + # Instead, use select_from() then chain .join() + first_table = from_tables[0].name + + if not self._is_select_star_query(expr=expr) and not self._is_cte_alias(first_table): + # Single table without joins (and not SELECT * or CTE which are handled elsewhere) + query_expr = self._chain_method( + query_expr, + "select_from", + [cst.Arg(self._cst_table_name(first_table))], + ) + # WHERE, GROUP BY, HAVING, ORDER BY using shared helpers query_expr = self._apply_where_clause(query_expr, expr) query_expr = self._apply_group_by_clause(query_expr, expr) query_expr = self._apply_having_clause(query_expr, expr) - # FROM / JOIN (ORM-specific handling) - if include_from_join: - from_tables, join_nodes = self._extract_from_and_joins(expr) - - if from_tables: - # For ORM, don't call .join() on class names - # Instead, use select_from() then chain .join() - first_table = from_tables[0].name - - if join_nodes or len(from_tables) > 1: - # We have joins - build the FROM with proper ORM semantics - # select().select_from(Table).join(...) - query_expr = self._chain_method( - query_expr, "select_from", [cst.Arg(self._cst_table_name(first_table))] - ) - - if join_nodes: - for j in join_nodes: - jt_name, args_list, join_method = self._process_join_node(j) - query_expr = self._chain_method(query_expr, join_method, args_list) - else: - for t in from_tables[1:]: - jt_name = t.name - query_expr = self._chain_method( - query_expr, "join", [cst.Arg(self._cst_table_name(jt_name))] - ) - elif not self._is_select_star_query(expr=expr) and not self._is_cte_alias(first_table): - # Single table without joins (and not SELECT * or CTE which are handled elsewhere) - query_expr = self._chain_method( - query_expr, "select_from", [cst.Arg(self._cst_table_name(first_table))] - ) - # ORDER BY using shared helper query_expr = self._apply_order_by_clause(query_expr, expr) # Apply LIMIT and OFFSET using shared helper - query_expr = self._apply_limit_offset_clauses(query_expr, expr) - - return query_expr + return self._apply_limit_offset_clauses(query_expr, expr) def _query_select(self, expr: exp.Select) -> list[cst.SimpleStatementLine]: - """Override parent's _query_select to use ORM-compatible method chaining with best practices. + """Override parent's _query_select to use ORM-compatible method chaining. Method order follows SQLAlchemy 2.0 best practices: - 1. select() - 2. .select_from() early if needed (to set FROM context) - 3. .distinct() - 4. .where(), .group_by(), .having() - 5. .join() - 6. .order_by() - 7. .limit(), .offset() + - select() + - .select_from() early if needed (to set FROM context) + - .join() + - .distinct() + - .where(), .group_by(), .having() + - .order_by() + - .limit(), .offset() """ body = [] - # Process CTEs - these remain as separate statements - body.extend(self._process_ctes(expr)) + # Emit CTEs using shared helper + self._build_and_emit_ctes(expr, body) # Build the initial select() call using parent's column processor final_select_args = self._process_select_columns(expr) query_expr = cst_call_func(func="select", args=final_select_args) - # Determine if we need explicit .select_from() and add it early (best practice) + # Determine if we need explicit .select_from() and add it early needs_select_from, table_name, is_cte = self._needs_explicit_select_from(expr, final_select_args) if needs_select_from and table_name: # For CTEs, use the CTE variable name directly; for tables, use _cst_table_name - if is_cte: - select_from_arg = cst.Name(table_name) - else: - select_from_arg = self._cst_table_name(table_name) + select_from_arg = cst.Name(table_name) if is_cte else self._cst_table_name(table_name) query_expr = self._chain_method(query_expr, "select_from", [cst.Arg(select_from_arg)]) # DISTINCT if expr.args.get("distinct"): query_expr = self._chain_method(query_expr, "distinct", []) - # Use shared helpers for standard clause application - query_expr = self._apply_where_clause(query_expr, expr) - query_expr = self._apply_group_by_clause(query_expr, expr) - query_expr = self._apply_having_clause(query_expr, expr) - - # JOINs (ORM-specific) + # JOINs (ORM specific) from_tables, join_nodes = self._extract_from_and_joins(expr) - if from_tables and join_nodes: # Chain each JOIN for j in join_nodes: - jt_name, args_list, join_method = self._process_join_node(j) - query_expr = self._chain_method(query_expr, join_method, args_list) - elif from_tables and len(from_tables) > 1: - # Multiple tables without explicit joins - for t in from_tables[1:]: - jt_name = t.name - query_expr = self._chain_method(query_expr, "join", [cst.Arg(self._cst_table_name(jt_name))]) + args_list = self._process_join_node(j) + query_expr = self._chain_method(query_expr, "join", args_list) + + # Use shared helpers for WHERE, GROUP BY, HAVING clauses + query_expr = self._apply_where_clause(query_expr, expr) + query_expr = self._apply_group_by_clause(query_expr, expr) + query_expr = self._apply_having_clause(query_expr, expr) # Use shared helper for ORDER BY query_expr = self._apply_order_by_clause(query_expr, expr) @@ -480,7 +427,8 @@ def _query_select(self, expr: exp.Select) -> list[cst.SimpleStatementLine]: return body def _generate_self_contained_test_function(self) -> str: - """Generate a test function that creates an in-memory SQLite database, + """Generate a test function that creates an in-memory SQLite database,. + populates it with faker-generated test data, and runs basic queries. This function is appended to the generated module when include_self_contained_test=True. @@ -489,31 +437,3 @@ def _generate_self_contained_test_function(self) -> str: String containing just the test function definition (no module imports) """ return generate_self_contained_test_function(self.tables, self.dialect) - - def _process_ctes(self, expr: exp.Expression) -> list[cst.SimpleStatementLine]: - """Process Common Table Expressions (CTEs) with ORM tracking. - - Builds CTE variable assignments and populates internal alias tracking - (_cte_aliases, _recursive_cte_aliases) so that column references within - CTE definitions can be properly resolved. - - Returns: - List of libcst statement lines for CTE definitions. - """ - body = [] - - # Populate recursive CTE tracking BEFORE building CTEs - # This is needed so _emit_column can check during CTE definition building - with_node = expr.args.get("with_") - if with_node and getattr(with_node, "expressions", None): - recursive = getattr(with_node, "recursive", False) - for cte in with_node.expressions: - alias_name = self._extract_cte_alias_name(cte) - self._cte_aliases.add(alias_name) - if recursive: - self._recursive_cte_aliases.add(alias_name) - - # Use shared helper to build and emit CTEs (now with tracking already populated) - self._build_and_emit_ctes(expr, body) - - return body diff --git a/sqlalchemy_cauldron/generator/orm_test_generator.py b/sqlalchemy_cauldron/generator/orm_test_generator.py index 6297d34..e80350f 100644 --- a/sqlalchemy_cauldron/generator/orm_test_generator.py +++ b/sqlalchemy_cauldron/generator/orm_test_generator.py @@ -2,22 +2,24 @@ from __future__ import annotations +from typing import TYPE_CHECKING + from sqlalchemy_cauldron.config import Dialect -from sqlalchemy_cauldron.generator.core_generator import TablesType from sqlalchemy_cauldron.generator.helpers import safe_column_name_for_orm, safe_model_name, safe_table_name from sqlalchemy_cauldron.generator.sqlalchemy_runtime import AliasedTable from sqlalchemy_cauldron.generator.type_inference import ColumnType, TypeInference +if TYPE_CHECKING: + from sqlalchemy_cauldron.generator.core_generator import TablesType + def generate_self_contained_test_function( tables: TablesType, dialect: Dialect = Dialect.POSTGRES, ) -> str: - """Generate a test function that creates an in-memory SQLite database, - populates it with faker-generated test data, and runs basic queries. + """Generate a test function that creates an in-memory SQLite database,. - This function assumes the module already has the necessary imports: - - sqlite3, datetime, sqlalchemy, faker (these should be added to the ORM imports) + populates it with faker-generated test data, and runs basic queries. Args: tables: Dictionary of table objects with column information @@ -26,11 +28,9 @@ def generate_self_contained_test_function( Returns: String containing the test function definition (without module-level imports) """ - # Map inferred types to faker methods type_faker_map = { ColumnType.VARCHAR: "faker.unique.word()", - ColumnType.STRING: "faker.unique.word()", ColumnType.TEXT: "faker.unique.paragraph()", ColumnType.INTEGER: "faker.unique.random_int(min=1, max=1000)", ColumnType.FLOAT: "faker.unique.pyfloat(min_value=0.0, max_value=1000.0)", @@ -42,8 +42,7 @@ def generate_self_contained_test_function( lines = [ "", 'if __name__ == "__main__":', - " # Self-contained test: create in-memory SQLite DB and populate with faker data", - " import sqlite3", + f" # Self-contained test: ({dialect.value}) create in-memory SQLite DB and populate with faker data", " from datetime import datetime", " from sqlalchemy import create_engine, select, func", " from sqlalchemy.orm import Session", @@ -56,6 +55,7 @@ def generate_self_contained_test_function( " Base.metadata.create_all(engine)", "", " # Initialize faker", + " Faker.seed(42)", " faker = Faker()", "", " # Create a session", @@ -71,9 +71,13 @@ def generate_self_contained_test_function( model_name = safe_model_name(table_name) var_name = safe_table_name(table_name) - lines.append(f" # Populate {table_name} with test data") - lines.append(" for _ in range(10):") - lines.append(f" {var_name} = {model_name}(") + lines.extend( + ( + f" # Populate {table_name} with test data", + " for _ in range(10):", + f" {var_name} = {model_name}(", + ), + ) for col in table.columns: col_name = col.name @@ -82,9 +86,13 @@ def generate_self_contained_test_function( faker_expr = type_faker_map.get(inferred_type, "faker.word()") lines.append(f" {safe_attr_name}={faker_expr},") - lines.append(" )") - lines.append(f" session.add({var_name})") - lines.append("") + lines.extend( + ( + " )", + f" session.add({var_name})", + "", + ), + ) lines.extend( [ @@ -92,7 +100,7 @@ def generate_self_contained_test_function( " ", " # Run basic sanity checks", "", - ] + ], ) # Add basic queries for each table @@ -107,18 +115,22 @@ def generate_self_contained_test_function( lines.extend( [ f" # Verify {table_name} was populated", - f" {var_name}_count = session.scalar(select(func.count()).select_from({model_name}))", + f" {var_name}_count = session.scalar(select(func.count()).select_from({model_name})) or 0", f' assert {var_name}_count > 0, f"Expected {table_name} to have rows"', f' print(f"OK - {table_name}: {{{var_name}_count}} rows")', "", - ] + ], ) # execute the original query and print results - lines.append(" result = session.execute(query).fetchall()") - lines.append(' print(f"OK - query returned:\\n{result}")') - lines.append("") - lines.append(' print("All tests passed!")') - lines.append("") + lines.extend( + ( + " result = session.execute(query).fetchall()", + ' print(f"OK - query returned:\\n{result}")', + "", + ' print("All tests passed!")', + "", + ), + ) return "\n".join(lines) diff --git a/sqlalchemy_cauldron/generator/schema_extractor.py b/sqlalchemy_cauldron/generator/schema_extractor.py index 382ce2b..0b617c4 100644 --- a/sqlalchemy_cauldron/generator/schema_extractor.py +++ b/sqlalchemy_cauldron/generator/schema_extractor.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections import defaultdict from dataclasses import dataclass, field from sqlglot import exp @@ -19,12 +20,6 @@ class InfoItem: name: str alias: str | None = None - def __str__(self) -> str: - return f"{self.name}" + (f" AS {self.alias}" if self.alias else "") - - def __hash__(self) -> int: - return hash((self.name, self.alias)) - @dataclass class ColumnInfo(InfoItem): @@ -66,8 +61,8 @@ def unique(self) -> None: class SchemaExtractor: - """ - Extracts table and column metadata from a parsed sqlglot expression. + """Extracts table and column metadata from a parsed sqlglot expression. + Returns a list of TableInfo objects containing table names, aliases, and columns. Handles SELECT * by collecting all referenced columns from the full query. @@ -75,7 +70,14 @@ class SchemaExtractor: @staticmethod def _get_column_alias(col: exp.Column) -> str | None: - """Get the column alias if it's wrapped in an Alias node.""" + """Get the column alias if it's wrapped in an Alias node. + + Args: + col: The column expression to check. + + Returns: + The alias name if present, otherwise None. + """ alias_node = col.find_ancestor(exp.Alias) if alias_node and hasattr(alias_node, "alias"): return alias_node.alias or None @@ -91,28 +93,23 @@ def _add_column( ) -> None: """Add a column to a table in table_info_dict if the table exists.""" key = (table_name, alias) - if key in table_info_dict: - table_info_dict[key].columns.append(ColumnInfo(name=col_name, alias=col_alias)) + table_info_dict[key].columns.append(ColumnInfo(name=col_name, alias=col_alias)) - def _get_select_tables(self, select_parent: exp.Select | None) -> dict[str, set[str | None]]: + def _get_select_tables(self, node: exp.Select | None) -> dict[str, set[str | None]]: """Extract tables from a SELECT's FROM clause and JOINs.""" - select_tables = {} - if not select_parent: + select_tables = defaultdict(set) + if not node: return select_tables - from_clause = select_parent.args.get("from_") or select_parent.args.get("from") + from_clause = node.args.get("from_") if from_clause: for table in from_clause.find_all(exp.Table): - if table.name not in select_tables: - select_tables[table.name] = set() select_tables[table.name].add(table.alias or None) - joins = select_parent.args.get("joins") + joins = node.args.get("joins") if joins: for join in joins: for table in join.find_all(exp.Table): - if table.name not in select_tables: - select_tables[table.name] = set() select_tables[table.name].add(table.alias or None) return select_tables @@ -136,17 +133,15 @@ def _find_column_tables( if table_ref in alias_to_table: # It's an alias key = (real_table, table_ref) - if key in table_info_dict: - results.append(key) - else: + results.append(key) + elif (real_table, None) in table_info_dict: # It's a table name - try None first, then pick first alias - if (real_table, None) in table_info_dict: - results.append((real_table, None)) - else: - aliases_for_table = table_to_aliases.get(real_table, set()) - if aliases_for_table: - first_alias = next(iter(aliases_for_table)) - results.append((real_table, first_alias)) + results.append((real_table, None)) + else: + aliases_for_table = table_to_aliases.get(real_table, set()) + if aliases_for_table: + first_alias = next(iter(aliases_for_table)) + results.append((real_table, first_alias)) else: # Unqualified column: try SELECT context first, then global fallback select_tables = self._get_select_tables(select_parent) @@ -154,13 +149,11 @@ def _find_column_tables( if len(select_tables) == 1: # Single table in this SELECT only_table = next(iter(select_tables)) - for alias in select_tables[only_table]: - results.append((only_table, alias)) + results.extend((only_table, alias) for alias in select_tables[only_table]) elif len(table_to_aliases) == 1: # Single table globally only_table = next(iter(table_to_aliases)) - for alias in table_to_aliases[only_table]: - results.append((only_table, alias)) + results.extend((only_table, alias) for alias in table_to_aliases[only_table]) return results @@ -168,7 +161,7 @@ def extract(self, expr: exp.Expression, add_missing_pk: bool = True) -> list[Tab """Extract table and column information from an expression. Args: - expr: A sqlglot Expression (e.g. parsed from a SQL query) + expr: SQLGlot Expression (e.g. parsed from a SQL query) add_missing_pk: Whether to add a fallback primary key if none is found Returns: @@ -177,9 +170,11 @@ def extract(self, expr: exp.Expression, add_missing_pk: bool = True) -> list[Tab """ # Use (table_name, alias) tuple as key to handle self-joins # where the same table appears with different aliases - table_info_dict: dict[tuple[str, str | None], TableInfo] = {} + table_info_dict: dict[tuple[str, str | None], TableInfo] = defaultdict( + lambda: TableInfo(name="", alias=None, columns=[]), + ) alias_to_table: dict[str, str] = {} # key: alias, value: table name - table_to_aliases: dict[str, set[str | None]] = {} # key: table name, value: set of aliases + table_to_aliases: dict[str, set[str | None]] = defaultdict(set) # key: table name, value: set of aliases # First pass: collect all tables and their aliases for table in expr.find_all(exp.Table): @@ -187,8 +182,6 @@ def extract(self, expr: exp.Expression, add_missing_pk: bool = True) -> list[Tab alias = table.alias or None # Normalize empty string to None # Track which aliases exist for each table (None means no alias) - if table_name not in table_to_aliases: - table_to_aliases[table_name] = set() table_to_aliases[table_name].add(alias) # Create a TableInfo for this (table, alias) pair if we haven't seen it @@ -202,7 +195,7 @@ def extract(self, expr: exp.Expression, add_missing_pk: bool = True) -> list[Tab # Second pass: collect columns for col in expr.find_all(exp.Column): - # Skip wildcard columns like * or table.* + # Skip asterisk columns like * or table.* if col.name == "*": continue @@ -222,19 +215,17 @@ def extract(self, expr: exp.Expression, add_missing_pk: bool = True) -> list[Tab if isinstance(expr, exp.Insert) and isinstance(expr.this, exp.Schema): table_name = expr.this.this.name # INSERT never has aliases, so use (table_name, None) - if (table_name, None) in table_info_dict: - for col_id in expr.this.expressions: - if hasattr(col_id, "name"): - self._add_column(table_info_dict, table_name, None, col_id.name, None) + for col in expr.this.expressions: + self._add_column(table_info_dict, table_name, None, col.name, None) # For CREATE TABLE: extract columns from the Schema node if isinstance(expr, exp.Create) and isinstance(expr.this, exp.Schema): table_name = expr.this.this.name # CREATE TABLE never has aliases, so use (table_name, None) - if (table_name, None) in table_info_dict: - for col_def in expr.this.expressions: - if isinstance(col_def, exp.ColumnDef) and hasattr(col_def, "name"): - self._add_column(table_info_dict, table_name, None, col_def.name, None) + for col in expr.this.expressions: + if not isinstance(col, exp.ColumnDef): + continue + self._add_column(table_info_dict, table_name, None, col.name, None) # For any table with no extracted columns, add a default FALLBACK_PK primary key # This handles: SELECT * FROM users, SELECT 1 FROM users, SELECT COUNT(*) FROM users, etc. @@ -245,14 +236,12 @@ def extract(self, expr: exp.Expression, add_missing_pk: bool = True) -> list[Tab # Merge columns across all aliases of the same table # For example, if we have "employees e" and "employees m", both should have all columns used - table_to_all_columns: dict[str, list[ColumnInfo]] = {} - for (table_name, alias), table_info in table_info_dict.items(): - if table_name not in table_to_all_columns: - table_to_all_columns[table_name] = [] + table_to_all_columns: dict[str, list[ColumnInfo]] = defaultdict(list) + for (table_name, _alias), table_info in table_info_dict.items(): table_to_all_columns[table_name].extend(table_info.columns) # Apply merged columns back to all aliases of each table - for (table_name, alias), table_info in table_info_dict.items(): + for (table_name, _alias), table_info in table_info_dict.items(): table_info.columns = list(table_to_all_columns[table_name]) table_info.unique() # Deduplicate again after merge diff --git a/sqlalchemy_cauldron/generator/sqlalchemy_runtime.py b/sqlalchemy_cauldron/generator/sqlalchemy_runtime.py index 2544c3f..e42a176 100644 --- a/sqlalchemy_cauldron/generator/sqlalchemy_runtime.py +++ b/sqlalchemy_cauldron/generator/sqlalchemy_runtime.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any from sqlalchemy import ( Column, @@ -9,13 +9,15 @@ ) from sqlalchemy_cauldron.config import FALLBACK_PK -from sqlalchemy_cauldron.generator.schema_extractor import TableInfo from sqlalchemy_cauldron.generator.type_inference import ( TYPE_MAP, ColumnType, TypeInference, ) +if TYPE_CHECKING: + from sqlalchemy_cauldron.generator.schema_extractor import TableInfo + class AliasedTable: """Wrapper around a SQLAlchemy Table that carries an explicit alias. @@ -24,7 +26,7 @@ class AliasedTable: user-defined alias names during code generation. """ - def __init__(self, table: Table, alias: str): + def __init__(self, table: Table, alias: str) -> None: """Initialize an AliasedTable. Args: @@ -32,12 +34,12 @@ def __init__(self, table: Table, alias: str): alias: The alias name to use (e.g., 'e', 'm') """ self._table = table - self.table_alias = alias + self._alias = alias @property def alias(self) -> str: """Get the alias for this table.""" - return self.table_alias + return self._alias @property def name(self) -> str: @@ -45,16 +47,16 @@ def name(self) -> str: return self._table.name @property - def columns(self) -> Any: + def columns(self) -> list[Column]: """Get the columns (delegate to underlying table).""" - return self._table.columns + return list(self._table.columns) @property def metadata(self) -> MetaData: """Get the metadata (delegate to underlying table).""" return self._table.metadata - def __getattr__(self, name: str) -> Any: + def __getattr__(self, name: str) -> Any: # noqa: ANN401 """Delegate unknown attributes to the underlying table. Args: @@ -71,8 +73,7 @@ def __getattr__(self, name: str) -> Any: class SQLAlchemyRuntimeBuilder: - """ - Builds SQLAlchemy Core Table objects from a list of TableInfo objects. + """Builds SQLAlchemy Core Table objects from a list of TableInfo objects. Input: list[TableInfo] (table name, alias, columns) Output: (MetaData, {table_ref: Table | AliasedTable}) @@ -130,8 +131,7 @@ def build(self, table_infos: list[TableInfo], add_missing_pk: bool = True) -> Ru return metadata, tables def _find_primary_key(self, table_name: str, columns: set[str]) -> str | None: - """ - Find the primary key column name from existing columns. + """Find the primary key column name from existing columns. Returns: str: The name of the primary key column, or None if not found diff --git a/sqlalchemy_cauldron/generator/type_inference.py b/sqlalchemy_cauldron/generator/type_inference.py index a98b7f4..b96584d 100644 --- a/sqlalchemy_cauldron/generator/type_inference.py +++ b/sqlalchemy_cauldron/generator/type_inference.py @@ -1,6 +1,7 @@ from __future__ import annotations from enum import Enum +from typing import TYPE_CHECKING, Any from sqlalchemy import ( Boolean, @@ -9,17 +10,23 @@ Integer, Numeric, String, - Table, Text, ) from sqlalchemy_cauldron.config import FALLBACK_PK +if TYPE_CHECKING: + from sqlalchemy import Table + from sqlalchemy.sql.schema import _TypeEngineArgument + + from sqlalchemy_cauldron.generator.sqlalchemy_runtime import AliasedTable + class ColumnType(Enum): + """Enum representing inferred column types for SQLAlchemy code generation.""" + INTEGER = "Integer" FLOAT = "Float" - STRING = "String" TEXT = "Text" VARCHAR = "String" BOOLEAN = "Boolean" @@ -28,10 +35,9 @@ class ColumnType(Enum): # Map ColumnType enum members to SQLAlchemy column types -TYPE_MAP = { +TYPE_MAP: dict[ColumnType, _TypeEngineArgument[Any]] = { ColumnType.INTEGER: Integer, ColumnType.FLOAT: Float, - ColumnType.STRING: String, ColumnType.TEXT: Text, ColumnType.VARCHAR: String, ColumnType.BOOLEAN: Boolean, @@ -40,10 +46,9 @@ class ColumnType(Enum): } # Map ColumnType to Python type hints for type annotations -TYPEHINT_MAP = { +TYPEHINT_MAP: dict[ColumnType, str] = { ColumnType.INTEGER: "int", ColumnType.FLOAT: "float", - ColumnType.STRING: "str", ColumnType.TEXT: "str", ColumnType.VARCHAR: "str", ColumnType.BOOLEAN: "bool", @@ -62,8 +67,11 @@ class ColumnType(Enum): "timestamp", "cdate", "udate", - } + }, ) +LIKELY_DATETIME_PREFIXES = ("created_", "updated_", "modified_", "deleted_", "date_") +LIKELY_DATETIME_SUFFIXES = ("_at", "_on", "_date", "_time", "_timestamp", "_ts") + LIKELY_BOOLEAN = frozenset( { "active", @@ -72,8 +80,10 @@ class ColumnType(Enum): "archived", "deleted", "verified", - } + }, ) +LIKELY_BOOLEAN_PREFIXES = ("is_", "has_", "can_", "should_") + LIKELY_NUMERIC = frozenset( { "balance", @@ -83,8 +93,10 @@ class ColumnType(Enum): "decimal", "numeric", "value", - } + }, ) +LIKELY_NUMERIC_SUFFIXES = ("_decimal", "_numeric", "_rate", "_value") + LIKELY_INTEGER = frozenset( { "id", @@ -96,15 +108,17 @@ class ColumnType(Enum): "year", "datumn", "count", - } + }, ) +LIKELY_INTEGER_SUFFIXES = ("_id", "_count", "_number") + LIKELY_FLOAT = frozenset( { "price", "amount", "total", "cost", - } + }, ) LIKELY_TEXT = frozenset( { @@ -119,8 +133,10 @@ class ColumnType(Enum): "datuz", "memo", "body", - } + }, ) +LIKELY_TEXT_PREFIXES = ("text_", "content_") + LIKELY_VARCHAR = frozenset( { "name", @@ -144,13 +160,12 @@ class ColumnType(Enum): "url", "path", "slug", - } + }, ) def get_python_type_hint_for_sqlalchemy_type(type_name: str) -> str: - """ - Map SQLAlchemy type name string to Python type hint. + """Map SQLAlchemy type name string to Python type hint. Args: type_name: SQLAlchemy type name (e.g., "Integer", "String", "DateTime") @@ -159,18 +174,39 @@ def get_python_type_hint_for_sqlalchemy_type(type_name: str) -> str: Python type hint string (e.g., "int", "str", "datetime") """ for col_type, sa_type in TYPE_MAP.items(): - if sa_type.__name__ == type_name: + if sa_type.__name__ == type_name: # ty: ignore[unresolved-attribute] return TYPEHINT_MAP[col_type] return "object" def get_python_type_hint(column_type: ColumnType, fallback: str = "str") -> str: - """Map ColumnType to Python type hint string""" + """Map ColumnType to Python type hint string.""" return TYPEHINT_MAP.get(column_type, fallback) +TYPE_CHECKS = [ + (LIKELY_DATETIME, LIKELY_DATETIME_PREFIXES, LIKELY_DATETIME_SUFFIXES, ColumnType.DATETIME), + (LIKELY_BOOLEAN, LIKELY_BOOLEAN_PREFIXES, (), ColumnType.BOOLEAN), + (LIKELY_NUMERIC, (), LIKELY_NUMERIC_SUFFIXES, ColumnType.NUMERIC), + (LIKELY_INTEGER, (), LIKELY_INTEGER_SUFFIXES, ColumnType.INTEGER), + (LIKELY_FLOAT, (), (), ColumnType.FLOAT), + (LIKELY_TEXT, LIKELY_TEXT_PREFIXES, (), ColumnType.TEXT), + (LIKELY_VARCHAR, (), (), ColumnType.VARCHAR), +] + + +def _matches_any( + name: str, + exact_set: set[str] | frozenset[str], + suffixes: tuple[str, ...] = (), + prefixes: tuple[str, ...] = (), +) -> bool: + return name in exact_set or name.endswith(suffixes) or name.startswith(prefixes) + + class TypeInference: - @staticmethod + """Utility class for inferring SQLAlchemy column types from column names.""" + @staticmethod def infer(column_name: str) -> ColumnType: """Infer SQLAlchemy column type from column name using heuristics. @@ -196,105 +232,65 @@ def infer(column_name: str) -> ColumnType: """ name = column_name.lower() - # DateTime detection: names containing temporal indicators - if ( - name in LIKELY_DATETIME - or name.endswith("_at") - or name.endswith("_date") - or name.endswith("_on") - or name.endswith("_time") - or name.startswith("created_") - or name.startswith("updated_") - or name.startswith("modified_") - or name.startswith("deleted_") - or name.startswith("date_") - ): - return ColumnType.DATETIME - - # Boolean detection: names starting with 'is_', 'has_', 'can_', 'should_', or specific patterns - if ( - name in LIKELY_BOOLEAN - or name.startswith("is_") - or name.startswith("has_") - or name.startswith("can_") - or name.startswith("should_") - or name.endswith("_enabled") - or name.endswith("_active") - or name.endswith("_deleted") - or name.endswith("_verified") - ): - return ColumnType.BOOLEAN - - # Numeric (DECIMAL/NUMERIC) detection: names with financial/percentage indicators - if ( - name in LIKELY_NUMERIC - or name.endswith("_decimal") - or name.endswith("_numeric") - or name.endswith("_rate") - or name.endswith("_value") - ): - return ColumnType.NUMERIC - - # INTEGER detection - check for specific matches before 'count' substring - if name in LIKELY_INTEGER or name.endswith("_id") or name.endswith("_count") or name.endswith("_number"): - return ColumnType.INTEGER - - if name in LIKELY_FLOAT: - return ColumnType.FLOAT - - # TEXT detection: names with large text content indicators - # Check this BEFORE VARCHAR since some patterns might overlap - if ( - name in LIKELY_TEXT - or name.endswith("_text") - or name.endswith("_content") - or name.endswith("_description") - or name.endswith("_note") - or name.endswith("_notes") - or name.endswith("_memo") - or name.endswith("_comment") - or name.endswith("_body") - or name.endswith("_message") - or name.startswith("text_") - or name.startswith("content_") - ): - return ColumnType.TEXT - - # VARCHAR detection: shorter named string columns - if ( - name in LIKELY_VARCHAR - or name.endswith("_name") - or name.endswith("_title") - or name.endswith("_code") - or name.endswith("_status") - or name.endswith("_type") - ): - return ColumnType.VARCHAR + for exact_set, prefixes, suffixes, col_type in TYPE_CHECKS: + suffixes_ = suffixes + if not suffixes: + suffixes_ = tuple(f"_{t}" for t in exact_set) + if _matches_any(name, exact_set, suffixes_, prefixes): + return col_type # Default to VARCHAR for generic string columns return ColumnType.VARCHAR @staticmethod - def infer_primary_key_from_table(table_name: str, table: Table, fallback: str | None = FALLBACK_PK) -> str | None: - """ - Infer primary key column name using these conventions (in order): + def infer_primary_key_from_table(table: Table | AliasedTable, fallback: str = FALLBACK_PK) -> str: + """Infer primary key column name. + + Using these conventions (in order): 1. "recid" 2. "id" 3. "{table_name}_id" (e.g., "users_id") 4. First Integer column (inferred by column name heuristics) 5. First Integer column (based on SQLAlchemy type) - 6. Return fallback if no primary key candidate is found (default "recid") + 6. Return fallback if no primary key candidate is found (default "recid"). + + Args: + table: SQLAlchemy Table or AliasedTable object containing column info. + fallback: Fallback column name to use if no primary key candidate is found. + + Returns: + Name of the inferred primary key column, or fallback if none found. """ col_names = {col.name for col in table.columns} - candidate = TypeInference.infer_primary_key_from_columns(table_name, col_names, fallback=None) + candidate = TypeInference.infer_primary_key_from_columns(table.name, col_names, fallback=None) if candidate: return candidate return fallback @staticmethod def infer_primary_key_from_columns( - table_name: str, columns: set[str], fallback: str | None = FALLBACK_PK + table_name: str, + columns: set[str], + fallback: str | None = FALLBACK_PK, ) -> str | None: + """Infer primary key column name from a set of column names. + + Using these conventions (in order): + 1. "recid" + 2. "id" + 3. "{table_name}_id" (e.g., "users_id") + 4. First Integer column (inferred by column name heuristics) + 5. First Integer column (based on SQLAlchemy type) + 6. Return fallback if no primary key candidate is found (default "recid"). + + Args: + table_name: Name of the table to infer primary key for. + columns: Set of column names to consider. + fallback: Fallback column name to use if no primary key candidate is found. + + Returns: + Name of the inferred primary key column, or fallback if none found. + """ # Check for "recid" first if "recid" in columns: return "recid" diff --git a/sqlalchemy_cauldron/main.py b/sqlalchemy_cauldron/main.py index 99c778a..036f6a6 100644 --- a/sqlalchemy_cauldron/main.py +++ b/sqlalchemy_cauldron/main.py @@ -37,8 +37,6 @@ from dataclasses import dataclass from typing import TYPE_CHECKING -from sqlglot.expressions import Expression - from sqlalchemy_cauldron.config import Dialect, Generator, GeneratorConfig from sqlalchemy_cauldron.generator.core_generator import CodeGenerator from sqlalchemy_cauldron.generator.orm_generator import ORMCodeGenerator @@ -48,6 +46,7 @@ if TYPE_CHECKING: from sqlalchemy import Table + from sqlglot.expressions import Expression from sqlalchemy_cauldron.generator.sqlalchemy_runtime import AliasedTable @@ -121,7 +120,8 @@ def brew(sql: str, config: GeneratorConfig | None = None) -> str: include_self_contained_test=config.include_self_contained_test, ) else: - raise ValueError(f"Unsupported generator type: {config.generator}") + msg = f"Unsupported generator type: {config.generator}" + raise ValueError(msg) # optionally write to file if output path is specified if config.output: @@ -145,4 +145,3 @@ def brew(sql: str, config: GeneratorConfig | None = None) -> str: LIMIT 10 """ code = brew(sql) - print(code) diff --git a/sqlalchemy_cauldron/parser/normalizer.py b/sqlalchemy_cauldron/parser/normalizer.py index 299f6e1..c23f38f 100644 --- a/sqlalchemy_cauldron/parser/normalizer.py +++ b/sqlalchemy_cauldron/parser/normalizer.py @@ -43,10 +43,10 @@ simplify, ) -# Mapping from dialect names to sqlglot dialect names -# None means no specific dialect (use sqlglot default) +# Mapping from dialect names to SQLGlot dialect names +# None means no specific dialect (use SQLGlot default) DIALECT_MAPPING: dict[str | None, str | None] = { - None: None, # No dialect - use sqlglot default + None: None, # No dialect - use SQLGlot default "postgres": "postgres", "oracle": "oracle", "sqlite": "sqlite", @@ -82,8 +82,7 @@ def normalize_and_parse( ParseResult containing normalized SQL and parsed expression """ # Convert string to Dialect enum if needed - if isinstance(dialect, str): - dialect = Dialect(dialect) + dialect = Dialect(dialect) normalizer = SQLNormalizer( dialect=dialect, @@ -123,7 +122,7 @@ class SQLNormalizer: def __init__( self, - dialect: Dialect | None = Dialect.POSTGRES, + dialect: Dialect | str | None = Dialect.POSTGRES, optimize_enabled: bool = False, resolve_aliases: bool = False, ) -> None: @@ -131,17 +130,17 @@ def __init__( Args: dialect: SQL dialect to use (default: POSTGRES). - optimize_enabled: Enable sqlglot optimization (default: False). + optimize_enabled: Enable SQLGlot optimization (default: False). resolve_aliases: Resolve aliases to table names (default: False). """ - self.sqlglot_dialect = DIALECT_MAPPING.get(dialect.lower(), None) if dialect else None + self.sqlglot_dialect = DIALECT_MAPPING.get(dialect.lower()) if dialect else None self.optimize_enabled = optimize_enabled self.resolve_aliases = resolve_aliases def normalize(self, sql: str) -> exp.Expression: """Normalize SQL statement. - Transpiles SQL to target dialect, optionally optimizes, and + Transpile SQL to target dialect, optionally optimizes, and optionally resolves aliases. Returns normalized SQL expression. Args: @@ -160,10 +159,10 @@ def normalize(self, sql: str) -> exp.Expression: return expr def expression_to_sql(self, expression: exp.Expression) -> str: - """Convert a sqlglot expression back to SQL string using the normalizer's dialect. + """Convert a SQLGlot expression back to SQL string using the normalizer's dialect. Args: - expression: sqlglot Expression to convert. + expression: SQLGlot Expression to convert. Returns: SQL string representation of the expression. From 04c23ca661a931bd7e78415f77f9de6528494f47 Mon Sep 17 00:00:00 2001 From: mpf82 Date: Wed, 17 Jun 2026 01:09:34 +0200 Subject: [PATCH 2/2] test: improve type hints, added tests, pushed code coverage beyond 95% Changelog: test --- tests/caching.py | 162 +-- tests/conftest.py | 20 + tests/generator/test_core_generator.py | 1226 +++++++++++------ tests/generator/test_helpers.py | 52 +- tests/generator/test_orm_advanced.py | 216 +-- tests/generator/test_orm_comprehensive.py | 175 ++- tests/generator/test_orm_data_all_cases.py | 206 ++- tests/generator/test_orm_generator.py | 178 ++- tests/generator/test_orm_self_contained.py | 25 +- tests/generator/test_orm_sql_comparison.py | 79 +- tests/generator/test_orm_string_literals.py | 79 +- tests/generator/test_orm_test_generator.py | 10 +- .../generator/test_placeholder_generation.py | 10 +- tests/generator/test_schema_extractor.py | 125 +- tests/generator/test_self_join_alias.py | 10 +- tests/generator/test_sqlalchemy_runtime.py | 7 +- tests/generator/test_type_inference.py | 53 +- tests/parser/test_normalizer.py | 72 +- tests/test_cases.py | 56 +- tests/test_cli.py | 102 +- tests/test_main.py | 79 +- tests/test_roundtrip.py | 136 +- tests/test_simple_case.py | 26 +- 23 files changed, 1810 insertions(+), 1294 deletions(-) create mode 100644 tests/conftest.py diff --git a/tests/caching.py b/tests/caching.py index a06c040..8924556 100644 --- a/tests/caching.py +++ b/tests/caching.py @@ -15,6 +15,7 @@ from __future__ import annotations import functools +import keyword import re import sqlglot @@ -40,7 +41,11 @@ # Pattern for auto-generated alias names that should be stripped # Matches: function names (count_1, sum_1, etc.), type names (varchar_1, int_1, etc.), # and other generated names (anon_1, _col_0, etc.) -AUTO_GENERATED_ALIAS_PATTERN = r"^(_col|count|sum|avg|min|max|abs|round|ceil|floor|varchar|int|bigint|text|numeric|substring|concat|case|coalesce|anon|upper|lower|trim|ltrim|rtrim|length|sqrt|exp|log|ln|power|mod|sign|year|month|day|hour|minute|second|date|time|timestamp|nullif|least|greatest)_\d+$" +AUTO_GENERATED_ALIAS_PATTERN = ( + r"^(_col|count|sum|avg|min|max|abs|round|ceil|floor|varchar|int|bigint|text|numeric|substring|" + r"concat|case|coalesce|anon|upper|lower|trim|ltrim|rtrim|length|sqrt|exp|log|ln|power|mod|sign|" + r"year|month|day|hour|minute|second|date|time|timestamp|nullif|least|greatest)_\d+$" +) def _apply_post_normalizations_via_transform(expr: sqlglot.exp.Expression) -> sqlglot.exp.Expression: @@ -50,14 +55,6 @@ def _apply_post_normalizations_via_transform(expr: sqlglot.exp.Expression) -> sq """ def normalize_node(node: sqlglot.exp.Expression) -> sqlglot.exp.Expression: - # Convert DATE(expr) function back to CAST(expr AS DATE) - # This preserves the semantic intention better for code generation - # sqlglot converts CAST(x AS DATE) to DATE(x) for sqlite, but we want to - # preserve CAST for consistent code generation across dialects - # if isinstance(node, sqlglot.exp.Date): - # # Convert to CAST(expr AS DATE) - # return sqlglot.exp.Cast(this=node.this.copy() if node.this else None, to=sqlglot.exp.DataType.build("DATE")) - # Normalize FETCH to LIMIT (semantically equivalent) # FETCH FIRST/NEXT n ROWS ONLY -> LIMIT n if isinstance(node, sqlglot.exp.Fetch): @@ -82,7 +79,7 @@ def normalize_node(node: sqlglot.exp.Expression) -> sqlglot.exp.Expression: if is_zero_offset: # Return None to remove this node from the tree - return None + return None # ty: ignore[invalid-return-type] return node # Normalize LIMIT with OFFSET 0 to just LIMIT (OFFSET 0 is no-op) @@ -90,13 +87,11 @@ def normalize_node(node: sqlglot.exp.Expression) -> sqlglot.exp.Expression: if isinstance(node, sqlglot.exp.Limit): offset = node.args.get("offset") # Check if offset is 0 (literal or expression) - if offset: - if isinstance(offset, sqlglot.exp.Literal): - if offset.this == "0": - # Remove the OFFSET 0 - node = node.copy() - node.set("offset", None) - return node + if offset and isinstance(offset, sqlglot.exp.Literal) and offset.this == "0": + # Remove the OFFSET 0 + node = node.copy() + node.set("offset", None) + return node return node # Normalize CTE: remove column list parameter specs @@ -151,8 +146,6 @@ def _normalize_output_formatting(sql: str, dialect: str | None = None) -> str: sql: SQL string to normalize dialect: sqlglot dialect name (e.g., 'mysql', 'postgres') for correct quote style """ - import keyword - # SQL keywords that MUST be quoted when used as identifiers # (keywords that are reserved in SQL standard and are problematic) sql_reserved_keywords = SQL_KEYWORDS @@ -162,14 +155,13 @@ def _normalize_output_formatting(sql: str, dialect: str | None = None) -> str: quote_char = "`" if dialect == "mysql" else '"' # Remove unnecessary quotes from identifiers (but keep them for reserved keywords) - def remove_quotes(match): + def remove_quotes(match: re.Match) -> str: identifier = match.group(1) if identifier.lower() in sql_reserved_keywords or identifier.lower() in reserved: # Keep quotes for reserved keywords with correct style return f"{quote_char}{identifier}{quote_char}" - else: - # Remove quotes for non-reserved identifiers - return identifier + # Remove quotes for non-reserved identifiers + return identifier sql = re.sub(r'"([a-zA-Z_][a-zA-Z0-9_]*)"', remove_quotes, sql) sql = re.sub(r"`([a-zA-Z_][a-zA-Z0-9_]*)`", remove_quotes, sql) # MySQL-style backticks @@ -178,8 +170,6 @@ def remove_quotes(match): sql = re.sub(r"\bINNER\s+JOIN\b", "JOIN", sql, flags=re.IGNORECASE) sql = re.sub(r"\bLEFT\s+OUTER\s+JOIN\b", "LEFT JOIN", sql, flags=re.IGNORECASE) sql = re.sub(r"\bRIGHT\s+OUTER\s+JOIN\b", "RIGHT JOIN", sql, flags=re.IGNORECASE) - sql = re.sub(r"\bRIGHT\s+JOIN\b", "JOIN", sql, flags=re.IGNORECASE) - sql = re.sub(r"\bFULL\s+(?:OUTER\s+)?JOIN\b", "JOIN", sql, flags=re.IGNORECASE) # Normalize ORDER BY: remove explicit ASC (it's the default) # Match spaces before ASC keyword and remove both @@ -189,40 +179,38 @@ def remove_quotes(match): # Normalize NOT = to <> (WHERE NOT x = y -> WHERE x <> y) # This must be done carefully to avoid matching NOT IN, NOT EXISTS, etc. - sql = re.sub(r"\bWHERE\s+NOT\s+(\w+(?:\.\w+)?)\s*=\s*", r"WHERE \1 <> ", sql, flags=re.IGNORECASE) - - return sql + return re.sub(r"\bWHERE\s+NOT\s+(\w+(?:\.\w+)?)\s*=\s*", r"WHERE \1 <> ", sql, flags=re.IGNORECASE) @functools.cache -def cached_normalizer(dialect: str, optimize_enabled=False, resolve_aliases=True): +def cached_normalizer(dialect: str, optimize_enabled: bool = False, resolve_aliases: bool = True) -> SQLNormalizer: """Return a cached SQLNormalizer instance for the given configuration.""" return SQLNormalizer(dialect=dialect, optimize_enabled=optimize_enabled, resolve_aliases=resolve_aliases) @functools.cache def _parse_cached(sql: str, dialect: str | None = None) -> sqlglot.exp.Expression: - """Cache sqlglot parsing results by SQL string and dialect. + """Cache SQLGlot parsing results by SQL string and dialect. Parsing is expensive (~30-40% of normalize time), so caching by SQL and dialect eliminates redundant parsing across normalization passes and schema extraction. Args: sql: SQL string to parse - dialect: Optional sqlglot dialect name to use for parsing + dialect: Optional SQLGlot dialect name to use for parsing Returns: - Parsed sqlglot Expression + Parsed SQLGlot Expression """ return transpile_and_parse(sql, sqlglot_dialect=dialect) @functools.cache -def normalize( +def normalize( # noqa: PLR0915 sql: str, dialect: Dialect | None = None, ) -> sqlglot.exp.Expression: - """Parse and re-emit SQL through sqlglot for canonical comparison. + """Parse and re-emit SQL through SQLGlot for canonical comparison. Normalizes via: unescape %%, standard normalization, remove aliases, simplify qualifications, add duplicate labels, normalize JOINs via transform, @@ -236,31 +224,26 @@ def normalize( dialect: SQL dialect to use (default: Dialect.POSTGRES) Returns: - Normalized sqlglot Expression (use .sql() to get SQL string). + Normalized SQLGlot Expression (use .sql() to get SQL string). """ dialect = dialect or Dialect.POSTGRES - sqlglot_dialect = DIALECT_MAPPING.get(dialect.name.lower(), None) if dialect else None + sqlglot_dialect = DIALECT_MAPPING.get(dialect.lower(), None) if dialect else None sql = sql.replace("%%", "%") # Unescape early for modulo operators # Check if this is a self-join (same table appears multiple times with different aliases) # If so, we must preserve aliases to avoid losing disambiguation information - has_self_join = False - try: - # temp_expr = sqlglot.parse_one(sql, dialect=sqlglot_dialect) - temp_expr = _parse_cached(sql, dialect=sqlglot_dialect) - table_aliases: dict[str, list[str]] = {} # table_name -> [aliases] - for table in temp_expr.find_all(sqlglot.exp.Table): - table_name = table.name.lower() - alias = table.alias.lower() if table.alias else table_name - if table_name not in table_aliases: - table_aliases[table_name] = [] - if alias not in table_aliases[table_name]: - table_aliases[table_name].append(alias) - - # If any table appears multiple times, it's a self-join - has_self_join = any(len(aliases) > 1 for aliases in table_aliases.values()) - except Exception: - pass + temp_expr = _parse_cached(sql, dialect=sqlglot_dialect) + table_aliases: dict[str, list[str]] = {} # table_name -> [aliases] + for table in temp_expr.find_all(sqlglot.exp.Table): + table_name = table.name.lower() + alias = table.alias.lower() if table.alias else table_name + if table_name not in table_aliases: + table_aliases[table_name] = [] + if alias not in table_aliases[table_name]: + table_aliases[table_name].append(alias) + + # If any table appears multiple times, it's a self-join + has_self_join = any(len(aliases) > 1 for aliases in table_aliases.values()) # First pass: standard normalization # For self-joins, don't resolve aliases (we need them for disambiguation) @@ -273,8 +256,7 @@ def normalize( # Extract all table names to determine if qualifications are needed table_names = set() - for table in expr.find_all(sqlglot.exp.Table): - table_names.add(table.name.lower()) + table_names.update(table.name.lower() for table in expr.find_all(sqlglot.exp.Table)) # If there's only one table, we can remove qualifications single_table = len(table_names) == 1 @@ -321,7 +303,7 @@ def normalize( # First occurrence: no label, subsequent: add _2, _3, etc. if duplicate_position[col_name] > 1: label_name = f"{select_expr.name}_{duplicate_position[col_name]}" - select_expr = sqlglot.exp.Alias(this=select_expr.copy(), alias=label_name) + select_expr = sqlglot.exp.Alias(this=select_expr.copy(), alias=label_name) # noqa: PLW2901 new_expressions.append(select_expr) @@ -361,12 +343,16 @@ def normalize_aliases_qualifications_and_keywords(node: sqlglot.exp.Expression) return node # Remove table qualification if single table (for all Column nodes, not just SELECT) - if single_table and isinstance(node, sqlglot.exp.Column): - if node.table and node.table.lower() == single_table_name: - # Remove unnecessary table qualification - node = node.copy() - node.set("table", None) - return node + if ( + single_table + and isinstance(node, sqlglot.exp.Column) + and node.table + and node.table.lower() == single_table_name + ): + # Remove unnecessary table qualification + node = node.copy() + node.set("table", None) + return node # Normalize LIKE patterns: unescape double %% back to single % if isinstance(node, sqlglot.exp.Like): @@ -380,41 +366,6 @@ def normalize_aliases_qualifications_and_keywords(node: sqlglot.exp.Expression) node.set("expression", sqlglot.exp.Literal.string(unescaped)) return node - # Normalize ORDER BY: make ASC explicit if no DESC - if isinstance(node, sqlglot.exp.Ordered): - # If there's no DESC and no explicit ASC, add ASC - if not node.desc: - node = node.copy() - node.set("desc", False) # Explicitly set to False (ASC) - return node - - # Normalize JOIN keywords - if isinstance(node, sqlglot.exp.Join): - join_kind = node.kind - if join_kind: - join_upper = join_kind.upper() if isinstance(join_kind, str) else "" - # INNER JOIN -> JOIN (remove INNER prefix) - if join_upper == "INNER": - node = node.copy() - node.set("kind", None) - return node - # LEFT OUTER JOIN -> LEFT JOIN - elif "LEFT" in join_upper and "OUTER" in join_upper: - node = node.copy() - node.set("kind", "LEFT") - return node - # RIGHT OUTER JOIN -> RIGHT JOIN - elif "RIGHT" in join_upper and "OUTER" in join_upper: - node = node.copy() - node.set("kind", "RIGHT") - return node - # RIGHT JOIN, FULL OUTER JOIN, FULL JOIN -> JOIN - # These are semantically different but generators may not support them - elif join_upper in ("RIGHT", "FULL OUTER", "FULL"): - node = node.copy() - node.set("kind", None) - return node - return node # Transform to normalize @@ -424,11 +375,10 @@ def normalize_aliases_qualifications_and_keywords(node: sqlglot.exp.Expression) if single_table: def remove_table_qualifications(node: sqlglot.exp.Expression) -> sqlglot.exp.Expression: - if isinstance(node, sqlglot.exp.Column): - if node.table and node.table.lower() == single_table_name: - node = node.copy() - node.set("table", None) - return node + if isinstance(node, sqlglot.exp.Column) and node.table and node.table.lower() == single_table_name: + node = node.copy() + node.set("table", None) + return node return node expr = expr.transform(remove_table_qualifications) @@ -450,11 +400,7 @@ def remove_table_qualifications(node: sqlglot.exp.Expression) -> sqlglot.exp.Exp # Re-parse the normalized SQL to get the final expression tree # This ensures we return an Expression instead of a string, eliminating the redundant # parse-unparse-parse cycle that was happening when callers called normalize() then _parse_cached() - final_expr = _parse_cached(result, dialect=sqlglot_dialect) - - # Return expression directly instead of converting to SQL - # Callers can use normalizer.expression_to_sql(expr) when they need SQL string with dialect support - return final_expr + return _parse_cached(result, dialect=sqlglot_dialect) # normalize booleans to 0/1 for dialects that don't support true/false literals @@ -463,7 +409,7 @@ def normalize_booleans(node: sqlglot.exp.Expression) -> sqlglot.exp.Expression: value = node.this if value is True: return sqlglot.exp.Literal.number(1) - elif value is False: + if value is False: return sqlglot.exp.Literal.number(0) return node diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..0b58021 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from sqlalchemy import Engine, create_engine + +if TYPE_CHECKING: + from collections.abc import Iterator + + +@pytest.fixture +def sa_engine() -> Iterator[Engine]: + """Fixture providing an in-memory SQLite engine.""" + engine = create_engine("sqlite:///:memory:", echo=False) + try: + yield engine + finally: + engine.dispose() diff --git a/tests/generator/test_core_generator.py b/tests/generator/test_core_generator.py index 6390d30..8e0dacb 100644 --- a/tests/generator/test_core_generator.py +++ b/tests/generator/test_core_generator.py @@ -1,9 +1,17 @@ from __future__ import annotations +# ruff: noqa: PLR2004, S608 +import libcst as cst import pytest +from sqlglot import ParseError, exp + from sqlalchemy_cauldron.config import Dialect -from sqlalchemy_cauldron.generator.core_generator import DDL_MSG_MAP, CodeGenerator +from sqlalchemy_cauldron.generator.core_generator import ( + DDL_MSG_MAP, + CodeGenerator, + cst_bool, +) from sqlalchemy_cauldron.parser.normalizer import SQLNormalizer from tests.caching import ( RUNTIME_BUILDER, @@ -31,149 +39,173 @@ def generate( expr = normalizer.normalize(sql) schema = SCHEMA_EXTRACTOR.extract(expr) _, tables = RUNTIME_BUILDER.build(schema) - return CodeGenerator().generate_module(expr=expr, tables=tables, original_sql=sql if with_original_sql else "") + return CodeGenerator().generate_module( + expr=expr, + tables=tables, + original_sql=sql if with_original_sql else "", + ) class TestSelect: - def test_simple_column(self): + def test_simple_column(self) -> None: + code = generate("SELECT id FROM users") + assert "select(users.c.id).select_from(users)" in code + + def test_simple_column_qualified(self) -> None: code = generate("SELECT users.id FROM users") - assert "users.c.id" in code + assert "select(users.c.id).select_from(users)" in code - def test_column_with_alias(self): + def test_column_with_alias(self) -> None: code = generate("SELECT users.id AS user_id FROM users") - assert 'users.c.id.label("user_id")' in code + assert 'select(users.c.id.label("user_id")).select_from(users)' in code - def test_table_alias_resolved(self): + def test_table_alias_resolved(self) -> None: code = generate("SELECT u.id FROM users u") - assert "users.c.id" in code - assert "u.c.id" not in code + assert "select(users.c.id).select_from(users)" in code + assert "select(u.c.id).select_from(users)" not in code - def test_count_star(self): + def test_count_star(self) -> None: code = generate("SELECT COUNT(*) FROM users") assert "func.count()" in code - def test_count_column(self): + def test_count_column(self) -> None: + code = generate("SELECT COUNT(users.id) FROM users") + assert "func.count(users.c.id)" in code + + def test_count_column_alias(self) -> None: code = generate("SELECT COUNT(orders.id) AS order_count FROM users JOIN orders ON users.id = orders.user_id") - assert 'func.count(orders.c.id).label("order_count")' in code + assert ( + 'select(func.count(orders.c.id).label("order_count"))' + ".select_from(users.join(orders, users.c.id == orders.c.user_id))" + ) in code - def test_multiple_columns(self): + def test_multiple_columns(self) -> None: + code = generate("SELECT id, name FROM users") + assert "select(users.c.id, users.c.name).select_from(users)" in code + + def test_multiple_columns_qualified(self) -> None: code = generate("SELECT users.id, users.name FROM users") - assert "users.c.id" in code - assert "users.c.name" in code + assert "select(users.c.id, users.c.name).select_from(users)" in code - def test_select_star(self): + def test_select_star(self) -> None: code = generate("SELECT * FROM users") assert 'literal_column("*")' in code assert "select_from(users)" in code - def test_select_table_star(self): + def test_select_table_star(self) -> None: code = generate("SELECT users.* FROM users") assert 'literal_column("users.*")' in code assert "select_from(users)" in code - def test_select_alias_table_star(self): + def test_select_alias_table_star(self) -> None: code = generate("SELECT u.* FROM users u") assert 'literal_column("users.*")' in code assert "select_from(users)" in code - def test_select_star_unqualified_where(self): + def test_select_star_unqualified_where(self) -> None: code = generate("SELECT * FROM users WHERE userid = 1") assert 'literal_column("*")' in code assert "select_from(users)" in code assert "users.c.userid == 1" in code - def test_select_star_unqualified_where_ambiguous(self): + def test_select_star_unqualified_where_ambiguous(self) -> None: code = generate("SELECT * FROM users JOIN orders ON users.id = orders.user_id WHERE userid = 1") assert 'literal_column("*")' in code assert "select_from(users.join(orders" in code assert 'where(literal_column("userid") == 1)' in code - def test_select_star_unqualified_multiple_tables(self): + def test_select_star_unqualified_multiple_tables(self) -> None: code = generate("SELECT users.*, orders.* FROM users user JOIN orders ON user.id = orders.user_id") assert 'literal_column("users.*")' in code assert 'literal_column("orders.*")' in code assert "users.join(orders, users.c.id == orders.c.user_id)" in code - def test_select_distinct(self): + def test_select_distinct(self) -> None: code = generate("SELECT DISTINCT users.id FROM users") - assert ".distinct()" in code + assert "select(users.c.id).distinct().select_from(users)" in code + + def test_select_from_multiple_tables(self) -> None: + code = generate("SELECT * FROM users, orders") + assert "select_from(users.join(orders" in code class TestWhere: - def test_string_literal(self): + def test_string_literal(self) -> None: code = generate("SELECT orders.id FROM orders WHERE orders.material = 'steel'") assert 'orders.c.material == "steel"' in code - def test_integer_literal(self): + def test_integer_literal(self) -> None: code = generate("SELECT users.id FROM users WHERE users.active = 1") assert "users.c.active == 1" in code - def test_column_reference_rhs(self): + def test_column_reference_rhs(self) -> None: code = generate( - "SELECT orders.id FROM orders JOIN users ON orders.user_id = users.id WHERE orders.cid = users.cid" + "SELECT orders.id FROM orders JOIN users ON orders.user_id = users.id WHERE orders.cid = users.cid", ) assert "orders.c.cid == users.c.cid" in code assert '"users.cid"' not in code - def test_multiple_and_conditions(self): + def test_multiple_and_conditions(self) -> None: code = generate( "SELECT orders.id FROM orders JOIN users ON orders.user_id = users.id " - "WHERE orders.cid = users.cid AND orders.material = 'steel' AND users.active = 1" + "WHERE orders.cid = users.cid AND orders.material = 'steel' AND users.active = 1", ) assert "orders.c.cid == users.c.cid" in code assert 'orders.c.material == "steel"' in code assert "users.c.active == 1" in code - def test_table_alias_in_where(self): + def test_table_alias_in_where(self) -> None: code = generate("SELECT u.id FROM users u WHERE u.active = 1") assert "users.c.active == 1" in code assert "u.c.active" not in code class TestOrderBy: - def test_order_by_desc(self): + def test_order_by_desc(self) -> None: code = generate("SELECT users.id FROM users ORDER BY users.id DESC") assert "users.c.id.desc()" in code - def test_order_by_asc(self): + def test_order_by_asc(self) -> None: code = generate("SELECT users.id FROM users ORDER BY users.id ASC") assert "users.c.id.asc()" in code - def test_order_by_alias_column(self): + def test_order_by_alias_column(self) -> None: code = generate( "SELECT users.id, COUNT(orders.id) AS order_count " "FROM users JOIN orders ON users.id = orders.user_id " "GROUP BY users.id " - "ORDER BY order_count DESC" + "ORDER BY order_count DESC", ) assert 'literal_column("order_count").desc()' in code - def test_order_by_position(self): + def test_order_by_position(self) -> None: code = generate("SELECT users.name, users.age FROM users ORDER BY 1 ASC, 2 DESC") assert 'literal_column("1").asc()' in code assert 'literal_column("2").desc()' in code class TestGroupBy: - def test_group_by_single(self): + def test_group_by_single(self) -> None: code = generate("SELECT users.id FROM users GROUP BY users.id") assert ".group_by(users.c.id)" in code - def test_group_by_multiple(self): + def test_group_by_multiple(self) -> None: code = generate("SELECT users.id, users.name FROM users GROUP BY users.id, users.name") assert "users.c.id" in code assert "users.c.name" in code - def test_having_single(self): + def test_having_single(self) -> None: code = generate( - "SELECT users.id, COUNT(orders.id) AS order_count FROM users JOIN orders ON users.id = orders.user_id GROUP BY users.id HAVING COUNT(orders.id) > 1" + "SELECT users.id, COUNT(orders.id) AS order_count FROM users " + "JOIN orders ON users.id = orders.user_id GROUP BY users.id HAVING COUNT(orders.id) > 1", ) assert ".having(" in code assert "func.count(orders.c.id) > 1" in code - def test_having_multiple_conditions(self): + def test_having_multiple_conditions(self) -> None: code = generate( - "SELECT u.id FROM users u JOIN orders o ON u.id = o.user_id GROUP BY u.id HAVING SUM(o.amount) > 100 AND COUNT(o.id) > 1" + "SELECT u.id FROM users u JOIN orders o ON u.id = o.user_id " + "GROUP BY u.id HAVING SUM(o.amount) > 100 AND COUNT(o.id) > 1", ) assert ".having(" in code assert "func.sum(orders.c.amount) > 100" in code @@ -181,22 +213,22 @@ def test_having_multiple_conditions(self): class TestLimitOffset: - def test_limit(self): + def test_limit(self) -> None: code = generate("SELECT users.id FROM users LIMIT 10") assert ".limit(10)" in code - def test_offset(self): + def test_offset(self) -> None: code = generate("SELECT users.id FROM users OFFSET 5") assert ".offset(5)" in code - def test_limit_offset(self): + def test_limit_offset(self) -> None: code = generate("SELECT users.id FROM users LIMIT 10 OFFSET 5") assert ".limit(10)" in code assert ".offset(5)" in code class TestFullQuery: - def test_complex_query(self): + def test_complex_query(self) -> None: sql = """ SELECT u.id, u.name, COUNT(o.id) AS order_count FROM users u @@ -207,25 +239,76 @@ def test_complex_query(self): LIMIT 10 """ code = generate(sql) - assert "users.c.id" in code - assert "users.c.name" in code - assert 'func.count(orders.c.id).label("order_count")' in code - assert "users.c.active == 1" in code - assert 'orders.c.material == "steel"' in code - assert ".group_by" in code - assert 'literal_column("order_count").desc()' in code - assert ".limit(10)" in code + assert ( + 'select(users.c.id, users.c.name, func.count(orders.c.id).label("order_count"))' + ".select_from(users.join(orders, users.c.id == orders.c.user_id, isouter = True))" + '.where(users.c.active == 1, orders.c.material == "steel")' + '.group_by(users.c.id, users.c.name).order_by(literal_column("order_count").desc()).limit(10)' + ) in code + + +class TestJoinsTypes: + @pytest.mark.parametrize( + "sql", + [ + "SELECT user.id FROM users AS user JOIN orders ON user.id = orders.user_id", + "SELECT user.id FROM users AS user INNER JOIN orders ON user.id = orders.user_id", + ], + ) + def test_inner_join(self, sql: str) -> None: + code = generate(sql) + assert "join(orders, users.c.id == orders.c.user_id)" in code + + def test_left_join(self) -> None: + sql = "SELECT user.id FROM users AS user LEFT JOIN orders ON user.id = orders.user_id" + code = generate(sql) + assert "join(orders, users.c.id == orders.c.user_id, isouter = True)" in code + + def test_left_outer_join(self) -> None: + sql = "SELECT user.id FROM users AS user LEFT OUTER JOIN orders ON user.id = orders.user_id" + code = generate(sql) + assert "join(orders, users.c.id == orders.c.user_id, isouter = True)" in code + + def test_right_join(self) -> None: + sql = "SELECT user.id FROM users AS user RIGHT JOIN orders ON user.id = orders.user_id" + code = generate(sql) + assert "join(orders, users.c.id == orders.c.user_id, isouter = True)" in code + + def test_right_outer_join(self) -> None: + sql = "SELECT user.id FROM users AS user RIGHT OUTER JOIN orders ON user.id = orders.user_id" + code = generate(sql) + assert "join(orders, users.c.id == orders.c.user_id, isouter = True)" in code + + def test_full_join(self) -> None: + sql = "SELECT user.id FROM users AS user FULL JOIN orders ON user.id = orders.user_id" + code = generate(sql) + assert "join(orders, users.c.id == orders.c.user_id, full = True)" in code + + def test_full_outer_join(self) -> None: + sql = "SELECT user.id FROM users AS user FULL OUTER JOIN orders ON user.id = orders.user_id" + code = generate(sql) + assert "join(orders, users.c.id == orders.c.user_id, isouter = True, full = True)" in code + + def test_cross_join(self) -> None: + sql = "SELECT user.id FROM users AS user CROSS JOIN orders" + code = generate(sql) + assert "select(users.c.id).select_from(users.join(orders))" in code + + def test_self_join(self) -> None: + sql = "SELECT u1.id FROM users AS u1 JOIN users AS u2 ON u1.referrer_id = u2.id" + code = generate(sql, resolve_aliases=False) + assert "select(u1.c.id).select_from(u1.join(u2, u1.c.referrer_id == u2.c.id))" in code class TestJoins: - def test_join_on_preserved(self): + def test_join_on_preserved(self) -> None: sql = "SELECT user.*, orders.* FROM users AS user JOIN orders ON user.id = orders.user_id" code = generate(sql) assert 'literal_column("users.*")' in code # join should include ON condition expressed as a comparison assert "join(orders, users.c.id == orders.c.user_id)" in code - def test_comma_from_star(self): + def test_comma_from_star(self) -> None: sql = "SELECT * FROM users, orders WHERE users.id = orders.user_id" code = generate(sql) assert 'literal_column("*")' in code @@ -233,57 +316,93 @@ def test_comma_from_star(self): assert "select_from(users.join(orders))" in code assert "users.c.id == orders.c.user_id" in code - def test_multi_join_chain(self): + def test_multi_join_chain(self) -> None: sql = "SELECT * FROM a JOIN b ON a.id = b.aid JOIN c ON b.id = c.bid" code = generate(sql) # ensure both join conditions are present in the nested join assert "a.c.id == b.c.aid" in code assert "b.c.id == c.c.bid" in code - def test_join_on_multiple_conditions(self): + def test_join_on_multiple_conditions(self) -> None: sql = "SELECT * FROM a JOIN b ON a.id = b.aid AND a.type = b.type" code = generate(sql) assert "and_(" in code assert "a.c.id == b.c.aid" in code assert "a.c.type == b.c.type" in code - def test_unqualified_column_with_multiple_tables_is_literal(self): - sql = "SELECT users.id FROM users, orders WHERE userid = 1" + def test_join_on_multiple_conditions_with_or(self) -> None: + sql = "SELECT * FROM a JOIN b ON a.id != b.aid OR a.type = b.type" code = generate(sql) - assert 'literal_column("userid")' in code + assert "or_(" in code + assert "a.c.id != b.c.aid" in code + assert "a.c.type == b.c.type" in code - def test_left_join_emitted_as_outerjoin(self): - sql = "SELECT users.id FROM users LEFT JOIN orders ON users.id = orders.user_id" + def test_unqualified_column_with_multiple_tables_is_literal(self) -> None: + sql = "SELECT users.id FROM users, orders WHERE userid = 1" code = generate(sql) - # LEFT JOIN should map to SQLAlchemy's `outerjoin` helper - assert "outerjoin(orders" in code + assert 'literal_column("userid")' in code class TestInsert: - def test_insert_basic(self): + def test_insert_basic(self) -> None: code = generate("INSERT INTO users (id, name) VALUES (1, 'Alice')") assert "insert(users)" in code assert "id = 1" in code assert 'name = "Alice"' in code - def test_insert_multiple_columns(self): + def test_insert_multiple_columns(self) -> None: code = generate("INSERT INTO users (id, name, active) VALUES (1, 'Bob', 0)") assert "insert(users)" in code assert "id = 1" in code assert 'name = "Bob"' in code assert "active = 0" in code - def test_insert_produces_values_call(self): + def test_insert_produces_values_call(self) -> None: code = generate("INSERT INTO users (id, name) VALUES (42, 'Charlie')") assert "insert(users).values(" in code - def test_insert_table_defined(self): + def test_insert_table_defined(self) -> None: code = generate("INSERT INTO users (id, name) VALUES (1, 'Alice')") assert 'Table("users"' in code + def test_insert_no_columns(self) -> None: + # no columns in SQL should trigger fallback PK column to be added to schema, + # allowing code generation to proceed without error, + # however since there are no columns in the SQL, the generated code won't have any + # column references and will just insert a default row with the fallback PK value + code = generate("INSERT INTO users VALUES (1, 'Alice')") + assert "insert(users).values(" in code + assert ( + "Alice" not in code + ) # since there are no columns, the string literal should not be present in the generated code + assert 'Column("recid", Integer)' in code # fallback PK column should be added to schema + + def test_insert_plain(self) -> None: + with pytest.raises(ParseError, match="Expected table"): + generate("INSERT") + + def test_insert_into_plain(self) -> None: + with pytest.raises(ParseError, match="Expected table"): + generate("INSERT INTO") + + def test_insert_into_table_no_values(self) -> None: + code = generate("INSERT INTO users") + assert "insert(users).values()" in code + + @pytest.mark.xfail( + reason="Currently the generator does not handle RETURNING clauses on INSERT statements", + strict=True, + ) + def test_insert_returning(self) -> None: + code = generate("INSERT INTO users (id, name) VALUES (1, 'Alice') RETURNING id") + assert "insert(users)" in code + assert "id = 1" in code + assert 'name = "Alice"' in code + assert ".returning(users.c.id)" in code + class TestCTE: - def test_simple_cte(self): + def test_simple_cte(self) -> None: sql = "WITH cte AS (SELECT id, name FROM users WHERE active = 1) SELECT cte.id, cte.name FROM cte" code = generate(sql) # CTE should be emitted as a variable assigned to