instructure-pysqlsync 0.8.5.dev0__py3-none-any.whl

pysqlsync/dialect/redshift/connection.py

@@ -0,0 +1,156 @@
+import dataclasses
+import logging
+import re
+import typing
+from io import BytesIO
+from typing import Iterable, Optional, TypeVar
+
+import redshift_connector
+from strong_typing.inspection import DataclassInstance, is_dataclass_type
+
+from pysqlsync.base import (
+    BaseConnection,
+    BaseContext,
+    ClassRef,
+    DataSource,
+    QueryException,
+    RecordType,
+)
+from pysqlsync.formation.object_types import Table
+from pysqlsync.model.id_types import LocalId, SupportsQualifiedId
+from pysqlsync.model.properties import is_identity_type
+from pysqlsync.resultset import resultset_unwrap_object, resultset_unwrap_tuple
+from pysqlsync.util.dispatch import thread_dispatch
+from pysqlsync.util.typing import override
+
+D = TypeVar("D", bound=DataclassInstance)
+T = TypeVar("T")
+
+LOGGER = logging.getLogger("pysqlsync.redshift")
+
+redshift_connector.paramstyle = "numeric"
+
+
+class RedshiftConnection(BaseConnection):
+    native: redshift_connector.Connection
+
+    @override
+    @thread_dispatch
+    def open(self) -> BaseContext:
+        LOGGER.info("connecting to %s", self.params)
+        conn = redshift_connector.connect(
+            iam=True,
+            host=self.params.host,
+            port=self.params.port,
+            database=self.params.database,
+            db_user=self.params.username,
+        )
+        conn.autocommit = True
+
+        self.native = conn
+        return RedshiftContext(self)
+
+    @override
+    @thread_dispatch
+    def close(self) -> None:
+        self.native.close()
+
+
+class RedshiftContext(BaseContext):
+    def __init__(self, connection: RedshiftConnection) -> None:
+        super().__init__(connection)
+
+    @property
+    def native_connection(self) -> redshift_connector.Connection:
+        return typing.cast(RedshiftConnection, self.connection).native
+
+    @override
+    @thread_dispatch
+    def _execute(self, statement: str) -> None:
+        with self.native_connection.cursor() as cursor:
+            # The library `redshift_connector` executes all statements as prepared statements, which locks out
+            # the possibility to execute multiple statements in one go. We artificially split SQL scripts into
+            # multiple statements to work around this limitation in the connector.
+            statement = statement.rstrip("\r\n\t\v ;")
+            for s in re.split(r";$", statement, flags=re.MULTILINE):
+                try:
+                    cursor.execute(s)
+                except Exception as e:
+                    raise QueryException(s) from e
+
+    @override
+    async def _execute_all(self, statement: str, source: DataSource) -> None:
+        async for batch in source.batches():
+            await self._internal_execute_all(statement, batch)
+
+    @thread_dispatch
+    def _internal_execute_all(self, statement: str, records: Iterable[RecordType]) -> None:
+        with self.native_connection.cursor() as cursor:
+            cursor.executemany(statement, records)
+
+    @override
+    @thread_dispatch
+    def _query_all(self, signature: type[T], statement: str) -> list[T]:
+        with self.native_connection.cursor() as cursor:
+            records = cursor.execute(statement).fetchall()
+
+            if is_dataclass_type(signature):
+                return resultset_unwrap_object(signature, records)  # type: ignore
+            else:
+                return resultset_unwrap_tuple(signature, records)
+
+    @override
+    @thread_dispatch
+    def insert_data(self, table: type[D], data: Iterable[D]) -> None:
+        if not is_dataclass_type(table):
+            raise TypeError(f"expected dataclass type, got: {table}")
+        generator = self.connection.generator
+        table_name = generator.get_qualified_id(ClassRef(table))
+        columns = [LocalId(field.name) for field in dataclasses.fields(table) if not is_identity_type(field.type)]
+        records = generator.get_dataclasses_as_records(table, data, skip_identity=True)
+        self._insert_copy_stream(table_name, columns, records)
+
+    @override
+    async def _insert_rows(
+        self,
+        table: Table,
+        source: DataSource,
+        *,
+        field_types: tuple[type, ...],
+        field_names: Optional[tuple[str, ...]] = None,
+    ) -> None:
+        order = tuple(name for name in field_names if name) if field_names else None
+        columns = [col.name for col in table.get_columns(order)]
+        record_generator = await self._generate_records(table, source, field_types=field_types, field_names=field_names)
+        async for batch in record_generator.batches():
+            await self._insert_copy_stream_async(table.name, columns, batch)
+
+    @thread_dispatch
+    def _insert_copy_stream_async(
+        self,
+        table_name: SupportsQualifiedId,
+        columns: list[LocalId],
+        records: Iterable[RecordType],
+    ) -> None:
+        self._insert_copy_stream(table_name, columns, records)
+
+    def _insert_copy_stream(
+        self,
+        table_name: SupportsQualifiedId,
+        columns: list[LocalId],
+        records: Iterable[RecordType],
+    ) -> None:
+        column_list = ", ".join(str(column) for column in columns)
+        copy_query = f"COPY {table_name} ({column_list}) FROM STDIN"
+
+        with BytesIO() as f:
+            for record in records:
+                f.write(b"|".join(str(value).encode("utf-8") for value in record))
+                f.write(b"\n")
+            stream_data = f.getvalue()
+
+        cursor = self.native_connection.cursor()
+        try:
+            cursor.execute(copy_query, stream=BytesIO(stream_data))
+        finally:
+            cursor.close()

