snowflake-connector-python/test/test_transaction.py at master · byule/snowflake-connector-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2012-2019 Snowflake Computing Inc. All right reserved.
#

import pytest

import snowflake.connector


def test_transaction(conn_cnx, db_parameters):
    u"""
    Transaction API
    """
    with conn_cnx() as cnx:
        cnx.cursor().execute(u"create table {name} (c1 int)".format(
            name=db_parameters['name']))
        cnx.cursor().execute(u"insert into {name}(c1) "
                             u"values(1234),(3456)".format(
            name=db_parameters['name']))
        c = cnx.cursor()
        c.execute(u"select * from {name}".format(name=db_parameters['name']))
        total = 0
        for rec in c:
            total += rec[0]
        assert total == 4690, u'total integer'

        #
        cnx.cursor().execute(u"begin")
        cnx.cursor().execute(
            u"insert into {name}(c1) values(5678),(7890)".format(
                name=db_parameters['name']))
        c = cnx.cursor()
        c.execute(u"select * from {name}".format(name=db_parameters['name']))
        total = 0
        for rec in c:
            total += rec[0]
        assert total == 18258, u'total integer'
        cnx.rollback()

        c.execute(u"select * from {name}".format(name=db_parameters['name']))
        total = 0
        for rec in c:
            total += rec[0]
        assert total == 4690, u'total integer'

        #
        cnx.cursor().execute(u"begin")
        cnx.cursor().execute(
            u"insert into {name}(c1) values(2345),(6789)".format(
                name=db_parameters['name']))
        c = cnx.cursor()
        c.execute(u"select * from {name}".format(name=db_parameters['name']))
        total = 0
        for rec in c:
            total += rec[0]
        assert total == 13824, u'total integer'
        cnx.commit()
        cnx.rollback()
        c = cnx.cursor()
        c.execute(u"select * from {name}".format(name=db_parameters['name']))
        total = 0
        for rec in c:
            total += rec[0]
        assert total == 13824, u'total integer'


def test_connection_context_manager(request, db_parameters):
    db_config = {
        'protocol': db_parameters['protocol'],
        'account': db_parameters['account'],
        'user': db_parameters['user'],
        'password': db_parameters['password'],
        'host': db_parameters['host'],
        'port': db_parameters['port'],
        'database': db_parameters['database'],
        'schema': db_parameters['schema'],
        'timezone': 'UTC',
    }

    def fin():
        with snowflake.connector.connect(**db_config) as cnx:
            cnx.cursor().execute("""
DROP TABLE IF EXISTS {name}
""".format(name=db_parameters['name']))

    request.addfinalizer(fin)

    try:
        with snowflake.connector.connect(**db_config) as cnx:
            cnx.autocommit(False)
            cnx.cursor().execute("""
CREATE OR REPLACE TABLE {name} (cc1 int)
""".format(name=db_parameters['name']))
            cnx.cursor().execute("""
INSERT INTO {name} VALUES(1),(2),(3)
""".format(name=db_parameters['name']))
            ret = cnx.cursor().execute("""
SELECT SUM(cc1) FROM {name}
""".format(name=db_parameters['name'])).fetchone()
            assert ret[0] == 6
            cnx.commit()
            cnx.cursor().execute("""
INSERT INTO {name} VALUES(4),(5),(6)
""".format(name=db_parameters['name']))
            ret = cnx.cursor().execute("""
SELECT SUM(cc1) FROM {name}
""".format(name=db_parameters['name'])).fetchone()
            assert ret[0] == 21
            cnx.cursor().execute("""
SELECT WRONG SYNTAX QUERY
""".format(name=db_parameters['name']))
            raise Exception("Failed to cause the syntax error")
    except snowflake.connector.Error as e:
        # syntax error should be caught here
        # and the last change must have been rollbacked
        with snowflake.connector.connect(**db_config) as cnx:
            ret = cnx.cursor().execute("""
SELECT SUM(cc1) FROM {name}
""".format(name=db_parameters['name'])).fetchone()
            assert ret[0] == 6