Feat: Add Postgres Support (built-in, Airflow, dbt) (#677)

* add postgres support

* remove single line returns

* cleanup

* fix type

* fix test

* fix primary key

* update test

Author: eakmanrq

docs/integrations/engines.md

@@ -142,6 +142,43 @@ sqlmesh_airflow = SQLMeshAirflow(
 ## DuckDB - Airflow
 DuckDB only works when running locally; therefore it does not support Airflow. 
 
+# Postgres
+## Postgres - Local/Built-in Scheduler
+| Option            | Description                                                                                                                          |  Type  | Required |
+|-------------------|--------------------------------------------------------------------------------------------------------------------------------------|:------:|:--------:|
+| `host`            | The hostname of the Postgres server                                                                                                  | string |    Y     |
+| `user`            | The username to use for authentication with the Postgres server                                                                      | string |    Y     |
+| `password`        | The password to use for authentication with the Postgres server                                                                      | string |    Y     |
+| `port`            | The port number of the Postgres server                                                                                               |  int   |    Y     |
+| `database`        | The name of the database instance to connect to                                                                                      | string |    Y     |
+| `keepalives_idle` | The number of seconds between each keepalive packet sent to the server. If set to 0, the system default will be used. (Default: `0`) |  int   |    N     |
+| `connect_timeout` | The number of seconds to wait for the connection to the server. (Default: `10`)                                                      |  int   |    N     |
+| `role`            | The role to use for authentication with the Postgres server                                                                          | string |    N     |
+| `sslmode`         | The security of the connection to the Postgres server.                                                                               | string |    N     |
+
+## Postgres - Airflow
+**Engine Name:** `postgres`
+
+The SQLMesh Postgres Operator is similar to the [PostgresOperator](https://airflow.apache.org/docs/apache-airflow-providers-postgres/stable/_api/airflow/providers/postgres/operators/postgres/index.html), and relies on the same [PostgresHook](https://airflow.apache.org/docs/apache-airflow-providers-postgres/stable/_api/airflow/providers/postgres/hooks/postgres/index.html) implementation.
+
+To enable support for this operator, the Airflow Postgres provider package should be installed on the target Airflow cluster along with SQLMesh with the Postgres extra:
+```
+pip install "apache-airflow-providers-postgres"
+pip install "sqlmesh[postgres]"
+```
+
+The operator requires an [Airflow connection](https://airflow.apache.org/docs/apache-airflow/stable/howto/connection.html) to determine the target Postgres account. Refer to [Postgres connection](https://airflow.apache.org/docs/apache-airflow-providers-postgres/stable/connections/postgres.html) for more details.
+
+By default, the connection ID is set to `postgres_default`, but can be overridden using the `engine_operator_args` parameter to the `SQLMeshAirflow` instance as in the example below:
+```python linenums="1"
+sqlmesh_airflow = SQLMeshAirflow(
+    "postgres",
+    engine_operator_args={
+        "postgres_conn_id": "<Connection ID>"
+    },
+)
+```
+
 # Redshift
 ## Redshift - Local/Built-in Scheduler
 | Option                  | Description                                                                                                 |  Type  | Required |

setup.cfg

@@ -63,6 +63,9 @@ ignore_missing_imports = True
 [mypy-fsspec]
 ignore_missing_imports = True
 
+[mypy-psycopg2.*]
+ignore_missing_imports = True
+
 [autoflake]
 in-place = True
 expand-star-imports = True

setup.py

@@ -47,6 +47,13 @@
         "fsspec",
     ],
     extras_require={
+        "bigquery": [
+            "google-cloud-bigquery[pandas]",
+        ],
+        "databricks": [
+            "databricks-sql-connector",
+            "databricks-cli",
+        ],
         "dev": [
             f"apache-airflow=={os.environ.get('AIRFLOW_VERSION', '2.3.3')}",
             "autoflake==1.7.7",
@@ -76,27 +83,23 @@
             "types-pytz",
             "types-requests==2.28.8",
         ],
-        "web": [
-            "fastapi==0.95.0",
-            "hyperscript==0.0.1",
-            "pyarrow==11.0.0",
-            "uvicorn==0.21.1",
-        ],
-        "snowflake": [
-            "snowflake-connector-python[pandas]",
-        ],
-        "bigquery": [
-            "google-cloud-bigquery[pandas]",
+        "dbt": [
+            "dbt-core",
         ],
-        "databricks": [
-            "databricks-sql-connector",
-            "databricks-cli",
+        "postgres": [
+            "psycopg2",
         ],
         "redshift": [
             "redshift_connector",
         ],
-        "dbt": [
-            "dbt-core",
+        "snowflake": [
+            "snowflake-connector-python[pandas]",
+        ],
+        "web": [
+            "fastapi==0.95.0",
+            "hyperscript==0.0.1",
+            "pyarrow==11.0.0",
+            "uvicorn==0.21.1",
         ],
     },
     classifiers=[

sqlmesh/core/config/__init__.py

@@ -6,6 +6,7 @@
     DatabricksSparkSessionConnectionConfig,
     DatabricksSQLConnectionConfig,
     DuckDBConnectionConfig,
+    PostgresConnectionConfig,
     RedshiftConnectionConfig,
     SnowflakeConnectionConfig,
 )

sqlmesh/core/config/connection.py

@@ -516,15 +516,56 @@ def _connection_factory(self) -> t.Callable:
         return connect
 
 
+class PostgresConnectionConfig(_ConnectionConfig):
+    host: str
+    user: str
+    password: str
+    port: int
+    database: str
+    keepalives_idle: int = 0
+    connect_timeout: int = 10
+    role: t.Optional[str] = None
+    sslmode: t.Optional[str] = None
+
+    concurrent_tasks: int = 4
+
+    type_: Literal["postgres"] = Field(alias="type", default="postgres")
+
+    @property
+    def _connection_kwargs_keys(self) -> t.Set[str]:
+        return {
+            "host",
+            "user",
+            "password",
+            "port",
+            "database",
+            "keepalives_idle",
+            "connect_timeout",
+            "role",
+            "sslmode",
+        }
+
+    @property
+    def _engine_adapter(self) -> t.Type[EngineAdapter]:
+        return engine_adapter.PostgresEngineAdapter
+
+    @property
+    def _connection_factory(self) -> t.Callable:
+        from psycopg2 import connect
+
+        return connect
+
+
 ConnectionConfig = Annotated[
     t.Union[
-        DuckDBConnectionConfig,
-        SnowflakeConnectionConfig,
+        BigQueryConnectionConfig,
         DatabricksSQLConnectionConfig,
         DatabricksSparkSessionConnectionConfig,
         DatabricksConnectionConfig,
-        BigQueryConnectionConfig,
+        DuckDBConnectionConfig,
+        PostgresConnectionConfig,
         RedshiftConnectionConfig,
+        SnowflakeConnectionConfig,
     ],
     Field(discriminator="type_"),
 ]

sqlmesh/core/engine_adapter/__init__.py

@@ -27,13 +27,16 @@
     "mssql": EngineAdapterWithIndexSupport,
 }
 
+DIALECT_ALIASES = {
+    "postgresql": "postgres",
+}
+
 
 def create_engine_adapter(
     connection_factory: t.Callable[[], t.Any], dialect: str, multithreaded: bool = False
 ) -> EngineAdapter:
     dialect = dialect.lower()
-    if dialect == "postgresql":
-        dialect = "postgres"
+    dialect = DIALECT_ALIASES.get(dialect, dialect)
     if dialect == "databricks":
         try:
             from pyspark.sql import SparkSession