moto/tests/test_glue/test_datacatalog.py

from __future__ import unicode_literals

import sure  # noqa
import re
from nose.tools import assert_raises
import boto3
from botocore.client import ClientError


from datetime import datetime
import pytz

from moto import mock_glue
from . import helpers


@mock_glue
def test_create_database():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    helpers.create_database(client, database_name)

    response = helpers.get_database(client, database_name)
    database = response["Database"]

    database.should.equal({"Name": database_name})


@mock_glue
def test_create_database_already_exists():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "cantcreatethisdatabasetwice"
    helpers.create_database(client, database_name)

    with assert_raises(ClientError) as exc:
        helpers.create_database(client, database_name)

    exc.exception.response["Error"]["Code"].should.equal("AlreadyExistsException")


@mock_glue
def test_get_database_not_exits():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "nosuchdatabase"

    with assert_raises(ClientError) as exc:
        helpers.get_database(client, database_name)

    exc.exception.response["Error"]["Code"].should.equal("EntityNotFoundException")
    exc.exception.response["Error"]["Message"].should.match(
        "Database nosuchdatabase not found"
    )


@mock_glue
def test_get_databases_empty():
    client = boto3.client("glue", region_name="us-east-1")
    response = client.get_databases()
    response["DatabaseList"].should.have.length_of(0)


@mock_glue
def test_get_databases_several_items():
    client = boto3.client("glue", region_name="us-east-1")
    database_name_1, database_name_2 = "firstdatabase", "seconddatabase"

    helpers.create_database(client, database_name_1)
    helpers.create_database(client, database_name_2)

    database_list = sorted(client.get_databases()["DatabaseList"], key=lambda x: x["Name"])
    database_list.should.have.length_of(2)
    database_list[0].should.equal({"Name": database_name_1})
    database_list[1].should.equal({"Name": database_name_2})


@mock_glue
def test_create_table():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    helpers.create_database(client, database_name)

    table_name = "myspecialtable"
    table_input = helpers.create_table_input(database_name, table_name)
    helpers.create_table(client, database_name, table_name, table_input)

    response = helpers.get_table(client, database_name, table_name)
    table = response["Table"]

    table["Name"].should.equal(table_input["Name"])
    table["StorageDescriptor"].should.equal(table_input["StorageDescriptor"])
    table["PartitionKeys"].should.equal(table_input["PartitionKeys"])


@mock_glue
def test_create_table_already_exists():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    helpers.create_database(client, database_name)

    table_name = "cantcreatethistabletwice"
    helpers.create_table(client, database_name, table_name)

    with assert_raises(ClientError) as exc:
        helpers.create_table(client, database_name, table_name)

    exc.exception.response["Error"]["Code"].should.equal("AlreadyExistsException")


@mock_glue
def test_get_tables():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    helpers.create_database(client, database_name)

    table_names = ["myfirsttable", "mysecondtable", "mythirdtable"]
    table_inputs = {}

    for table_name in table_names:
        table_input = helpers.create_table_input(database_name, table_name)
        table_inputs[table_name] = table_input
        helpers.create_table(client, database_name, table_name, table_input)

    response = helpers.get_tables(client, database_name)

    tables = response["TableList"]

    tables.should.have.length_of(3)

    for table in tables:
        table_name = table["Name"]
        table_name.should.equal(table_inputs[table_name]["Name"])
        table["StorageDescriptor"].should.equal(
            table_inputs[table_name]["StorageDescriptor"]
        )
        table["PartitionKeys"].should.equal(table_inputs[table_name]["PartitionKeys"])


