sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/dialects/postgres/_pgvector.py

@@ -0,0 +1,36 @@
+"""PGVector dialect extending Postgres with vector distance operators.
+
+Adds support for pgvector distance operators:
+- <-> : L2 (Euclidean) distance (already in base Postgres)
+- <#> : Negative inner product
+- <=> : Cosine distance
+- <+> : L1 (Taxicab/Manhattan) distance
+- <~> : Hamming distance (binary vectors)
+- <%> : Jaccard distance (binary vectors)
+"""
+
+from sqlglot.dialects.dialect import Dialect
+from sqlglot.dialects.postgres import Postgres
+
+from sqlspec.dialects.postgres._generators import PGVectorGenerator
+from sqlspec.dialects.postgres._operators import PGVECTOR_OPERATOR_TOKENS, register_postgres_extension_operators
+
+__all__ = ("PGVector",)
+
+register_postgres_extension_operators()
+
+
+class PGVectorTokenizer(Postgres.Tokenizer):
+    """Tokenizer with pgvector distance operators."""
+
+    KEYWORDS = {**Postgres.Tokenizer.KEYWORDS, **PGVECTOR_OPERATOR_TOKENS}
+
+
+class PGVector(Postgres):
+    """PostgreSQL dialect with pgvector extension support."""
+
+    Tokenizer = PGVectorTokenizer
+    Generator = PGVectorGenerator
+
+
+Dialect.classes["pgvector"] = PGVector

sqlspec/dialects/spanner/__init__.py

@@ -0,0 +1,6 @@
+"""Google Cloud Spanner sqlglot dialects."""
+
+from sqlspec.dialects.spanner._spangres import Spangres
+from sqlspec.dialects.spanner._spanner import Spanner
+
+__all__ = ("Spangres", "Spanner")

sqlspec/dialects/spanner/_generators.py

