diff --git a/docs/changelog.rst b/docs/changelog.rst index 68cfc01be..301faaac2 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -189,11 +189,11 @@ v0.63.0 - Transactions, table fixtures, SQL fragments, storage, and kwargs param probe branches. (`#782 `_) -* Driver execution methods (:meth:`~sqlspec.driver.SyncDriverAdapterBase.execute`, - :meth:`~sqlspec.driver.SyncDriverAdapterBase.select`, etc.) enforce keyword argument parameter passing - (``execute(sql, a=1, b=2)`` or ``execute(sql, **params)``). Passing positional dictionary literals - is prohibited across documentation, examples, and internal extensions to take advantage of the driver - fast-path parameter dispatch. +* Docs, examples, and built-in extensions now pass named query values as keyword arguments + (``execute(sql, a=1, b=2)`` or ``execute(sql, **params)``). Drivers still accept a dict, list, + or tuple as a positional argument. Existing calls need no changes. + ``execute_many`` still accepts a collection of rows. + (`#769 `_) * The Litestar extension now requires ``litestar>=2.23.0``. @@ -212,11 +212,22 @@ v0.63.0 - Transactions, table fixtures, SQL fragments, storage, and kwargs param consumption without intermediate tuple relays. (`#771 `_) -* ``sql.values`` creates a :class:`~sqlspec.builder.Values` builder for parameterized bulk row lists rather than resolving as a column named ``values``. Use ``sql.column("values")`` to construct column expressions referencing that identifier. +* ``sql.values(...)`` creates a :class:`~sqlspec.builder.Values` builder that binds values for bulk row lists. + Use ``sql.column("values")`` to refer to a column named ``values``. (`#779 `_) +* Removed unused private helpers and shared repeated code across builders, drivers, and migrations. + Loader, service, and ADK artifact modules now group public methods before private helpers. + Supported public APIs and query behavior stay the same. + (`#784 `_) + **Fixed:** +* Cached statements no longer retry SQL when a query or result conversion fails. This prevents + duplicate writes. Cached dict and record rows keep their values. With pymssql, statement stacks + leave the caller's open transaction in place. + (`#742 `_) + * Preserve JSON objects and arrays as individual query parameters after placeholder conversion, instead of reinterpreting them as batches during parameter validation. diff --git a/sqlspec/adapters/adbc/config.py b/sqlspec/adapters/adbc/config.py index d4ff0c5eb..5b4c64be7 100644 --- a/sqlspec/adapters/adbc/config.py +++ b/sqlspec/adapters/adbc/config.py @@ -8,21 +8,23 @@ from sqlspec.adapters.adbc.core import ( apply_driver_features, build_connection_config, - build_postgres_extension_probe_names, detect_postgres_extensions, get_statement_config, is_postgres_dialect, - is_postgres_extension_active, resolve_dialect_from_config, resolve_dialect_name, resolve_driver_connect_func, - resolve_postgres_extension_state, - resolve_runtime_statement_config, ) from sqlspec.adapters.adbc.driver import AdbcDriver, AdbcExceptionHandler from sqlspec.config import ExtensionConfigs, NoPoolSyncConfig from sqlspec.core import StatementConfig from sqlspec.core.capabilities import TypeCoercionCapabilities +from sqlspec.core.config_runtime import ( + build_postgres_extension_probe_names, + is_postgres_extension_active, + resolve_postgres_extension_state, + resolve_runtime_statement_config, +) from sqlspec.driver._sync import SyncPoolConnectionContext, SyncPoolSessionFactory from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.events import EventRuntimeHints @@ -297,23 +299,6 @@ def create_connection(self) -> AdbcConnection: msg = f"Could not configure connection using driver '{err_driver_name}'. Error: {e}" raise ImproperConfigurationError(msg) from e - def _update_dialect_for_extensions(self) -> None: - """Update statement_config dialect based on detected extensions. - - Priority: paradedb > pg_textsearch > pgvector > postgres (default). - Only switches when current dialect is ``postgres``. - """ - current_dialect = self.statement_config.dialect or "postgres" - if current_dialect != "postgres": - return - - if self._paradedb_available: - self.statement_config = self.statement_config.replace(dialect="paradedb") - elif self._pg_textsearch_available: - self.statement_config = self.statement_config.replace(dialect="pg_textsearch") - elif self._pgvector_available: - self.statement_config = self.statement_config.replace(dialect="pgvector") - @property def pg_textsearch_available(self) -> bool: """Return True if the pg_textsearch extension is available.""" diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index 712537c08..16b0f7369 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -544,19 +544,13 @@ def select_to_arrow( arrow_schema=arrow_schema, ) - # Use ADBC cursor for native Arrow with self.with_cursor(self.connection) as cursor, exc_handler: if cursor is None: msg = "Failed to create cursor" raise DatabaseConnectionError(msg) - # Get compiled SQL and parameters sql, driver_params = self._compiled_sql(prepared_statement, config) - - # Execute query cursor.execute(sql, driver_params or ()) - - # Fetch as Arrow table (zero-copy!) arrow_table = cursor.fetch_arrow_table() arrow_result = build_arrow_result_from_table( @@ -594,7 +588,6 @@ def select_to_storage( ) -> "StorageBridgeJob": """Stream query results to storage via the Arrow fast path.""" - _ = kwargs self._require_capability("arrow_export_enabled") arrow_result = self.select_to_arrow(statement, *parameters, statement_config=statement_config, **kwargs) sync_pipeline = self._storage_pipeline() diff --git a/sqlspec/adapters/asyncpg/config.py b/sqlspec/adapters/asyncpg/config.py index b0c0c6c14..1870bd7a2 100644 --- a/sqlspec/adapters/asyncpg/config.py +++ b/sqlspec/adapters/asyncpg/config.py @@ -19,17 +19,19 @@ from sqlspec.adapters.asyncpg.core import ( apply_driver_features, build_connection_config, - build_postgres_extension_probe_names, default_statement_config, - is_postgres_extension_active, register_json_codecs, register_pgvector_support, - resolve_postgres_extension_state, - resolve_runtime_statement_config, ) from sqlspec.adapters.asyncpg.driver import AsyncpgDriver, AsyncpgExceptionHandler from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs from sqlspec.core.capabilities import TypeCoercionCapabilities +from sqlspec.core.config_runtime import ( + build_postgres_extension_probe_names, + is_postgres_extension_active, + resolve_postgres_extension_state, + resolve_runtime_statement_config, +) from sqlspec.driver._async import AsyncPoolConnectionContext, AsyncPoolSessionFactory from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError from sqlspec.extensions.events import EventRuntimeHints diff --git a/sqlspec/adapters/bigquery/adk/store.py b/sqlspec/adapters/bigquery/adk/store.py index 0758288b4..67ae6c3eb 100644 --- a/sqlspec/adapters/bigquery/adk/store.py +++ b/sqlspec/adapters/bigquery/adk/store.py @@ -20,7 +20,6 @@ from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk._config_utils import _adk_config_from_extension from sqlspec.utils.serializers import from_json, to_json -from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: from collections.abc import Iterable @@ -647,10 +646,6 @@ def _decode_json(value: Any) -> "dict[str, Any] | None": msg = f"Unsupported JSON column representation from BigQuery: {type(value).__name__}" raise TypeError(msg) - @staticmethod - def _new_id() -> str: - return str(uuid4()) - def _session_record_from_row(row: "dict[str, Any]") -> StoredSession: return { diff --git a/sqlspec/adapters/cockroach_asyncpg/config.py b/sqlspec/adapters/cockroach_asyncpg/config.py index 51b1eeb9a..69ccd90b5 100644 --- a/sqlspec/adapters/cockroach_asyncpg/config.py +++ b/sqlspec/adapters/cockroach_asyncpg/config.py @@ -12,7 +12,6 @@ default_statement_config, register_json_codecs, register_pgvector_support, - resolve_runtime_statement_config, ) from sqlspec.adapters.cockroach_asyncpg._typing import ( CockroachAsyncpgConnection, @@ -22,6 +21,7 @@ from sqlspec.adapters.cockroach_asyncpg.driver import CockroachAsyncpgDriver, CockroachAsyncpgExceptionHandler from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs from sqlspec.core.capabilities import TypeCoercionCapabilities +from sqlspec.core.config_runtime import resolve_runtime_statement_config from sqlspec.driver._async import AsyncPoolConnectionContext, AsyncPoolSessionFactory from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.events import EventRuntimeHints diff --git a/sqlspec/adapters/cockroach_psycopg/config.py b/sqlspec/adapters/cockroach_psycopg/config.py index 866e96a6d..5c60030b5 100644 --- a/sqlspec/adapters/cockroach_psycopg/config.py +++ b/sqlspec/adapters/cockroach_psycopg/config.py @@ -19,9 +19,9 @@ CockroachPsycopgSyncDriver, CockroachPsycopgSyncExceptionHandler, ) -from sqlspec.adapters.psycopg.core import resolve_runtime_statement_config from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs, SyncDatabaseConfig from sqlspec.core.capabilities import TypeCoercionCapabilities +from sqlspec.core.config_runtime import resolve_runtime_statement_config from sqlspec.driver._async import AsyncPoolConnectionContext, AsyncPoolSessionFactory from sqlspec.driver._sync import SyncPoolConnectionContext, SyncPoolSessionFactory from sqlspec.exceptions import ImproperConfigurationError diff --git a/sqlspec/adapters/psqlpy/config.py b/sqlspec/adapters/psqlpy/config.py index 30966be60..f6c8975ac 100644 --- a/sqlspec/adapters/psqlpy/config.py +++ b/sqlspec/adapters/psqlpy/config.py @@ -6,18 +6,16 @@ from typing_extensions import NotRequired from sqlspec.adapters.psqlpy._typing import PsqlpyConnection, PsqlpyCursor, PsqlpySessionContext -from sqlspec.adapters.psqlpy.core import ( - apply_driver_features, - build_connection_config, +from sqlspec.adapters.psqlpy.core import apply_driver_features, build_connection_config, default_statement_config +from sqlspec.adapters.psqlpy.driver import PsqlpyDriver, PsqlpyExceptionHandler +from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs +from sqlspec.core.capabilities import TypeCoercionCapabilities +from sqlspec.core.config_runtime import ( build_postgres_extension_probe_names, - default_statement_config, is_postgres_extension_active, resolve_postgres_extension_state, resolve_runtime_statement_config, ) -from sqlspec.adapters.psqlpy.driver import PsqlpyDriver, PsqlpyExceptionHandler -from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs -from sqlspec.core.capabilities import TypeCoercionCapabilities from sqlspec.driver._async import AsyncPoolConnectionContext, AsyncPoolSessionFactory from sqlspec.extensions.events import EventRuntimeHints from sqlspec.utils.config_tools import normalize_connection_config diff --git a/sqlspec/adapters/psycopg/config.py b/sqlspec/adapters/psycopg/config.py index 14529adce..45dd91dd7 100644 --- a/sqlspec/adapters/psycopg/config.py +++ b/sqlspec/adapters/psycopg/config.py @@ -15,14 +15,7 @@ PsycopgSyncCursor, PsycopgSyncSessionContext, ) -from sqlspec.adapters.psycopg.core import ( - apply_driver_features, - build_postgres_extension_probe_names, - default_statement_config, - is_postgres_extension_active, - resolve_postgres_extension_state, - resolve_runtime_statement_config, -) +from sqlspec.adapters.psycopg.core import apply_driver_features, default_statement_config from sqlspec.adapters.psycopg.driver import ( PsycopgAsyncDriver, PsycopgAsyncExceptionHandler, @@ -32,6 +25,12 @@ from sqlspec.adapters.psycopg.type_converter import register_pgvector_async, register_pgvector_sync from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs, SyncDatabaseConfig from sqlspec.core.capabilities import TypeCoercionCapabilities +from sqlspec.core.config_runtime import ( + build_postgres_extension_probe_names, + is_postgres_extension_active, + resolve_postgres_extension_state, + resolve_runtime_statement_config, +) from sqlspec.driver._async import AsyncPoolConnectionContext, AsyncPoolSessionFactory from sqlspec.driver._sync import SyncPoolConnectionContext, SyncPoolSessionFactory from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError diff --git a/sqlspec/adapters/sqlite/driver.py b/sqlspec/adapters/sqlite/driver.py index 104653470..c080d40ac 100644 --- a/sqlspec/adapters/sqlite/driver.py +++ b/sqlspec/adapters/sqlite/driver.py @@ -183,14 +183,13 @@ def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionRe statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) successful_count = 0 - last_cursor = cursor for stmt in statements: cursor.execute(stmt, normalize_execute_parameters(prepared_parameters)) successful_count += 1 return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True ) def execute_many( @@ -414,12 +413,12 @@ def _execute_cache_hit( returns_rows = cached.operation_profile.returns_rows self._invalidate_rowid_target_cache(cached.operation_type) try: - if not returns_rows: - try: - cursor = self.connection.execute(cached.compiled_sql, params) - except sqlite3.Error as exc: - raise create_mapped_exception(exc) from exc + try: + cursor = self.connection.execute(cached.compiled_sql, params) + except sqlite3.Error as exc: + raise create_mapped_exception(exc) from exc + if not returns_rows: rowcount = cursor.rowcount affected_rows = rowcount if isinstance(rowcount, int) and rowcount > 0 else 0 last_inserted_id = resolve_lastrowid( @@ -432,11 +431,6 @@ def _execute_cache_hit( ) return DMLResult(cached.operation_type, affected_rows, last_inserted_id) - try: - cursor = self.connection.execute(cached.compiled_sql, params) - except sqlite3.Error as exc: - raise create_mapped_exception(exc) from exc - fetched_data = cursor.fetchall() affected_rows = resolve_rowcount(cursor) last_inserted_id = resolve_lastrowid( diff --git a/sqlspec/builder/__init__.py b/sqlspec/builder/__init__.py index eef646bbc..81b0fb52c 100644 --- a/sqlspec/builder/__init__.py +++ b/sqlspec/builder/__init__.py @@ -30,6 +30,7 @@ InsertFromSelectMixin, InsertIntoClauseMixin, InsertValuesMixin, + ReturningClauseMixin, UpdateFromClauseMixin, UpdateSetClauseMixin, UpdateTableClauseMixin, @@ -80,7 +81,6 @@ LimitOffsetClauseMixin, OrderByClauseMixin, PivotClauseMixin, - ReturningClauseMixin, Select, SelectClauseMixin, SetOperationMixin, diff --git a/sqlspec/builder/_base.py b/sqlspec/builder/_base.py index 80eb74eb6..761a08156 100644 --- a/sqlspec/builder/_base.py +++ b/sqlspec/builder/_base.py @@ -29,7 +29,6 @@ SQL, ParameterStyle, ParameterStyleConfig, - SQLResult, StatementConfig, get_cache, get_cache_config, @@ -237,15 +236,6 @@ def _create_base_expression(self) -> exp.Expr: A new sqlglot expression appropriate for the query type. """ - @property - @abstractmethod - def _expected_result_type(self) -> "type[SQLResult]": - """The expected result type for the query being built. - - Returns: - type[ResultT]: The type of the result. - """ - @staticmethod def _raise_builder_error(message: str, cause: BaseException | None = None) -> NoReturn: """Helper to raise SQLBuilderError, potentially with a cause. @@ -1254,10 +1244,6 @@ def _create_base_expression(self) -> exp.Expr: self._raise_builder_error(msg) return self._expression - @property - def _expected_result_type(self) -> "type[SQLResult]": - return SQLResult - class _BuilderCacheEntry: __slots__ = ("dialect", "expression") diff --git a/sqlspec/builder/_ddl.py b/sqlspec/builder/_ddl.py index 99293bc1c..7b7735d5d 100644 --- a/sqlspec/builder/_ddl.py +++ b/sqlspec/builder/_ddl.py @@ -14,7 +14,7 @@ from sqlspec.builder._base import BuiltQuery, QueryBuilder from sqlspec.builder._parsing_utils import _normalize_dialect from sqlspec.builder._select import Select -from sqlspec.core import SQL, SQLResult, StatementConfig +from sqlspec.core import SQL, StatementConfig from sqlspec.exceptions import SQLBuilderError from sqlspec.utils.type_guards import has_sqlglot_expression, has_with_method @@ -248,10 +248,6 @@ def _resolve_select_query(self, query: object, context: str, *, require_select_t return select_expr - @property - def _expected_result_type(self) -> "type[SQLResult]": - return SQLResult - def _prepare_expression(self, dialect: "DialectType" = None) -> None: target_dialect = _normalize_dialect(dialect or self.dialect) if self._expression is not None and target_dialect != self._expression_dialect: diff --git a/sqlspec/builder/_delete.py b/sqlspec/builder/_delete.py index e46936559..6d8bd6f6b 100644 --- a/sqlspec/builder/_delete.py +++ b/sqlspec/builder/_delete.py @@ -9,10 +9,9 @@ from sqlglot import exp from sqlspec.builder._base import BuiltQuery, QueryBuilder -from sqlspec.builder._dml import DeleteFromClauseMixin +from sqlspec.builder._dml import DeleteFromClauseMixin, ReturningClauseMixin from sqlspec.builder._explain import ExplainMixin -from sqlspec.builder._select import ReturningClauseMixin, WhereClauseMixin -from sqlspec.core import SQLResult +from sqlspec.builder._select import WhereClauseMixin from sqlspec.exceptions import SQLBuilderError if TYPE_CHECKING: @@ -44,15 +43,6 @@ def __init__(self, table: str | None = None, **kwargs: Any) -> None: if table: self.from_(table) - @property - def _expected_result_type(self) -> "type[SQLResult]": - """Get the expected result type for DELETE operations. - - Returns: - The ExecuteResult type for DELETE statements. - """ - return SQLResult - def _create_base_expression(self) -> "exp.Delete": """Create a new sqlglot Delete expression. diff --git a/sqlspec/builder/_dml.py b/sqlspec/builder/_dml.py index 52ff5b28b..ef5506ee2 100644 --- a/sqlspec/builder/_dml.py +++ b/sqlspec/builder/_dml.py @@ -1,24 +1,30 @@ """Reusable mixins for INSERT/UPDATE/DELETE builders.""" from collections.abc import Mapping, Sequence -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from mypy_extensions import trait from sqlglot import exp from typing_extensions import Self from sqlspec.builder._base import BuiltQuery, QueryBuilder -from sqlspec.builder._parsing_utils import extract_sql_object_expression +from sqlspec.builder._parsing_utils import extract_expression, extract_sql_object_expression from sqlspec.exceptions import SQLBuilderError from sqlspec.protocols import SQLBuilderProtocol from sqlspec.utils.serializers import schema_dump from sqlspec.utils.type_guards import has_expression_and_sql, has_parameter_builder, is_dict +if TYPE_CHECKING: + from sqlspec.builder._column import Column + from sqlspec.builder._expression_wrappers import ExpressionWrapper + from sqlspec.builder._select import Case + __all__ = ( "DeleteFromClauseMixin", "InsertFromSelectMixin", "InsertIntoClauseMixin", "InsertValuesMixin", + "ReturningClauseMixin", "UpdateFromClauseMixin", "UpdateSetClauseMixin", "UpdateTableClauseMixin", @@ -444,3 +450,34 @@ def from_(self, table: str | exp.Expr | Any, alias: str | None = None) -> Self: from_table.append("joins", exp.Join(this=table_expr)) return self + + +@trait +class ReturningClauseMixin: + """Mixin providing RETURNING clause support for DML builders.""" + + __slots__ = () + + _expression: exp.Expr | None + + def returning(self, *columns: "str | exp.Expr | Column | ExpressionWrapper | Case") -> Self: + """Add RETURNING clause to the DML statement. + + Args: + *columns: Columns or expressions to return. + + Returns: + The builder instance for method chaining. + + Raises: + SQLBuilderError: If expression not initialized or not DML. + """ + if self._expression is None: + msg = "Cannot add RETURNING: expression not initialized." + raise SQLBuilderError(msg) + if not isinstance(self._expression, (exp.Insert, exp.Update, exp.Delete)): + msg = "RETURNING only supported for INSERT, UPDATE, DELETE statements." + raise SQLBuilderError(msg) + returning_exprs = [extract_expression(col) for col in columns] + self._expression.set("returning", exp.Returning(expressions=returning_exprs)) + return self diff --git a/sqlspec/builder/_explain.py b/sqlspec/builder/_explain.py index 2b5f3d4f3..dd3275958 100644 --- a/sqlspec/builder/_explain.py +++ b/sqlspec/builder/_explain.py @@ -4,7 +4,7 @@ dialect-aware SQL generation. """ -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final from mypy_extensions import trait from sqlglot import Dialect, exp @@ -42,7 +42,12 @@ DUCKDB_DIALECTS = frozenset({"duckdb"}) ORACLE_DIALECTS = frozenset({"oracle"}) BIGQUERY_DIALECTS = frozenset({"bigquery"}) -SPANNER_DIALECTS = frozenset({"spanner"}) +_MYSQL_FORMAT_MAP: Final[dict[ExplainFormat, str]] = { + ExplainFormat.JSON: "JSON", + ExplainFormat.TREE: "TREE", + ExplainFormat.TRADITIONAL: "TRADITIONAL", + ExplainFormat.TEXT: "TRADITIONAL", +} def normalize_dialect_name(dialect: "DialectType | None") -> str | None: @@ -124,13 +129,7 @@ def build_mysql_explain(statement_sql: str, options: "ExplainOptions") -> str: return f"EXPLAIN ANALYZE {statement_sql}" if options.format is not None: - format_map = { - ExplainFormat.JSON: "JSON", - ExplainFormat.TREE: "TREE", - ExplainFormat.TRADITIONAL: "TRADITIONAL", - ExplainFormat.TEXT: "TRADITIONAL", - } - fmt = format_map.get(options.format, "TRADITIONAL") + fmt = _MYSQL_FORMAT_MAP.get(options.format, "TRADITIONAL") return f"EXPLAIN FORMAT = {fmt} {statement_sql}" return f"EXPLAIN {statement_sql}" @@ -255,8 +254,6 @@ def build_explain_sql(statement_sql: str, options: "ExplainOptions", dialect: "D return build_oracle_explain(statement_sql, options) if dialect_name in BIGQUERY_DIALECTS: return build_bigquery_explain(statement_sql, options) - if dialect_name in SPANNER_DIALECTS: - return build_generic_explain(statement_sql, options) return build_generic_explain(statement_sql, options) diff --git a/sqlspec/builder/_factory.py b/sqlspec/builder/_factory.py index e60149d6c..f27d83a27 100644 --- a/sqlspec/builder/_factory.py +++ b/sqlspec/builder/_factory.py @@ -69,38 +69,7 @@ from sqlspec.protocols import SQLBuilderProtocol -__all__ = ( - "AlterTable", - "Case", - "Column", - "CommentOn", - "CreateIndex", - "CreateMaterializedView", - "CreateSchema", - "CreateTable", - "CreateTableAsSelect", - "CreateView", - "Delete", - "DropIndex", - "DropMaterializedView", - "DropSchema", - "DropTable", - "DropView", - "Explain", - "Insert", - "Merge", - "RenameTable", - "SQLFactory", - "Select", - "Truncate", - "Update", - "Values", - "WindowFunctionBuilder", - "build_copy_from_statement", - "build_copy_statement", - "build_copy_to_statement", - "sql", -) +__all__ = ("SQLFactory", "build_copy_from_statement", "build_copy_statement", "build_copy_to_statement", "sql") logger = get_logger("sqlspec.builder.factory") @@ -137,59 +106,6 @@ } -def _fingerprint_sql(sql: str) -> str: - digest = hashlib.sha256(sql.encode("utf-8", errors="replace")).hexdigest() - return digest[:12] - - -def _normalize_copy_dialect(dialect: DialectType | None) -> str: - if dialect is None: - return "postgres" - if isinstance(dialect, str): - return dialect - return str(dialect) - - -def _to_copy_schema(table: str, columns: "Sequence[str] | None") -> exp.Expr: - base = exp.table_(table) - if not columns: - return base - column_nodes = [exp.column(column_name) for column_name in columns] - return exp.Schema(this=base, expressions=column_nodes) - - -def _build_copy_expression( - *, direction: str, table: str, location: str, columns: "Sequence[str] | None", options: "Mapping[str, Any] | None" -) -> exp.Copy: - copy_args: dict[str, Any] = {"this": _to_copy_schema(table, columns), "files": [exp.Literal.string(location)]} - - if direction == "from": - copy_args["kind"] = True - elif direction == "to": - copy_args["kind"] = False - - if options: - params: list[exp.CopyParameter] = [] - for key, value in options.items(): - identifier = exp.Var(this=str(key).upper()) - value_expression: exp.Expr - if isinstance(value, bool): - value_expression = exp.Boolean(this=value) - elif value is None: - value_expression = exp.null() - elif isinstance(value, (int, float)): - value_expression = exp.Literal.number(value) - elif isinstance(value, (list, tuple)): - elements = [exp.Literal.string(str(item)) for item in value] - value_expression = exp.Array(expressions=elements) - else: - value_expression = exp.Literal.string(str(value)) - params.append(exp.CopyParameter(this=identifier, expression=value_expression)) - copy_args["params"] = params - - return exp.Copy(**copy_args) - - def build_copy_statement( *, direction: str, @@ -295,15 +211,7 @@ def __call__(self, statement: str, dialect: DialectType = None) -> "Any": msg = f"Failed to parse SQL: {e}" raise SQLBuilderError(msg) from e actual_type = type(parsed_expr).__name__.upper() - expr_type_map = { - "SELECT": "SELECT", - "INSERT": "INSERT", - "UPDATE": "UPDATE", - "DELETE": "DELETE", - "MERGE": "MERGE", - "WITH": "WITH", - } - actual_type_str = expr_type_map.get(actual_type, actual_type) + actual_type_str = actual_type if actual_type_str == "SELECT" or ( actual_type_str == "WITH" and parsed_expr.this and isinstance(parsed_expr.this, exp.Select) ): @@ -1608,4 +1516,57 @@ def _create_window_function( return FunctionExpression(exp.Window(this=func_expr, **over_args)) +def _fingerprint_sql(sql: str) -> str: + digest = hashlib.sha256(sql.encode("utf-8", errors="replace")).hexdigest() + return digest[:12] + + +def _normalize_copy_dialect(dialect: DialectType | None) -> str: + if dialect is None: + return "postgres" + if isinstance(dialect, str): + return dialect + return str(dialect) + + +def _to_copy_schema(table: str, columns: "Sequence[str] | None") -> exp.Expr: + base = exp.table_(table) + if not columns: + return base + column_nodes = [exp.column(column_name) for column_name in columns] + return exp.Schema(this=base, expressions=column_nodes) + + +def _build_copy_expression( + *, direction: str, table: str, location: str, columns: "Sequence[str] | None", options: "Mapping[str, Any] | None" +) -> exp.Copy: + copy_args: dict[str, Any] = {"this": _to_copy_schema(table, columns), "files": [exp.Literal.string(location)]} + + if direction == "from": + copy_args["kind"] = True + elif direction == "to": + copy_args["kind"] = False + + if options: + params: list[exp.CopyParameter] = [] + for key, value in options.items(): + identifier = exp.Var(this=str(key).upper()) + value_expression: exp.Expr + if isinstance(value, bool): + value_expression = exp.Boolean(this=value) + elif value is None: + value_expression = exp.null() + elif isinstance(value, (int, float)): + value_expression = exp.Literal.number(value) + elif isinstance(value, (list, tuple)): + elements = [exp.Literal.string(str(item)) for item in value] + value_expression = exp.Array(expressions=elements) + else: + value_expression = exp.Literal.string(str(value)) + params.append(exp.CopyParameter(this=identifier, expression=value_expression)) + copy_args["params"] = params + + return exp.Copy(**copy_args) + + sql = SQLFactory() diff --git a/sqlspec/builder/_insert.py b/sqlspec/builder/_insert.py index f38266174..82071bdfc 100644 --- a/sqlspec/builder/_insert.py +++ b/sqlspec/builder/_insert.py @@ -10,11 +10,9 @@ from typing_extensions import Self from sqlspec.builder._base import QueryBuilder -from sqlspec.builder._dml import InsertFromSelectMixin, InsertIntoClauseMixin, InsertValuesMixin +from sqlspec.builder._dml import InsertFromSelectMixin, InsertIntoClauseMixin, InsertValuesMixin, ReturningClauseMixin from sqlspec.builder._explain import ExplainMixin from sqlspec.builder._parsing_utils import extract_sql_object_expression -from sqlspec.builder._select import ReturningClauseMixin -from sqlspec.core import SQLResult from sqlspec.exceptions import SQLBuilderError from sqlspec.utils.serializers import schema_dump, serialize_collection from sqlspec.utils.type_guards import has_expression_and_sql @@ -67,15 +65,6 @@ def _create_base_expression(self) -> exp.Insert: """ return exp.Insert() - @property - def _expected_result_type(self) -> "type[SQLResult]": - """Specifies the expected result type for an INSERT query. - - Returns: - The type of result expected for INSERT operations. - """ - return SQLResult - def _insert_expression(self) -> exp.Insert: """Safely gets and casts the internal expression to exp.Insert. diff --git a/sqlspec/builder/_locking.py b/sqlspec/builder/_locking.py index 3c4ce0d4a..db9d5da30 100644 --- a/sqlspec/builder/_locking.py +++ b/sqlspec/builder/_locking.py @@ -3,7 +3,7 @@ from collections.abc import Callable, MutableMapping from typing import TYPE_CHECKING, ClassVar, Protocol, cast -from sqlglot import exp +from sqlglot import Dialect, exp from sqlglot.generator import Generator from sqlspec.builder._generation import invalidate_generator_dispatch @@ -66,8 +66,6 @@ def _generator_class_for_dialect(dialect: "DialectType | str | None") -> "type[_ if dialect is None: return cast("type[_GeneratorClass]", Generator) - from sqlglot import Dialect - dialect_class: type[Dialect] if isinstance(dialect, str): dialect_class = type(Dialect.get_or_raise(dialect)) diff --git a/sqlspec/builder/_merge.py b/sqlspec/builder/_merge.py index a427b8b4c..15043faf0 100644 --- a/sqlspec/builder/_merge.py +++ b/sqlspec/builder/_merge.py @@ -20,7 +20,6 @@ from sqlspec.builder._explain import ExplainMixin from sqlspec.builder._parsing_utils import _coerce_column, _resolve_dialect, extract_sql_object_expression from sqlspec.builder._select import is_explicitly_quoted -from sqlspec.core import SQLResult from sqlspec.exceptions import DialectNotSupportedError, SQLBuilderError from sqlspec.utils.dispatch import TypeDispatcher from sqlspec.utils.serializers import to_json @@ -668,15 +667,6 @@ def __init__(self, target_table: str | None = None, **kwargs: Any) -> None: if target_table: self.into(target_table) - @property - def _expected_result_type(self) -> "type[SQLResult]": - """Return the expected result type for this builder. - - Returns: - The SQLResult type for MERGE statements. - """ - return SQLResult - def _create_base_expression(self) -> "exp.Merge": """Create a base MERGE expression. diff --git a/sqlspec/builder/_select.py b/sqlspec/builder/_select.py index 8437a9984..0a58f0b67 100644 --- a/sqlspec/builder/_select.py +++ b/sqlspec/builder/_select.py @@ -31,7 +31,8 @@ parse_table_expression, to_expression, ) -from sqlspec.core import SQL, ParameterStyle, SQLResult +from sqlspec.core import SQL, ParameterStyle +from sqlspec.core.query_modifiers import expr_eq, expr_gt, expr_gte, expr_lt, expr_lte, expr_neq, expr_not_like from sqlspec.exceptions import SQLBuilderError from sqlspec.utils.type_guards import ( has_expression_and_parameters, @@ -59,7 +60,6 @@ "LimitOffsetClauseMixin", "OrderByClauseMixin", "PivotClauseMixin", - "ReturningClauseMixin", "Select", "SelectClauseMixin", "SetOperationMixin", @@ -84,30 +84,6 @@ def is_explicitly_quoted(identifier: Any) -> bool: ) -def _expr_eq(col: "exp.Expr", placeholder: "exp.Placeholder") -> "exp.Expr": - return exp.EQ(this=col, expression=placeholder) - - -def _expr_neq(col: "exp.Expr", placeholder: "exp.Placeholder") -> "exp.Expr": - return exp.NEQ(this=col, expression=placeholder) - - -def _expr_gt(col: "exp.Expr", placeholder: "exp.Placeholder") -> "exp.Expr": - return exp.GT(this=col, expression=placeholder) - - -def _expr_gte(col: "exp.Expr", placeholder: "exp.Placeholder") -> "exp.Expr": - return exp.GTE(this=col, expression=placeholder) - - -def _expr_lt(col: "exp.Expr", placeholder: "exp.Placeholder") -> "exp.Expr": - return exp.LT(this=col, expression=placeholder) - - -def _expr_lte(col: "exp.Expr", placeholder: "exp.Placeholder") -> "exp.Expr": - return exp.LTE(this=col, expression=placeholder) - - def _expr_like_exp(col: "exp.Expr", placeholder: "exp.Placeholder") -> "exp.Expr": return exp.Like(this=col, expression=placeholder) @@ -116,25 +92,21 @@ def _expr_like_method(col: "exp.Expr", placeholder: "exp.Placeholder") -> "exp.E return cast("exp.Expr", col.like(placeholder)) -def _expr_not_like(col: "exp.Expr", placeholder: "exp.Placeholder") -> "exp.Expr": - return exp.Not(this=exp.Like(this=col, expression=placeholder)) - - def _expr_ilike(col: "exp.Expr", placeholder: "exp.Placeholder") -> "exp.Expr": return cast("exp.Expr", col.ilike(placeholder)) _SIMPLE_OPERATOR_MAP: dict[str, Any] = { - "=": _expr_eq, - "==": _expr_eq, - "!=": _expr_neq, - "<>": _expr_neq, - ">": _expr_gt, - ">=": _expr_gte, - "<": _expr_lt, - "<=": _expr_lte, + "=": expr_eq, + "==": expr_eq, + "!=": expr_neq, + "<>": expr_neq, + ">": expr_gt, + ">=": expr_gte, + "<": expr_lt, + "<=": expr_lte, "LIKE": _expr_like_exp, - "NOT LIKE": _expr_not_like, + "NOT LIKE": expr_not_like, } @@ -439,24 +411,6 @@ def offset(self, value: int) -> Self: return cast("Self", builder) -@trait -class ReturningClauseMixin: - __slots__ = () - - _expression: exp.Expr | None - - def returning(self, *columns: Union[str, exp.Expr, "Column", "ExpressionWrapper", Case]) -> Self: - if self._expression is None: - msg = "Cannot add RETURNING: expression not initialized." - raise SQLBuilderError(msg) - if not isinstance(self._expression, (exp.Insert, exp.Update, exp.Delete)): - msg = "RETURNING only supported for INSERT, UPDATE, DELETE statements." - raise SQLBuilderError(msg) - returning_exprs = [extract_expression(col) for col in columns] - self._expression.set("returning", exp.Returning(expressions=returning_exprs)) - return self - - @trait class WhereClauseMixin: __slots__ = () @@ -524,22 +478,7 @@ def _handle_in_operator(self, column_exp: exp.Expr, value: Any, column_name: str return exp.In(this=column_exp, expressions=[exp.Placeholder(this=param_name)]) def _handle_not_in_operator(self, column_exp: exp.Expr, value: Any, column_name: str = "column") -> exp.Expr: - builder = cast("SQLBuilderProtocol", self) - if has_parameter_builder(value) or isinstance(value, exp.Expr): - subquery_expr = self._normalize_subquery_expression(value, builder) - return exp.Not(this=exp.In(this=column_exp, expressions=[subquery_expr])) - if is_iterable_parameters(value): - placeholders = [] - for index, element in enumerate(value): - name_seed = column_name if len(value) == 1 else f"{column_name}_{index + 1}" - param_name = builder._next_parameter_name(name_seed) - _, param_name = builder.add_parameter(element, name=param_name) - placeholders.append(exp.Placeholder(this=param_name)) - return exp.Not(this=exp.In(this=column_exp, expressions=placeholders)) - - param_name = builder._next_parameter_name(column_name) - _, param_name = builder.add_parameter(value, name=param_name) - return exp.Not(this=exp.In(this=column_exp, expressions=[exp.Placeholder(this=param_name)])) + return exp.Not(this=self._handle_in_operator(column_exp, value, column_name=column_name)) def _handle_is_operator(self, column_exp: exp.Expr, value: Any) -> exp.Expr: value_expr = exp.Null() if value is None else exp.convert(value) @@ -689,7 +628,7 @@ def _create_or_expression(self, conditions: "list[exp.Expr]") -> exp.Expr: def _process_tuple_condition(self, condition: "tuple[Any, ...]") -> exp.Expr: if len(condition) == PAIR_LENGTH: column, value = condition - return self._create_parameterized_condition(column, value, _expr_eq) + return self._create_parameterized_condition(column, value, expr_eq) if len(condition) != TRIPLE_LENGTH: msg = f"Condition tuple must have 2 or 3 elements, got {len(condition)}" @@ -873,22 +812,22 @@ def _build_not_exists_condition(self, subquery: Any) -> exp.Expr: return exp.Not(this=self._build_exists_condition(subquery)) def where_eq(self, column: str | exp.Column, value: Any) -> Self: - return self.where(self._build_comparison_condition(column, value, _expr_eq)) + return self.where(self._build_comparison_condition(column, value, expr_eq)) def where_neq(self, column: str | exp.Column, value: Any) -> Self: - return self.where(self._build_comparison_condition(column, value, _expr_neq)) + return self.where(self._build_comparison_condition(column, value, expr_neq)) def where_lt(self, column: str | exp.Column, value: Any) -> Self: - return self.where(self._build_comparison_condition(column, value, _expr_lt)) + return self.where(self._build_comparison_condition(column, value, expr_lt)) def where_lte(self, column: str | exp.Column, value: Any) -> Self: - return self.where(self._build_comparison_condition(column, value, _expr_lte)) + return self.where(self._build_comparison_condition(column, value, expr_lte)) def where_gt(self, column: str | exp.Column, value: Any) -> Self: - return self.where(self._build_comparison_condition(column, value, _expr_gt)) + return self.where(self._build_comparison_condition(column, value, expr_gt)) def where_gte(self, column: str | exp.Column, value: Any) -> Self: - return self.where(self._build_comparison_condition(column, value, _expr_gte)) + return self.where(self._build_comparison_condition(column, value, expr_gte)) def where_between(self, column: str | exp.Column, low: Any, high: Any) -> Self: return self.where(self._build_between_condition(column, low, high)) @@ -897,7 +836,7 @@ def where_like(self, column: str | exp.Column, pattern: str, escape: str | None return self.where(self._build_like_condition(column, pattern, escape)) def where_not_like(self, column: str | exp.Column, pattern: str) -> Self: - return self.where(self._build_comparison_condition(column, pattern, _expr_not_like)) + return self.where(self._build_comparison_condition(column, pattern, expr_not_like)) def where_ilike(self, column: str | exp.Column, pattern: str) -> Self: return self.where(self._build_comparison_condition(column, pattern, _expr_ilike)) @@ -932,22 +871,22 @@ def where_like_any(self, column: str | exp.Column, patterns: list[str]) -> Self: return self.where(or_condition) def or_where_eq(self, column: str | exp.Column, value: Any) -> Self: - return self._combine_with_or(self._build_comparison_condition(column, value, _expr_eq)) + return self._combine_with_or(self._build_comparison_condition(column, value, expr_eq)) def or_where_neq(self, column: str | exp.Column, value: Any) -> Self: - return self._combine_with_or(self._build_comparison_condition(column, value, _expr_neq)) + return self._combine_with_or(self._build_comparison_condition(column, value, expr_neq)) def or_where_lt(self, column: str | exp.Column, value: Any) -> Self: - return self._combine_with_or(self._build_comparison_condition(column, value, _expr_lt)) + return self._combine_with_or(self._build_comparison_condition(column, value, expr_lt)) def or_where_lte(self, column: str | exp.Column, value: Any) -> Self: - return self._combine_with_or(self._build_comparison_condition(column, value, _expr_lte)) + return self._combine_with_or(self._build_comparison_condition(column, value, expr_lte)) def or_where_gt(self, column: str | exp.Column, value: Any) -> Self: - return self._combine_with_or(self._build_comparison_condition(column, value, _expr_gt)) + return self._combine_with_or(self._build_comparison_condition(column, value, expr_gt)) def or_where_gte(self, column: str | exp.Column, value: Any) -> Self: - return self._combine_with_or(self._build_comparison_condition(column, value, _expr_gte)) + return self._combine_with_or(self._build_comparison_condition(column, value, expr_gte)) def or_where_between(self, column: str | exp.Column, low: Any, high: Any) -> Self: return self._combine_with_or(self._build_between_condition(column, low, high)) @@ -956,7 +895,7 @@ def or_where_like(self, column: str | exp.Column, pattern: str, escape: str | No return self._combine_with_or(self._build_like_condition(column, pattern, escape)) def or_where_not_like(self, column: str | exp.Column, pattern: str) -> Self: - return self._combine_with_or(self._build_comparison_condition(column, pattern, _expr_not_like)) + return self._combine_with_or(self._build_comparison_condition(column, pattern, expr_not_like)) def or_where_ilike(self, column: str | exp.Column, pattern: str) -> Self: return self._combine_with_or(self._build_comparison_condition(column, pattern, _expr_ilike)) @@ -1271,15 +1210,6 @@ def __init__(self, *columns: str, **kwargs: Any) -> None: if columns: self.select(*columns) - @property - def _expected_result_type(self) -> "type[SQLResult]": - """Get the expected result type for SELECT operations. - - Returns: - type: The SelectResult type. - """ - return SQLResult - def _create_base_expression(self) -> exp.Select: """Create base SELECT expression.""" if self._expression is None or not isinstance(self._expression, exp.Select): diff --git a/sqlspec/builder/_update.py b/sqlspec/builder/_update.py index 6f76d860f..e8a0e77e5 100644 --- a/sqlspec/builder/_update.py +++ b/sqlspec/builder/_update.py @@ -10,11 +10,15 @@ from typing_extensions import Self from sqlspec.builder._base import BuiltQuery, QueryBuilder -from sqlspec.builder._dml import UpdateFromClauseMixin, UpdateSetClauseMixin, UpdateTableClauseMixin +from sqlspec.builder._dml import ( + ReturningClauseMixin, + UpdateFromClauseMixin, + UpdateSetClauseMixin, + UpdateTableClauseMixin, +) from sqlspec.builder._explain import ExplainMixin from sqlspec.builder._join import build_join_clause -from sqlspec.builder._select import ReturningClauseMixin, WhereClauseMixin -from sqlspec.core import SQLResult +from sqlspec.builder._select import WhereClauseMixin from sqlspec.exceptions import SQLBuilderError if TYPE_CHECKING: @@ -56,11 +60,6 @@ def __init__(self, table: str | None = None, **kwargs: Any) -> None: if table: self.table(table) - @property - def _expected_result_type(self) -> "type[SQLResult]": - """Return the expected result type for this builder.""" - return SQLResult - def _create_base_expression(self) -> exp.Update: """Create a base UPDATE expression. diff --git a/sqlspec/builder/_values.py b/sqlspec/builder/_values.py index 6b23aa45e..cb0f22de2 100644 --- a/sqlspec/builder/_values.py +++ b/sqlspec/builder/_values.py @@ -12,7 +12,6 @@ from sqlspec.builder._base import QueryBuilder from sqlspec.builder._parsing_utils import extract_sql_object_expression -from sqlspec.core import SQLResult from sqlspec.exceptions import SQLBuilderError from sqlspec.utils.type_guards import has_expression_and_sql @@ -57,11 +56,6 @@ def _create_base_expression(self) -> exp.Values: """Create initial empty VALUES expression.""" return exp.Values() - @property - def _expected_result_type(self) -> type[SQLResult]: - """Return expected result type for VALUES queries.""" - return SQLResult - @property def alias_name(self) -> str | None: """Get the table alias name if defined. diff --git a/sqlspec/builder/_vector_distance.py b/sqlspec/builder/_vector_distance.py index f5f628a7c..8bd175cae 100644 --- a/sqlspec/builder/_vector_distance.py +++ b/sqlspec/builder/_vector_distance.py @@ -26,6 +26,24 @@ ) _VECTOR_DISTANCE_META_KEY: Final[str] = "sqlspec_vector_distance_metric" +_POSTGRES_OPERATOR_MAP: Final[dict[str, str]] = {"euclidean": "<->", "cosine": "<=>", "inner_product": "<#>"} +_MYSQL_METRIC_MAP: Final[dict[str, str]] = {"euclidean": "EUCLIDEAN", "cosine": "COSINE", "inner_product": "DOT"} +_ORACLE_METRIC_MAP: Final[dict[str, str]] = { + "euclidean": "EUCLIDEAN", + "cosine": "COSINE", + "inner_product": "DOT", + "euclidean_squared": "EUCLIDEAN_SQUARED", +} +_BIGQUERY_FUNCTION_MAP: Final[dict[str, str]] = { + "euclidean": "EUCLIDEAN_DISTANCE", + "cosine": "COSINE_DISTANCE", + "inner_product": "DOT_PRODUCT", +} +_DUCKDB_FUNCTION_MAP: Final[dict[str, str]] = { + "euclidean": "array_distance", + "cosine": "array_cosine_distance", + "inner_product": "array_negative_inner_product", +} _OperatorTransform = Callable[[Any, exp.Operator], str] _SQLGLOT_VECTOR_DISTANCE_REGISTERED = False _BASE_OPERATOR_TRANSFORM: _OperatorTransform | None = None @@ -90,9 +108,7 @@ def VectorDistance(*, this: exp.Expr, expression: exp.Expr, metric: Any = "eucli def render_vector_distance_postgres(left: str, right: str, metric: str) -> str: """Render PostgreSQL pgvector operator syntax.""" - operator_map = {"euclidean": "<->", "cosine": "<=>", "inner_product": "<#>"} - - operator = operator_map.get(metric) + operator = _POSTGRES_OPERATOR_MAP.get(metric) if operator: return f"{left} {operator} {right}" @@ -101,9 +117,7 @@ def render_vector_distance_postgres(left: str, right: str, metric: str) -> str: def render_vector_distance_mysql(left: str, right: str, metric: str) -> str: """Render MySQL DISTANCE function syntax.""" - metric_map = {"euclidean": "EUCLIDEAN", "cosine": "COSINE", "inner_product": "DOT"} - - mysql_metric = metric_map.get(metric, "EUCLIDEAN") + mysql_metric = _MYSQL_METRIC_MAP.get(metric, "EUCLIDEAN") if ("ARRAY" in right or "[" in right) and "STRING_TO_VECTOR" not in right: right = f"STRING_TO_VECTOR({right})" @@ -113,14 +127,7 @@ def render_vector_distance_mysql(left: str, right: str, metric: str) -> str: def render_vector_distance_oracle(left: str, right: str, metric: str) -> str: """Render Oracle VECTOR_DISTANCE function syntax.""" - metric_map = { - "euclidean": "EUCLIDEAN", - "cosine": "COSINE", - "inner_product": "DOT", - "euclidean_squared": "EUCLIDEAN_SQUARED", - } - - oracle_metric = metric_map.get(metric, "EUCLIDEAN") + oracle_metric = _ORACLE_METRIC_MAP.get(metric, "EUCLIDEAN") if ("[" in right or "ARRAY" in right) and "TO_VECTOR" not in right: right = f"TO_VECTOR({right})" @@ -130,9 +137,7 @@ def render_vector_distance_oracle(left: str, right: str, metric: str) -> str: def render_vector_distance_bigquery(left: str, right: str, metric: str) -> str: """Render BigQuery vector distance function syntax.""" - function_map = {"euclidean": "EUCLIDEAN_DISTANCE", "cosine": "COSINE_DISTANCE", "inner_product": "DOT_PRODUCT"} - - function_name = function_map.get(metric) + function_name = _BIGQUERY_FUNCTION_MAP.get(metric) if function_name: return f"{function_name}({left}, {right})" @@ -141,12 +146,7 @@ def render_vector_distance_bigquery(left: str, right: str, metric: str) -> str: def render_vector_distance_duckdb(left: str, right: str, metric: str, *, dimension: int | None = None) -> str: """Render DuckDB VSS extension function syntax.""" - function_map = { - "euclidean": "array_distance", - "cosine": "array_cosine_distance", - "inner_product": "array_negative_inner_product", - } - function_name = function_map.get(metric) + function_name = _DUCKDB_FUNCTION_MAP.get(metric) if function_name: target_type = f"DOUBLE[{dimension}]" if dimension is not None else "DOUBLE[]" return f"{function_name}({left}, CAST({right} AS {target_type}))" diff --git a/sqlspec/cli.py b/sqlspec/cli.py index 5db95bc2e..da3531628 100644 --- a/sqlspec/cli.py +++ b/sqlspec/cli.py @@ -518,6 +518,32 @@ async def _async_show(config: Any) -> None: sqlspec_config = get_config_by_bind_key(ctx, bind_key) _show_for_config(sqlspec_config) + def _create_echo_settings( + no_echo: bool, summary_only: bool, use_logger: bool + ) -> tuple[bool | None, bool | None, bool, Callable[[Any], bool]]: + """Create echo and summary settings with echo predicate closure. + + Args: + no_echo: Flag indicating console output should be suppressed. + summary_only: Flag indicating only summary logs should be emitted. + use_logger: Flag indicating logger is used rather than console. + + Returns: + Tuple of (echo_setting, summary_setting, echo_enabled, should_echo_predicate). + """ + effective_no_echo = no_echo or summary_only + echo_setting = False if effective_no_echo else None + summary_setting = True if summary_only else None + echo_enabled = not effective_no_echo and not use_logger + + def _should_echo(config: Any) -> bool: + if not echo_enabled: + return False + migration_config = cast("dict[str, Any]", getattr(config, "migration_config", None)) or {} + return not bool(migration_config.get("use_logger", False)) + + return echo_setting, summary_setting, echo_enabled, _should_echo + @database_group.command(name="downgrade", help="Downgrade database to a specific revision.") @bind_key_option @no_prompt_option @@ -540,18 +566,10 @@ def downgrade_database( # pyright: ignore[reportUnusedFunction] summary_only: bool, ) -> None: """Downgrade the database to the latest revision.""" - ctx = _ensure_click_context() - effective_no_echo = no_echo or summary_only - echo_setting = False if effective_no_echo else None - summary_setting = True if summary_only else None - echo_enabled = not effective_no_echo and not use_logger - - def _should_echo(config: Any) -> bool: - if not echo_enabled: - return False - migration_config = cast("dict[str, Any]", getattr(config, "migration_config", None)) or {} - return not bool(migration_config.get("use_logger", False)) + echo_setting, summary_setting, echo_enabled, _should_echo = _create_echo_settings( + no_echo, summary_only, use_logger + ) def _downgrade_for_config(config: Any) -> None: """Downgrade a single config with sync/async dispatch.""" @@ -646,16 +664,9 @@ def upgrade_database( # pyright: ignore[reportUnusedFunction] ) -> None: """Upgrade the database to the latest revision.""" ctx = _ensure_click_context() - effective_no_echo = no_echo or summary_only - echo_setting = False if effective_no_echo else None - summary_setting = True if summary_only else None - echo_enabled = not effective_no_echo and not use_logger - - def _should_echo(config: Any) -> bool: - if not echo_enabled: - return False - migration_config = cast("dict[str, Any]", getattr(config, "migration_config", None)) or {} - return not bool(migration_config.get("use_logger", False)) + echo_setting, summary_setting, echo_enabled, _should_echo = _create_echo_settings( + no_echo, summary_only, use_logger + ) def _upgrade_for_config(config: Any) -> None: """Upgrade a single config with sync/async dispatch.""" diff --git a/sqlspec/config.py b/sqlspec/config.py index 3de76a9eb..b4cbfee47 100644 --- a/sqlspec/config.py +++ b/sqlspec/config.py @@ -273,6 +273,23 @@ class MigrationConfig(TypedDict): } +def validate_migration_config_keys(migration_config: "Mapping[str, Any]") -> None: + """Reject migration configuration keys that SQLSpec does not read. + + Args: + migration_config: Migration configuration mapping to check. + + Raises: + ImproperConfigurationError: If the mapping contains an unrecognized key. + """ + lines = _report_unknown_keys(migration_config, MIGRATION_CONFIG_KEYS, "", "") + templates = migration_config.get("templates") + if templates is not None: + lines.extend(_report_template_keys(templates)) + if lines: + raise ImproperConfigurationError(" ".join(lines)) + + def _report_unknown_keys( mapping: "Mapping[str, Any]", valid_keys: "frozenset[str]", prefix: str, scope: str ) -> "list[str]": @@ -325,23 +342,6 @@ def _report_template_keys(templates: Any) -> "list[str]": return lines -def validate_migration_config_keys(migration_config: "Mapping[str, Any]") -> None: - """Reject migration configuration keys that SQLSpec does not read. - - Args: - migration_config: Migration configuration mapping to check. - - Raises: - ImproperConfigurationError: If the mapping contains an unrecognized key. - """ - lines = _report_unknown_keys(migration_config, MIGRATION_CONFIG_KEYS, "", "") - templates = migration_config.get("templates") - if templates is not None: - lines.extend(_report_template_keys(templates)) - if lines: - raise ImproperConfigurationError(" ".join(lines)) - - class FlaskConfig(TypedDict): """Configuration options for Flask SQLSpec extension. diff --git a/sqlspec/core/__init__.py b/sqlspec/core/__init__.py index 9d2199c4c..8d21e56b0 100644 --- a/sqlspec/core/__init__.py +++ b/sqlspec/core/__init__.py @@ -1,87 +1,24 @@ -"""SQLSpec Core Module - SQL Processing System. +"""SQLSpec Core Module - SQL Processing Infrastructure. -This module provides the core SQL processing infrastructure for SQLSpec, implementing -a complete pipeline for SQL statement compilation, parameter processing, caching, -and result management. All components are optimized for MyPyC compilation to -reduce overhead. +This module provides the core SQL compilation, parameter processing, +caching, and result management layer for SQLSpec. -Architecture Overview: - The core module implements a single-pass processing pipeline where SQL statements - are parsed once, transformed once, and validated once. The SQL object serves as - the single source of truth throughout the system. - -Key Components: - statement.py: SQL statement representation and configuration management - - SQL class for statement encapsulation with lazy compilation - - StatementConfig for processing pipeline configuration - - ProcessedState for cached compilation results - - Support for execute_many and script execution modes - - parameters.py: Type-safe parameter processing and style conversion - - Automatic parameter style detection and conversion - - Support for QMARK (?), NAMED (:name), NUMERIC ($1), FORMAT (%s) styles - - Parameter validation and type coercion - - Batch parameter handling for execute_many operations - - compiler.py: SQL compilation with validation and optimization - - SQLProcessor for statement compilation and validation - - Operation type detection (SELECT, INSERT, UPDATE, DELETE, etc.) - - AST-based SQL analysis using SQLGlot - - Support for multiple SQL dialects - - Compiled result caching for performance - - result.py: Comprehensive result handling for all SQL operations - - SQLResult for standard query results with metadata - - ArrowResult for Apache Arrow format integration - - Support for DML operations with RETURNING clauses - - Script execution result aggregation - - Iterator protocol support for result rows - - filters.py: Composable SQL statement filters - - BeforeAfterFilter for date range filtering - - InCollectionFilter for IN clause generation - - LimitOffsetFilter for pagination - - OrderByFilter for dynamic sorting - - SearchFilter for text search operations - - Parameter conflict resolution - - cache.py: Caching system with LRU eviction - - LRUCache with configurable TTL and size limits - - NamespacedCache for statement, expression, optimized, builder, and file caching - - Thread-safe operations with fine-grained locking - - Cache statistics and monitoring - - splitter.py: Dialect-aware SQL script splitting - - Support for Oracle PL/SQL, T-SQL, PostgreSQL, MySQL - - Proper handling of block structures (BEGIN/END) - - Dollar-quoted string support for PostgreSQL - - Batch separator recognition (GO for T-SQL) - - Comment and string literal preservation - - hashing.py: Efficient cache key generation - - SQL statement hashing with parameter consideration - - Expression tree hashing for AST caching - - Parameter set hashing for batch operations - - Optimized hash computation with caching - -Performance Optimizations: - - MyPyC compilation support with proper annotations - - __slots__ usage for memory efficiency - - Final annotations for constant folding - - Lazy evaluation and compilation - - Comprehensive result caching - - Minimal object allocation in hot paths - -Thread Safety: - All caching components are thread-safe with RLock protection. - The processing pipeline is stateless and safe for concurrent use. +Responsibilities: + - statement: SQL statement encapsulation, configuration, and execution state. + - parameters: Parameter extraction, validation, coercion, and style conversion. + - compiler: Dialect translation, AST analysis, and SQLGlot compilation. + - result: Query execution result encapsulation and Arrow integration. + - filters: Composable query filters and pagination. + - cache: Statement, expression, and compilation caching. + - splitter: Dialect-aware SQL script splitting. + - hashing: Cache key hashing for statements and expressions. Example Usage: >>> from sqlspec.core import SQL, StatementConfig >>> config = StatementConfig(dialect="postgresql") >>> stmt = SQL( - ... "SELECT * FROM users WHERE id = ?", - ... 1, + ... "SELECT * FROM users WHERE id = :id", + ... {"id": 1}, ... statement_config=config, ... ) >>> compiled_sql, params = stmt.compile() diff --git a/sqlspec/core/cache.py b/sqlspec/core/cache.py index fc88ac3af..099e264ec 100644 --- a/sqlspec/core/cache.py +++ b/sqlspec/core/cache.py @@ -1001,6 +1001,13 @@ class FiltersView: Provides zero-copy access to filters with methods for querying, iteration, and canonical representation generation. + + Note: + `FiltersView` and its `to_canonical` method operate on lightweight + `Filter` records (field_name, operation, value) for cache key generation. + This is distinct from `StatementFilter` canonicalization in + `sqlspec.core.filters.canonicalize_filters`, which handles full SQL + AST statement filter objects. """ __slots__ = ("_filters_ref",) @@ -1055,7 +1062,6 @@ def to_canonical(self) -> "tuple[Any, ...]": Returns: Canonical tuple representation of filters """ - # Convert to Filter objects if needed, then canonicalize filter_objects = [] for f in self._filters_ref: if isinstance(f, Filter): diff --git a/sqlspec/core/compiler.py b/sqlspec/core/compiler.py index ea4d55407..3b422ed12 100644 --- a/sqlspec/core/compiler.py +++ b/sqlspec/core/compiler.py @@ -455,7 +455,7 @@ def _materialize_cached_result(self, cached_result: CompiledSQL, parameters: Any # Structural fingerprinting means same SQL structure = same cache entry, # but we must still process the caller's actual parameter values. - processed_params = self._parameter_processor._transform_cached_parameters( # pyright: ignore[reportPrivateUsage] + processed_params = self._parameter_processor.transform_cached_parameters( parameters, cached_result.parameter_profile, self._parameter_config, diff --git a/sqlspec/core/parameters/_converter.py b/sqlspec/core/parameters/_converter.py index c1d395b88..589eef367 100644 --- a/sqlspec/core/parameters/_converter.py +++ b/sqlspec/core/parameters/_converter.py @@ -6,6 +6,7 @@ from mypy_extensions import mypyc_attr from sqlspec.core.parameters._types import ( + _EXPANDING_POSITIONAL_STYLES, _NAMED_STYLES, _POSITIONAL_STYLES, ConvertedParameters, @@ -27,11 +28,6 @@ ParameterStyle.QMARK, ParameterStyle.POSITIONAL_PYFORMAT, }) -_EXPANDING_POSITIONAL_STYLES: Final[frozenset[ParameterStyle]] = frozenset({ - ParameterStyle.QMARK, - ParameterStyle.POSITIONAL_PYFORMAT, - ParameterStyle.POSITIONAL_COLON, -}) def _placeholder_qmark(_: Any) -> str: diff --git a/sqlspec/core/parameters/_processor.py b/sqlspec/core/parameters/_processor.py index e8a6568df..4e3e953c4 100644 --- a/sqlspec/core/parameters/_processor.py +++ b/sqlspec/core/parameters/_processor.py @@ -8,6 +8,7 @@ from sqlspec.core.parameters._converter import ParameterConverter from sqlspec.core.parameters._types import ( + _EXPANDING_POSITIONAL_STYLES, _NAMED_STYLE_VALUES, _NAMED_STYLES, _POSITIONAL_STYLE_VALUES, @@ -22,6 +23,7 @@ wrap_with_type, ) from sqlspec.core.parameters._validator import ParameterValidator +from sqlspec.exceptions import SQLSpecError from sqlspec.utils.dispatch import TypeDispatcher __all__ = ( @@ -37,11 +39,6 @@ _EXECUTE_MANY_SAMPLE_THRESHOLD: Final[int] = 10 _EXECUTE_MANY_SAMPLE_SIZE: Final[int] = 3 -_OCCURRENCE_BASED_POSITIONAL_STYLES: Final[frozenset[ParameterStyle]] = frozenset({ - ParameterStyle.QMARK, - ParameterStyle.POSITIONAL_COLON, - ParameterStyle.POSITIONAL_PYFORMAT, -}) _TYPE_COERCION_DISPATCHERS: Final[dict[tuple[TypeCoercionFallback, ...], TypeDispatcher[Callable[[Any], Any]]]] = {} @@ -308,7 +305,7 @@ def _process_internal( ) # Return cached SQL transformation with NEW parameters transformed # to match the cached SQL's placeholder format - transformed_params = self._transform_cached_parameters( + transformed_params = self.transform_cached_parameters( parameters, cached_result.parameter_profile, config, @@ -524,18 +521,7 @@ def _coerce_parameter_types( is_many: bool = False, ) -> "ConvertedParameters": fallback_items = type_coercion_fallbacks(type_coercion_map) - result = _coerce_parameters_payload(parameters, type_coercion_map, fallback_items, is_many) - # Fast type narrowing - _coerce_parameters_payload returns object but produces concrete types - if result is None: - return None - result_type = type(result) - if result_type is dict: - return result - if result_type is list: - return result - if result_type is tuple: - return result - return result + return _coerce_parameters_payload(parameters, type_coercion_map, fallback_items, is_many) def _store_cached_result( self, cache_key: Any | None, result: "ParameterProcessingResult" @@ -548,7 +534,7 @@ def _store_cached_result( self._cache.popitem(last=False) return result - def _transform_cached_parameters( + def transform_cached_parameters( self, parameters: "ParameterPayload", cached_profile: "ParameterProfile", @@ -682,12 +668,7 @@ def _map_named_to_positional( for idx, row in enumerate(parameter_rows): if type(row) is dict or isinstance(row, Mapping): if strict: - missing = [name for name in named_order if name not in row] - if missing: - from sqlspec.exceptions import SQLSpecError - - msg = f"Missing required parameters: {missing}" - raise SQLSpecError(msg) + _validate_missing_parameters(named_order, row) mapped_row: Any = tuple(row.get(name) for name in named_order) else: mapped_row = row @@ -706,12 +687,7 @@ def _map_named_to_positional( if isinstance(parameters, Mapping): if strict: - missing = [name for name in named_order if name not in parameters] - if missing: - from sqlspec.exceptions import SQLSpecError - - msg = f"Missing required parameters: {missing}" - raise SQLSpecError(msg) + _validate_missing_parameters(named_order, parameters) return tuple(parameters.get(name) for name in named_order) return parameters @@ -1141,6 +1117,22 @@ def _named_parameters_for_style( param_info: "list[ParameterInfo]", target_style: "ParameterStyle | None" ) -> "tuple[str, ...]": names = tuple(p.name for p in param_info if p.name is not None) - if target_style in _OCCURRENCE_BASED_POSITIONAL_STYLES: + if target_style in _EXPANDING_POSITIONAL_STYLES: return names return tuple(dict.fromkeys(names)) + + +def _validate_missing_parameters(named_order: Sequence[str], parameters: Mapping[str, Any]) -> None: + """Validate that all required parameters are present in parameters mapping. + + Args: + named_order: Expected parameter names in order. + parameters: Parameter mapping to validate. + + Raises: + SQLSpecError: If any required parameters are missing. + """ + missing = [name for name in named_order if name not in parameters] + if missing: + msg = f"Missing required parameters: {missing}" + raise SQLSpecError(msg) diff --git a/sqlspec/core/parameters/_types.py b/sqlspec/core/parameters/_types.py index 888930ddc..e6fbadb8b 100644 --- a/sqlspec/core/parameters/_types.py +++ b/sqlspec/core/parameters/_types.py @@ -10,6 +10,7 @@ from mypy_extensions import mypyc_attr __all__ = ( + "_EXPANDING_POSITIONAL_STYLES", "ConvertedParameters", "DriverParameterProfile", "NamedParameterOutput", @@ -162,6 +163,11 @@ class ParameterStyle(str, Enum): }) _NAMED_STYLE_VALUES: Final[frozenset[str]] = frozenset(style.value for style in _NAMED_STYLES) _POSITIONAL_STYLE_VALUES: Final[frozenset[str]] = frozenset(style.value for style in _POSITIONAL_STYLES) +_EXPANDING_POSITIONAL_STYLES: Final[frozenset[ParameterStyle]] = frozenset({ + ParameterStyle.QMARK, + ParameterStyle.POSITIONAL_PYFORMAT, + ParameterStyle.POSITIONAL_COLON, +}) @mypyc_attr(allow_interpreted_subclasses=False) diff --git a/sqlspec/core/statement.py b/sqlspec/core/statement.py index 2fe3631fb..dcf84f334 100644 --- a/sqlspec/core/statement.py +++ b/sqlspec/core/statement.py @@ -818,7 +818,7 @@ def _rebind_cached_parameters(self, state: "ProcessedState") -> "tuple[str, Any] validator_cache_max_size=0, ) processor = self._rebind_processor - rebound_params = processor._transform_cached_parameters( # pyright: ignore[reportPrivateUsage] + rebound_params = processor.transform_cached_parameters( params, state.parameter_profile, self._statement_config.parameter_config, diff --git a/sqlspec/data_dictionary/_loader.py b/sqlspec/data_dictionary/_loader.py index 8cb07a9ba..212c73f9f 100644 --- a/sqlspec/data_dictionary/_loader.py +++ b/sqlspec/data_dictionary/_loader.py @@ -30,10 +30,6 @@ __all__ = ("DataDictionaryLoader", "get_data_dictionary_loader") -SQL_RESOURCE_PACKAGE = "sqlspec.data_dictionary" -SQL_RESOURCE_NAME = "sql" - - @mypyc_attr(allow_interpreted_subclasses=False) class DataDictionaryLoader: """Loads and manages data dictionary SQL for all dialects.""" @@ -330,10 +326,6 @@ def get_data_dictionary_loader() -> DataDictionaryLoader: return _loader_instance -def _sql_resource_root() -> "Traversable": - return resources.files(SQL_RESOURCE_PACKAGE).joinpath(SQL_RESOURCE_NAME) - - def _normalize_domain_key(name: str) -> str: return slugify(name, separator="_") diff --git a/sqlspec/dialects/__init__.py b/sqlspec/dialects/__init__.py index e03dcbe7a..7f6dd47e4 100644 --- a/sqlspec/dialects/__init__.py +++ b/sqlspec/dialects/__init__.py @@ -10,6 +10,7 @@ ``import sqlspec`` does not pay the sqlglot dialect-machinery cost. """ +import importlib from typing import TYPE_CHECKING, Any if TYPE_CHECKING: @@ -32,6 +33,5 @@ def __getattr__(name: str) -> Any: if module_name is None: msg = f"module {__name__!r} has no attribute {name!r}" raise AttributeError(msg) - import importlib return getattr(importlib.import_module(module_name), name) diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index e3f326b78..7710aad51 100644 --- a/sqlspec/driver/_common.py +++ b/sqlspec/driver/_common.py @@ -86,7 +86,7 @@ if TYPE_CHECKING: from collections import abc - from collections.abc import Awaitable, Callable + from collections.abc import Awaitable from types import TracebackType from sqlspec.core import ArrowResult, FilterTypeT, StatementFilter @@ -558,10 +558,6 @@ def sort_tables_topologically(self, tables: "list[str]", foreign_keys: "list[For result.raise_for_cycles() return [identity.name for identity in result.ordered] - def _resolve_log_adapter(self) -> str: - """Resolve adapter identifier for logging.""" - return str(type(self).dialect) - def _log_version_detected(self, adapter: str, version: VersionInfo) -> None: """Log detected database version with db.system context.""" logger.debug( @@ -712,7 +708,7 @@ def stmt_cache_rebind( and not needs_style_remap ): return params - return self._stmt_cache_rebind_processor._transform_cached_parameters( + return self._stmt_cache_rebind_processor.transform_cached_parameters( params, cached.parameter_profile, config, @@ -987,15 +983,15 @@ def prepare_driver_parameters( for param_set in parameters: if isinstance(param_set, dict): for value in param_set.values(): - if self._needs_coercion_candidate(value, type_coercion_map, fallback_items): + if parameter_value_needs_processing(value, type_coercion_map, fallback_items): needs_transform = True break elif isinstance(param_set, (list, tuple)): for value in param_set: - if self._needs_coercion_candidate(value, type_coercion_map, fallback_items): + if parameter_value_needs_processing(value, type_coercion_map, fallback_items): needs_transform = True break - elif self._needs_coercion_candidate(param_set, type_coercion_map, fallback_items): + elif parameter_value_needs_processing(param_set, type_coercion_map, fallback_items): needs_transform = True if needs_transform: break @@ -1402,14 +1398,6 @@ def _apply_filters(self, sql_statement: "SQL", filters: "list[StatementFilter]") sql_statement = filter_obj.append_to_statement(sql_statement) return sql_statement - def _needs_coercion_candidate( - self, - value: object, - type_coercion_map: "dict[type, Callable[[Any], Any]] | None", - fallback_items: "tuple[tuple[type, Any], ...]", - ) -> bool: - return parameter_value_needs_processing(value, type_coercion_map, fallback_items) - def _batch_parameters( self, parameters: "StatementParameters", statement_config: "StatementConfig" ) -> "ConvertedParameters": @@ -1631,7 +1619,7 @@ def _compiled_statement( # but we must still use the caller's actual parameter values. parameter_profile = cached_result.parameter_profile if parameter_profile is not None: - prepared_parameters = self._stmt_cache_rebind_processor._transform_cached_parameters( + prepared_parameters = self._stmt_cache_rebind_processor.transform_cached_parameters( statement.parameters, parameter_profile, statement_config.parameter_config, diff --git a/sqlspec/extensions/adk/artifact/service.py b/sqlspec/extensions/adk/artifact/service.py index 74a0ba08b..3678cf65a 100644 --- a/sqlspec/extensions/adk/artifact/service.py +++ b/sqlspec/extensions/adk/artifact/service.py @@ -13,6 +13,7 @@ import json import logging import re +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any from google.adk.artifacts.base_artifact_service import BaseArtifactService @@ -22,8 +23,6 @@ from sqlspec.utils.logging import get_logger, log_with_context if TYPE_CHECKING: - from datetime import datetime - from google.adk.artifacts.base_artifact_service import ArtifactVersion from google.genai import types @@ -34,148 +33,6 @@ logger = get_logger("sqlspec.extensions.adk.artifact.service") -# Matches path traversal and absolute path components -_UNSAFE_PATH_CHARS = re.compile(r"(?:^|/)\.\.(?:/|$)|[\x00]") - - -def _sanitize_path_component(value: str) -> str: - """Sanitize a path component to prevent directory traversal. - - Removes leading/trailing slashes, rejects ``..`` traversals, and - replaces NUL bytes. - - Args: - value: Raw path component. - - Returns: - Sanitized path component. - - Raises: - ValueError: If the value contains path traversal sequences. - """ - value = value.strip("/") - if _UNSAFE_PATH_CHARS.search(value): - msg = f"Unsafe path component: {value!r}" - raise ValueError(msg) - return value - - -def _build_content_path( - app_name: str, user_id: str, filename: str, version: int, session_id: "str | None" = None -) -> str: - """Build the storage path for artifact content. - - Pattern: - ``apps/{app_name}/users/{user_id}/[sessions/{session_id}/]artifacts/{filename}/v{version}`` - - All path components are sanitized to prevent directory traversal. - - Args: - app_name: Application name. - user_id: User identifier. - filename: Artifact filename. - version: Version number. - session_id: Optional session identifier. - - Returns: - Sanitized storage path. - """ - parts = ["apps", _sanitize_path_component(app_name), "users", _sanitize_path_component(user_id)] - if session_id is not None: - parts.extend(["sessions", _sanitize_path_component(session_id)]) - parts.extend(["artifacts", _sanitize_path_component(filename), f"v{version}"]) - return "/".join(parts) - - -def _extract_mime_type(artifact: "types.Part | dict[str, Any]") -> "str | None": - """Extract MIME type from an artifact Part. - - Checks ``inline_data.mime_type`` and ``file_data.mime_type`` on the Part. - - Args: - artifact: ADK Part or dict representation. - - Returns: - MIME type string, or None if not determinable. - """ - if isinstance(artifact, dict): - # Handle camelCase and snake_case keys - inline = artifact.get("inline_data") or artifact.get("inlineData") - if isinstance(inline, dict): - return inline.get("mime_type") or inline.get("mimeType") - file_data = artifact.get("file_data") or artifact.get("fileData") - if isinstance(file_data, dict): - return file_data.get("mime_type") or file_data.get("mimeType") - return None - - # types.Part object - if hasattr(artifact, "inline_data") and artifact.inline_data is not None: - return getattr(artifact.inline_data, "mime_type", None) - if hasattr(artifact, "file_data") and artifact.file_data is not None: - return getattr(artifact.file_data, "mime_type", None) - return None - - -def _serialize_artifact(artifact: "types.Part | dict[str, Any]") -> bytes: - """Serialize an artifact Part to bytes for content storage. - - The artifact is serialized as JSON via ``model_dump(exclude_none=True)``. - This preserves the full Part structure including text, inline_data, - file_data, and any future Part fields. - - Args: - artifact: ADK Part or dict representation. - - Returns: - JSON-encoded bytes. - """ - if isinstance(artifact, dict): - return json.dumps(artifact, default=str).encode("utf-8") - - # Use Pydantic model serialization - if hasattr(artifact, "model_dump"): - data = artifact.model_dump(exclude_none=True) - return json.dumps(data, default=str).encode("utf-8") - - # Fallback for unexpected types - return json.dumps({"text": str(artifact)}).encode("utf-8") - - -def _deserialize_artifact(data: bytes) -> "types.Part": - """Deserialize bytes back into an ADK Part. - - Args: - data: JSON-encoded bytes from content storage. - - Returns: - Reconstructed Part object. - """ - from google.genai import types - - parsed = json.loads(data.decode("utf-8")) - return types.Part.model_validate(parsed) - - -def _record_to_artifact_version(record: "StoredArtifact") -> "ArtifactVersion": - """Convert a database artifact record to an ADK ArtifactVersion. - - Args: - record: Database artifact record. - - Returns: - ArtifactVersion model instance. - """ - from google.adk.artifacts.base_artifact_service import ArtifactVersion - - return ArtifactVersion( - version=record["version"], - canonical_uri=record["canonical_uri"], - custom_metadata=record["custom_metadata"] or {}, - create_time=record["created_at"].timestamp(), - mime_type=record["mime_type"], - ) - - class SQLSpecArtifactService(BaseArtifactService): """SQLSpec-backed implementation of BaseArtifactService. @@ -241,33 +98,18 @@ async def save_artifact( """ from google.adk.artifacts.base_artifact_service import ensure_part - # Normalize artifact to Part artifact_part: types.Part = ensure_part(artifact) - - # Determine the next version version = await self._store.get_next_version( app_name=app_name, user_id=user_id, filename=filename, session_id=session_id ) - - # Build the content path and canonical URI content_path = _build_content_path( app_name=app_name, user_id=user_id, filename=filename, version=version, session_id=session_id ) canonical_uri = f"{self._artifact_storage_uri}/{content_path}" - - # Serialize content content_bytes = _serialize_artifact(artifact_part) - - # Extract MIME type mime_type = _extract_mime_type(artifact_part) - - # Write content first (fail-fast before metadata) backend = self._registry.get(self._artifact_storage_uri) await _call_storage_backend(backend, "write_bytes_async", "write_bytes_sync", content_path, content_bytes) - - # Insert metadata row - from datetime import datetime, timezone - record = StoredArtifact( app_name=app_name, user_id=user_id, @@ -527,6 +369,143 @@ async def _delete_content(self, records: "list[StoredArtifact]") -> None: ) +_UNSAFE_PATH_CHARS = re.compile(r"(?:^|/)\.\.(?:/|$)|[\x00]") + + +def _sanitize_path_component(value: str) -> str: + """Sanitize a path component to prevent directory traversal. + + Removes leading/trailing slashes, rejects ``..`` traversals, and + replaces NUL bytes. + + Args: + value: Raw path component. + + Returns: + Sanitized path component. + + Raises: + ValueError: If the value contains path traversal sequences. + """ + value = value.strip("/") + if _UNSAFE_PATH_CHARS.search(value): + msg = f"Unsafe path component: {value!r}" + raise ValueError(msg) + return value + + +def _build_content_path( + app_name: str, user_id: str, filename: str, version: int, session_id: "str | None" = None +) -> str: + """Build the storage path for artifact content. + + Pattern: + ``apps/{app_name}/users/{user_id}/[sessions/{session_id}/]artifacts/{filename}/v{version}`` + + All path components are sanitized to prevent directory traversal. + + Args: + app_name: Application name. + user_id: User identifier. + filename: Artifact filename. + version: Version number. + session_id: Optional session identifier. + + Returns: + Sanitized storage path. + """ + parts = ["apps", _sanitize_path_component(app_name), "users", _sanitize_path_component(user_id)] + if session_id is not None: + parts.extend(["sessions", _sanitize_path_component(session_id)]) + parts.extend(["artifacts", _sanitize_path_component(filename), f"v{version}"]) + return "/".join(parts) + + +def _extract_mime_type(artifact: "types.Part | dict[str, Any]") -> "str | None": + """Extract MIME type from an artifact Part. + + Checks ``inline_data.mime_type`` and ``file_data.mime_type`` on the Part. + + Args: + artifact: ADK Part or dict representation. + + Returns: + MIME type string, or None if not determinable. + """ + if isinstance(artifact, dict): + inline = artifact.get("inline_data") or artifact.get("inlineData") + if isinstance(inline, dict): + return inline.get("mime_type") or inline.get("mimeType") + file_data = artifact.get("file_data") or artifact.get("fileData") + if isinstance(file_data, dict): + return file_data.get("mime_type") or file_data.get("mimeType") + return None + + if hasattr(artifact, "inline_data") and artifact.inline_data is not None: + return getattr(artifact.inline_data, "mime_type", None) + if hasattr(artifact, "file_data") and artifact.file_data is not None: + return getattr(artifact.file_data, "mime_type", None) + return None + + +def _serialize_artifact(artifact: "types.Part | dict[str, Any]") -> bytes: + """Serialize an artifact Part to bytes for content storage. + + The artifact is serialized as JSON via ``model_dump(exclude_none=True)``. + This preserves the full Part structure including text, inline_data, + file_data, and any future Part fields. + + Args: + artifact: ADK Part or dict representation. + + Returns: + JSON-encoded bytes. + """ + if isinstance(artifact, dict): + return json.dumps(artifact, default=str).encode("utf-8") + + if hasattr(artifact, "model_dump"): + data = artifact.model_dump(exclude_none=True) + return json.dumps(data, default=str).encode("utf-8") + + return json.dumps({"text": str(artifact)}).encode("utf-8") + + +def _deserialize_artifact(data: bytes) -> "types.Part": + """Deserialize bytes back into an ADK Part. + + Args: + data: JSON-encoded bytes from content storage. + + Returns: + Reconstructed Part object. + """ + from google.genai import types + + parsed = json.loads(data.decode("utf-8")) + return types.Part.model_validate(parsed) + + +def _record_to_artifact_version(record: "StoredArtifact") -> "ArtifactVersion": + """Convert a database artifact record to an ADK ArtifactVersion. + + Args: + record: Database artifact record. + + Returns: + ArtifactVersion model instance. + """ + from google.adk.artifacts.base_artifact_service import ArtifactVersion + + return ArtifactVersion( + version=record["version"], + canonical_uri=record["canonical_uri"], + custom_metadata=record["custom_metadata"] or {}, + create_time=record["created_at"].timestamp(), + mime_type=record["mime_type"], + ) + + async def _call_storage_backend( backend: Any, async_method_name: str, sync_method_name: str, *args: Any, **kwargs: Any ) -> Any: diff --git a/sqlspec/extensions/adk/memory/store.py b/sqlspec/extensions/adk/memory/store.py index 6e15376e9..4e72c6b8d 100644 --- a/sqlspec/extensions/adk/memory/store.py +++ b/sqlspec/extensions/adk/memory/store.py @@ -163,24 +163,6 @@ def _reset_drop_memory_table_sql(self) -> list[str]: statements.extend(self._drop_sql_for_table(cand)) return unique_statements(statements) - def _require_enabled(self) -> None: - if not self._enabled: - msg = "ADK memory store is disabled for this database configuration" - raise RuntimeError(msg) - - def _effective_limit(self, limit: int | None) -> int: - return limit if limit is not None else self._max_results - - def _log_operation(self, event: str, **kwargs: Any) -> None: - log_with_context( - logger, - logging.DEBUG, - event, - table_name=self._memory_table, - db_system=resolve_db_system(type(self).__name__), - **kwargs, - ) - class BaseAsyncADKMemoryStore(_ADKMemoryStoreCommon[ConfigT], ABC): """Base class for async SQLSpec-backed ADK memory stores. diff --git a/sqlspec/loader.py b/sqlspec/loader.py index 6eba12b99..8b34e37b7 100644 --- a/sqlspec/loader.py +++ b/sqlspec/loader.py @@ -84,18 +84,6 @@ } -def _parse_parameter_declaration(param_match: "re.Match[str]") -> ParameterDeclaration: - """Build a parameter declaration from a matched ``-- param:`` line.""" - description = param_match.group("desc") - required = param_match.group("optional") != "?" - if description is not None and PARAM_OPTIONAL_DESCRIPTION_PATTERN.search(description): - required = False - description = PARAM_OPTIONAL_DESCRIPTION_PATTERN.sub("", description).strip() or None - return ParameterDeclaration( - name=param_match.group("name"), type_str=param_match.group("type"), description=description, required=required - ) - - class SlotDeclaration: """A fill point declared by a ``/* slot: name */`` marker. @@ -221,6 +209,18 @@ def __init__( self.statement_names = tuple(parsed_statements.keys()) +def _parse_parameter_declaration(param_match: "re.Match[str]") -> ParameterDeclaration: + """Build a parameter declaration from a matched ``-- param:`` line.""" + description = param_match.group("desc") + required = param_match.group("optional") != "?" + if description is not None and PARAM_OPTIONAL_DESCRIPTION_PATTERN.search(description): + required = False + description = PARAM_OPTIONAL_DESCRIPTION_PATTERN.sub("", description).strip() or None + return ParameterDeclaration( + name=param_match.group("name"), type_str=param_match.group("type"), description=description, required=required + ) + + class SQLFileLoader: """Loads and parses SQL files with named SQL queries. @@ -277,999 +277,996 @@ def set_observability_runtime(self, runtime: "ObservabilityRuntime | None") -> N self._runtime = runtime - def _raise_file_not_found(self, path: str) -> None: - """Raise SQLFileNotFoundError for nonexistent file. + def load_sql(self, *paths: str | Path) -> None: + """Load SQL files and parse named queries. Args: - path: File path that was not found. - - Raises: - SQLFileNotFoundError: Always raised. + *paths: One or more file paths or directory paths to load. """ - raise SQLFileNotFoundError(path) + runtime = self._runtime + span = None + error: Exception | None = None + start_time = time.perf_counter() + path_count = len(paths) + previous_correlation_id = CorrelationContext.get() + if runtime is not None: + runtime.increment_metric("loader.load.invocations") + runtime.increment_metric("loader.paths.requested", path_count) + span = runtime.start_span( + "sqlspec.loader.load", + attributes={"sqlspec.loader.path_count": path_count, "sqlspec.loader.encoding": self.encoding}, + ) - def _raise_statement_not_found(self, name: str, normalized_name: str) -> None: - """Raise SQLStatementNotFoundError for nonexistent statements. + try: + for path in paths: + path_str = str(path) + if "://" in path_str or self.storage_registry.is_alias_registered(path_str.split("/", maxsplit=1)[0]): + self._load_single_file(path, None) + continue + + path_obj = Path(path) + if path_obj.is_dir(): + self._load_directory(path_obj) + elif path_obj.exists(): + self._load_single_file(path_obj, None) + elif path_obj.suffix: + self._raise_file_not_found(str(path)) + + except Exception as exc: + error = exc + if runtime is not None: + runtime.increment_metric("loader.load.errors") + raise + finally: + duration_ms = (time.perf_counter() - start_time) * 1000 + if runtime is not None: + runtime.record_metric("loader.last_load_ms", duration_ms) + runtime.increment_metric("loader.load.duration_ms", duration_ms) + runtime.end_span(span, error=error) + CorrelationContext.set(previous_correlation_id) + + def add_named_sql( + self, + name: str, + sql: str, + dialect: "str | None" = None, + parameters: "Sequence[ParameterDeclaration] | None" = None, + ) -> None: + """Add a named SQL query directly without loading from a file. + + The SQL may contain ``/* include: name */`` and ``/* slot: name */`` markers; + slots added this way have no defaults. Args: - name: Name requested by the caller. - normalized_name: Normalized statement name used for lookup. + name: Name for the SQL query. + sql: Raw SQL content. + dialect: Optional dialect for the SQL statement. + parameters: Optional declared parameter metadata for the query. Raises: - SQLStatementNotFoundError: Always raised. + ValueError: If query name already exists. """ - raise SQLStatementNotFoundError(name=name, normalized_name=normalized_name, query_count=len(self._queries)) - def _file_cache_key(self, path: str | Path) -> str: - """Generate cache key for a file path. + normalized_name = _normalize_query_name(name) - Args: - path: File path to generate key for. + if normalized_name in self._queries: + existing_source = self._query_to_file.get(normalized_name, "") + msg = f"Query name '{name}' already exists (source: {existing_source})" + raise ValueError(msg) - Returns: - Cache key string for the file. - """ - path_str = str(path) - path_hash = hashlib.md5(path_str.encode(), usedforsecurity=False).hexdigest() - return f"file:{path_hash[:16]}" + if dialect is not None: + dialect = _normalize_dialect(dialect) - @staticmethod - def _compute_checksum(content: str) -> str: - """Compute MD5 checksum from already-read file content.""" - return hashlib.md5(content.encode(), usedforsecurity=False).hexdigest() + declared = tuple(parameters) if parameters else () + clean_sql = sql.strip() + slots = _merge_slot_markers((), clean_sql) + has_includes = bool(_include_markers(clean_sql)) + if not slots and not has_includes: + self._check_declared_parameters(clean_sql, declared, name, "") - def _calculate_file_checksum(self, path: str | Path) -> str: - """Calculate checksum for file content validation. + statement = NamedStatement( + name=normalized_name, + sql=clean_sql, + dialect=dialect, + start_line=0, + parameters=declared, + slots=slots, + has_includes=has_includes, + ) + self._queries[normalized_name] = statement + self._query_to_file[normalized_name] = "" - Args: - path: File path to calculate checksum for. + def add_fragment(self, name: str, sql: str) -> None: + """Add a reusable SQL fragment directly without loading from a file. - Returns: - MD5 checksum of file content. + Args: + name: Name for the fragment, referenced by ``/* include: name */`` markers. + sql: Fragment SQL text; may contain include and slot markers. Raises: - SQLFileParseError: If file cannot be read. + ValueError: If the fragment name already exists. """ - try: - return self._compute_checksum(self._read_file_content(path)) - except Exception as e: - raise SQLFileParseError(str(path), str(path), e) from e + normalized_name = _normalize_query_name(name) + if normalized_name in self._fragments: + existing_source = self._fragment_to_file.get(normalized_name, "") + msg = f"Fragment name '{name}' already exists (source: {existing_source})" + raise ValueError(msg) + self._fragments[normalized_name] = SQLFragment(name=normalized_name, sql=sql.strip()) + self._fragment_to_file[normalized_name] = "" + self._invalidate_resolved() - def _is_file_unchanged(self, path: str | Path, cached_file: SQLFile) -> bool: - """Check if file has changed since caching. + def has_fragment(self, name: str) -> bool: + """Check if a fragment exists. Args: - path: File path to check. - cached_file: Cached file data. + name: Fragment name to check. Returns: - True if file is unchanged, False otherwise. + True if the fragment exists. """ - try: - current_checksum = self._calculate_file_checksum(path) - except Exception: - return False - else: - return current_checksum == cached_file.checksum - - def _reload_changed_files(self) -> "list[str]": - """Reload tracked SQL files whose content checksum changed. + return _normalize_query_name(name) in self._fragments - Every changed file's queries and fragments are removed before any changed - file is loaded again, so names may move between changed files. + def list_fragments(self) -> "list[str]": + """List all available fragment names. Returns: - Paths of files that were reloaded. + Sorted list of fragment names. """ - pending_reloads: list[tuple[str, str | None]] = [] - for path, sql_file in list(self._files.items()): - if self._is_file_unchanged(path, sql_file): - if self._runtime is not None: - self._runtime.increment_metric("loader.reload.skipped") - continue - - query_names = [name for name, source in self._query_to_file.items() if source == path] - fragment_names = [name for name, source in self._fragment_to_file.items() if source == path] - namespaces = { - name[: -(len(statement.name) + 1)] - for name in query_names - if (statement := self._queries.get(name)) is not None and name.endswith(f".{statement.name}") - } - namespaces.update( - name[: -(len(fragment.name) + 1)] - for name in fragment_names - if (fragment := self._fragments.get(name)) is not None and name.endswith(f".{fragment.name}") - ) - namespace = next(iter(namespaces)) if len(namespaces) == 1 else None - for name in query_names: - self._queries.pop(name, None) - self._query_to_file.pop(name, None) - self._compiled_statements.pop(name, None) - for name in fragment_names: - self._fragments.pop(name, None) - self._fragment_to_file.pop(name, None) - self._files.pop(path, None) - pending_reloads.append((path, namespace)) - - if pending_reloads: - self._invalidate_resolved() - changed_paths: list[str] = [] - for path, namespace in pending_reloads: - self._load_single_file(path, namespace) - changed_paths.append(path) - if self._runtime is not None: - self._runtime.increment_metric("loader.reload.changed") - return changed_paths + return sorted(self._fragments.keys()) - def _content_matches_cache(self, content: str, cached_file: SQLFileCacheEntry) -> bool: - """Check if already-read file content matches cached checksum.""" - return self._compute_checksum(content) == cached_file.sql_file.checksum + def get_fragment_text(self, name: str) -> str: + """Get a fragment's SQL text with its includes resolved. - def _read_file_content(self, path: str | Path) -> str: - """Read file content using storage backend. + Slot markers are left in place. Args: - path: File path (can be local path or URI). + name: Fragment name. Returns: - File content as string. + Fragment SQL text with ``/* include: */`` markers replaced. Raises: - SQLFileNotFoundError: If file does not exist. - SQLFileParseError: If file cannot be read or parsed. + SQLFragmentNotFoundError: If the fragment or an included fragment does not exist. + SQLFileParseError: If the includes form a cycle. """ - path_str = str(path) + safe_name = self._find_fragment_name(name, None) + return self._resolve_includes( + self._fragments[safe_name].sql, namespace=_namespace_of(safe_name), stack=(safe_name,) + ) - try: - backend = self.storage_registry.get(path) + def get_query_slots(self, name: str) -> "tuple[SlotDeclaration, ...]": + """Get the slots of a query, including slots contributed by included fragments. - # If path_str contains a '/', we check if the first part is a registered alias. - # This is specifically for when a path is provided relative to an alias. - parts = path_str.split("/", 1) - if len(parts) > 1 and self.storage_registry.is_alias_registered(parts[0]): - return backend.read_text_sync(parts[1], encoding=self.encoding) + Slots declared with ``-- slot:`` come first in declaration order, then + undeclared (required) markers in the query's own text, then markers + contributed by included fragments; markers keep their order of appearance. - if path_str.startswith("file://"): - parsed = urlparse(path_str) - file_path = unquote(parsed.path) - if file_path and len(file_path) > 2 and file_path[2] == ":": # noqa: PLR2004 - file_path = file_path[1:] - return backend.read_text_sync(Path(file_path).name, encoding=self.encoding) - - if isinstance(path, Path) or is_local_path(path_str): - return backend.read_text_sync(Path(path_str).name, encoding=self.encoding) + Args: + name: Query name (hyphens are converted to underscores). - return backend.read_text_sync(path_str, encoding=self.encoding) - except KeyError as e: - raise SQLFileNotFoundError(path_str) from e - except FileNotFoundInStorageError as e: - raise SQLFileNotFoundError(path_str) from e - except FileNotFoundError as e: - raise SQLFileNotFoundError(path_str) from e - except StorageOperationFailedError as e: - raise SQLFileParseError(path_str, path_str, e) from e - except Exception as e: - raise SQLFileParseError(path_str, path_str, e) from e + Returns: + Tuple of slot declarations; empty if the query has none. - @staticmethod - def _strip_leading_comments(sql_text: str) -> str: - """Remove leading comment lines from a SQL string.""" - lines = sql_text.strip().split("\n") - first_sql_line_index = -1 - for i, line in enumerate(lines): - if line.strip() and not line.strip().startswith("--"): - first_sql_line_index = i - break - if first_sql_line_index == -1: - return "" - return "\n".join(lines[first_sql_line_index:]).strip() + Raises: + SQLStatementNotFoundError: If the query or an included fragment does not exist. + SQLFileParseError: If a declared slot has no marker or the includes form a cycle. + """ + safe_name = _normalize_query_name(name) + if safe_name not in self._queries: + self._raise_statement_not_found(name, safe_name) + statement = self._queries[safe_name] + if not statement.slots and not statement.has_includes: + return () + return _merge_slot_markers(statement.slots, self._resolve_statement_text(safe_name)) - @staticmethod - def _parse_directive_block( - statement_section: str, file_path: str, strict: bool, base_line: int = 0 - ) -> "tuple[str | None, tuple[ParameterDeclaration, ...], str, tuple[SlotDeclaration, ...]]": - """Scan a section's leading comment block for ``dialect``/``param``/``slot`` directives. + def get_query_parameters(self, name: str) -> "tuple[ParameterDeclaration, ...]": + """Get declared parameter metadata for a query. Args: - statement_section: The statement body including any leading directive lines. - file_path: File path for error reporting. - strict: When True, a malformed ``-- param:`` line raises instead of warning. - base_line: 0-based line offset of ``statement_section`` within the file. + name: Query name (hyphens are converted to underscores). Returns: - The resolved dialect, the declared parameters, the SQL body with the - leading directive/comment lines removed, and the declared slots. + Tuple of declared parameters; empty if the query declares none. Raises: - SQLFileParseError: If ``strict`` and a ``-- param:`` line is malformed, or a - ``-- slot:`` line is malformed, duplicated, or placed after the SQL body begins. + SQLStatementNotFoundError: If the query does not exist. """ - dialect: str | None = None - params: list[ParameterDeclaration] = [] - slots: list[SlotDeclaration] = [] - raw_lines = statement_section.split("\n") - body_start = len(raw_lines) - for idx, raw in enumerate(raw_lines): - stripped = raw.strip() - if not stripped: - continue - if not stripped.startswith("--"): - body_start = idx - break - dialect_match = DIALECT_PATTERN.match(stripped) - if dialect_match: - dialect = _normalize_dialect(dialect_match.group("dialect").lower()) - continue - param_match = PARAM_PATTERN.match(stripped) - if param_match: - params.append(_parse_parameter_declaration(param_match)) - continue - slot_match = SLOT_DIRECTIVE_PATTERN.match(stripped) - if slot_match: - slot_name = slot_match.group("name") - if any(slot.name == slot_name for slot in slots): - raise SQLFileParseError( - file_path, - file_path, - ValueError(f"Duplicate -- slot: directive for slot '{slot_name}'"), - line=base_line + idx + 1, - ) - slots.append(SlotDeclaration(slot_name, slot_match.group("default"))) - continue - if SLOT_COMMENT_PATTERN.match(stripped): - raise SQLFileParseError( - file_path, - file_path, - ValueError(f"Malformed -- slot: directive: {stripped}"), - line=base_line + idx + 1, - ) - if PARAM_PREFIX_PATTERN.match(stripped): - line_number = base_line + idx + 1 - if strict: - raise SQLFileParseError( - file_path, file_path, ValueError(f"Malformed -- param: directive: {stripped}"), line=line_number - ) - log_with_context( - logger, - logging.WARNING, - f"sql.parse.param: malformed parameter directive in {file_path} at line {line_number}: {stripped}", - file_path=file_path, - line_number=line_number, - directive=stripped, - status="malformed", - ) - body_text = "\n".join(raw_lines[body_start:]) - if SLOT_COMMENT_PATTERN.search(body_text) is not None: - for start, end, is_block in _scan_sql_comments(body_text): - if is_block or SLOT_COMMENT_PATTERN.match(body_text, start) is None: - continue - line_start = body_text.rfind("\n", 0, start) + 1 - if body_text[line_start:start].strip(): - continue - raise SQLFileParseError( - file_path, - file_path, - ValueError( - f"-- slot: directive must appear in the leading directive block: {body_text[start:end].strip()}" - ), - line=base_line + body_start + body_text.count("\n", 0, start) + 1, - ) - return dialect, tuple(params), body_text, tuple(slots) - - @staticmethod - def _check_declared_parameters( - clean_sql: str, - declared: "tuple[ParameterDeclaration, ...]", - statement_name: str, - file_path: str, - start_line: "int | None" = None, - ) -> None: - """Validate declared parameters against the query's actual placeholders. + safe_name = _normalize_query_name(name) + if safe_name not in self._queries: + self._raise_statement_not_found(name, safe_name) + return self._queries[safe_name].parameters - For named binding, every declared name must appear among the SQL placeholders - (declared names may be a subset; filters and undeclared params are allowed). For - positional binding, the declared count must equal the placeholder count. + def get_file(self, path: str | Path) -> "SQLFile | None": + """Get a loaded SQLFile object by path. Args: - clean_sql: The SQL body with directives/comments stripped. - declared: Declared parameters for the query. - statement_name: Raw query name for error messages. - file_path: File path for error reporting. - start_line: Optional 0-based line of the statement within the file. + path: Path of the file. - Raises: - SQLFileParseError: On name drift (named) or count mismatch (positional). + Returns: + SQLFile object if loaded, None otherwise. """ - if not declared: - return - error_line = start_line + 1 if start_line is not None else None - infos = ParameterValidator().extract_parameters(clean_sql) - named = {info.name for info in infos if info.name and not info.name.isdigit()} - if named: - for decl in declared: - if decl.name not in named: - raise SQLFileParseError( - file_path, - file_path, - ValueError( - f"Declared parameter '{decl.name}' for query '{statement_name}' is not present in the " - f"SQL placeholders {sorted(named)}" - ), - line=error_line, - ) - elif len(declared) != len(infos): - raise SQLFileParseError( - file_path, - file_path, - ValueError( - f"Query '{statement_name}' declares {len(declared)} parameter(s) but the SQL has " - f"{len(infos)} positional placeholder(s)" - ), - line=error_line, - ) - - @staticmethod - def _parse_statements( - content: str, file_path: str, strict_parameter_annotations: bool = False - ) -> "tuple[dict[str, NamedStatement], dict[str, SQLFragment]]": - """Parse SQL content and extract named statements and fragments. + return self._files.get(str(path)) - A section starts at a ``-- name:`` or ``-- fragment:`` marker and ends at the - next marker of either kind. Files without any markers are gracefully skipped - by returning empty dictionaries. The caller is responsible for handling empty - results appropriately. + def get_file_for_query(self, name: str) -> "SQLFile | None": + """Get the SQLFile object containing a query. Args: - content: Raw SQL file content to parse. - file_path: File path for error reporting. - strict_parameter_annotations: Raise on malformed parameter declarations instead of skipping them. + name: Query name (hyphens are converted to underscores). Returns: - Dictionaries mapping normalized names to NamedStatement and SQLFragment - objects. Both are empty if no markers are found in the content. - - Raises: - SQLFileParseError: If sections are malformed (duplicate names, directives - on fragments, invalid slot directives, or no content after parsing). + SQLFile object if query exists, None otherwise. """ - statements: dict[str, NamedStatement] = {} - fragments: dict[str, SQLFragment] = {} + safe_name = _normalize_query_name(name) + if safe_name in self._query_to_file: + file_path = self._query_to_file[safe_name] + return self._files.get(file_path) + return None - section_matches = list(SECTION_MARKER_PATTERN.finditer(content)) - if not section_matches: - return {}, {} + def list_queries(self) -> "list[str]": + """List all available query names. - for i, match in enumerate(section_matches): - fragment_name = match.group("fragment") - is_fragment = fragment_name is not None - raw_statement_name = (fragment_name if is_fragment else match.group("name")).strip() - statement_start_line = content[: match.start()].count("\n") + Returns: + Sorted list of query names. + """ + return sorted(self._queries.keys()) - start_pos = match.end() - end_pos = section_matches[i + 1].start() if i + 1 < len(section_matches) else len(content) + def list_files(self) -> "list[str]": + """List all loaded file paths. - section_raw = content[start_pos:end_pos] - statement_section = section_raw.strip() - if not raw_statement_name or not statement_section: - continue + Returns: + Sorted list of file paths. + """ + return sorted(self._files.keys()) - section_lead = len(section_raw) - len(section_raw.lstrip()) - section_start_line = content[:start_pos].count("\n") + section_raw[:section_lead].count("\n") + def has_query(self, name: str) -> bool: + """Check if a query exists. - dialect, declared_params, statement_sql, declared_slots = SQLFileLoader._parse_directive_block( - statement_section, file_path, strict_parameter_annotations or is_fragment, base_line=section_start_line - ) + Args: + name: Query name to check. - clean_sql = SQLFileLoader._strip_leading_comments(statement_sql) - if not clean_sql: - continue - normalized_name = _normalize_query_name(raw_statement_name) + Returns: + True if query exists. + """ + safe_name = _normalize_query_name(name) + return safe_name in self._queries - if is_fragment: - if dialect is not None or declared_params or declared_slots: - raise SQLFileParseError( - file_path, - file_path, - ValueError( - f"Fragment '{raw_statement_name}' cannot declare -- dialect:, -- param:, or -- slot: " - "directives" - ), - line=statement_start_line + 1, - ) - if normalized_name in fragments: - raise SQLFileParseError( - file_path, - file_path, - ValueError(f"Duplicate fragment name: {raw_statement_name}"), - line=statement_start_line + 1, - ) - fragments[normalized_name] = SQLFragment( - name=normalized_name, sql=clean_sql, start_line=statement_start_line - ) - continue + def clear_cache(self) -> None: + """Clear all cached files and queries.""" + self._compiled_statements.clear() + self._files.clear() + self._queries.clear() + self._query_to_file.clear() + self._fragments.clear() + self._fragment_to_file.clear() + self._resolved_text.clear() - if normalized_name in statements: - raise SQLFileParseError( - file_path, - file_path, - ValueError(f"Duplicate statement name: {raw_statement_name}"), - line=statement_start_line + 1, - ) + cache_config = get_cache_config() + if cache_config.compiled_cache_enabled: + cache = get_cache() + cache.clear() - slots = _merge_slot_markers(declared_slots, clean_sql) - has_includes = bool(_include_markers(clean_sql)) - if not slots and not has_includes: - SQLFileLoader._check_declared_parameters( - clean_sql, declared_params, raw_statement_name, file_path, start_line=statement_start_line - ) + def clear_file_cache(self) -> None: + """Clear the file cache only, keeping loaded queries.""" + cache_config = get_cache_config() + if cache_config.compiled_cache_enabled: + cache = get_cache() + cache.clear() - statements[normalized_name] = NamedStatement( - name=normalized_name, - sql=clean_sql, - dialect=dialect, - start_line=statement_start_line, - parameters=declared_params, - slots=slots, - has_includes=has_includes, - ) - log_with_context( - logger, logging.DEBUG, "sql.parse", file_path=file_path, query_name=normalized_name, dialect=dialect - ) + def get_query_text(self, name: str) -> str: + """Get raw SQL text for a query. - if not statements and not fragments: - raise SQLFileParseError(file_path, file_path, ValueError("No valid SQL statements found after parsing")) + Includes are resolved; slot markers are left in place. - return statements, fragments + Args: + name: Query name. - def load_sql(self, *paths: str | Path) -> None: - """Load SQL files and parse named queries. + Returns: + Raw SQL text. + + Raises: + SQLStatementNotFoundError: If the query or an included fragment does not exist. + SQLFileParseError: If a declared slot has no marker or the includes form a cycle. + """ + safe_name = _normalize_query_name(name) + if safe_name not in self._queries: + self._raise_statement_not_found(name, safe_name) + return self._resolve_statement_text(safe_name) + + def get_sql(self, name: str, **slots: Any) -> "SQL": + """Get a SQL object by statement name, filling its slots. + + Each ``/* slot: name */`` marker is replaced by the matching keyword value, + or by the slot's ``-- slot:`` default when no value is given. A value may be + a ``str`` (spliced verbatim), a sqlglot expression (rendered with the + statement's dialect), or a ``SQL`` object (its text is spliced and its named + parameters are bound on the returned statement). Slot values are SQL, not + data: pass user input as parameters of a ``SQL`` value. + + The statement is cached only when no slot values are given. Args: - *paths: One or more file paths or directory paths to load. + name: Name of the statement (from -- name: in SQL file). + Hyphens in names are converted to underscores. + **slots: Values for the statement's slots, keyed by slot name. + + Returns: + SQL object ready for execution. + + Raises: + SQLSlotError: If a required slot is missing, a slot name is unknown, a + ``SQL`` value uses positional parameters, or slot parameter names collide + with each other or with the statement's placeholders. + TypeError: If a slot value is not a ``str``, sqlglot expression, or ``SQL``. + SQLFileParseError: If declared parameters do not match the filled SQL or the + SQL cannot be compiled. """ - runtime = self._runtime - span = None - error: Exception | None = None - start_time = time.perf_counter() - path_count = len(paths) - previous_correlation_id = CorrelationContext.get() - if runtime is not None: - runtime.increment_metric("loader.load.invocations") - runtime.increment_metric("loader.paths.requested", path_count) - span = runtime.start_span( - "sqlspec.loader.load", - attributes={"sqlspec.loader.path_count": path_count, "sqlspec.loader.encoding": self.encoding}, - ) + safe_name = _normalize_query_name(name) - try: - for path in paths: - path_str = str(path) - # If it looks like a URI or a potential alias (contains no path separators, or is in registry) - if "://" in path_str or self.storage_registry.is_alias_registered(path_str.split("/", maxsplit=1)[0]): - self._load_single_file(path, None) - continue + if safe_name not in self._queries: + self._raise_statement_not_found(name, safe_name) + if not slots and safe_name in self._compiled_statements: + return self._compiled_statements[safe_name] - path_obj = Path(path) - if path_obj.is_dir(): - self._load_directory(path_obj) - elif path_obj.exists(): - self._load_single_file(path_obj, None) - elif path_obj.suffix: - self._raise_file_not_found(str(path)) + parsed_statement = self._queries[safe_name] + sqlglot_dialect = None + if parsed_statement.dialect: + sqlglot_dialect = _normalize_dialect(parsed_statement.dialect) + statement_text = parsed_statement.sql + slot_parameters: dict[str, Any] = {} + if slots or parsed_statement.slots or parsed_statement.has_includes: + statement_text, slot_parameters = self._fill_slots( + safe_name, self._resolve_statement_text(safe_name), slots, sqlglot_dialect + ) + self._check_declared_parameters( + statement_text, + parsed_statement.parameters, + name, + self._query_to_file.get(safe_name, ""), + start_line=parsed_statement.start_line, + ) + + sql = SQL(statement_text, dialect=sqlglot_dialect, declared_parameters=parsed_statement.parameters) + try: + sql.compile() except Exception as exc: - error = exc - if runtime is not None: - runtime.increment_metric("loader.load.errors") - raise - finally: - duration_ms = (time.perf_counter() - start_time) * 1000 - if runtime is not None: - runtime.record_metric("loader.last_load_ms", duration_ms) - runtime.increment_metric("loader.load.duration_ms", duration_ms) - runtime.end_span(span, error=error) - CorrelationContext.set(previous_correlation_id) + raise SQLFileParseError(name=name, path="", original_error=exc) from exc + if slot_parameters: + return SQL( + statement_text, + slot_parameters, + dialect=sqlglot_dialect, + declared_parameters=parsed_statement.parameters, + ) + if not slots: + self._compiled_statements[safe_name] = sql + return sql - def _load_directory(self, dir_path: Path) -> None: - """Load all SQL files from a directory. + def _raise_file_not_found(self, path: str) -> None: + """Raise SQLFileNotFoundError for nonexistent file. Args: - dir_path: Directory path to load SQL files from. + path: File path that was not found. + + Raises: + SQLFileNotFoundError: Always raised. """ - runtime = self._runtime - if runtime is not None: - runtime.increment_metric("loader.directories.scanned") + raise SQLFileNotFoundError(path) - sql_files = list(dir_path.rglob("*.sql")) - if not sql_files: - return + def _raise_statement_not_found(self, name: str, normalized_name: str) -> None: + """Raise SQLStatementNotFoundError for nonexistent statements. - for file_path in sql_files: - relative_path = file_path.relative_to(dir_path) - namespace_parts = relative_path.parent.parts - self._load_single_file(file_path, ".".join(namespace_parts) if namespace_parts else None) + Args: + name: Name requested by the caller. + normalized_name: Normalized statement name used for lookup. - def _load_single_file(self, file_path: str | Path, namespace: str | None) -> bool: - """Load a single SQL file with optional namespace. + Raises: + SQLStatementNotFoundError: Always raised. + """ + raise SQLStatementNotFoundError(name=name, normalized_name=normalized_name, query_count=len(self._queries)) + + def _file_cache_key(self, path: str | Path) -> str: + """Generate cache key for a file path. Args: - file_path: Path to the SQL file. - namespace: Optional namespace prefix for queries. + path: File path to generate key for. Returns: - True if file was newly loaded, False if already cached. + Cache key string for the file. """ - path_str = str(file_path) - runtime = self._runtime - if runtime is not None: - runtime.increment_metric("loader.files.considered") + path_str = str(path) + path_hash = hashlib.md5(path_str.encode(), usedforsecurity=False).hexdigest() + return f"file:{path_hash[:16]}" - if path_str in self._files: - if runtime is not None: - runtime.increment_metric("loader.cache.hit") - return False + @staticmethod + def _compute_checksum(content: str) -> str: + """Compute MD5 checksum from already-read file content.""" + return hashlib.md5(content.encode(), usedforsecurity=False).hexdigest() - cache_config = get_cache_config() - if not cache_config.compiled_cache_enabled: - self._load_uncached_file(file_path, namespace) - if runtime is not None: - runtime.increment_metric("loader.cache.miss") - return True + def _calculate_file_checksum(self, path: str | Path) -> str: + """Calculate checksum for file content validation. - cache_key_str = self._file_cache_key(file_path) - cache = get_cache() - cached_file = cache.get_file(cache_key_str) + Args: + path: File path to calculate checksum for. - if cached_file is not None and isinstance(cached_file, SQLFileCacheEntry): - try: - file_content = self._read_file_content(file_path) - except Exception: - file_content = None + Returns: + MD5 checksum of file content. - if file_content is not None and self._content_matches_cache(file_content, cached_file): - self._files[path_str] = cached_file.sql_file - for name, statement in cached_file.parsed_statements.items(): - namespaced_name = f"{namespace}.{name}" if namespace else name - if namespaced_name in self._queries: - existing_file = self._query_to_file.get(namespaced_name, "unknown") - if existing_file != path_str: - raise SQLFileParseError( - path_str, - path_str, - ValueError(f"Query name '{namespaced_name}' already exists in file: {existing_file}"), - line=statement.start_line + 1, - ) - self._queries[namespaced_name] = statement - self._query_to_file[namespaced_name] = path_str - self._register_fragments(cached_file.parsed_fragments, path_str, namespace) - if runtime is not None: - runtime.increment_metric("loader.cache.hit") - return True + Raises: + SQLFileParseError: If file cannot be read. + """ + try: + return self._compute_checksum(self._read_file_content(path)) + except Exception as e: + raise SQLFileParseError(str(path), str(path), e) from e - loaded_statements, loaded_fragments = self._load_uncached_file(file_path, namespace, content=file_content) - else: - loaded_statements, loaded_fragments = self._load_uncached_file(file_path, namespace) + def _is_file_unchanged(self, path: str | Path, cached_file: SQLFile) -> bool: + """Check if file has changed since caching. - if path_str in self._files: - sql_file = self._files[path_str] - cached_file_data = SQLFileCacheEntry( - sql_file=sql_file, parsed_statements=loaded_statements, parsed_fragments=loaded_fragments - ) - cache.put_file(cache_key_str, cached_file_data) - if runtime is not None: - runtime.increment_metric("loader.cache.miss") - runtime.increment_metric("loader.files.loaded") - runtime.increment_metric("loader.statements.loaded", len(loaded_statements)) + Args: + path: File path to check. + cached_file: Cached file data. - return True + Returns: + True if file is unchanged, False otherwise. + """ + try: + current_checksum = self._calculate_file_checksum(path) + except Exception: + return False + else: + return current_checksum == cached_file.checksum - def _load_uncached_file( - self, file_path: str | Path, namespace: "str | None", content: "str | None" = None - ) -> "tuple[dict[str, NamedStatement], dict[str, SQLFragment]]": - """Load a single SQL file without using cache. + def _reload_changed_files(self) -> "list[str]": + """Reload tracked SQL files whose content checksum changed. - Args: - file_path: Path to the SQL file. - namespace: Optional namespace prefix for queries and fragments. - content: Pre-read file content. If provided, skips the disk read. + Every changed file's queries and fragments are removed before any changed + file is loaded again, so names may move between changed files. Returns: - The file's parsed statements and fragments keyed by un-namespaced name; - both are empty when the file contains no named sections. + Paths of files that were reloaded. """ - path_str = str(file_path) - runtime = self._runtime - if content is None: - content = self._read_file_content(file_path) - statements, fragments = self._parse_statements(content, path_str, self.strict_parameter_annotations) + pending_reloads: list[tuple[str, str | None]] = [] + for path, sql_file in list(self._files.items()): + if self._is_file_unchanged(path, sql_file): + if self._runtime is not None: + self._runtime.increment_metric("loader.reload.skipped") + continue - if not statements and not fragments: - log_with_context( - logger, logging.DEBUG, "sql.load", file_path=path_str, status="skipped", reason="no_named_statements" + query_names = [name for name, source in self._query_to_file.items() if source == path] + fragment_names = [name for name, source in self._fragment_to_file.items() if source == path] + namespaces = { + name[: -(len(statement.name) + 1)] + for name in query_names + if (statement := self._queries.get(name)) is not None and name.endswith(f".{statement.name}") + } + namespaces.update( + name[: -(len(fragment.name) + 1)] + for name in fragment_names + if (fragment := self._fragments.get(name)) is not None and name.endswith(f".{fragment.name}") ) - return {}, {} - - sql_file = SQLFile(content=content, path=path_str) - self._files[path_str] = sql_file + namespace = next(iter(namespaces)) if len(namespaces) == 1 else None + for name in query_names: + self._queries.pop(name, None) + self._query_to_file.pop(name, None) + self._compiled_statements.pop(name, None) + for name in fragment_names: + self._fragments.pop(name, None) + self._fragment_to_file.pop(name, None) + self._files.pop(path, None) + pending_reloads.append((path, namespace)) - for name, statement in statements.items(): - namespaced_name = f"{namespace}.{name}" if namespace else name - if namespaced_name in self._queries: - existing_file = self._query_to_file.get(namespaced_name, "unknown") - if existing_file != path_str: - raise SQLFileParseError( - path_str, - path_str, - ValueError(f"Query name '{namespaced_name}' already exists in file: {existing_file}"), - line=statement.start_line + 1, - ) - self._queries[namespaced_name] = statement - self._query_to_file[namespaced_name] = path_str - self._register_fragments(fragments, path_str, namespace) - log_with_context( - logger, logging.DEBUG, "sql.load", file_path=path_str, statement_count=len(statements), status="loaded" - ) - if runtime is not None: - runtime.increment_metric("loader.files.loaded") - runtime.increment_metric("loader.statements.loaded", len(statements)) - return statements, fragments + if pending_reloads: + self._invalidate_resolved() + changed_paths: list[str] = [] + for path, namespace in pending_reloads: + self._load_single_file(path, namespace) + changed_paths.append(path) + if self._runtime is not None: + self._runtime.increment_metric("loader.reload.changed") + return changed_paths - def _register_fragments(self, fragments: "dict[str, SQLFragment]", path_str: str, namespace: "str | None") -> None: - """Register a file's fragments under their namespaced names. + def _content_matches_cache(self, content: str, cached_file: SQLFileCacheEntry) -> bool: + """Check if already-read file content matches cached checksum.""" + return self._compute_checksum(content) == cached_file.sql_file.checksum - Registering any fragment invalidates include-resolved statement text. + def _read_file_content(self, path: str | Path) -> str: + """Read file content using storage backend. Args: - fragments: Fragments keyed by un-namespaced name. - path_str: Source file path. - namespace: Optional namespace prefix for the fragment names. + path: File path (can be local path or URI). + + Returns: + File content as string. Raises: - SQLFileParseError: If a fragment name is already registered from another file. + SQLFileNotFoundError: If file does not exist. + SQLFileParseError: If file cannot be read or parsed. """ - if not fragments: - return - for name, fragment in fragments.items(): - namespaced_name = f"{namespace}.{name}" if namespace else name - existing_file = self._fragment_to_file.get(namespaced_name) - if existing_file is not None and existing_file != path_str: - raise SQLFileParseError( - path_str, - path_str, - ValueError(f"Fragment name '{namespaced_name}' already exists in file: {existing_file}"), - line=fragment.start_line + 1, - ) - self._fragments[namespaced_name] = fragment - self._fragment_to_file[namespaced_name] = path_str - self._invalidate_resolved() + path_str = str(path) - def _invalidate_resolved(self) -> None: - """Drop include-resolved text and compiled statements that depend on fragments.""" - self._resolved_text.clear() - stale_names = [ - name - for name in self._compiled_statements - if (statement := self._queries.get(name)) is not None and statement.has_includes - ] - for name in stale_names: - del self._compiled_statements[name] + try: + backend = self.storage_registry.get(path) - def add_named_sql( - self, - name: str, - sql: str, - dialect: "str | None" = None, - parameters: "Sequence[ParameterDeclaration] | None" = None, - ) -> None: - """Add a named SQL query directly without loading from a file. + parts = path_str.split("/", 1) + if len(parts) > 1 and self.storage_registry.is_alias_registered(parts[0]): + return backend.read_text_sync(parts[1], encoding=self.encoding) - The SQL may contain ``/* include: name */`` and ``/* slot: name */`` markers; - slots added this way have no defaults. + if path_str.startswith("file://"): + parsed = urlparse(path_str) + file_path = unquote(parsed.path) + if file_path and len(file_path) > 2 and file_path[2] == ":": # noqa: PLR2004 + file_path = file_path[1:] + return backend.read_text_sync(Path(file_path).name, encoding=self.encoding) - Args: - name: Name for the SQL query. - sql: Raw SQL content. - dialect: Optional dialect for the SQL statement. - parameters: Optional declared parameter metadata for the query. + if isinstance(path, Path) or is_local_path(path_str): + return backend.read_text_sync(Path(path_str).name, encoding=self.encoding) - Raises: - ValueError: If query name already exists. - """ + return backend.read_text_sync(path_str, encoding=self.encoding) + except KeyError as e: + raise SQLFileNotFoundError(path_str) from e + except FileNotFoundInStorageError as e: + raise SQLFileNotFoundError(path_str) from e + except FileNotFoundError as e: + raise SQLFileNotFoundError(path_str) from e + except StorageOperationFailedError as e: + raise SQLFileParseError(path_str, path_str, e) from e + except Exception as e: + raise SQLFileParseError(path_str, path_str, e) from e - normalized_name = _normalize_query_name(name) + @staticmethod + def _strip_leading_comments(sql_text: str) -> str: + """Remove leading comment lines from a SQL string.""" + lines = sql_text.strip().split("\n") + first_sql_line_index = -1 + for i, line in enumerate(lines): + if line.strip() and not line.strip().startswith("--"): + first_sql_line_index = i + break + if first_sql_line_index == -1: + return "" + return "\n".join(lines[first_sql_line_index:]).strip() - if normalized_name in self._queries: - existing_source = self._query_to_file.get(normalized_name, "") - msg = f"Query name '{name}' already exists (source: {existing_source})" - raise ValueError(msg) + @staticmethod + def _parse_directive_block( + statement_section: str, file_path: str, strict: bool, base_line: int = 0 + ) -> "tuple[str | None, tuple[ParameterDeclaration, ...], str, tuple[SlotDeclaration, ...]]": + """Scan a section's leading comment block for ``dialect``/``param``/``slot`` directives. - if dialect is not None: - dialect = _normalize_dialect(dialect) + Args: + statement_section: The statement body including any leading directive lines. + file_path: File path for error reporting. + strict: When True, a malformed ``-- param:`` line raises instead of warning. + base_line: 0-based line offset of ``statement_section`` within the file. - declared = tuple(parameters) if parameters else () - clean_sql = sql.strip() - slots = _merge_slot_markers((), clean_sql) - has_includes = bool(_include_markers(clean_sql)) - if not slots and not has_includes: - self._check_declared_parameters(clean_sql, declared, name, "") + Returns: + The resolved dialect, the declared parameters, the SQL body with the + leading directive/comment lines removed, and the declared slots. + + Raises: + SQLFileParseError: If ``strict`` and a ``-- param:`` line is malformed, or a + ``-- slot:`` line is malformed, duplicated, or placed after the SQL body begins. + """ + dialect: str | None = None + params: list[ParameterDeclaration] = [] + slots: list[SlotDeclaration] = [] + raw_lines = statement_section.split("\n") + body_start = len(raw_lines) + for idx, raw in enumerate(raw_lines): + stripped = raw.strip() + if not stripped: + continue + if not stripped.startswith("--"): + body_start = idx + break + dialect_match = DIALECT_PATTERN.match(stripped) + if dialect_match: + dialect = _normalize_dialect(dialect_match.group("dialect").lower()) + continue + param_match = PARAM_PATTERN.match(stripped) + if param_match: + params.append(_parse_parameter_declaration(param_match)) + continue + slot_match = SLOT_DIRECTIVE_PATTERN.match(stripped) + if slot_match: + slot_name = slot_match.group("name") + if any(slot.name == slot_name for slot in slots): + raise SQLFileParseError( + file_path, + file_path, + ValueError(f"Duplicate -- slot: directive for slot '{slot_name}'"), + line=base_line + idx + 1, + ) + slots.append(SlotDeclaration(slot_name, slot_match.group("default"))) + continue + if SLOT_COMMENT_PATTERN.match(stripped): + raise SQLFileParseError( + file_path, + file_path, + ValueError(f"Malformed -- slot: directive: {stripped}"), + line=base_line + idx + 1, + ) + if PARAM_PREFIX_PATTERN.match(stripped): + line_number = base_line + idx + 1 + if strict: + raise SQLFileParseError( + file_path, file_path, ValueError(f"Malformed -- param: directive: {stripped}"), line=line_number + ) + log_with_context( + logger, + logging.WARNING, + f"sql.parse.param: malformed parameter directive in {file_path} at line {line_number}: {stripped}", + file_path=file_path, + line_number=line_number, + directive=stripped, + status="malformed", + ) + body_text = "\n".join(raw_lines[body_start:]) + if SLOT_COMMENT_PATTERN.search(body_text) is not None: + for start, end, is_block in _scan_sql_comments(body_text): + if is_block or SLOT_COMMENT_PATTERN.match(body_text, start) is None: + continue + line_start = body_text.rfind("\n", 0, start) + 1 + if body_text[line_start:start].strip(): + continue + raise SQLFileParseError( + file_path, + file_path, + ValueError( + f"-- slot: directive must appear in the leading directive block: {body_text[start:end].strip()}" + ), + line=base_line + body_start + body_text.count("\n", 0, start) + 1, + ) + return dialect, tuple(params), body_text, tuple(slots) - statement = NamedStatement( - name=normalized_name, - sql=clean_sql, - dialect=dialect, - start_line=0, - parameters=declared, - slots=slots, - has_includes=has_includes, - ) - self._queries[normalized_name] = statement - self._query_to_file[normalized_name] = "" + @staticmethod + def _check_declared_parameters( + clean_sql: str, + declared: "tuple[ParameterDeclaration, ...]", + statement_name: str, + file_path: str, + start_line: "int | None" = None, + ) -> None: + """Validate declared parameters against the query's actual placeholders. - def add_fragment(self, name: str, sql: str) -> None: - """Add a reusable SQL fragment directly without loading from a file. + For named binding, every declared name must appear among the SQL placeholders + (declared names may be a subset; filters and undeclared params are allowed). For + positional binding, the declared count must equal the placeholder count. Args: - name: Name for the fragment, referenced by ``/* include: name */`` markers. - sql: Fragment SQL text; may contain include and slot markers. + clean_sql: The SQL body with directives/comments stripped. + declared: Declared parameters for the query. + statement_name: Raw query name for error messages. + file_path: File path for error reporting. + start_line: Optional 0-based line of the statement within the file. Raises: - ValueError: If the fragment name already exists. - """ - normalized_name = _normalize_query_name(name) - if normalized_name in self._fragments: - existing_source = self._fragment_to_file.get(normalized_name, "") - msg = f"Fragment name '{name}' already exists (source: {existing_source})" - raise ValueError(msg) - self._fragments[normalized_name] = SQLFragment(name=normalized_name, sql=sql.strip()) - self._fragment_to_file[normalized_name] = "" - self._invalidate_resolved() - - def has_fragment(self, name: str) -> bool: - """Check if a fragment exists. - - Args: - name: Fragment name to check. - - Returns: - True if the fragment exists. - """ - return _normalize_query_name(name) in self._fragments - - def list_fragments(self) -> "list[str]": - """List all available fragment names. - - Returns: - Sorted list of fragment names. + SQLFileParseError: On name drift (named) or count mismatch (positional). """ - return sorted(self._fragments.keys()) + if not declared: + return + error_line = start_line + 1 if start_line is not None else None + infos = ParameterValidator().extract_parameters(clean_sql) + named = {info.name for info in infos if info.name and not info.name.isdigit()} + if named: + for decl in declared: + if decl.name not in named: + raise SQLFileParseError( + file_path, + file_path, + ValueError( + f"Declared parameter '{decl.name}' for query '{statement_name}' is not present in the " + f"SQL placeholders {sorted(named)}" + ), + line=error_line, + ) + elif len(declared) != len(infos): + raise SQLFileParseError( + file_path, + file_path, + ValueError( + f"Query '{statement_name}' declares {len(declared)} parameter(s) but the SQL has " + f"{len(infos)} positional placeholder(s)" + ), + line=error_line, + ) - def get_fragment_text(self, name: str) -> str: - """Get a fragment's SQL text with its includes resolved. + @staticmethod + def _parse_statements( + content: str, file_path: str, strict_parameter_annotations: bool = False + ) -> "tuple[dict[str, NamedStatement], dict[str, SQLFragment]]": + """Parse SQL content and extract named statements and fragments. - Slot markers are left in place. + A section starts at a ``-- name:`` or ``-- fragment:`` marker and ends at the + next marker of either kind. Files without any markers are gracefully skipped + by returning empty dictionaries. The caller is responsible for handling empty + results appropriately. Args: - name: Fragment name. + content: Raw SQL file content to parse. + file_path: File path for error reporting. + strict_parameter_annotations: Raise on malformed parameter declarations instead of skipping them. Returns: - Fragment SQL text with ``/* include: */`` markers replaced. + Dictionaries mapping normalized names to NamedStatement and SQLFragment + objects. Both are empty if no markers are found in the content. Raises: - SQLFragmentNotFoundError: If the fragment or an included fragment does not exist. - SQLFileParseError: If the includes form a cycle. + SQLFileParseError: If sections are malformed (duplicate names, directives + on fragments, invalid slot directives, or no content after parsing). """ - safe_name = self._find_fragment_name(name, None) - return self._resolve_includes( - self._fragments[safe_name].sql, namespace=_namespace_of(safe_name), stack=(safe_name,) - ) - - def get_query_slots(self, name: str) -> "tuple[SlotDeclaration, ...]": - """Get the slots of a query, including slots contributed by included fragments. - - Slots declared with ``-- slot:`` come first in declaration order, then - undeclared (required) markers in the query's own text, then markers - contributed by included fragments; markers keep their order of appearance. + statements: dict[str, NamedStatement] = {} + fragments: dict[str, SQLFragment] = {} - Args: - name: Query name (hyphens are converted to underscores). + section_matches = list(SECTION_MARKER_PATTERN.finditer(content)) + if not section_matches: + return {}, {} - Returns: - Tuple of slot declarations; empty if the query has none. + for i, match in enumerate(section_matches): + fragment_name = match.group("fragment") + is_fragment = fragment_name is not None + raw_statement_name = (fragment_name if is_fragment else match.group("name")).strip() + statement_start_line = content[: match.start()].count("\n") - Raises: - SQLStatementNotFoundError: If the query or an included fragment does not exist. - SQLFileParseError: If a declared slot has no marker or the includes form a cycle. - """ - safe_name = _normalize_query_name(name) - if safe_name not in self._queries: - self._raise_statement_not_found(name, safe_name) - statement = self._queries[safe_name] - if not statement.slots and not statement.has_includes: - return () - return _merge_slot_markers(statement.slots, self._resolve_statement_text(safe_name)) + start_pos = match.end() + end_pos = section_matches[i + 1].start() if i + 1 < len(section_matches) else len(content) - def get_query_parameters(self, name: str) -> "tuple[ParameterDeclaration, ...]": - """Get declared parameter metadata for a query. + section_raw = content[start_pos:end_pos] + statement_section = section_raw.strip() + if not raw_statement_name or not statement_section: + continue - Args: - name: Query name (hyphens are converted to underscores). + section_lead = len(section_raw) - len(section_raw.lstrip()) + section_start_line = content[:start_pos].count("\n") + section_raw[:section_lead].count("\n") - Returns: - Tuple of declared parameters; empty if the query declares none. + dialect, declared_params, statement_sql, declared_slots = SQLFileLoader._parse_directive_block( + statement_section, file_path, strict_parameter_annotations or is_fragment, base_line=section_start_line + ) - Raises: - SQLStatementNotFoundError: If the query does not exist. - """ - safe_name = _normalize_query_name(name) - if safe_name not in self._queries: - self._raise_statement_not_found(name, safe_name) - return self._queries[safe_name].parameters + clean_sql = SQLFileLoader._strip_leading_comments(statement_sql) + if not clean_sql: + continue + normalized_name = _normalize_query_name(raw_statement_name) - def get_file(self, path: str | Path) -> "SQLFile | None": - """Get a loaded SQLFile object by path. + if is_fragment: + if dialect is not None or declared_params or declared_slots: + raise SQLFileParseError( + file_path, + file_path, + ValueError( + f"Fragment '{raw_statement_name}' cannot declare -- dialect:, -- param:, or -- slot: " + "directives" + ), + line=statement_start_line + 1, + ) + if normalized_name in fragments: + raise SQLFileParseError( + file_path, + file_path, + ValueError(f"Duplicate fragment name: {raw_statement_name}"), + line=statement_start_line + 1, + ) + fragments[normalized_name] = SQLFragment( + name=normalized_name, sql=clean_sql, start_line=statement_start_line + ) + continue - Args: - path: Path of the file. + if normalized_name in statements: + raise SQLFileParseError( + file_path, + file_path, + ValueError(f"Duplicate statement name: {raw_statement_name}"), + line=statement_start_line + 1, + ) - Returns: - SQLFile object if loaded, None otherwise. - """ - return self._files.get(str(path)) + slots = _merge_slot_markers(declared_slots, clean_sql) + has_includes = bool(_include_markers(clean_sql)) + if not slots and not has_includes: + SQLFileLoader._check_declared_parameters( + clean_sql, declared_params, raw_statement_name, file_path, start_line=statement_start_line + ) - def get_file_for_query(self, name: str) -> "SQLFile | None": - """Get the SQLFile object containing a query. + statements[normalized_name] = NamedStatement( + name=normalized_name, + sql=clean_sql, + dialect=dialect, + start_line=statement_start_line, + parameters=declared_params, + slots=slots, + has_includes=has_includes, + ) + log_with_context( + logger, logging.DEBUG, "sql.parse", file_path=file_path, query_name=normalized_name, dialect=dialect + ) - Args: - name: Query name (hyphens are converted to underscores). + if not statements and not fragments: + raise SQLFileParseError(file_path, file_path, ValueError("No valid SQL statements found after parsing")) - Returns: - SQLFile object if query exists, None otherwise. - """ - safe_name = _normalize_query_name(name) - if safe_name in self._query_to_file: - file_path = self._query_to_file[safe_name] - return self._files.get(file_path) - return None + return statements, fragments - def list_queries(self) -> "list[str]": - """List all available query names. + def _load_directory(self, dir_path: Path) -> None: + """Load all SQL files from a directory. - Returns: - Sorted list of query names. + Args: + dir_path: Directory path to load SQL files from. """ - return sorted(self._queries.keys()) + runtime = self._runtime + if runtime is not None: + runtime.increment_metric("loader.directories.scanned") - def list_files(self) -> "list[str]": - """List all loaded file paths. + sql_files = list(dir_path.rglob("*.sql")) + if not sql_files: + return - Returns: - Sorted list of file paths. - """ - return sorted(self._files.keys()) + for file_path in sql_files: + relative_path = file_path.relative_to(dir_path) + namespace_parts = relative_path.parent.parts + self._load_single_file(file_path, ".".join(namespace_parts) if namespace_parts else None) - def has_query(self, name: str) -> bool: - """Check if a query exists. + def _load_single_file(self, file_path: str | Path, namespace: str | None) -> bool: + """Load a single SQL file with optional namespace. Args: - name: Query name to check. + file_path: Path to the SQL file. + namespace: Optional namespace prefix for queries. Returns: - True if query exists. + True if file was newly loaded, False if already cached. """ - safe_name = _normalize_query_name(name) - return safe_name in self._queries + path_str = str(file_path) + runtime = self._runtime + if runtime is not None: + runtime.increment_metric("loader.files.considered") - def clear_cache(self) -> None: - """Clear all cached files and queries.""" - self._compiled_statements.clear() - self._files.clear() - self._queries.clear() - self._query_to_file.clear() - self._fragments.clear() - self._fragment_to_file.clear() - self._resolved_text.clear() + if path_str in self._files: + if runtime is not None: + runtime.increment_metric("loader.cache.hit") + return False cache_config = get_cache_config() - if cache_config.compiled_cache_enabled: - cache = get_cache() - cache.clear() + if not cache_config.compiled_cache_enabled: + self._load_uncached_file(file_path, namespace) + if runtime is not None: + runtime.increment_metric("loader.cache.miss") + return True - def clear_file_cache(self) -> None: - """Clear the file cache only, keeping loaded queries.""" - cache_config = get_cache_config() - if cache_config.compiled_cache_enabled: - cache = get_cache() - cache.clear() + cache_key_str = self._file_cache_key(file_path) + cache = get_cache() + cached_file = cache.get_file(cache_key_str) - def get_query_text(self, name: str) -> str: - """Get raw SQL text for a query. + if cached_file is not None and isinstance(cached_file, SQLFileCacheEntry): + try: + file_content = self._read_file_content(file_path) + except Exception: + file_content = None - Includes are resolved; slot markers are left in place. + if file_content is not None and self._content_matches_cache(file_content, cached_file): + self._files[path_str] = cached_file.sql_file + for name, statement in cached_file.parsed_statements.items(): + namespaced_name = f"{namespace}.{name}" if namespace else name + if namespaced_name in self._queries: + existing_file = self._query_to_file.get(namespaced_name, "unknown") + if existing_file != path_str: + raise SQLFileParseError( + path_str, + path_str, + ValueError(f"Query name '{namespaced_name}' already exists in file: {existing_file}"), + line=statement.start_line + 1, + ) + self._queries[namespaced_name] = statement + self._query_to_file[namespaced_name] = path_str + self._register_fragments(cached_file.parsed_fragments, path_str, namespace) + if runtime is not None: + runtime.increment_metric("loader.cache.hit") + return True + + loaded_statements, loaded_fragments = self._load_uncached_file(file_path, namespace, content=file_content) + else: + loaded_statements, loaded_fragments = self._load_uncached_file(file_path, namespace) + + if path_str in self._files: + sql_file = self._files[path_str] + cached_file_data = SQLFileCacheEntry( + sql_file=sql_file, parsed_statements=loaded_statements, parsed_fragments=loaded_fragments + ) + cache.put_file(cache_key_str, cached_file_data) + if runtime is not None: + runtime.increment_metric("loader.cache.miss") + runtime.increment_metric("loader.files.loaded") + runtime.increment_metric("loader.statements.loaded", len(loaded_statements)) + + return True + + def _load_uncached_file( + self, file_path: str | Path, namespace: "str | None", content: "str | None" = None + ) -> "tuple[dict[str, NamedStatement], dict[str, SQLFragment]]": + """Load a single SQL file without using cache. Args: - name: Query name. + file_path: Path to the SQL file. + namespace: Optional namespace prefix for queries and fragments. + content: Pre-read file content. If provided, skips the disk read. Returns: - Raw SQL text. - - Raises: - SQLStatementNotFoundError: If the query or an included fragment does not exist. - SQLFileParseError: If a declared slot has no marker or the includes form a cycle. + The file's parsed statements and fragments keyed by un-namespaced name; + both are empty when the file contains no named sections. """ - safe_name = _normalize_query_name(name) - if safe_name not in self._queries: - self._raise_statement_not_found(name, safe_name) - return self._resolve_statement_text(safe_name) + path_str = str(file_path) + runtime = self._runtime + if content is None: + content = self._read_file_content(file_path) + statements, fragments = self._parse_statements(content, path_str, self.strict_parameter_annotations) - def get_sql(self, name: str, **slots: Any) -> "SQL": - """Get a SQL object by statement name, filling its slots. + if not statements and not fragments: + log_with_context( + logger, logging.DEBUG, "sql.load", file_path=path_str, status="skipped", reason="no_named_statements" + ) + return {}, {} - Each ``/* slot: name */`` marker is replaced by the matching keyword value, - or by the slot's ``-- slot:`` default when no value is given. A value may be - a ``str`` (spliced verbatim), a sqlglot expression (rendered with the - statement's dialect), or a ``SQL`` object (its text is spliced and its named - parameters are bound on the returned statement). Slot values are SQL, not - data: pass user input as parameters of a ``SQL`` value. + sql_file = SQLFile(content=content, path=path_str) + self._files[path_str] = sql_file - The statement is cached only when no slot values are given. + for name, statement in statements.items(): + namespaced_name = f"{namespace}.{name}" if namespace else name + if namespaced_name in self._queries: + existing_file = self._query_to_file.get(namespaced_name, "unknown") + if existing_file != path_str: + raise SQLFileParseError( + path_str, + path_str, + ValueError(f"Query name '{namespaced_name}' already exists in file: {existing_file}"), + line=statement.start_line + 1, + ) + self._queries[namespaced_name] = statement + self._query_to_file[namespaced_name] = path_str + self._register_fragments(fragments, path_str, namespace) + log_with_context( + logger, logging.DEBUG, "sql.load", file_path=path_str, statement_count=len(statements), status="loaded" + ) + if runtime is not None: + runtime.increment_metric("loader.files.loaded") + runtime.increment_metric("loader.statements.loaded", len(statements)) + return statements, fragments - Args: - name: Name of the statement (from -- name: in SQL file). - Hyphens in names are converted to underscores. - **slots: Values for the statement's slots, keyed by slot name. + def _register_fragments(self, fragments: "dict[str, SQLFragment]", path_str: str, namespace: "str | None") -> None: + """Register a file's fragments under their namespaced names. - Returns: - SQL object ready for execution. + Registering any fragment invalidates include-resolved statement text. + + Args: + fragments: Fragments keyed by un-namespaced name. + path_str: Source file path. + namespace: Optional namespace prefix for the fragment names. Raises: - SQLSlotError: If a required slot is missing, a slot name is unknown, a - ``SQL`` value uses positional parameters, or slot parameter names collide - with each other or with the statement's placeholders. - TypeError: If a slot value is not a ``str``, sqlglot expression, or ``SQL``. - SQLFileParseError: If declared parameters do not match the filled SQL or the - SQL cannot be compiled. + SQLFileParseError: If a fragment name is already registered from another file. """ - safe_name = _normalize_query_name(name) - - if safe_name not in self._queries: - self._raise_statement_not_found(name, safe_name) - if not slots and safe_name in self._compiled_statements: - return self._compiled_statements[safe_name] - - parsed_statement = self._queries[safe_name] - sqlglot_dialect = None - if parsed_statement.dialect: - sqlglot_dialect = _normalize_dialect(parsed_statement.dialect) - - statement_text = parsed_statement.sql - slot_parameters: dict[str, Any] = {} - if slots or parsed_statement.slots or parsed_statement.has_includes: - statement_text, slot_parameters = self._fill_slots( - safe_name, self._resolve_statement_text(safe_name), slots, sqlglot_dialect - ) - self._check_declared_parameters( - statement_text, - parsed_statement.parameters, - name, - self._query_to_file.get(safe_name, ""), - start_line=parsed_statement.start_line, - ) + if not fragments: + return + for name, fragment in fragments.items(): + namespaced_name = f"{namespace}.{name}" if namespace else name + existing_file = self._fragment_to_file.get(namespaced_name) + if existing_file is not None and existing_file != path_str: + raise SQLFileParseError( + path_str, + path_str, + ValueError(f"Fragment name '{namespaced_name}' already exists in file: {existing_file}"), + line=fragment.start_line + 1, + ) + self._fragments[namespaced_name] = fragment + self._fragment_to_file[namespaced_name] = path_str + self._invalidate_resolved() - sql = SQL(statement_text, dialect=sqlglot_dialect, declared_parameters=parsed_statement.parameters) - try: - sql.compile() - except Exception as exc: - raise SQLFileParseError(name=name, path="", original_error=exc) from exc - if slot_parameters: - return SQL( - statement_text, - slot_parameters, - dialect=sqlglot_dialect, - declared_parameters=parsed_statement.parameters, - ) - if not slots: - self._compiled_statements[safe_name] = sql - return sql + def _invalidate_resolved(self) -> None: + """Drop include-resolved text and compiled statements that depend on fragments.""" + self._resolved_text.clear() + stale_names = [ + name + for name in self._compiled_statements + if (statement := self._queries.get(name)) is not None and statement.has_includes + ] + for name in stale_names: + del self._compiled_statements[name] def _fill_slots( self, safe_name: str, resolved_text: str, provided: "dict[str, Any]", dialect: "str | None" diff --git a/sqlspec/migrations/base.py b/sqlspec/migrations/base.py index 0dbe9dbb5..58479b9ad 100644 --- a/sqlspec/migrations/base.py +++ b/sqlspec/migrations/base.py @@ -367,17 +367,6 @@ def _record_squashed_migration_statement( ) ) - def _column_exists_query(self) -> Select: - """Get SQL to check what columns exist in the tracking table. - - Returns a query that will fail gracefully if the table doesn't exist, - and returns column names if it does. - - Returns: - SQL builder object for column check query. - """ - return sql.select("*").from_(self.version_table).limit(0) - def _detect_missing_columns(self, existing_columns: "set[str]") -> "set[str]": """Detect which columns are missing from the current schema. diff --git a/sqlspec/migrations/commands.py b/sqlspec/migrations/commands.py index 2704f1d08..65c175efe 100644 --- a/sqlspec/migrations/commands.py +++ b/sqlspec/migrations/commands.py @@ -165,8 +165,6 @@ def __init__(self, config: "SyncConfigT") -> None: """ super().__init__(config) self.tracker = self._create_tracker() - - # Create context with extension configurations context = MigrationContext.from_config(config) context.extension_config = self.extension_configs @@ -340,7 +338,7 @@ def record_version(exec_time: int, migration: "LoadedMigrationMetadata" = migrat except Exception as exc: use_txn = self.runner.should_use_transaction(migration, self.config) rollback_msg = " (transaction rolled back)" if use_txn else "" - _output_exception( + _output_error( use_logger, echo, summary_only, @@ -394,7 +392,7 @@ def remove_version(exec_time: int, version: str = version) -> None: except Exception as exc: use_txn = self.runner.should_use_transaction(migration, self.config) rollback_msg = " (transaction rolled back)" if use_txn else "" - _output_exception( + _output_error( use_logger, echo, summary_only, @@ -563,7 +561,6 @@ def upgrade( self._validate_migration_schema(driver) self.tracker.ensure_tracking_table(driver) - # config auto_sync=False cannot be overridden by the call-site flag. if auto_sync and self.config.migration_config.get("auto_sync", True): self._synchronize_version_records( driver, use_logger=ul, echo=echo_value, summary_only=summary_value @@ -819,27 +816,7 @@ def revision(self, message: str, file_type: str | None = None) -> None: message: Description for the migration. file_type: Type of migration file to create ('sql' or 'py'). """ - version = generate_timestamp_version() - selected_format = file_type or self._template_settings.default_format - file_path = create_migration_file( - self.migrations_path, - version, - message, - selected_format, - config=self.config, - template_settings=self._template_settings, - ) - log_with_context( - logger, - logging.DEBUG, - "migration.create", - db_system=resolve_db_system(type(self.config).__name__), - version=version, - file_path=str(file_path), - file_type=selected_format, - description=message, - ) - console.print(f"[green]Created migration:[/] {file_path}") + _create_revision(self.migrations_path, self.config, self._template_settings, message, file_type) def squash( self, @@ -873,7 +850,6 @@ def squash( """ squasher = MigrationSquasher(self.migrations_path, self.runner, self._template_settings) - # Infer start/end from all sequential migrations when not provided if start_version is None or end_version is None: all_migrations = self.runner.get_migration_files() sequential = [(v, p) for v, p in all_migrations if v.isdigit() or v.lstrip("0").isdigit()] @@ -886,7 +862,6 @@ def squash( end_version = sequential[-1][0] console.print(f"[cyan]Squashing range: {start_version} to {end_version}[/]") - # Prompt for description when not provided if description is None: from rich.prompt import Prompt @@ -896,7 +871,6 @@ def squash( start_version, end_version, description, allow_gaps=allow_gaps, output_format=output_format ) - # Display plan for each squash group table = Table(title="Squash Plan") table.add_column("Version", style="cyan") table.add_column("File") @@ -1042,8 +1016,6 @@ def __init__(self, config: "AsyncConfigT") -> None: """ super().__init__(config) self.tracker = self._create_tracker() - - # Create context with extension configurations context = MigrationContext.from_config(config) context.extension_config = self.extension_configs @@ -1216,7 +1188,7 @@ async def record_version(exec_time: int, migration: "LoadedMigrationMetadata" = except Exception as exc: use_txn = self.runner.should_use_transaction(migration, self.config) rollback_msg = " (transaction rolled back)" if use_txn else "" - _output_exception( + _output_error( use_logger, echo, summary_only, @@ -1270,7 +1242,7 @@ async def remove_version(exec_time: int, version: str = version) -> None: except Exception as exc: use_txn = self.runner.should_use_transaction(migration, self.config) rollback_msg = " (transaction rolled back)" if use_txn else "" - _output_exception( + _output_error( use_logger, echo, summary_only, @@ -1439,7 +1411,6 @@ async def upgrade( await self._validate_migration_schema(driver) await self.tracker.ensure_tracking_table(driver) - # config auto_sync=False cannot be overridden by the call-site flag. if auto_sync and self.config.migration_config.get("auto_sync", True): await self._synchronize_version_records( driver, use_logger=ul, echo=echo_value, summary_only=summary_value @@ -1702,27 +1673,7 @@ async def revision(self, message: str, file_type: str | None = None) -> None: message: Description for the migration. file_type: Type of migration file to create ('sql' or 'py'). """ - version = generate_timestamp_version() - selected_format = file_type or self._template_settings.default_format - file_path = create_migration_file( - self.migrations_path, - version, - message, - selected_format, - config=self.config, - template_settings=self._template_settings, - ) - log_with_context( - logger, - logging.DEBUG, - "migration.create", - db_system=resolve_db_system(type(self.config).__name__), - version=version, - file_path=str(file_path), - file_type=selected_format, - description=message, - ) - console.print(f"[green]Created migration:[/] {file_path}") + _create_revision(self.migrations_path, self.config, self._template_settings, message, file_type) async def squash( self, @@ -1768,7 +1719,6 @@ async def squash( squasher = MigrationSquasher(self.migrations_path, sync_runner, self._template_settings) - # Infer start/end from all sequential migrations when not provided if start_version is None or end_version is None: all_migrations = sync_runner.get_migration_files() sequential = [(v, p) for v, p in all_migrations if v.isdigit() or v.lstrip("0").isdigit()] @@ -1781,7 +1731,6 @@ async def squash( end_version = sequential[-1][0] console.print(f"[cyan]Squashing range: {start_version} to {end_version}[/]") - # Prompt for description when not provided if description is None: import anyio from rich.prompt import Prompt @@ -1798,7 +1747,6 @@ async def squash( output_format=output_format, ) - # Display plan for each squash group table = Table(title="Squash Plan") table.add_column("Version", style="cyan") table.add_column("File") @@ -1967,16 +1915,34 @@ def _output_info( console.print(rich_message or message % args if args else message) -def _output_warning( - use_logger: bool, echo: bool, summary_only: bool, message: str, *args: Any, rich_message: str | None = None +def _create_revision( + migrations_path: "Path", config: Any, template_settings: Any, message: str, file_type: str | None = None ) -> None: - """Output a warning message to logger or console.""" - if use_logger: - logger.warning(message, *args) - else: - if not echo: - return - console.print(rich_message or message % args if args else message) + """Create a new migration file with timestamp-based versioning and log it. + + Args: + migrations_path: Path to migrations directory. + config: Database configuration. + template_settings: Migration template settings. + message: Description for the migration. + file_type: Type of migration file to create ('sql' or 'py'). + """ + version = generate_timestamp_version() + selected_format = file_type or template_settings.default_format + file_path = create_migration_file( + migrations_path, version, message, selected_format, config=config, template_settings=template_settings + ) + log_with_context( + logger, + logging.DEBUG, + "migration.create", + db_system=resolve_db_system(type(config).__name__), + version=version, + file_path=str(file_path), + file_type=selected_format, + description=message, + ) + console.print(f"[green]Created migration:[/] {file_path}") def _output_error( @@ -1991,18 +1957,6 @@ def _output_error( console.print(rich_message or message % args if args else message) -def _output_exception( - use_logger: bool, echo: bool, summary_only: bool, message: str, *args: Any, rich_message: str | None = None -) -> None: - """Output an exception message to logger or console.""" - if use_logger: - logger.error(message, *args) - else: - if not echo: - return - console.print(rich_message or message % args if args else message) - - def _log_command_summary( *, use_logger: bool, diff --git a/sqlspec/service.py b/sqlspec/service.py index ab318657d..4bb78d066 100644 --- a/sqlspec/service.py +++ b/sqlspec/service.py @@ -37,20 +37,6 @@ logger = get_logger("sqlspec.service") -class _TransactionState: - __slots__ = ("driver", "origin", "owner") - - def __init__(self, driver: AsyncDriverAdapterBase | SyncDriverAdapterBase) -> None: - self.driver: AsyncDriverAdapterBase | SyncDriverAdapterBase | None = driver - self.owner = _execution_owner() - self.origin = _owner_identity(self.owner) - - -_TRANSACTIONS: ContextVar[dict[object, _TransactionState] | None] = ContextVar( - "sqlspec_service_transactions", default=None -) - - def _execution_owner() -> tuple[int, object]: try: task = asyncio.current_task() @@ -63,78 +49,18 @@ def _owner_identity(owner: tuple[int, object]) -> tuple[int, int]: return owner[0], id(owner[1]) -def _active_transaction(key: object) -> _TransactionState | None: - state = (_TRANSACTIONS.get() or {}).get(key) - if state is None: - return None - if state.driver is None: - if state.origin == _owner_identity(_execution_owner()): - _discard_transaction(state) - return None - msg = "The inherited service transaction is no longer active." - raise ImproperConfigurationError(msg) - if state.owner != _execution_owner(): - msg = "A service transaction cannot be implicitly reused by another task or thread; pass session= explicitly." - raise ImproperConfigurationError(msg) - return state - - -def _live_transaction(key: object) -> _TransactionState | None: - current = _TRANSACTIONS.get() - state = None if current is None else current.get(key) - if state is None: - return None - if state.driver is None: - _discard_transaction(state) - return None - return _active_transaction(key) - - -def _discard_transaction(state: _TransactionState) -> None: - current = _TRANSACTIONS.get() - if current is not None and any(value is state for value in current.values()): - _TRANSACTIONS.set({key: value for key, value in current.items() if value is not state} or None) - - -def _session_transaction(state: _TransactionState | None) -> _TransactionState | None: - if state is None or state.driver is None: - return None - if state.owner != _execution_owner(): - msg = ( - "A service transaction is active in another task or thread; nested begin_transaction() blocks must " - "run in the task or thread that entered the outer block." - ) - raise ImproperConfigurationError(msg) - return state - - -def _connection_in_transaction(driver: AsyncDriverAdapterBase | SyncDriverAdapterBase) -> bool: - try: - return driver._connection_in_transaction() - except NotImplementedError: - return False - - -def _transaction_session(key: object) -> AsyncDriverAdapterBase | SyncDriverAdapterBase | None: - state = _active_transaction(key) - return None if state is None else state.driver - +class _TransactionState: + __slots__ = ("driver", "origin", "owner") -def _bind_transaction( - key: object, driver: AsyncDriverAdapterBase | SyncDriverAdapterBase -) -> tuple[_TransactionState, Token[dict[object, _TransactionState] | None]]: - state = _TransactionState(driver) - token = _TRANSACTIONS.set({**(_TRANSACTIONS.get() or {}), key: state}) - return state, token + def __init__(self, driver: AsyncDriverAdapterBase | SyncDriverAdapterBase) -> None: + self.driver: AsyncDriverAdapterBase | SyncDriverAdapterBase | None = driver + self.owner = _execution_owner() + self.origin = _owner_identity(self.owner) -def _release_transaction(state: _TransactionState, token: Token[dict[object, _TransactionState] | None]) -> None: - state.driver = None - state.owner = (0, None) - try: - _TRANSACTIONS.reset(token) - except ValueError: - _discard_transaction(state) +_TRANSACTIONS: ContextVar[dict[object, _TransactionState] | None] = ContextVar( + "sqlspec_service_transactions", default=None +) @mypyc_attr(allow_interpreted_subclasses=True) @@ -703,6 +629,80 @@ def begin_transaction(self) -> "_SyncBeginTransactionContext[SyncDriverT]": return _SyncBeginTransactionContext(self) +def _active_transaction(key: object) -> _TransactionState | None: + state = (_TRANSACTIONS.get() or {}).get(key) + if state is None: + return None + if state.driver is None: + if state.origin == _owner_identity(_execution_owner()): + _discard_transaction(state) + return None + msg = "The inherited service transaction is no longer active." + raise ImproperConfigurationError(msg) + if state.owner != _execution_owner(): + msg = "A service transaction cannot be implicitly reused by another task or thread; pass session= explicitly." + raise ImproperConfigurationError(msg) + return state + + +def _live_transaction(key: object) -> _TransactionState | None: + current = _TRANSACTIONS.get() + state = None if current is None else current.get(key) + if state is None: + return None + if state.driver is None: + _discard_transaction(state) + return None + return _active_transaction(key) + + +def _discard_transaction(state: _TransactionState) -> None: + current = _TRANSACTIONS.get() + if current is not None and any(value is state for value in current.values()): + _TRANSACTIONS.set({key: value for key, value in current.items() if value is not state} or None) + + +def _session_transaction(state: _TransactionState | None) -> _TransactionState | None: + if state is None or state.driver is None: + return None + if state.owner != _execution_owner(): + msg = ( + "A service transaction is active in another task or thread; nested begin_transaction() blocks must " + "run in the task or thread that entered the outer block." + ) + raise ImproperConfigurationError(msg) + return state + + +def _connection_in_transaction(driver: AsyncDriverAdapterBase | SyncDriverAdapterBase) -> bool: + try: + return driver._connection_in_transaction() + except NotImplementedError: + return False + + +def _transaction_session(key: object) -> AsyncDriverAdapterBase | SyncDriverAdapterBase | None: + state = _active_transaction(key) + return None if state is None else state.driver + + +def _bind_transaction( + key: object, driver: AsyncDriverAdapterBase | SyncDriverAdapterBase +) -> tuple[_TransactionState, Token[dict[object, _TransactionState] | None]]: + state = _TransactionState(driver) + token = _TRANSACTIONS.set({**(_TRANSACTIONS.get() or {}), key: state}) + return state, token + + +def _release_transaction(state: _TransactionState, token: Token[dict[object, _TransactionState] | None]) -> None: + state.driver = None + state.owner = (0, None) + try: + _TRANSACTIONS.reset(token) + except ValueError: + _discard_transaction(state) + + class _AsyncBeginTransactionContext(Generic[AsyncDriverT]): __slots__ = ("_nested", "_service", "_stack", "_state") diff --git a/sqlspec/typing.py b/sqlspec/typing.py index 9ac3378dd..ed5807cb8 100644 --- a/sqlspec/typing.py +++ b/sqlspec/typing.py @@ -45,11 +45,9 @@ StructStub, UnsetType, convert, - import_optional, - import_optional_attr, - module_available, msgspec_fields, ) +from sqlspec.utils.module_loader import import_optional, import_optional_attr, module_available if TYPE_CHECKING: from sqlspec._typing import ( diff --git a/sqlspec/utils/arrow_helpers.py b/sqlspec/utils/arrow_helpers.py index 4c2754f1f..217030da7 100644 --- a/sqlspec/utils/arrow_helpers.py +++ b/sqlspec/utils/arrow_helpers.py @@ -170,30 +170,6 @@ def arrow_type_from_token(token: str) -> Any: } -def _resolve_null_columns(table: "ArrowTable", column_types: "Mapping[str, str] | None" = None) -> "ArrowTable": - """Give value-inferred Arrow ``null`` columns their real type. - - ``pa.Table.from_pylist`` infers column types from values alone, so a column - that is ``NULL`` in every row arrives as Arrow ``null`` and loses the type - the database declared. When the adapter reported that column's type, use it. - Otherwise fall back to ``string``, which holds the nulls without breaking - string operations downstream. - """ - import pyarrow as pa - - if not any(pa.types.is_null(field.type) for field in table.schema): - return table - - hints = column_types or {} - fields = [ - field.with_type(arrow_type_from_token(hints[field.name])) - if pa.types.is_null(field.type) and field.name in hints - else (field.with_type(pa.string()) if pa.types.is_null(field.type) else field) - for field in table.schema - ] - return table.cast(pa.schema(fields)) - - def convert_dict_to_arrow_with_schema( data: "list[dict[str, Any]]", return_format: Literal["table", "reader", "batch", "batches"] = "table", @@ -302,29 +278,6 @@ def coerce_arrow_table(source: "ArrowResult | Any") -> "ArrowTable": raise TypeError(msg) -def _coerce_arrow_table_identity(source: Any) -> Any: - return source - - -def _coerce_arrow_record_batch(source: Any) -> Any: - import pyarrow as pa - - return pa.Table.from_batches([source]) - - -def _get_arrow_table_coercer() -> "TypeDispatcher[Any]": - global _ARROW_TABLE_COERCER - if _ARROW_TABLE_COERCER is None: - ensure_pyarrow() - import pyarrow as pa - - dispatcher = TypeDispatcher[Any]() - dispatcher.register(pa.Table, _coerce_arrow_table_identity) - dispatcher.register(pa.RecordBatch, _coerce_arrow_record_batch) - _ARROW_TABLE_COERCER = dispatcher - return _ARROW_TABLE_COERCER - - def ensure_arrow_table(data: Any) -> "ArrowTable": """Ensure data is a PyArrow Table.""" ensure_pyarrow() @@ -400,48 +353,14 @@ def arrow_table_to_rows( msg = "Arrow table has no columns to import" raise ValueError(msg) - # Use column-oriented access with zip transpose (O(n) vs O(n*m) row iteration) - # Extract columns as Python lists and transpose to rows col_data = [table.column(col).to_pylist() for col in resolved_columns] - - # Handle empty table case if not col_data or not col_data[0]: return resolved_columns, [] - # Transpose columns to rows using zip records: list[tuple[Any, ...]] = [tuple(row) for row in zip(*col_data, strict=False)] return resolved_columns, records -def _arrow_type_needs_preparation(data_type: Any) -> bool: - ensure_pyarrow() - import pyarrow as pa - - if pa.types.is_dictionary(data_type): - return _arrow_type_needs_preparation(data_type.value_type) - - is_extension = getattr(pa.types, "is_extension", None) - if is_extension is not None and is_extension(data_type): - return _arrow_type_needs_preparation(data_type.storage_type) - - nested_type_checks = ( - "is_struct", - "is_list", - "is_large_list", - "is_fixed_size_list", - "is_map", - "is_union", - "is_list_view", - "is_large_list_view", - ) - return any(getattr(pa.types, check, lambda _: False)(data_type) for check in nested_type_checks) - - -@lru_cache(maxsize=_ARROW_SCHEMA_DECISION_CACHE_SIZE) -def _arrow_schema_needs_preparation(schema: Any) -> bool: - return any(_arrow_type_needs_preparation(field.type) for field in schema) - - def arrow_table_needs_parameter_preparation(table: "ArrowTable") -> bool: """Return whether Arrow rows may emit nested values needing driver preparation.""" return _arrow_schema_needs_preparation(table.schema) @@ -510,36 +429,6 @@ def build_ingest_telemetry(table: "ArrowTable", *, format_label: str = "arrow") return {"rows_processed": rows, "bytes_processed": bytes_processed, "format": format_label} -class _DeferredCloseBatchIterator: - """Iterate RecordBatches from a reader and invoke a callback on exhaustion or error.""" - - __slots__ = ("_close_callback", "_closed", "_reader") - - def __init__(self, reader: Any, close_callback: "Callable[[], None]") -> None: - self._reader = reader - self._close_callback = close_callback - self._closed = False - - def __iter__(self) -> "_DeferredCloseBatchIterator": - return self - - def _finalize(self) -> None: - if not self._closed: - self._closed = True - with contextlib.suppress(Exception): - self._close_callback() - - def __next__(self) -> Any: - try: - return self._reader.read_next_batch() - except StopIteration: - self._finalize() - raise - except Exception: - self._finalize() - raise - - def arrow_reader_with_deferred_close(reader: Any, close_callback: "Callable[[], None]") -> "ArrowRecordBatchReader": """Wrap a RecordBatchReader so close_callback fires when it is exhausted or errors.""" ensure_pyarrow() @@ -615,3 +504,109 @@ def _arrow_uuid_column_to_pylist(column: Any, data_type: Any) -> "list[Any]": for value in nested_values ] return cast("list[Any]", column.to_pylist()) + + +def _resolve_null_columns(table: "ArrowTable", column_types: "Mapping[str, str] | None" = None) -> "ArrowTable": + """Give value-inferred Arrow ``null`` columns their real type. + + ``pa.Table.from_pylist`` infers column types from values alone, so a column + that is ``NULL`` in every row arrives as Arrow ``null`` and loses the type + the database declared. When the adapter reported that column's type, use it. + Otherwise fall back to ``string``, which holds the nulls without breaking + string operations downstream. + """ + import pyarrow as pa + + if not any(pa.types.is_null(field.type) for field in table.schema): + return table + + hints = column_types or {} + fields = [ + field.with_type(arrow_type_from_token(hints[field.name])) + if pa.types.is_null(field.type) and field.name in hints + else (field.with_type(pa.string()) if pa.types.is_null(field.type) else field) + for field in table.schema + ] + return table.cast(pa.schema(fields)) + + +def _coerce_arrow_table_identity(source: Any) -> Any: + return source + + +def _coerce_arrow_record_batch(source: Any) -> Any: + import pyarrow as pa + + return pa.Table.from_batches([source]) + + +def _get_arrow_table_coercer() -> "TypeDispatcher[Any]": + global _ARROW_TABLE_COERCER + if _ARROW_TABLE_COERCER is None: + ensure_pyarrow() + import pyarrow as pa + + dispatcher = TypeDispatcher[Any]() + dispatcher.register(pa.Table, _coerce_arrow_table_identity) + dispatcher.register(pa.RecordBatch, _coerce_arrow_record_batch) + _ARROW_TABLE_COERCER = dispatcher + return _ARROW_TABLE_COERCER + + +def _arrow_type_needs_preparation(data_type: Any) -> bool: + ensure_pyarrow() + import pyarrow as pa + + if pa.types.is_dictionary(data_type): + return _arrow_type_needs_preparation(data_type.value_type) + + is_extension = getattr(pa.types, "is_extension", None) + if is_extension is not None and is_extension(data_type): + return _arrow_type_needs_preparation(data_type.storage_type) + + nested_type_checks = ( + "is_struct", + "is_list", + "is_large_list", + "is_fixed_size_list", + "is_map", + "is_union", + "is_list_view", + "is_large_list_view", + ) + return any(getattr(pa.types, check, lambda _: False)(data_type) for check in nested_type_checks) + + +@lru_cache(maxsize=_ARROW_SCHEMA_DECISION_CACHE_SIZE) +def _arrow_schema_needs_preparation(schema: Any) -> bool: + return any(_arrow_type_needs_preparation(field.type) for field in schema) + + +class _DeferredCloseBatchIterator: + """Iterate RecordBatches from a reader and invoke a callback on exhaustion or error.""" + + __slots__ = ("_close_callback", "_closed", "_reader") + + def __init__(self, reader: Any, close_callback: "Callable[[], None]") -> None: + self._reader = reader + self._close_callback = close_callback + self._closed = False + + def __iter__(self) -> "_DeferredCloseBatchIterator": + return self + + def _finalize(self) -> None: + if not self._closed: + self._closed = True + with contextlib.suppress(Exception): + self._close_callback() + + def __next__(self) -> Any: + try: + return self._reader.read_next_batch() + except StopIteration: + self._finalize() + raise + except Exception: + self._finalize() + raise diff --git a/sqlspec/utils/logging.py b/sqlspec/utils/logging.py index 5cd8487b3..eecc66283 100644 --- a/sqlspec/utils/logging.py +++ b/sqlspec/utils/logging.py @@ -95,6 +95,8 @@ def format(self, record: LogRecord) -> str: Returns: JSON formatted log entry """ + from sqlspec.utils.serializers import to_json + record_dict = record.__dict__ log_entry = { "timestamp": self.formatTime(record, self.datefmt), @@ -132,8 +134,6 @@ def format(self, record: LogRecord) -> str: if record.exc_info: log_entry["exception"] = self.formatException(record.exc_info) - from sqlspec.utils.serializers import to_json - return to_json(log_entry) diff --git a/sqlspec/utils/schema.py b/sqlspec/utils/schema.py index 03ac34c9f..2fbda47b2 100644 --- a/sqlspec/utils/schema.py +++ b/sqlspec/utils/schema.py @@ -136,6 +136,47 @@ def _transform_list(data: list, converter: Callable[[str], str]) -> list: # ============================================================================= +@overload +def to_schema(data: "list[DataT]", *, schema_type: "type[SchemaT]") -> "list[SchemaT]": ... +@overload +def to_schema(data: "list[DataT]", *, schema_type: None = None) -> "list[DataT]": ... +@overload +def to_schema(data: "DataT", *, schema_type: "type[SchemaT]") -> "SchemaT": ... +@overload +def to_schema(data: "DataT", *, schema_type: None = None) -> "DataT": ... + + +def to_schema(data: Any, *, schema_type: Any = None) -> Any: + """Convert data to a specified schema type. + + Supports transformation to various schema types including: + - TypedDict + - dataclasses + - msgspec Structs + - Pydantic models + - attrs classes + + Args: + data: Input data to convert (dict, list of dicts, or other) + schema_type: Target schema type for conversion. If None, returns data unchanged. + + Returns: + Converted data in the specified schema type, or original data if schema_type is None + + Raises: + SQLSpecError: If schema_type is not a supported type + """ + if schema_type is None: + return data + + conv = _get_schema_converter(schema_type) + if conv is None: + msg = "`schema_type` should be a valid Dataclass, Pydantic model, Msgspec struct, Attrs class, or TypedDict" + raise SQLSpecError(msg) + + return conv(data, schema_type) + + def _is_list_type_target(target_type: Any) -> "TypeGuard[list[object]]": """Check if target type is a list type.""" try: @@ -486,8 +527,6 @@ def _convert_attrs(data: Any, schema_type: Any) -> Any: return schema_type(**data) if is_dict(data) else data -# Cache for schema converters - maps type directly to converter callable (or None if unsupported) -# Manual dict cache is faster than lru_cache for mypyc: direct dict[type] lookup vs decorated call _SCHEMA_CONVERTER_CACHE: "dict[type, Callable[[Any, Any], Any] | None]" = {} @@ -506,7 +545,6 @@ def _get_schema_converter(schema_type: type) -> "Callable[[Any, Any], Any] | Non try: return _SCHEMA_CONVERTER_CACHE[schema_type] except KeyError: - # Determine converter - order by expected frequency if is_typed_dict(schema_type): conv: Callable[[Any, Any], Any] | None = _convert_typed_dict elif is_dataclass(schema_type): @@ -525,51 +563,80 @@ def _get_schema_converter(schema_type: type) -> "Callable[[Any, Any], Any] | Non return conv -@overload -def to_schema(data: "list[DataT]", *, schema_type: "type[SchemaT]") -> "list[SchemaT]": ... -@overload -def to_schema(data: "list[DataT]", *, schema_type: None = None) -> "list[DataT]": ... -@overload -def to_schema(data: "DataT", *, schema_type: "type[SchemaT]") -> "SchemaT": ... -@overload -def to_schema(data: "DataT", *, schema_type: None = None) -> "DataT": ... +# ============================================================================= +# Scalar Type Conversion +# ============================================================================= -def to_schema(data: Any, *, schema_type: Any = None) -> Any: - """Convert data to a specified schema type. +def to_value_type(value: Any, value_type: "type[ValueT]") -> "ValueT": + """Convert a database value to the specified Python type. - Supports transformation to various schema types including: - - TypedDict - - dataclasses - - msgspec Structs - - Pydantic models - - attrs classes + This function handles type conversion for common database return values, + providing runtime type safety for scalar queries. When the value is already + the correct type, it is returned as-is without conversion overhead. Strict + type identity check handles subclass gotchas (bool is subclass of int, + datetime is subclass of date). + + Also supports schema types (Pydantic models, dataclasses, msgspec Structs, + attrs classes, and TypedDict). For schema types, JSON strings are automatically + parsed before conversion. Args: - data: Input data to convert (dict, list of dicts, or other) - schema_type: Target schema type for conversion. If None, returns data unchanged. + value: The value to convert. + value_type: The target Python type. Supported types include: + + - Primitives: int, float, str, bool + - Temporal: datetime, date, time + - Numeric: Decimal + - Identifiers: UUID, Path + - Collections: dict, list (for JSON/JSONB columns) + - Schema types: Pydantic models, dataclasses, msgspec Structs, + attrs classes, TypedDict (for JSONB columns) Returns: - Converted data in the specified schema type, or original data if schema_type is None + The converted value of the specified type. Raises: - SQLSpecError: If schema_type is not a supported type + TypeError: If the value cannot be converted to the specified type. """ - if schema_type is None: - return data - - # Get cached converter - single dict lookup, no string indirection - conv = _get_schema_converter(schema_type) - if conv is None: - msg = "`schema_type` should be a valid Dataclass, Pydantic model, Msgspec struct, Attrs class, or TypedDict" - raise SQLSpecError(msg) + if type(value) is value_type: + return value - return conv(data, schema_type) + if value_type is int: + return cast("ValueT", _convert_to_int(value)) + if value_type is str: + return cast("ValueT", str(value)) + if value_type is float: + return cast("ValueT", _convert_to_float(value)) + if value_type is bool: + return cast("ValueT", _convert_to_bool(value)) + if value_type is datetime.datetime: + return cast("ValueT", _convert_to_datetime(value)) + if value_type is datetime.date: + return cast("ValueT", _convert_to_date(value)) + if value_type is datetime.time: + return cast("ValueT", _convert_to_time(value)) + if value_type is Decimal: + return cast("ValueT", _convert_to_decimal(value)) + if value_type is UUID: + return cast("ValueT", _convert_to_uuid(value)) + if value_type is Path: + return cast("ValueT", _convert_to_path(value)) + if value_type is dict: + return cast("ValueT", _convert_to_dict(value)) + if value_type is list: + return cast("ValueT", _convert_to_list(value)) + schema_converter = _get_schema_converter(value_type) + if schema_converter is not None: + parsed = _ensure_json_parsed(value) + return cast("ValueT", schema_converter(parsed, value_type)) -# ============================================================================= -# Scalar Type Conversion -# ============================================================================= + try: + return value_type(value) # type: ignore[call-arg] + except (TypeError, ValueError) as e: + msg = f"Cannot convert {type(value).__name__} to {value_type.__name__}" + raise TypeError(msg) from e def _ensure_json_parsed(value: Any) -> Any: @@ -607,12 +674,7 @@ def _try_parse_json(value: str) -> Any: return None -# Boolean true values for string conversion _BOOL_TRUE_VALUES: Final[frozenset[str]] = frozenset({"true", "1", "yes", "y", "t", "on"}) - -# Types requiring strict type() identity check due to subclass gotchas: -# - bool is subclass of int (isinstance(True, int) is True) -# - datetime is subclass of date (isinstance(datetime(...), date) is True) _STRICT_IDENTITY_TYPES: Final[tuple[type, ...]] = (int, bool, datetime.date, datetime.time) @@ -636,7 +698,6 @@ def _convert_to_int(value: Any) -> int: try: return int(value) except ValueError: - # Try parsing as float first for values like "42.0" try: return int(float(value)) except ValueError: @@ -735,10 +796,8 @@ def _convert_to_date(value: Any) -> datetime.date: return value if isinstance(value, str): try: - # Try ISO format first return datetime.date.fromisoformat(value) except ValueError: - # Try parsing as datetime and extracting date try: return datetime.datetime.fromisoformat(value).date() except ValueError: @@ -899,80 +958,3 @@ def _convert_to_list(value: Any) -> list[Any]: return list(value) msg = f"Cannot convert {type(value).__name__} to list" raise TypeError(msg) - - -def to_value_type(value: Any, value_type: "type[ValueT]") -> "ValueT": - """Convert a database value to the specified Python type. - - This function handles type conversion for common database return values, - providing runtime type safety for scalar queries. When the value is already - the correct type, it is returned as-is without conversion overhead. - - Also supports schema types (Pydantic models, dataclasses, msgspec Structs, - attrs classes, and TypedDict). For schema types, JSON strings are automatically - parsed before conversion. - - Args: - value: The value to convert. - value_type: The target Python type. Supported types include: - - - Primitives: int, float, str, bool - - Temporal: datetime, date, time - - Numeric: Decimal - - Identifiers: UUID, Path - - Collections: dict, list (for JSON/JSONB columns) - - Schema types: Pydantic models, dataclasses, msgspec Structs, - attrs classes, TypedDict (for JSONB columns) - - Returns: - The converted value of the specified type. - - Raises: - TypeError: If the value cannot be converted to the specified type. - """ - # Fast path: already correct type (handles ~90% of cases) - # Uses strict type() identity which correctly handles subclass gotchas: - # - type(True) is int → False (bool is subclass of int) - # - type(datetime(...)) is date → False (datetime is subclass of date) - if type(value) is value_type: - return value - - # Scalar type conversions - most common cases, fast pointer comparisons under mypyc - if value_type is int: - return cast("ValueT", _convert_to_int(value)) - if value_type is str: - return cast("ValueT", str(value)) - if value_type is float: - return cast("ValueT", _convert_to_float(value)) - if value_type is bool: - return cast("ValueT", _convert_to_bool(value)) - if value_type is datetime.datetime: - return cast("ValueT", _convert_to_datetime(value)) - if value_type is datetime.date: - return cast("ValueT", _convert_to_date(value)) - if value_type is datetime.time: - return cast("ValueT", _convert_to_time(value)) - if value_type is Decimal: - return cast("ValueT", _convert_to_decimal(value)) - if value_type is UUID: - return cast("ValueT", _convert_to_uuid(value)) - if value_type is Path: - return cast("ValueT", _convert_to_path(value)) - if value_type is dict: - return cast("ValueT", _convert_to_dict(value)) - if value_type is list: - return cast("ValueT", _convert_to_list(value)) - - # Schema types (Pydantic, dataclass, msgspec, attrs, TypedDict) - # Deferred after scalar checks to avoid overhead for common scalar queries - schema_converter = _get_schema_converter(value_type) - if schema_converter is not None: - parsed = _ensure_json_parsed(value) - return cast("ValueT", schema_converter(parsed, value_type)) - - # Fallback: try direct construction for custom types - try: - return value_type(value) # type: ignore[call-arg] - except (TypeError, ValueError) as e: - msg = f"Cannot convert {type(value).__name__} to {value_type.__name__}" - raise TypeError(msg) from e diff --git a/sqlspec/utils/type_converters.py b/sqlspec/utils/type_converters.py index 73e45a6be..4db8df24d 100644 --- a/sqlspec/utils/type_converters.py +++ b/sqlspec/utils/type_converters.py @@ -205,12 +205,10 @@ def build_uuid_coercions(*, native: bool = False) -> "dict[type[Any], Callable[[ When ``False`` (default), convert to ``str`` (for drivers that need a plain string. """ - import uuid as _uuid_mod - coercions: dict[type[Any], Callable[[Any], Any]] = {} if not native: - coercions[_uuid_mod.UUID] = _uuid_to_string + coercions[UUID] = _uuid_to_string uuid_utils_uuid = import_optional_attr("uuid_utils", "UUID") if uuid_utils_uuid is not None: diff --git a/tests/unit/adapters/test_adbc/test_extension_detection.py b/tests/unit/adapters/test_adbc/test_extension_detection.py index 050acb361..a00306af4 100644 --- a/tests/unit/adapters/test_adbc/test_extension_detection.py +++ b/tests/unit/adapters/test_adbc/test_extension_detection.py @@ -144,56 +144,57 @@ def test_resolve_postgres_extension_state_promotes_paradedb() -> None: assert paradedb_available is True -def test_adbc_config_update_dialect_for_extensions_pgvector() -> None: +def test_adbc_config_resolve_dialect_for_extensions_pgvector() -> None: """Dialect switches to pgvector when pgvector is available.""" config = AdbcConfig(connection_config={"uri": "postgresql://localhost/test"}) - config._pgvector_available = True # pyright: ignore[reportPrivateUsage] - config._paradedb_available = False # pyright: ignore[reportPrivateUsage] - config._pg_textsearch_available = False # pyright: ignore[reportPrivateUsage] - config._update_dialect_for_extensions() # pyright: ignore[reportPrivateUsage] - assert config.statement_config.dialect == "pgvector" + statement_config, pgvector, paradedb = resolve_postgres_extension_state( + config.statement_config, config.driver_features, {"vector"} + ) + assert statement_config.dialect == "pgvector" + assert pgvector is True + assert paradedb is False -def test_adbc_config_update_dialect_for_extensions_pg_textsearch() -> None: +def test_adbc_config_resolve_dialect_for_extensions_pg_textsearch() -> None: """Dialect switches to pg_textsearch when pg_textsearch is available.""" - config = AdbcConfig(connection_config={"uri": "postgresql://localhost/test"}) - config._pgvector_available = True # pyright: ignore[reportPrivateUsage] - config._paradedb_available = False # pyright: ignore[reportPrivateUsage] - config._pg_textsearch_available = True # pyright: ignore[reportPrivateUsage] - config._update_dialect_for_extensions() # pyright: ignore[reportPrivateUsage] - assert config.statement_config.dialect == "pg_textsearch" - assert config.pg_textsearch_available is True + config = AdbcConfig( + connection_config={"uri": "postgresql://localhost/test"}, driver_features={"enable_pg_textsearch": True} + ) + statement_config, _, _ = resolve_postgres_extension_state( + config.statement_config, config.driver_features, {"vector", "pg_textsearch"} + ) + assert statement_config.dialect == "pg_textsearch" -def test_adbc_config_update_dialect_for_extensions_paradedb() -> None: +def test_adbc_config_resolve_dialect_for_extensions_paradedb() -> None: """Dialect switches to paradedb when both extensions available (paradedb > pgvector).""" - config = AdbcConfig(connection_config={"uri": "postgresql://localhost/test"}) - config._pgvector_available = True # pyright: ignore[reportPrivateUsage] - config._paradedb_available = True # pyright: ignore[reportPrivateUsage] - config._pg_textsearch_available = True # pyright: ignore[reportPrivateUsage] - config._update_dialect_for_extensions() # pyright: ignore[reportPrivateUsage] - assert config.statement_config.dialect == "paradedb" + config = AdbcConfig( + connection_config={"uri": "postgresql://localhost/test"}, driver_features={"enable_pg_textsearch": True} + ) + statement_config, _, _ = resolve_postgres_extension_state( + config.statement_config, config.driver_features, {"vector", "pg_search", "pg_textsearch"} + ) + assert statement_config.dialect == "paradedb" -def test_adbc_config_update_dialect_skips_non_postgres() -> None: +def test_adbc_config_resolve_dialect_skips_non_postgres() -> None: """Dialect is not changed for non-postgres backends.""" config = AdbcConfig(connection_config={"uri": ":memory:", "driver_name": "sqlite"}) - original_dialect = config.statement_config.dialect - config._pgvector_available = True # pyright: ignore[reportPrivateUsage] - config._paradedb_available = True # pyright: ignore[reportPrivateUsage] - config._update_dialect_for_extensions() # pyright: ignore[reportPrivateUsage] - assert config.statement_config.dialect == original_dialect + statement_config, _, _ = resolve_postgres_extension_state( + config.statement_config, config.driver_features, {"vector", "pg_search"} + ) + assert statement_config.dialect == config.statement_config.dialect -def test_adbc_config_update_dialect_preserves_custom_dialect() -> None: +def test_adbc_config_resolve_dialect_preserves_custom_dialect() -> None: """If user explicitly set a non-postgres dialect, don't override it.""" config = AdbcConfig( connection_config={"uri": "postgresql://localhost/test"}, statement_config=StatementConfig(dialect="custom") ) - config._pgvector_available = True # pyright: ignore[reportPrivateUsage] - config._paradedb_available = True # pyright: ignore[reportPrivateUsage] - config._update_dialect_for_extensions() # pyright: ignore[reportPrivateUsage] - assert config.statement_config.dialect == "custom" + statement_config, _, _ = resolve_postgres_extension_state( + config.statement_config, config.driver_features, {"vector", "pg_search"} + ) + assert statement_config.dialect == "custom" def test_adbc_config_provide_session_skips_extension_probe_for_non_postgres(monkeypatch: MonkeyPatch) -> None: diff --git a/tests/unit/builder/test_values.py b/tests/unit/builder/test_values.py index 4cdcc23e4..7156e4eb1 100644 --- a/tests/unit/builder/test_values.py +++ b/tests/unit/builder/test_values.py @@ -130,14 +130,6 @@ def test_values_as_and_set_columns() -> None: val.set_columns("only_one") -def test_values_expected_result_type() -> None: - """Test expected result type property.""" - from sqlspec.core import SQLResult - - val = Values([(1, "a")]) - assert val._expected_result_type == SQLResult - - def test_values_alias_without_columns() -> None: """Test Values with alias but without column list.""" val = Values([(1, "a")], alias="v") diff --git a/tests/unit/driver/test_query_cache.py b/tests/unit/driver/test_query_cache.py index 558e84ffe..0dedeb24d 100644 --- a/tests/unit/driver/test_query_cache.py +++ b/tests/unit/driver/test_query_cache.py @@ -218,13 +218,13 @@ def test_stmt_cache_rebind_reuses_driver_owned_processor(sqlite_sync_driver: Any ) processor = sqlite_sync_driver._stmt_cache_rebind_processor calls: list[object] = [] - original_transform = ParameterProcessor._transform_cached_parameters + original_transform = ParameterProcessor.transform_cached_parameters def wrapped_transform(self: ParameterProcessor, *args: Any, **kwargs: Any) -> Any: calls.append(self) return original_transform(self, *args, **kwargs) - monkeypatch.setattr(ParameterProcessor, "_transform_cached_parameters", wrapped_transform) + monkeypatch.setattr(ParameterProcessor, "transform_cached_parameters", wrapped_transform) sqlite_sync_driver.stmt_cache_rebind({"id": 1}, cached) sqlite_sync_driver.stmt_cache_rebind({"id": 2}, cached) diff --git a/tests/unit/migrations/test_tracker_idempotency.py b/tests/unit/migrations/test_tracker_idempotency.py index 03a29550b..f86ade7ba 100644 --- a/tests/unit/migrations/test_tracker_idempotency.py +++ b/tests/unit/migrations/test_tracker_idempotency.py @@ -46,7 +46,6 @@ def test_sync_tracker_qualifies_table_sql_when_schema_is_configured() -> None: str( tracker._record_squashed_migration_statement("0002", "sequential", 2, "squash", 0, "def", "tester", "0001") ), # pyright: ignore[reportPrivateUsage] - str(tracker._column_exists_query()), # pyright: ignore[reportPrivateUsage] ] assert all('"history"."ddl_migrations"' in statement for statement in rendered_statements)