@mock_glue
def test_get_table_versions():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    helpers.create_database(client, database_name)

    table_name = "myfirsttable"
    version_inputs = {}

    table_input = helpers.create_table_input(database_name, table_name)
    helpers.create_table(client, database_name, table_name, table_input)
    version_inputs["1"] = table_input

    columns = [{"Name": "country", "Type": "string"}]
    table_input = helpers.create_table_input(database_name, table_name, columns=columns)
    helpers.update_table(client, database_name, table_name, table_input)
    version_inputs["2"] = table_input

    # Updateing with an identical input should still create a new version
    helpers.update_table(client, database_name, table_name, table_input)
    version_inputs["3"] = table_input

    response = helpers.get_table_versions(client, database_name, table_name)

    vers = response["TableVersions"]

    vers.should.have.length_of(3)
    vers[0]["Table"]["StorageDescriptor"]["Columns"].should.equal([])
    vers[-1]["Table"]["StorageDescriptor"]["Columns"].should.equal(columns)

    for n, ver in enumerate(vers):
        n = str(n + 1)
        ver["VersionId"].should.equal(n)
        ver["Table"]["Name"].should.equal(table_name)
        ver["Table"]["StorageDescriptor"].should.equal(
            version_inputs[n]["StorageDescriptor"]
        )
        ver["Table"]["PartitionKeys"].should.equal(version_inputs[n]["PartitionKeys"])

    response = helpers.get_table_version(client, database_name, table_name, "3")
    ver = response["TableVersion"]

    ver["VersionId"].should.equal("3")
    ver["Table"]["Name"].should.equal(table_name)
    ver["Table"]["StorageDescriptor"]["Columns"].should.equal(columns)


@mock_glue
def test_get_table_version_not_found():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    helpers.create_database(client, database_name)
    helpers.create_table(client, database_name, table_name)

    with assert_raises(ClientError) as exc:
        helpers.get_table_version(client, database_name, "myfirsttable", "20")

    exc.exception.response["Error"]["Code"].should.equal("EntityNotFoundException")
    exc.exception.response["Error"]["Message"].should.match("version", re.I)


@mock_glue
def test_get_table_version_invalid_input():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    helpers.create_database(client, database_name)
    helpers.create_table(client, database_name, table_name)

    with assert_raises(ClientError) as exc:
        helpers.get_table_version(client, database_name, "myfirsttable", "10not-an-int")

    exc.exception.response["Error"]["Code"].should.equal("InvalidInputException")


@mock_glue
def test_get_table_not_exits():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    helpers.create_database(client, database_name)

    with assert_raises(ClientError) as exc:
        helpers.get_table(client, database_name, "myfirsttable")

    exc.exception.response["Error"]["Code"].should.equal("EntityNotFoundException")
    exc.exception.response["Error"]["Message"].should.match(
        "Table myfirsttable not found"
    )


@mock_glue
def test_get_table_when_database_not_exits():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "nosuchdatabase"

    with assert_raises(ClientError) as exc:
        helpers.get_table(client, database_name, "myfirsttable")

    exc.exception.response["Error"]["Code"].should.equal("EntityNotFoundException")
    exc.exception.response["Error"]["Message"].should.match(
        "Database nosuchdatabase not found"
    )


@mock_glue
def test_delete_table():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    helpers.create_database(client, database_name)

    table_name = "myspecialtable"
    table_input = helpers.create_table_input(database_name, table_name)
    helpers.create_table(client, database_name, table_name, table_input)

    result = client.delete_table(DatabaseName=database_name, Name=table_name)
    result["ResponseMetadata"]["HTTPStatusCode"].should.equal(200)

    # confirm table is deleted
    with assert_raises(ClientError) as exc:
        helpers.get_table(client, database_name, table_name)

    exc.exception.response["Error"]["Code"].should.equal("EntityNotFoundException")
    exc.exception.response["Error"]["Message"].should.match(
        "Table myspecialtable not found"
    )


@mock_glue
def test_batch_delete_table():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    helpers.create_database(client, database_name)

    table_name = "myspecialtable"
    table_input = helpers.create_table_input(database_name, table_name)
    helpers.create_table(client, database_name, table_name, table_input)

    result = client.batch_delete_table(
        DatabaseName=database_name, TablesToDelete=[table_name]
    )
    result["ResponseMetadata"]["HTTPStatusCode"].should.equal(200)

    # confirm table is deleted
    with assert_raises(ClientError) as exc:
        helpers.get_table(client, database_name, table_name)

    exc.exception.response["Error"]["Code"].should.equal("EntityNotFoundException")
    exc.exception.response["Error"]["Message"].should.match(
        "Table myspecialtable not found"
    )


