codekingpro/portable-devtools

1# engine/mock.py
2# Copyright (C) 2005-2024 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7
8from __future__ import annotations
9
10from operator import attrgetter
11import typing
12from typing import Any
13from typing import Callable
14from typing import cast
15from typing import Optional
16from typing import Type
17from typing import Union
18
19from . import url as _url
20from .. import util
21
22
23if typing.TYPE_CHECKING:
24    from .base import Engine
25    from .interfaces import _CoreAnyExecuteParams
26    from .interfaces import CoreExecuteOptionsParameter
27    from .interfaces import Dialect
28    from .url import URL
29    from ..sql.base import Executable
30    from ..sql.ddl import SchemaDropper
31    from ..sql.ddl import SchemaGenerator
32    from ..sql.schema import HasSchemaAttr
33    from ..sql.schema import SchemaItem
34
35
36class MockConnection:
37    def __init__(self, dialect: Dialect, execute: Callable[..., Any]):
38        self._dialect = dialect
39        self._execute_impl = execute
40
41    engine: Engine = cast(Any, property(lambda s: s))
42    dialect: Dialect = cast(Any, property(attrgetter("_dialect")))
43    name: str = cast(Any, property(lambda s: s._dialect.name))
44
45    def connect(self, **kwargs: Any) -> MockConnection:
46        return self
47
48    def schema_for_object(self, obj: HasSchemaAttr) -> Optional[str]:
49        return obj.schema
50
51    def execution_options(self, **kw: Any) -> MockConnection:
52        return self
53
54    def _run_ddl_visitor(
55        self,
56        visitorcallable: Type[Union[SchemaGenerator, SchemaDropper]],
57        element: SchemaItem,
58        **kwargs: Any,
59    ) -> None:
60        kwargs["checkfirst"] = False
61        visitorcallable(self.dialect, self, **kwargs).traverse_single(element)
62
63    def execute(
64        self,
65        obj: Executable,
66        parameters: Optional[_CoreAnyExecuteParams] = None,
67        execution_options: Optional[CoreExecuteOptionsParameter] = None,
68    ) -> Any:
69        return self._execute_impl(obj, parameters)
70
71
72def create_mock_engine(
73    url: Union[str, URL], executor: Any, **kw: Any
74) -> MockConnection:
75    """Create a "mock" engine used for echoing DDL.
76
77    This is a utility function used for debugging or storing the output of DDL
78    sequences as generated by :meth:`_schema.MetaData.create_all`
79    and related methods.
80
81    The function accepts a URL which is used only to determine the kind of
82    dialect to be used, as well as an "executor" callable function which
83    will receive a SQL expression object and parameters, which can then be
84    echoed or otherwise printed.   The executor's return value is not handled,
85    nor does the engine allow regular string statements to be invoked, and
86    is therefore only useful for DDL that is sent to the database without
87    receiving any results.
88
89    E.g.::
90
91        from sqlalchemy import create_mock_engine
92
93        def dump(sql, *multiparams, **params):
94            print(sql.compile(dialect=engine.dialect))
95
96        engine = create_mock_engine('postgresql+psycopg2://', dump)
97        metadata.create_all(engine, checkfirst=False)
98
99    :param url: A string URL which typically needs to contain only the
100     database backend name.
101
102    :param executor: a callable which receives the arguments ``sql``,
103     ``*multiparams`` and ``**params``.  The ``sql`` parameter is typically
104     an instance of :class:`.ExecutableDDLElement`, which can then be compiled
105     into a string using :meth:`.ExecutableDDLElement.compile`.
106
107    .. versionadded:: 1.4 - the :func:`.create_mock_engine` function replaces
108       the previous "mock" engine strategy used with
109       :func:`_sa.create_engine`.
110
111    .. seealso::
112
113        :ref:`faq_ddl_as_string`
114
115    """
116
117    # create url.URL object
118    u = _url.make_url(url)
119
120    dialect_cls = u.get_dialect()
121
122    dialect_args = {}
123    # consume dialect arguments from kwargs
124    for k in util.get_cls_kwargs(dialect_cls):
125        if k in kw:
126            dialect_args[k] = kw.pop(k)
127
128    # create dialect
129    dialect = dialect_cls(**dialect_args)
130
131    return MockConnection(dialect, executor)
132