path: root/tests/test_pgexecute.py

from textwrap import dedent

import psycopg2
import pytest
from unittest.mock import patch, MagicMock
from pgspecial.main import PGSpecial, NO_QUERY
from utils import run, dbtest, requires_json, requires_jsonb

from pgcli.main import PGCli
from pgcli.packages.parseutils.meta import FunctionMetadata


def function_meta_data(
    func_name,
    schema_name="public",
    arg_names=None,
    arg_types=None,
    arg_modes=None,
    return_type=None,
    is_aggregate=False,
    is_window=False,
    is_set_returning=False,
    is_extension=False,
    arg_defaults=None,
):
    return FunctionMetadata(
        schema_name,
        func_name,
        arg_names,
        arg_types,
        arg_modes,
        return_type,
        is_aggregate,
        is_window,
        is_set_returning,
        is_extension,
        arg_defaults,
    )


@dbtest
def test_conn(executor):
    run(executor, """create table test(a text)""")
    run(executor, """insert into test values('abc')""")
    assert run(executor, """select * from test""", join=True) == dedent(
        """\
        +-----+
        | a   |
        |-----|
        | abc |
        +-----+
        SELECT 1"""
    )


@dbtest
def test_copy(executor):
    executor_copy = executor.copy()
    run(executor_copy, """create table test(a text)""")
    run(executor_copy, """insert into test values('abc')""")
    assert run(executor_copy, """select * from test""", join=True) == dedent(
        """\
        +-----+
        | a   |
        |-----|
        | abc |
        +-----+
        SELECT 1"""
    )


@dbtest
def test_bools_are_treated_as_strings(executor):
    run(executor, """create table test(a boolean)""")
    run(executor, """insert into test values(True)""")
    assert run(executor, """select * from test""", join=True) == dedent(
        """\
        +------+
        | a    |
        |------|
        | True |
        +------+
        SELECT 1"""
    )


@dbtest
def test_expanded_slash_G(executor, pgspecial):
    # Tests whether we reset the expanded output after a \G.
    run(executor, """create table test(a boolean)""")
    run(executor, """insert into test values(True)""")
    results = run(executor, r"""select * from test \G""", pgspecial=pgspecial)
    assert pgspecial.expanded_output == False


@dbtest
def test_schemata_table_views_and_columns_query(executor):
    run(executor, "create table a(x text, y text)")
    run(executor, "create table b(z text)")
    run(executor, "create view d as select 1 as e")
    run(executor, "create schema schema1")
    run(executor, "create table schema1.c (w text DEFAULT 'meow')")
    run(executor, "create schema schema2")

    # schemata
    # don't enforce all members of the schemas since they may include postgres
    # temporary schemas
    assert set(executor.schemata()) >= {
        "public",
        "pg_catalog",
        "information_schema",
        "schema1",
        "schema2",
    }
    assert executor.search_path() == ["pg_catalog", "public"]

    # tables
    assert set(executor.tables()) >= {
        ("public", "a"),
        ("public", "b"),
        ("schema1", "c"),
    }

    assert set(executor.table_columns()) >= {
        ("public", "a", "x", "text", False, None),
        ("public", "a", "y", "text", False, None),
        ("public", "b", "z", "text", False, None),
        ("schema1", "c", "w", "text", True, "'meow'::text"),
    }

    # views
    assert set(executor.views()) >= {("public", "d")}

    assert set(executor.view_columns()) >= {
        ("public", "d", "e", "integer", False, None)
    }


@dbtest
def test_foreign_key_query(executor):
    run(executor, "create schema schema1")
    run(executor, "create schema schema2")
    run(executor, "create table schema1.parent(parentid int PRIMARY KEY)")
    run(
        executor,
        "create table schema2.child(childid int PRIMARY KEY, motherid int REFERENCES schema1.parent)",
    )

    assert set(executor.foreignkeys()) >= {
        ("schema1", "parent", "parentid", "schema2", "child", "motherid")
    }