@mock_glue
def test_get_partitions_empty():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    helpers.create_database(client, database_name)

    helpers.create_table(client, database_name, table_name)

    response = client.get_partitions(DatabaseName=database_name, TableName=table_name)

    response["Partitions"].should.have.length_of(0)


@mock_glue
def test_create_partition():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    values = ["2018-10-01"]
    helpers.create_database(client, database_name)

    helpers.create_table(client, database_name, table_name)

    before = datetime.now(pytz.utc)

    part_input = helpers.create_partition_input(
        database_name, table_name, values=values
    )
    helpers.create_partition(client, database_name, table_name, part_input)

    after = datetime.now(pytz.utc)

    response = client.get_partitions(DatabaseName=database_name, TableName=table_name)

    partitions = response["Partitions"]

    partitions.should.have.length_of(1)

    partition = partitions[0]

    partition["TableName"].should.equal(table_name)
    partition["StorageDescriptor"].should.equal(part_input["StorageDescriptor"])
    partition["Values"].should.equal(values)
    partition["CreationTime"].should.be.greater_than(before)
    partition["CreationTime"].should.be.lower_than(after)


@mock_glue
def test_create_partition_already_exist():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    values = ["2018-10-01"]
    helpers.create_database(client, database_name)

    helpers.create_table(client, database_name, table_name)

    helpers.create_partition(client, database_name, table_name, values=values)

    with assert_raises(ClientError) as exc:
        helpers.create_partition(client, database_name, table_name, values=values)

    exc.exception.response["Error"]["Code"].should.equal("AlreadyExistsException")


@mock_glue
def test_get_partition_not_found():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    values = ["2018-10-01"]
    helpers.create_database(client, database_name)

    helpers.create_table(client, database_name, table_name)

    with assert_raises(ClientError) as exc:
        helpers.get_partition(client, database_name, table_name, values)

    exc.exception.response["Error"]["Code"].should.equal("EntityNotFoundException")
    exc.exception.response["Error"]["Message"].should.match("partition")


@mock_glue
def test_batch_create_partition():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    helpers.create_database(client, database_name)

    helpers.create_table(client, database_name, table_name)

    before = datetime.now(pytz.utc)

    partition_inputs = []
    for i in range(0, 20):
        values = ["2018-10-{:2}".format(i)]
        part_input = helpers.create_partition_input(
            database_name, table_name, values=values
        )
        partition_inputs.append(part_input)

    client.batch_create_partition(
        DatabaseName=database_name,
        TableName=table_name,
        PartitionInputList=partition_inputs,
    )

    after = datetime.now(pytz.utc)

    response = client.get_partitions(DatabaseName=database_name, TableName=table_name)

    partitions = response["Partitions"]

    partitions.should.have.length_of(20)

    for idx, partition in enumerate(partitions):
        partition_input = partition_inputs[idx]

        partition["TableName"].should.equal(table_name)
        partition["StorageDescriptor"].should.equal(
            partition_input["StorageDescriptor"]
        )
        partition["Values"].should.equal(partition_input["Values"])
        partition["CreationTime"].should.be.greater_than(before)
        partition["CreationTime"].should.be.lower_than(after)


@mock_glue
def test_batch_create_partition_already_exist():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    values = ["2018-10-01"]
    helpers.create_database(client, database_name)

    helpers.create_table(client, database_name, table_name)

    helpers.create_partition(client, database_name, table_name, values=values)

    partition_input = helpers.create_partition_input(
        database_name, table_name, values=values
    )

    response = client.batch_create_partition(
        DatabaseName=database_name,
        TableName=table_name,
        PartitionInputList=[partition_input],
    )

    response.should.have.key("Errors")
    response["Errors"].should.have.length_of(1)
    response["Errors"][0]["PartitionValues"].should.equal(values)
    response["Errors"][0]["ErrorDetail"]["ErrorCode"].should.equal(
        "AlreadyExistsException"
    )


@mock_glue
def test_get_partition():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    helpers.create_database(client, database_name)

    helpers.create_table(client, database_name, table_name)

    values = [["2018-10-01"], ["2018-09-01"]]

    helpers.create_partition(client, database_name, table_name, values=values[0])
    helpers.create_partition(client, database_name, table_name, values=values[1])

    response = client.get_partition(
        DatabaseName=database_name, TableName=table_name, PartitionValues=values[1]
    )

    partition = response["Partition"]

    partition["TableName"].should.equal(table_name)
    partition["Values"].should.equal(values[1])