pysqlsync/dialect/redshift/data_types.py

@@ -0,0 +1,10 @@
+from pysqlsync.model.data_types import SqlVariableBinaryType
+
+
+class RedshiftVariableBinaryType(SqlVariableBinaryType):
+    def __str__(self) -> str:
+        if self.storage is not None:
+            storage = f"({self.storage})"
+        else:
+            storage = ""
+        return f"varbyte{storage}"

pysqlsync/dialect/redshift/dependency.py

@@ -0,0 +1,9 @@
+"""
+pysqlsync: Synchronize schema and large volumes of data.
+
+This module defines dependencies required for AWS Redshift.
+
+:see: https://github.com/instructure-internal/pysqlsync
+"""
+
+import redshift_connector  # pyright: ignore[reportUnusedImport] # noqa: F401

pysqlsync/dialect/redshift/engine.py

@@ -0,0 +1,20 @@
+from pysqlsync.base import BaseConnection, BaseEngine, BaseGenerator, Explorer
+
+from ..postgresql.discovery import PostgreSQLExplorer
+from .connection import RedshiftConnection
+from .generator import RedshiftGenerator
+
+
+class RedshiftEngine(BaseEngine):
+    @property
+    def name(self) -> str:
+        return "redshift"
+
+    def get_generator_type(self) -> type[BaseGenerator]:
+        return RedshiftGenerator
+
+    def get_connection_type(self) -> type[BaseConnection]:
+        return RedshiftConnection
+
+    def get_explorer_type(self) -> type[Explorer]:
+        return PostgreSQLExplorer

pysqlsync/dialect/redshift/generator.py