@@ -0,0 +1,206 @@
+"""Spanner dialect generators with sqlglot[c] compatibility.
+
+When sqlglot[c] is installed, generators are compiled and reject interpreted
+subclasses. We detect this and use TRANSFORMS-based extension on the base
+generators. When running pure-Python sqlglot, we create real subclasses.
+"""
+
+from typing import Any, cast
+
+from sqlglot import exp
+from sqlglot.generator import _DISPATCH_CACHE  # pyright: ignore[reportPrivateUsage]
+from sqlglot.generators.bigquery import BigQueryGenerator
+from sqlglot.generators.postgres import PostgresGenerator
+
+__all__ = ("SpangresGenerator", "SpannerGenerator")
+
+_TTL_MIN_COMPONENTS = 2
+_ROW_DELETION_NAME = "ROW_DELETION_POLICY"
+_INTERLEAVE_NAME = "INTERLEAVE_IN_PARENT"
+
+# Capture originals before any patching
+_original_bq_property_sql = BigQueryGenerator.property_sql
+_original_bq_properties_sql = BigQueryGenerator.properties_sql
+_original_bq_locate_properties = BigQueryGenerator.locate_properties
+_original_pg_property_sql = PostgresGenerator.property_sql
+
+
+def _is_post_schema_spanner_property(expression: exp.Expr) -> bool:
+    if not isinstance(expression, exp.Property) or not isinstance(expression.this, exp.Literal):
+        return False
+    return expression.this.name.upper() in {_INTERLEAVE_NAME, _ROW_DELETION_NAME, "TTL"}
+
+
+def _render_interval_sql(generator: Any, expression: exp.Expr) -> str:
+    if isinstance(expression, exp.Interval):
+        unit = expression.args.get("unit")
+        if isinstance(expression.this, exp.Literal) and not expression.this.is_string and isinstance(unit, exp.Expr):
+            return f"INTERVAL {generator.sql(expression.this)} {generator.sql(unit)}"
+
+    interval_sql = cast("str", generator.sql(expression))
+    if not interval_sql.upper().startswith("INTERVAL"):
+        return f"INTERVAL {interval_sql}"
+    return interval_sql
+
+
+# ---------------------------------------------------------------------------
+# Spanner property rendering (BigQuery-based)
+# ---------------------------------------------------------------------------
+
+
+def _spanner_property_sql(self: Any, expression: exp.Property) -> str:
+    """Render Spanner-specific properties (INTERLEAVE, ROW_DELETION_POLICY, TTL)."""
+    if isinstance(expression.this, exp.Literal) and expression.this.name.upper() == _INTERLEAVE_NAME:
+        values = expression.args.get("value")
+        if isinstance(values, exp.Tuple) and values.expressions:
+            parent = self.sql(values.expressions[0])
+            sql = f"INTERLEAVE IN PARENT {parent}"
+            if len(values.expressions) >= _TTL_MIN_COMPONENTS:
+                on_delete_expr = values.expressions[1]
+                if isinstance(on_delete_expr, exp.Literal):
+                    sql = f"{sql} ON DELETE {on_delete_expr.this}"
+            return sql
+
+    if isinstance(expression.this, exp.Literal) and expression.this.name.upper() == _ROW_DELETION_NAME:
+        values = expression.args.get("value")
+        if isinstance(values, exp.Tuple) and len(values.expressions) >= _TTL_MIN_COMPONENTS:
+            column = self.sql(values.expressions[0])
+            interval_sql = _render_interval_sql(self, values.expressions[1])
+            return f"ROW DELETION POLICY (OLDER_THAN({column}, {interval_sql}))"
+
+    if isinstance(expression.this, exp.Literal) and expression.this.name.upper() == "TTL":
+        values = expression.args.get("value")
+        if isinstance(values, exp.Tuple) and len(values.expressions) >= _TTL_MIN_COMPONENTS:
+            interval = _render_interval_sql(self, values.expressions[0]).removeprefix("INTERVAL ")
+            column = self.sql(values.expressions[1])
+            return f"TTL INTERVAL {interval} ON {column}"
+
+    return str(_original_bq_property_sql(self, expression))
+
+
+def _spanner_properties_sql(self: Any, expression: exp.Properties) -> str:
+    """Render custom Spanner properties without BigQuery's OPTIONS wrapper."""
+    root_properties: list[exp.Expr] = []
+    with_properties: list[exp.Expr] = []
+
+    for property_expression in expression.expressions:
+        if _is_post_schema_spanner_property(property_expression):
+            root_properties.append(property_expression)
+            continue
+
+        property_location = self.PROPERTIES_LOCATION[property_expression.__class__]
+        if property_location == exp.Properties.Location.POST_WITH:
+            with_properties.append(property_expression)
+        elif property_location == exp.Properties.Location.POST_SCHEMA:
+            root_properties.append(property_expression)
+
+    root_props_ast = exp.Properties(expressions=root_properties)
+    root_props_ast.parent = expression.parent
+    with_props_ast = exp.Properties(expressions=with_properties)
+    with_props_ast.parent = expression.parent
+
+    root_props = self.root_properties(root_props_ast)
+    with_props = self.with_properties(with_props_ast)
+
+    if root_props and with_props and not self.pretty:
+        with_props = f" {with_props}"
+
+    return str(root_props) + str(with_props)
+
+
+# ---------------------------------------------------------------------------
+# Spangres property rendering (Postgres-based)
+# ---------------------------------------------------------------------------
+
+
+def _spangres_property_sql(self: Any, expression: exp.Property) -> str:
+    """Render Spangres row deletion policies."""
+    if isinstance(expression.this, exp.Literal) and expression.this.name.upper() == _ROW_DELETION_NAME:
+        values = expression.args.get("value")
+        if isinstance(values, exp.Tuple) and len(values.expressions) >= _TTL_MIN_COMPONENTS:
+            column = self.sql(values.expressions[0])
+            interval_sql = _render_interval_sql(self, values.expressions[1])
+            return f"ROW DELETION POLICY (OLDER_THAN({column}, {interval_sql}))"
+
+    return str(_original_pg_property_sql(self, expression))
+
+
+# ---------------------------------------------------------------------------
+# Unified extension logic (sqlglot[c] & Pure-Python)
+# ---------------------------------------------------------------------------
+
+# BigQuery / Spanner
+_original_bq_property_transform = BigQueryGenerator.TRANSFORMS.get(exp.Property)
+_original_bq_properties_transform = BigQueryGenerator.TRANSFORMS.get(exp.Properties)
+_original_bq_create_transform = BigQueryGenerator.TRANSFORMS.get(exp.Create)
+
+
+def _bq_property_transform(self: Any, expression: exp.Property) -> str:
+    dialect_class = getattr(self.dialect, "__class__", None)
+    dialect_name = dialect_class.__name__ if dialect_class else None
+    if dialect_name == "Spanner":
+        return _spanner_property_sql(self, expression)
+    if _original_bq_property_transform is not None:
+        return str(_original_bq_property_transform(self, expression))
+    return str(_original_bq_property_sql(self, expression))
+
+
+def _bq_properties_transform(self: Any, expression: exp.Properties) -> str:
+    dialect_class = getattr(self.dialect, "__class__", None)
+    dialect_name = dialect_class.__name__ if dialect_class else None
+    if dialect_name == "Spanner":
+        return _spanner_properties_sql(self, expression)
+    if _original_bq_properties_transform is not None:
+        return str(_original_bq_properties_transform(self, expression))
+    return str(_original_bq_properties_sql(self, expression))
+
+
+def _bq_create_transform(self: Any, expression: exp.Create) -> str:
+    dialect_class = getattr(self.dialect, "__class__", None)
+    dialect_name = dialect_class.__name__ if dialect_class else None
+    if dialect_name == "Spanner" and expression.this and expression.kind == "TABLE":
+        properties = expression.args.get("properties")
+        if properties:
+            # Re-order properties so Spanner ones stay at the schema boundary
+            new_expressions = []
+            for p in properties.expressions:
+                if _is_post_schema_spanner_property(p):
+                    # Force to POST_SCHEMA if it's a Spanner property
+                    new_expressions.append(p)
+                else:
+                    new_expressions.append(p)
+            properties.set("expressions", new_expressions)
+
+    if _original_bq_create_transform is not None:
+        return str(_original_bq_create_transform(self, expression))
+    return str(self.create_sql(expression))
+
+
+BigQueryGenerator.TRANSFORMS[exp.Property] = _bq_property_transform
+BigQueryGenerator.TRANSFORMS[exp.Properties] = _bq_properties_transform
+BigQueryGenerator.TRANSFORMS[exp.Create] = _bq_create_transform
+
+_DISPATCH_CACHE.pop(BigQueryGenerator, None)
+
+SpannerGenerator = BigQueryGenerator  # pyright: ignore[reportAssignmentType]
+
+
+# Postgres / Spangres
+_original_pg_property_transform = PostgresGenerator.TRANSFORMS.get(exp.Property)
+
+
+def _pg_property_transform(self: Any, expression: exp.Property) -> str:
+    dialect_class = getattr(self.dialect, "__class__", None)
+    dialect_name = dialect_class.__name__ if dialect_class else None
+    if dialect_name == "Spangres":
+        return _spangres_property_sql(self, expression)
+    if _original_pg_property_transform is not None:
+        return str(_original_pg_property_transform(self, expression))
+    return str(_original_pg_property_sql(self, expression))
+
+
+PostgresGenerator.TRANSFORMS[exp.Property] = _pg_property_transform
+
+_DISPATCH_CACHE.pop(PostgresGenerator, None)
+
+SpangresGenerator = PostgresGenerator  # pyright: ignore[reportAssignmentType]