@mock_glue
def test_batch_get_partition():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    helpers.create_database(client, database_name)

    helpers.create_table(client, database_name, table_name)

    values = [["2018-10-01"], ["2018-09-01"]]

    helpers.create_partition(client, database_name, table_name, values=values[0])
    helpers.create_partition(client, database_name, table_name, values=values[1])

    partitions_to_get = [{"Values": values[0]}, {"Values": values[1]}]
    response = client.batch_get_partition(
        DatabaseName=database_name,
        TableName=table_name,
        PartitionsToGet=partitions_to_get,
    )

    partitions = response["Partitions"]
    partitions.should.have.length_of(2)

    partition = partitions[1]
    partition["TableName"].should.equal(table_name)
    partition["Values"].should.equal(values[1])


@mock_glue
def test_batch_get_partition_missing_partition():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    helpers.create_database(client, database_name)

    helpers.create_table(client, database_name, table_name)

    values = [["2018-10-01"], ["2018-09-01"], ["2018-08-01"]]

    helpers.create_partition(client, database_name, table_name, values=values[0])
    helpers.create_partition(client, database_name, table_name, values=values[2])

    partitions_to_get = [
        {"Values": values[0]},
        {"Values": values[1]},
        {"Values": values[2]},
    ]
    response = client.batch_get_partition(
        DatabaseName=database_name,
        TableName=table_name,
        PartitionsToGet=partitions_to_get,
    )

    partitions = response["Partitions"]
    partitions.should.have.length_of(2)

    partitions[0]["Values"].should.equal(values[0])
    partitions[1]["Values"].should.equal(values[2])


@mock_glue
def test_update_partition_not_found_moving():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"

    helpers.create_database(client, database_name)
    helpers.create_table(client, database_name, table_name)

    with assert_raises(ClientError) as exc:
        helpers.update_partition(
            client,
            database_name,
            table_name,
            old_values=["0000-00-00"],
            values=["2018-10-02"],
        )

    exc.exception.response["Error"]["Code"].should.equal("EntityNotFoundException")
    exc.exception.response["Error"]["Message"].should.match("partition")


@mock_glue
def test_update_partition_not_found_change_in_place():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    values = ["2018-10-01"]

    helpers.create_database(client, database_name)
    helpers.create_table(client, database_name, table_name)

    with assert_raises(ClientError) as exc:
        helpers.update_partition(
            client, database_name, table_name, old_values=values, values=values
        )

    exc.exception.response["Error"]["Code"].should.equal("EntityNotFoundException")
    exc.exception.response["Error"]["Message"].should.match("partition")


@mock_glue
def test_update_partition_cannot_overwrite():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    helpers.create_database(client, database_name)

    helpers.create_table(client, database_name, table_name)

    values = [["2018-10-01"], ["2018-09-01"]]

    helpers.create_partition(client, database_name, table_name, values=values[0])
    helpers.create_partition(client, database_name, table_name, values=values[1])

    with assert_raises(ClientError) as exc:
        helpers.update_partition(
            client, database_name, table_name, old_values=values[0], values=values[1]
        )

    exc.exception.response["Error"]["Code"].should.equal("AlreadyExistsException")


@mock_glue
def test_update_partition():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    values = ["2018-10-01"]

    helpers.create_database(client, database_name)
    helpers.create_table(client, database_name, table_name)
    helpers.create_partition(client, database_name, table_name, values=values)

    response = helpers.update_partition(
        client,
        database_name,
        table_name,
        old_values=values,
        values=values,
        columns=[{"Name": "country", "Type": "string"}],
    )

    response = client.get_partition(
        DatabaseName=database_name, TableName=table_name, PartitionValues=values
    )
    partition = response["Partition"]

    partition["TableName"].should.equal(table_name)
    partition["StorageDescriptor"]["Columns"].should.equal(
        [{"Name": "country", "Type": "string"}]
    )


