def test_postgres_operator_test_multi(self):
sql = [
"TRUNCATE TABLE test_airflow",
"INSERT INTO test_airflow VALUES ('X')",
]
from airflow.operators.postgres_operator import PostgresOperator
t = PostgresOperator(
task_id='postgres_operator_test_multi', sql=sql, dag=self.dag)
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
def test_postgres_operator_test_multi(self):
sql = [
"CREATE TABLE IF NOT EXISTS test_airflow (dummy VARCHAR(50))",
"TRUNCATE TABLE test_airflow",
"INSERT INTO test_airflow VALUES ('X')",
]
from airflow.operators.postgres_operator import PostgresOperator
op = PostgresOperator(
task_id='postgres_operator_test_multi', sql=sql, dag=self.dag)
op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
def test_vacuum(self):
"""
Verifies the VACUUM operation runs well with the PostgresOperator
"""
from airflow.operators.postgres_operator import PostgresOperator
sql = "VACUUM ANALYZE;"
t = PostgresOperator(
task_id='postgres_operator_test_vacuum',
sql=sql,
dag=self.dag,
autocommit=True)
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
def test_postgres_operator_test(self):
sql = """
CREATE TABLE IF NOT EXISTS test_airflow (
dummy VARCHAR(50)
);
"""
from airflow.operators.postgres_operator import PostgresOperator
t = PostgresOperator(task_id='basic_postgres', sql=sql, dag=self.dag)
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
autocommitTask = operators.postgres_operator.PostgresOperator(
task_id='basic_postgres_with_autocommit',
sql=sql,
dag=self.dag,
autocommit=True)
autocommitTask.run(
start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE,
ignore_ti_state=True)
def test_overwrite_schema(self):
"""
Verifies option to overwrite connection schema
"""
from airflow.operators.postgres_operator import PostgresOperator
sql = "SELECT 1;"
t = PostgresOperator(
task_id='postgres_operator_test_schema_overwrite',
sql=sql,
dag=self.dag,
autocommit=True,
database="foobar",
)
from psycopg2._psycopg import OperationalError
try:
t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
ignore_ti_state=True)
except OperationalError as e:
assert 'database "foobar" does not exist' in str(e)