@@ -0,0 +1,82 @@
+import ipaddress
+import uuid
+from typing import Any, Callable, Optional
+
+from strong_typing.core import JsonType
+
+from pysqlsync.base import BaseGenerator, GeneratorOptions
+from pysqlsync.formation.inspection import is_ip_address_type
+from pysqlsync.formation.object_types import Column
+from pysqlsync.formation.py_to_sql import (
+    ArrayMode,
+    DataclassConverter,
+    DataclassConverterOptions,
+    EnumMode,
+    NamespaceMapping,
+    StructMode,
+)
+from pysqlsync.model.data_types import SqlVariableCharacterType
+from pysqlsync.util.typing import override
+
+from ..postgresql.mutation import PostgreSQLMutator
+from ..postgresql.object_types import PostgreSQLObjectFactory
+from .data_types import RedshiftVariableBinaryType
+
+
+class RedshiftGenerator(BaseGenerator):
+    "Generator for AWS Redshift."
+
+    converter: DataclassConverter
+
+    def __init__(self, options: GeneratorOptions) -> None:
+        super().__init__(
+            options,
+            PostgreSQLObjectFactory(),
+            PostgreSQLMutator(options.synchronization),
+        )
+
+        self.check_enum_mode(matches=EnumMode.CHECK)
+        self.check_struct_mode(matches=StructMode.JSON)
+        self.check_array_mode(include=[ArrayMode.JSON, ArrayMode.RELATION])
+
+        self.converter = DataclassConverter(
+            options=DataclassConverterOptions(
+                enum_mode=options.enum_mode or EnumMode.RELATION,
+                struct_mode=options.struct_mode or StructMode.JSON,
+                array_mode=options.array_mode or ArrayMode.JSON,
+                namespaces=NamespaceMapping(options.namespaces),
+                foreign_constraints=False,
+                initialize_tables=options.initialize_tables,
+                substitutions={
+                    uuid.UUID: RedshiftVariableBinaryType(16),
+                    JsonType: SqlVariableCharacterType(),
+                    ipaddress.IPv4Address: RedshiftVariableBinaryType(4),
+                    ipaddress.IPv6Address: RedshiftVariableBinaryType(16),
+                },
+                factory=self.factory,
+                skip_annotations=options.skip_annotations,
+                auto_default=options.auto_default,
+            )
+        )
+
+    @override
+    def placeholder(self, index: int) -> str:
+        return f":{index}"
+
+    @override
+    def get_field_extractor(self, column: Column, field_name: str, field_type: type) -> Callable[[Any], Any]:
+        if field_type is uuid.UUID:
+            return lambda obj: getattr(obj, field_name).bytes
+        elif is_ip_address_type(field_type):
+            return lambda obj: getattr(obj, field_name).packed
+
+        return super().get_field_extractor(column, field_name, field_type)
+
+    @override
+    def get_value_transformer(self, column: Column, field_type: type) -> Optional[Callable[[Any], Any]]:
+        if field_type is uuid.UUID:
+            return lambda field: field.bytes
+        elif is_ip_address_type(field_type):
+            return lambda field: field.packed
+
+        return super().get_value_transformer(column, field_type)

pysqlsync/dialect/snowflake/data_types.py

@@ -0,0 +1,18 @@
+from pysqlsync.model.data_types import SqlJsonType, SqlTimestampType
+
+
+class SnowflakeDateTimeType(SqlTimestampType):
+    "Timestamp without time zone, equivalent to TIMESTAMP_NTZ."
+
+    def __init__(self) -> None:
+        self.precision = 9
+
+    def __str__(self) -> str:
+        return "datetime"
+
+
+class SnowflakeJsonType(SqlJsonType):
+    "Represents JSON data extracted with PARSE_JSON."
+
+    def __str__(self) -> str:
+        return "variant"

pysqlsync/dialect/snowflake/dependency.py

@@ -0,0 +1,9 @@
+"""
+pysqlsync: Synchronize schema and large volumes of data.
+
+This module defines dependencies required for Snowflake.
+
+:see: https://github.com/instructure-internal/pysqlsync
+"""
+
+import snowflake  # pyright: ignore[reportUnusedImport] # noqa: F401

pysqlsync/dialect/snowflake/engine.py

@@ -0,0 +1,18 @@
+from pysqlsync.base import BaseConnection, BaseEngine, BaseGenerator, Explorer
+
+from .generator import SnowflakeGenerator
+
+
+class SnowflakeEngine(BaseEngine):
+    @property
+    def name(self) -> str:
+        return "snowflake"
+
+    def get_generator_type(self) -> type[BaseGenerator]:
+        return SnowflakeGenerator
+
+    def get_connection_type(self) -> type[BaseConnection]:
+        raise NotImplementedError()
+
+    def get_explorer_type(self) -> type[Explorer]:
+        raise NotImplementedError()