sqlspec/dialects/spanner/_spangres.py

@@ -0,0 +1,77 @@
+r"""Google Cloud Spanner PostgreSQL-interface dialect ("Spangres")."""
+
+from typing import Any, Final, cast
+
+from sqlglot import exp
+from sqlglot.dialects.postgres import Postgres
+from sqlglot.parsers.postgres import PostgresParser
+from sqlglot.tokenizer_core import TokenType
+
+from sqlspec.dialects.spanner._generators import SpangresGenerator
+
+__all__ = ("Spangres",)
+
+_ROW_DELETION_NAME = "ROW_DELETION_POLICY"
+_TTL_MIN_COMPONENTS = 2
+_ORIGINAL_PARSE_PROPERTY_ATTR: Final[str] = "_sqlspec_original_parse_property"
+_HOOKS_REGISTERED_ATTR: Final[str] = "_sqlspec_spangres_hooks_registered"
+
+
+def _normalize_interval_expression(expression: exp.Expr) -> exp.Expr:
+    if isinstance(expression, exp.Alias):
+        alias = expression.args.get("alias")
+        if isinstance(alias, exp.Identifier) and isinstance(expression.this, exp.Expr):
+            return exp.Interval(this=expression.this.copy(), unit=alias.copy())
+    return expression
+
+
+def _is_spangres_dialect(parser: Any) -> bool:
+    dialect = getattr(parser, "dialect", None)
+    return dialect is not None and dialect.__class__.__name__ == "Spangres"
+
+
+def _original_postgres_parse_property() -> Any:
+    original = getattr(PostgresParser, _ORIGINAL_PARSE_PROPERTY_ATTR, None)
+    if callable(original):
+        return original
+    original = PostgresParser._parse_property
+    setattr(PostgresParser, _ORIGINAL_PARSE_PROPERTY_ATTR, original)
+    return original
+
+
+def _spangres_parse_property(self: Any) -> exp.Expr:
+    if _is_spangres_dialect(self) and self._match_text_seq("ROW", "DELETION", "POLICY"):
+        self._match(TokenType.L_PAREN)
+        self._match_text_seq("OLDER_THAN")
+        self._match(TokenType.L_PAREN)
+        column = cast("exp.Expr", self._parse_id_var())
+        self._match(TokenType.COMMA)
+        self._match_text_seq("INTERVAL")
+        interval = _normalize_interval_expression(cast("exp.Expr", self._parse_expression()))
+        self._match(TokenType.R_PAREN)
+        self._match(TokenType.R_PAREN)
+
+        return exp.Property(
+            this=exp.Literal.string(_ROW_DELETION_NAME), value=exp.Tuple(expressions=[column, interval])
+        )
+
+    return cast("exp.Expr", _original_postgres_parse_property()(self))
+
+
+def _register_postgres_spangres_parser_hooks() -> None:
+    if getattr(PostgresParser, _HOOKS_REGISTERED_ATTR, False):
+        return
+
+    _original_postgres_parse_property()
+    setattr(PostgresParser, "_parse_property", _spangres_parse_property)
+    setattr(PostgresParser, _HOOKS_REGISTERED_ATTR, True)
+
+
+_register_postgres_spangres_parser_hooks()
+
+
+class Spangres(Postgres):
+    """Spanner PostgreSQL-compatible dialect."""
+
+    Parser = Postgres.Parser
+    Generator = SpangresGenerator