@dbtest
def test_functions_query(executor):
    run(
        executor,
        """create function func1() returns int
                     language sql as $$select 1$$""",
    )
    run(executor, "create schema schema1")
    run(
        executor,
        """create function schema1.func2() returns int
                     language sql as $$select 2$$""",
    )

    run(
        executor,
        """create function func3()
                     returns table(x int, y int) language sql
                     as $$select 1, 2 from generate_series(1,5)$$;""",
    )

    run(
        executor,
        """create function func4(x int) returns setof int language sql
                     as $$select generate_series(1,5)$$;""",
    )

    funcs = set(executor.functions())
    assert funcs >= {
        function_meta_data(func_name="func1", return_type="integer"),
        function_meta_data(
            func_name="func3",
            arg_names=["x", "y"],
            arg_types=["integer", "integer"],
            arg_modes=["t", "t"],
            return_type="record",
            is_set_returning=True,
        ),
        function_meta_data(
            schema_name="public",
            func_name="func4",
            arg_names=("x",),
            arg_types=("integer",),
            return_type="integer",
            is_set_returning=True,
        ),
        function_meta_data(
            schema_name="schema1", func_name="func2", return_type="integer"
        ),
    }


@dbtest
def test_datatypes_query(executor):
    run(executor, "create type foo AS (a int, b text)")

    types = list(executor.datatypes())
    assert types == [("public", "foo")]


@dbtest
def test_database_list(executor):
    databases = executor.databases()
    assert "_test_db" in databases


@dbtest
def test_invalid_syntax(executor, exception_formatter):
    result = run(executor, "invalid syntax!", exception_formatter=exception_formatter)
    assert 'syntax error at or near "invalid"' in result[0]


@dbtest
def test_invalid_column_name(executor, exception_formatter):
    result = run(
        executor, "select invalid command", exception_formatter=exception_formatter
    )
    assert 'column "invalid" does not exist' in result[0]


@pytest.fixture(params=[True, False])
def expanded(request):
    return request.param


@dbtest
def test_unicode_support_in_output(executor, expanded):
    run(executor, "create table unicodechars(t text)")
    run(executor, "insert into unicodechars (t) values ('é')")

    # See issue #24, this raises an exception without proper handling
    assert "é" in run(
        executor, "select * from unicodechars", join=True, expanded=expanded
    )


@dbtest
def test_not_is_special(executor, pgspecial):
    """is_special is set to false for database queries."""
    query = "select 1"
    result = list(executor.run(query, pgspecial=pgspecial))
    success, is_special = result[0][5:]
    assert success == True
    assert is_special == False


@dbtest
def test_execute_from_file_no_arg(executor, pgspecial):
    r"""\i without a filename returns an error."""
    result = list(executor.run(r"\i", pgspecial=pgspecial))
    status, sql, success, is_special = result[0][3:]
    assert "missing required argument" in status
    assert success == False
    assert is_special == True


@dbtest
@patch("pgcli.main.os")
def test_execute_from_file_io_error(os, executor, pgspecial):
    r"""\i with an os_error returns an error."""
    # Inject an OSError.
    os.path.expanduser.side_effect = OSError("test")

    # Check the result.
    result = list(executor.run(r"\i test", pgspecial=pgspecial))
    status, sql, success, is_special = result[0][3:]
    assert status == "test"
    assert success == False
    assert is_special == True


@dbtest
def test_multiple_queries_same_line(executor):
    result = run(executor, "select 'foo'; select 'bar'")
    assert len(result) == 12  # 2 * (output+status) * 3 lines
    assert "foo" in result[3]
    assert "bar" in result[9]


@dbtest
def test_multiple_queries_with_special_command_same_line(executor, pgspecial):
    result = run(executor, r"select 'foo'; \d", pgspecial=pgspecial)
    assert len(result) == 11  # 2 * (output+status) * 3 lines
    assert "foo" in result[3]
    # This is a lame check. :(
    assert "Schema" in result[7]


@dbtest
def test_multiple_queries_same_line_syntaxerror(executor, exception_formatter):
    result = run(
        executor,
        "select 'fooé'; invalid syntax é",
        exception_formatter=exception_formatter,
    )
    assert "fooé" in result[3]
    assert 'syntax error at or near "invalid"' in result[-1]


@pytest.fixture
def pgspecial():
    return PGCli().pgspecial


@dbtest
def test_special_command_help(executor, pgspecial):
    result = run(executor, "\\?", pgspecial=pgspecial)[1].split("|")
    assert "Command" in result[1]
    assert "Descript