Source code for operators.create_schema_postgres

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.providers.postgres.hooks.postgres import PostgresHook




class CreateSchemaPostgresOperator(BaseOperator):
    """
    Create a schema in a Postgres database.
    :param postgres_conn_id: the Postgres connection ID
    :type postgres_conn_id: str
    :param schema_name: the name of the schema to create
    :type schema_name: str
    :param if_not_exists: whether to add the ``IF NOT EXISTS`` clause to the SQL statement.
        Default is ``False``.
    :type env: bool
    """

    ui_color = "#e4e3ef"
    template_fields = (
        "postgres_conn_id",
        "schema_name"
    )
    @apply_defaults
    def __init__(
        self, postgres_conn_id, schema_name, if_not_exists=False, *args, **kwargs
    ):
        super(CreateSchemaPostgresOperator, self).__init__(*args, **kwargs)
        self.postgres_conn_id = postgres_conn_id
        self.schema_name = schema_name
        self.if_not_exists = if_not_exists



    def execute(self, context):
        ine = ""
        if self.if_not_exists:
            ine = "IF NOT EXISTS "

        sql = "CREATE SCHEMA " + ine + self.schema_name + ";"
        
        pg_hook = PostgresHook(postgres_conn_id=self.postgres_conn_id)
        pg_hook.run(sql, autocommit=True)