sqlspec/dialects/spanner/_spanner.py

@@ -0,0 +1,179 @@
+"""Google Cloud Spanner SQL dialect (GoogleSQL variant).
+
+Extends the BigQuery dialect with Spanner-only DDL features:
+`INTERLEAVE IN PARENT` for interleaved tables and `ROW DELETION POLICY`
+for row-level time-to-live policies (GoogleSQL).
+"""
+
+import re
+from typing import Any, Final, cast
+
+from sqlglot import exp
+from sqlglot.dialects.bigquery import BigQuery
+from sqlglot.parsers.bigquery import BigQueryParser
+from sqlglot.tokenizer_core import TokenType
+
+from sqlspec.dialects.spanner._generators import SpannerGenerator
+
+__all__ = ("Spanner",)
+
+_SPANNER_KEYWORDS: "dict[str, TokenType]" = {}
+interleave_token = cast("TokenType | None", TokenType.__dict__.get("INTERLEAVE"))
+if interleave_token is not None:
+    _SPANNER_KEYWORDS["INTERLEAVE"] = interleave_token
+ttl_token = cast("TokenType | None", TokenType.__dict__.get("TTL"))
+if ttl_token is not None:
+    _SPANNER_KEYWORDS["TTL"] = ttl_token
+
+_TTL_MIN_COMPONENTS = 2
+_ROW_DELETION_NAME = "ROW_DELETION_POLICY"
+_INTERLEAVE_NAME = "INTERLEAVE_IN_PARENT"
+_ORIGINAL_PARSE_PROPERTY_ATTR: Final[str] = "_sqlspec_original_parse_property"
+_HOOKS_REGISTERED_ATTR: Final[str] = "_sqlspec_spanner_hooks_registered"
+_INTERLEAVE_PATTERN: Final[re.Pattern[str]] = re.compile(
+    r"""
+    \bINTERLEAVE\s+IN\s+PARENT\s+
+    (?P<parent>.+?)
+    (?:\s+ON\s+DELETE\s+(?P<on_delete>CASCADE|NO\s+ACTION))?
+    (?=\s+(?:ROW\s+DELETION\s+POLICY|TTL)\b|\s*$)
+    """,
+    re.IGNORECASE | re.DOTALL | re.VERBOSE,
+)
+
+
+def _normalize_on_delete_value(on_delete: str) -> str:
+    return " ".join(on_delete.upper().split())
+
+
+def _build_interleave_property(parent: exp.Expr, on_delete: str | None = None) -> exp.Property:
+    values: list[exp.Expr] = [parent]
+    if on_delete is not None:
+        values.append(exp.Literal.string(_normalize_on_delete_value(on_delete)))
+    return exp.Property(this=exp.Literal.string(_INTERLEAVE_NAME), value=exp.Tuple(expressions=values))
+
+
+def _normalize_interval_expression(expression: exp.Expr) -> exp.Expr:
+    if isinstance(expression, exp.Alias):
+        alias = expression.args.get("alias")
+        if isinstance(alias, exp.Identifier) and isinstance(expression.this, exp.Expr):
+            return exp.Interval(this=expression.this.copy(), unit=alias.copy())
+    return expression
+
+
+def _extract_interleave_property(sql: str) -> tuple[str, exp.Property | None]:
+    match = _INTERLEAVE_PATTERN.search(sql)
+    if match is None:
+        return sql, None
+
+    parent = exp.to_table(match.group("parent").strip())
+    on_delete = match.group("on_delete")
+    interleave_property = _build_interleave_property(parent, on_delete)
+    repaired_sql = f"{sql[: match.start()]} {sql[match.end() :]}".strip()
+    return repaired_sql, interleave_property
+
+
+def _attach_create_property(create: exp.Create, property_expression: exp.Property) -> exp.Create:
+    properties = create.args.get("properties")
+    if isinstance(properties, exp.Properties):
+        expressions = list(properties.expressions)
+        expressions.insert(0, property_expression)
+        properties.set("expressions", expressions)
+    else:
+        create.set("properties", exp.Properties(expressions=[property_expression]))
+    return create
+
+
+def _is_spanner_dialect(parser: Any) -> bool:
+    dialect = getattr(parser, "dialect", None)
+    return dialect is not None and dialect.__class__.__name__ == "Spanner"
+
+
+def _original_bigquery_parse_property() -> Any:
+    original = getattr(BigQueryParser, _ORIGINAL_PARSE_PROPERTY_ATTR, None)
+    if callable(original):
+        return original
+    original = BigQueryParser._parse_property
+    setattr(BigQueryParser, _ORIGINAL_PARSE_PROPERTY_ATTR, original)
+    return original
+
+
+def _spanner_parse_property(self: Any) -> exp.Expr:
+    if _is_spanner_dialect(self):
+        if self._match_text_seq("INTERLEAVE", "IN", "PARENT"):
+            parent = cast("exp.Expr", self._parse_table(schema=True, is_db_reference=True))
+            on_delete: str | None = None
+
+            if self._match_text_seq("ON", "DELETE"):
+                if self._match_text_seq("CASCADE"):
+                    on_delete = "CASCADE"
+                elif self._match_text_seq("NO", "ACTION"):
+                    on_delete = "NO ACTION"
+
+            return _build_interleave_property(parent, on_delete)
+
+        if self._match_text_seq("ROW", "DELETION", "POLICY"):
+            self._match(TokenType.L_PAREN)
+            self._match_text_seq("OLDER_THAN")
+            self._match(TokenType.L_PAREN)
+            column = cast("exp.Expr", self._parse_id_var())
+            self._match(TokenType.COMMA)
+            self._match_text_seq("INTERVAL")
+            interval = _normalize_interval_expression(cast("exp.Expr", self._parse_expression()))
+            self._match(TokenType.R_PAREN)
+            self._match(TokenType.R_PAREN)
+
+            return exp.Property(
+                this=exp.Literal.string(_ROW_DELETION_NAME), value=exp.Tuple(expressions=[column, interval])
+            )
+
+        if self._match_text_seq("TTL"):
+            self._match_text_seq("INTERVAL")
+            interval = _normalize_interval_expression(cast("exp.Expr", self._parse_expression()))
+            self._match_text_seq("ON")
+            column = cast("exp.Expr", self._parse_id_var())
+
+            return exp.Property(this=exp.Literal.string("TTL"), value=exp.Tuple(expressions=[interval, column]))
+
+    return cast("exp.Expr", _original_bigquery_parse_property()(self))
+
+
+def _register_bigquery_spanner_parser_hooks() -> None:
+    if getattr(BigQueryParser, _HOOKS_REGISTERED_ATTR, False):
+        return
+
+    _original_bigquery_parse_property()
+    setattr(BigQueryParser, "_parse_property", _spanner_parse_property)
+    setattr(BigQueryParser, _HOOKS_REGISTERED_ATTR, True)
+
+
+class SpannerTokenizer(BigQuery.Tokenizer):
+    """Tokenizer adds Spanner-only keywords when supported by sqlglot."""
+
+    KEYWORDS = {**BigQuery.Tokenizer.KEYWORDS, **_SPANNER_KEYWORDS}
+
+
+_register_bigquery_spanner_parser_hooks()
+
+
+class Spanner(BigQuery):
+    """Google Cloud Spanner SQL dialect."""
+
+    Tokenizer = SpannerTokenizer
+    Parser = BigQuery.Parser
+    Generator = SpannerGenerator
+
+    def parse(self, sql: str, **opts: Any) -> list[exp.Expr | None]:
+        """Repair CREATE TABLE statements that sqlglot still falls back to Command for."""
+        expressions = super().parse(sql, **opts)
+        if len(expressions) != 1 or not isinstance(expressions[0], exp.Command):
+            return expressions
+
+        repaired_sql, interleave_property = _extract_interleave_property(sql)
+        if interleave_property is None:
+            return expressions
+
+        reparsed = BigQuery.parse(self, repaired_sql, **opts)
+        if len(reparsed) != 1 or not isinstance(reparsed[0], exp.Create):
+            return expressions
+
+        return [_attach_create_property(reparsed[0], interleave_property)]