pysqlsync/dialect/snowflake/generator.py

@@ -0,0 +1,64 @@
+import datetime
+import ipaddress
+import uuid
+
+from strong_typing.core import JsonType
+
+from pysqlsync.base import BaseGenerator, GeneratorOptions
+from pysqlsync.formation.mutation import Mutator
+from pysqlsync.formation.py_to_sql import (
+    ArrayMode,
+    DataclassConverter,
+    DataclassConverterOptions,
+    EnumMode,
+    NamespaceMapping,
+    StructMode,
+)
+from pysqlsync.model.data_types import SqlFixedBinaryType
+from pysqlsync.util.typing import override
+
+from .data_types import SnowflakeDateTimeType, SnowflakeJsonType
+from .object_types import SnowflakeObjectFactory
+
+
+class SnowflakeGenerator(BaseGenerator):
+    "Generator for Snowflake."
+
+    converter: DataclassConverter
+
+    def __init__(self, options: GeneratorOptions) -> None:
+        super().__init__(
+            options,
+            SnowflakeObjectFactory(),
+            Mutator(options.synchronization),
+        )
+
+        self.check_enum_mode(matches=EnumMode.CHECK)
+        self.check_struct_mode(matches=StructMode.JSON)
+        self.check_array_mode(include=[ArrayMode.JSON, ArrayMode.RELATION])
+
+        self.converter = DataclassConverter(
+            options=DataclassConverterOptions(
+                enum_mode=options.enum_mode or EnumMode.CHECK,
+                struct_mode=options.struct_mode or StructMode.JSON,
+                array_mode=options.array_mode or ArrayMode.JSON,
+                namespaces=NamespaceMapping(options.namespaces),
+                check_constraints=False,
+                foreign_constraints=False,
+                initialize_tables=options.initialize_tables,
+                substitutions={
+                    datetime.datetime: SnowflakeDateTimeType(),
+                    uuid.UUID: SqlFixedBinaryType(16),
+                    JsonType: SnowflakeJsonType(),
+                    ipaddress.IPv4Address: SqlFixedBinaryType(4),
+                    ipaddress.IPv6Address: SqlFixedBinaryType(16),
+                },
+                factory=self.factory,
+                skip_annotations=options.skip_annotations,
+                auto_default=options.auto_default,
+            )
+        )
+
+    @override
+    def placeholder(self, index: int) -> str:
+        return f":{index}"

pysqlsync/dialect/snowflake/object_types.py