@mock_glue
def test_update_partition_move():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    values = ["2018-10-01"]
    new_values = ["2018-09-01"]

    helpers.create_database(client, database_name)
    helpers.create_table(client, database_name, table_name)
    helpers.create_partition(client, database_name, table_name, values=values)

    response = helpers.update_partition(
        client,
        database_name,
        table_name,
        old_values=values,
        values=new_values,
        columns=[{"Name": "country", "Type": "string"}],
    )

    with assert_raises(ClientError) as exc:
        helpers.get_partition(client, database_name, table_name, values)

    # Old partition shouldn't exist anymore
    exc.exception.response["Error"]["Code"].should.equal("EntityNotFoundException")

    response = client.get_partition(
        DatabaseName=database_name, TableName=table_name, PartitionValues=new_values
    )
    partition = response["Partition"]

    partition["TableName"].should.equal(table_name)
    partition["StorageDescriptor"]["Columns"].should.equal(
        [{"Name": "country", "Type": "string"}]
    )


@mock_glue
def test_delete_partition():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    values = ["2018-10-01"]
    helpers.create_database(client, database_name)
    helpers.create_table(client, database_name, table_name)

    part_input = helpers.create_partition_input(
        database_name, table_name, values=values
    )
    helpers.create_partition(client, database_name, table_name, part_input)

    client.delete_partition(
        DatabaseName=database_name, TableName=table_name, PartitionValues=values
    )

    response = client.get_partitions(DatabaseName=database_name, TableName=table_name)
    partitions = response["Partitions"]
    partitions.should.be.empty


@mock_glue
def test_delete_partition_bad_partition():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    values = ["2018-10-01"]
    helpers.create_database(client, database_name)
    helpers.create_table(client, database_name, table_name)

    with assert_raises(ClientError) as exc:
        client.delete_partition(
            DatabaseName=database_name, TableName=table_name, PartitionValues=values
        )

    exc.exception.response["Error"]["Code"].should.equal("EntityNotFoundException")


@mock_glue
def test_batch_delete_partition():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    helpers.create_database(client, database_name)
    helpers.create_table(client, database_name, table_name)

    partition_inputs = []
    for i in range(0, 20):
        values = ["2018-10-{:2}".format(i)]
        part_input = helpers.create_partition_input(
            database_name, table_name, values=values
        )
        partition_inputs.append(part_input)

    client.batch_create_partition(
        DatabaseName=database_name,
        TableName=table_name,
        PartitionInputList=partition_inputs,
    )

    partition_values = [{"Values": p["Values"]} for p in partition_inputs]

    response = client.batch_delete_partition(
        DatabaseName=database_name,
        TableName=table_name,
        PartitionsToDelete=partition_values,
    )

    response.should_not.have.key("Errors")


@mock_glue
def test_batch_delete_partition_with_bad_partitions():
    client = boto3.client("glue", region_name="us-east-1")
    database_name = "myspecialdatabase"
    table_name = "myfirsttable"
    helpers.create_database(client, database_name)
    helpers.create_table(client, database_name, table_name)

    partition_inputs = []
    for i in range(0, 20):
        values = ["2018-10-{:2}".format(i)]
        part_input = helpers.create_partition_input(
            database_name, table_name, values=values
        )
        partition_inputs.append(part_input)

    client.batch_create_partition(
        DatabaseName=database_name,
        TableName=table_name,
        PartitionInputList=partition_inputs,
    )

    partition_values = [{"Values": p["Values"]} for p in partition_inputs]

    partition_values.insert(5, {"Values": ["2018-11-01"]})
    partition_values.insert(10, {"Values": ["2018-11-02"]})
    partition_values.insert(15, {"Values": ["2018-11-03"]})

    response = client.batch_delete_partition(
        DatabaseName=database_name,
        TableName=table_name,
        PartitionsToDelete=partition_values,
    )

    response.should.have.key("Errors")
    response["Errors"].should.have.length_of(3)
    error_partitions = map(lambda x: x["PartitionValues"], response["Errors"])
    ["2018-11-01"].should.be.within(error_partitions)
    ["2018-11-02"].should.be.within(error_partitions)
    ["2018-11-03"].should.be.within(error_partitions)