sqlspec/driver/__init__.py

@@ -0,0 +1,49 @@
+"""Driver protocols and base classes for database adapters."""
+
+from sqlspec.driver._async import (
+    AsyncDataDictionaryBase,
+    AsyncDriverAdapterBase,
+    AsyncPoolConnectionContext,
+    AsyncPoolSessionFactory,
+)
+from sqlspec.driver._common import (
+    CommonDriverAttributesMixin,
+    DataDictionaryDialectMixin,
+    DataDictionaryMixin,
+    ExecutionResult,
+    StackExecutionObserver,
+    describe_stack_statement,
+    hash_stack_operations,
+)
+from sqlspec.driver._exception_handler import BaseAsyncExceptionHandler, BaseSyncExceptionHandler
+from sqlspec.driver._sql_helpers import convert_to_dialect
+from sqlspec.driver._sync import (
+    SyncDataDictionaryBase,
+    SyncDriverAdapterBase,
+    SyncPoolConnectionContext,
+    SyncPoolSessionFactory,
+)
+
+__all__ = (
+    "AsyncDataDictionaryBase",
+    "AsyncDriverAdapterBase",
+    "AsyncPoolConnectionContext",
+    "AsyncPoolSessionFactory",
+    "BaseAsyncExceptionHandler",
+    "BaseSyncExceptionHandler",
+    "CommonDriverAttributesMixin",
+    "DataDictionaryDialectMixin",
+    "DataDictionaryMixin",
+    "DriverAdapterProtocol",
+    "ExecutionResult",
+    "StackExecutionObserver",
+    "SyncDataDictionaryBase",
+    "SyncDriverAdapterBase",
+    "SyncPoolConnectionContext",
+    "SyncPoolSessionFactory",
+    "convert_to_dialect",
+    "describe_stack_statement",
+    "hash_stack_operations",
+)
+
+DriverAdapterProtocol = SyncDriverAdapterBase | AsyncDriverAdapterBase