@@ -0,0 +1,84 @@
+import re
+from typing import Optional
+
+from pysqlsync.formation.object_types import Column, ObjectFactory, Table
+from pysqlsync.model.data_types import SqlTimestampType
+from pysqlsync.model.id_types import LocalId
+
+_sql_quoted_str_table = str.maketrans(
+    {
+        "\\": "\\\\",
+        "'": "\\'",
+        '"': '\\"',
+        "\0": "\\0",
+        "\b": "\\b",
+        "\f": "\\f",
+        "\n": "\\n",
+        "\r": "\\r",
+        "\t": "\\t",
+    }
+)
+
+
+def sql_quoted_string(text: str) -> str:
+    if re.search(r"[\\\0\b\f\n\r\t]", text):
+        text = text.translate(_sql_quoted_str_table)
+    elif "'" in text:
+        text = text.replace("'", "''")
+    return f"'{text}'"
+
+
+class SnowflakeTable(Table):
+    def create_stmt(self) -> str:
+        defs: list[str] = []
+        defs.extend(str(c) for c in self.columns.values())
+        defs.append(self.create_keys())
+        definitions = ",\n".join(defs)
+        comment = f"\nCOMMENT = {sql_quoted_string(self.description)}" if self.description else ""
+        return f"CREATE TABLE {self.name} (\n{definitions}\n){comment};"
+
+    @property
+    def primary_key_constraint_id(self) -> LocalId:
+        return LocalId(f"pk_{self.name.local_id.replace('.', '_')}")
+
+
+class SnowflakeColumn(Column):
+    @property
+    def default_expr(self) -> str:
+        if self.default is None:
+            raise ValueError("default value is NULL")
+
+        if isinstance(self.data_type, SqlTimestampType):
+            m = re.match(
+                r"^'(?P<year>\d{4})-(?P<month>\d{2})-(?P<day>\d{2}) (?P<hour>\d{2}):(?P<minute>\d{2}):(?P<second>\d{2})'$",
+                self.default,
+            )
+            if m:
+                return f"TIMESTAMP {self.default}"
+
+        return self.default
+
+    @property
+    def data_spec(self) -> str:
+        nullable = " NOT NULL" if not self.nullable and not self.identity else ""
+        default = f" DEFAULT {self.default_expr}" if self.default is not None else ""
+        identity = " IDENTITY" if self.identity else ""
+        description = f" COMMENT {self.comment}" if self.description is not None else ""
+        return f"{self.data_type}{nullable}{default}{identity}{description}"
+
+    @property
+    def comment(self) -> Optional[str]:
+        if self.description is not None:
+            return sql_quoted_string(self.description)
+        else:
+            return None
+
+
+class SnowflakeObjectFactory(ObjectFactory):
+    @property
+    def column_class(self) -> type[Column]:
+        return SnowflakeColumn
+
+    @property
+    def table_class(self) -> type[Table]:
+        return SnowflakeTable

pysqlsync/dialect/test/dependency.py

@@ -0,0 +1,10 @@
+"""
+pysqlsync: Synchronize schema and large volumes of data.
+
+This module is only required for integration tests.
+
+:see: https://github.com/instructure-internal/pysqlsync
+"""
+
+# a fake import to trigger a ModuleNotFoundError exception
+import module_not_found  # pyright: ignore[reportMissingImports, reportUnusedImport] # noqa: F401

pysqlsync/dialect/trino/connection.py

@@ -0,0 +1,86 @@
+import typing
+from typing import Iterable, Optional, TypeVar
+
+import aiotrino
+from strong_typing.inspection import is_dataclass_type
+
+from pysqlsync.base import BaseConnection, BaseContext, DataSource
+from pysqlsync.formation.object_types import Table
+from pysqlsync.resultset import resultset_unwrap_tuple
+from pysqlsync.util.typing import override
+
+T = TypeVar("T")
+
+
+class TrinoConnection(BaseConnection):
+    native: aiotrino.dbapi.Connection
+
+    @override
+    async def open(self) -> BaseContext:
+        self.native = aiotrino.dbapi.connect(
+            host=self.params.host,
+            port=self.params.port,
+            user=self.params.username,
+            auth=aiotrino.auth.BasicAuthentication(self.params.username, self.params.password),
+            http_scheme="https",
+            catalog=self.params.database,
+        )
+        return TrinoContext(self)
+
+    @override
+    async def close(self) -> None:
+        await self.native.close()
+
+
+class TrinoContext(BaseContext):
+    def __init__(self, connection: TrinoConnection) -> None:
+        super().__init__(connection)
+
+    @property
+    def native_connection(self) -> aiotrino.dbapi.Connection:
+        return typing.cast(TrinoConnection, self.connection).native
+
+    @override
+    async def _execute(self, statement: str) -> None:
+        cur = await self.native_connection.cursor()
+        await cur.execute(statement)
+
+    @override
+    async def _execute_all(self, statement: str, source: DataSource) -> None:
+        raise NotImplementedError("operation not supported for Trino")
+
+    @override
+    async def _query_all(self, signature: type[T], statement: str) -> list[T]:
+        if is_dataclass_type(signature):
+            raise NotImplementedError()
+
+        cur = await self.native_connection.cursor()
+        await cur.execute(statement)
+        records = await cur.fetchall()
+        return resultset_unwrap_tuple(signature, records)
+
+    @override
+    async def insert_data(self, table: type[T], data: Iterable[T]) -> None:
+        raise NotImplementedError("operation not supported for Trino")
+
+    @override
+    async def _insert_rows(
+        self,
+        table: Table,
+        records: DataSource,
+        *,
+        field_types: tuple[type, ...],
+        field_names: Optional[tuple[str, ...]] = None,
+    ) -> None:
+        raise NotImplementedError("operation not supported for Trino")
+
+    @override
+    async def _upsert_rows(
+        self,
+        table: Table,
+        records: DataSource,
+        *,
+        field_types: tuple[type, ...],
+        field_names: Optional[tuple[str, ...]] = None,
+    ) -> None:
+        raise NotImplementedError("operation not supported for Trino")

