sqlmesh.core.engine_adapter.base_postgres

  1from __future__ import annotations
  2
  3import typing as t
  4import logging
  5
  6from sqlglot import exp
  7
  8from sqlmesh.core.dialect import to_schema
  9from sqlmesh.core.engine_adapter.base import EngineAdapter, _get_data_object_cache_key
 10from sqlmesh.core.engine_adapter.shared import (
 11    CatalogSupport,
 12    CommentCreationTable,
 13    CommentCreationView,
 14    DataObject,
 15    DataObjectType,
 16)
 17from sqlmesh.utils.errors import SQLMeshError
 18
 19if t.TYPE_CHECKING:
 20    from sqlmesh.core._typing import SchemaName, TableName
 21    from sqlmesh.core.engine_adapter._typing import QueryOrDF
 22
 23
 24logger = logging.getLogger(__name__)
 25
 26
 27class BasePostgresEngineAdapter(EngineAdapter):
 28    DEFAULT_BATCH_SIZE = 400
 29    COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
 30    COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
 31    SUPPORTS_QUERY_EXECUTION_TRACKING = True
 32    SUPPORTED_DROP_CASCADE_OBJECT_KINDS = ["SCHEMA", "TABLE", "VIEW"]
 33
 34    def columns(
 35        self, table_name: TableName, include_pseudo_columns: bool = False
 36    ) -> t.Dict[str, exp.DataType]:
 37        """Fetches column names and types for the target table."""
 38        table = exp.to_table(table_name)
 39
 40        sql = (
 41            exp.select(
 42                "attname AS column_name",
 43                "pg_catalog.format_type(atttypid, atttypmod) AS data_type",
 44            )
 45            .from_("pg_catalog.pg_attribute")
 46            .join("pg_catalog.pg_class", on="pg_class.oid = attrelid")
 47            .join("pg_catalog.pg_namespace", on="pg_namespace.oid = relnamespace")
 48            .where(
 49                exp.and_(
 50                    "attnum > 0",
 51                    "NOT attisdropped",
 52                    exp.column("relname").eq(table.alias_or_name),
 53                )
 54            )
 55        )
 56        if table.args.get("db"):
 57            sql = sql.where(exp.column("nspname").eq(table.args["db"].name))
 58
 59        self.execute(sql)
 60        resp = self.cursor.fetchall()
 61        if not resp:
 62            raise SQLMeshError(
 63                f"Could not get columns for table '{table.sql(dialect=self.dialect)}'. Table not found."
 64            )
 65
 66        return {
 67            column_name: exp.DataType.build(data_type, dialect=self.dialect, udt=True)
 68            for column_name, data_type in resp
 69        }
 70
 71    @property
 72    def catalog_support(self) -> CatalogSupport:
 73        return CatalogSupport.SINGLE_CATALOG_ONLY
 74
 75    def table_exists(self, table_name: TableName) -> bool:
 76        """
 77        Postgres doesn't support describe so I'm using what the redshift cursor does to check if a table
 78        exists. We don't use this directly in order for this to work as a base class for other postgres
 79
 80        Reference: https://github.com/aws/amazon-redshift-python-driver/blob/master/redshift_connector/cursor.py#L528-L553
 81        """
 82        table = exp.to_table(table_name)
 83        data_object_cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name)
 84        if data_object_cache_key in self._data_object_cache:
 85            logger.debug("Table existence cache hit: %s", data_object_cache_key)
 86            return self._data_object_cache[data_object_cache_key] is not None
 87
 88        sql = (
 89            exp.select("1")
 90            .from_("information_schema.tables")
 91            .where(f"table_name = '{table.alias_or_name}'")
 92        )
 93        database_name = table.db
 94        if database_name:
 95            sql = sql.where(f"table_schema = '{database_name}'")
 96
 97        self.execute(sql)
 98
 99        result = self.cursor.fetchone()
100
101        return result[0] == 1 if result is not None else False
102
103    def create_view(
104        self,
105        view_name: TableName,
106        query_or_df: QueryOrDF,
107        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
108        replace: bool = True,
109        materialized: bool = False,
110        materialized_properties: t.Optional[t.Dict[str, t.Any]] = None,
111        table_description: t.Optional[str] = None,
112        column_descriptions: t.Optional[t.Dict[str, str]] = None,
113        view_properties: t.Optional[t.Dict[str, exp.Expr]] = None,
114        source_columns: t.Optional[t.List[str]] = None,
115        **create_kwargs: t.Any,
116    ) -> None:
117        """
118        Postgres has very strict rules around view replacement. For example the new query must generate an identical setFormatter
119        of columns, using the same column names and data types as the old one. We have to delete the old view instead of replacing it
120        to work around these constraints.
121
122        Reference: https://www.postgresql.org/docs/current/sql-createview.html
123        """
124        with self.transaction():
125            if replace:
126                self.drop_view(view_name, materialized=materialized)
127            super().create_view(
128                view_name,
129                query_or_df,
130                target_columns_to_types=target_columns_to_types,
131                replace=False,
132                materialized=materialized,
133                materialized_properties=materialized_properties,
134                table_description=table_description,
135                column_descriptions=column_descriptions,
136                view_properties=view_properties,
137                source_columns=source_columns,
138                **create_kwargs,
139            )
140
141    def drop_view(
142        self,
143        view_name: TableName,
144        ignore_if_not_exists: bool = True,
145        materialized: bool = False,
146        **kwargs: t.Any,
147    ) -> None:
148        kwargs["cascade"] = kwargs.get("cascade", True)
149        return super().drop_view(
150            view_name,
151            ignore_if_not_exists=ignore_if_not_exists,
152            materialized=materialized,
153            **kwargs,
154        )
155
156    def _get_data_objects(
157        self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
158    ) -> t.List[DataObject]:
159        """
160        Returns all the data objects that exist in the given schema and optionally catalog.
161        """
162        catalog = self.get_current_catalog()
163        table_query = exp.select(
164            exp.column("schemaname").as_("schema_name"),
165            exp.column("tablename").as_("name"),
166            exp.Literal.string("TABLE").as_("type"),
167        ).from_("pg_tables")
168        view_query = exp.select(
169            exp.column("schemaname").as_("schema_name"),
170            exp.column("viewname").as_("name"),
171            exp.Literal.string("VIEW").as_("type"),
172        ).from_("pg_views")
173        materialized_view_query = exp.select(
174            exp.column("schemaname").as_("schema_name"),
175            exp.column("matviewname").as_("name"),
176            exp.Literal.string("MATERIALIZED_VIEW").as_("type"),
177        ).from_("pg_matviews")
178        subquery = exp.union(
179            table_query,
180            exp.union(view_query, materialized_view_query, distinct=False),
181            distinct=False,
182        )
183        query = (
184            exp.select("*")
185            .from_(subquery.subquery(alias="objs"))
186            .where(exp.column("schema_name").eq(to_schema(schema_name).db))
187        )
188        if object_names:
189            query = query.where(exp.column("name").isin(*object_names))
190        df = self.fetchdf(query)
191        return [
192            DataObject(
193                catalog=catalog,
194                schema=row.schema_name,
195                name=row.name,
196                type=DataObjectType.from_str(row.type),  # type: ignore
197            )
198            for row in df.itertuples()
199        ]
200
201    def _get_current_schema(self) -> str:
202        """Returns the current default schema for the connection."""
203        result = self.fetchone(exp.select(exp.func("current_schema")))
204        if result and result[0]:
205            return result[0]
206        return "public"