pysqlsync/dialect/trino/dependency.py

@@ -0,0 +1,9 @@
+"""
+pysqlsync: Synchronize schema and large volumes of data.
+
+This module defines dependencies required to support the query engine Trino.
+
+:see: https://github.com/instructure-internal/pysqlsync
+"""
+
+import aiotrino  # pyright: ignore[reportUnusedImport] # noqa: F401

pysqlsync/dialect/trino/discovery.py

@@ -0,0 +1,21 @@
+from pysqlsync.base import BaseContext
+from pysqlsync.formation.data_types import SqlDiscovery, SqlDiscoveryOptions
+from pysqlsync.formation.discovery import AnsiExplorer
+from pysqlsync.formation.object_types import ObjectFactory
+
+
+class TrinoExplorer(AnsiExplorer):
+    """
+    Discovers objects in a database exposed via Trino.
+
+    We recommend setting the following configuration options in Trino:
+    * postgresql.array-mapping = AS_ARRAY
+    * postgresql.include-system-tables = true
+    """
+
+    def __init__(self, conn: BaseContext) -> None:
+        super().__init__(
+            conn,
+            SqlDiscovery(SqlDiscoveryOptions(substitutions={})),
+            ObjectFactory(),
+        )

pysqlsync/dialect/trino/engine.py

@@ -0,0 +1,20 @@
+from pysqlsync.base import BaseConnection, BaseEngine, BaseGenerator, Explorer
+
+from .connection import TrinoConnection
+from .discovery import TrinoExplorer
+from .generator import TrinoGenerator
+
+
+class TrinoEngine(BaseEngine):
+    @property
+    def name(self) -> str:
+        return "trino"
+
+    def get_generator_type(self) -> type[BaseGenerator]:
+        return TrinoGenerator
+
+    def get_connection_type(self) -> type[BaseConnection]:
+        return TrinoConnection
+
+    def get_explorer_type(self) -> type[Explorer]:
+        return TrinoExplorer

pysqlsync/dialect/trino/generator.py

@@ -0,0 +1,31 @@
+from pysqlsync.base import BaseGenerator, GeneratorOptions
+from pysqlsync.formation.py_to_sql import (
+    ArrayMode,
+    DataclassConverter,
+    DataclassConverterOptions,
+    EnumMode,
+    NamespaceMapping,
+    StructMode,
+)
+from pysqlsync.util.typing import override
+
+
+class TrinoGenerator(BaseGenerator):
+    def __init__(self, options: GeneratorOptions) -> None:
+        super().__init__(options)
+        self.converter = DataclassConverter(
+            options=DataclassConverterOptions(
+                enum_mode=options.enum_mode or EnumMode.CHECK,
+                struct_mode=options.struct_mode or StructMode.JSON,
+                array_mode=options.array_mode or ArrayMode.JSON,
+                namespaces=NamespaceMapping(options.namespaces),
+                foreign_constraints=options.foreign_constraints,
+                initialize_tables=options.initialize_tables,
+                skip_annotations=options.skip_annotations,
+                auto_default=options.auto_default,
+            )
+        )
+
+    @override
+    def placeholder(self, index: int) -> str:
+        raise NotImplementedError("operation not supported for Trino")