Beispiel #1
0
    def test_read_end_stream(self, m):
        '''reads the entire table from the database and validates that the
        stream returns None to indicate the end'''

        mock_recs = [{
            'id': 1,
            'col1': 'foo1',
            'col2': 'bar1'
        }, {
            'id': 2,
            'col1': 'foo2',
            'col2': 'bar2'
        }, {
            'id': 3,
            'col1': 'foo3',
            'col2': 'bar3'
        }]

        inst = Postgres(self.source, OPTIONS)
        inst.tables = [{'value': 'my_schema.foo_bar'}]
        result_order = [mock_recs, []]
        m.return_value.cursor.return_value.fetchall.side_effect = result_order

        rows = inst.read()
        self.assertEqual(len(rows), len(mock_recs))

        empty = inst.read()
        self.assertEqual(empty, [])
        end = inst.read()
        self.assertEqual(end, None)
Beispiel #2
0
    def test_recover_from_state(self, mock_connect, mock_execute, _, __):
        """continues to read a table from the saved state"""

        tables = [
            {
                'value': 'public.test1'
            },
            {
                'value': 'public.test2'
            },
            {
                'value': 'public.test3'
            },
        ]
        last_index = 1

        self.source['state'] = {'last_index': last_index, 'max_value': 100}
        inst = Postgres(self.source, OPTIONS)
        inst.tables = tables
        cursor_return_value = mock_connect.return_value.cursor.return_value
        cursor_return_value.fetchall.return_value = [{
            'id': 101
        }, {
            'id': 102
        }, {
            'id': 103
        }]

        inst.read()
        first_query = mock_execute.call_args_list[0][0][0]
        self.assertTrue(
            "\"inckey\" >= 'incval' AND \"inckey\" <= '100'" in first_query)
        self.assertTrue('FROM "public"."test2"' in first_query)
Beispiel #3
0
 def test_connect_other_error(self, mock_connect):
     inst = Postgres(self.source, OPTIONS)
     inst.tables = [{'value': 'schema.foo'}]
     msg = 'something unexpected'
     mock_connect.side_effect = psycopg2.OperationalError(msg)
     with self.assertRaises(psycopg2.OperationalError):
         inst.get_tables()
Beispiel #4
0
    def test_read(self, m):
        '''reads a table from the database and validates that each row
        has a __tablename and __schemaname column'''

        mock_recs = [{
            'id': 1,
            'col1': 'foo1',
            'col2': 'bar1'
        }, {
            'id': 2,
            'col1': 'foo2',
            'col2': 'bar2'
        }, {
            'id': 3,
            'col1': 'foo3',
            'col2': 'bar3'
        }]

        inst = Postgres(self.source, OPTIONS)
        inst.tables = [{'value': 'my_schema.foo_bar'}]
        m.return_value.cursor.return_value.fetchall.return_value = mock_recs

        rows = inst.read()
        self.assertEqual(len(rows), len(mock_recs))
        for x in range(0, len(rows)):
            self.assertEqual(rows[x]['__tablename'], 'foo_bar')
            self.assertEqual(rows[x]['__schemaname'], 'my_schema')
Beispiel #5
0
 def test_connect_auth_error(self, mock_connect):
     inst = Postgres(self.source, OPTIONS)
     inst.tables = [{'value': 'schema.foo'}]
     msg = 'authentication failed'
     mock_connect.side_effect = psycopg2.OperationalError(msg)
     with self.assertRaises(PanoplyException):
         inst.get_tables()
Beispiel #6
0
    def test_read_retries(self, mock_connect):
        inst = Postgres(self.source, OPTIONS)
        inst.tables = [{'value': 'my_schema.foo_bar'}]
        mock_connect.side_effect = psycopg2.DatabaseError('TestRetiresError')
        with self.assertRaises(psycopg2.DatabaseError):
            inst.read()

        self.assertEqual(mock_connect.call_count, postgres.source.MAX_RETRIES)
Beispiel #7
0
    def test_incremental(self, mock_connect):
        inst = Postgres(self.source, OPTIONS)
        inst.tables = [{'value': 'schema.foo'}]
        inst.read()

        q = ('DECLARE cur CURSOR FOR '
             'SELECT * FROM "schema"."foo" WHERE inckey > \'incval\'')
        execute_mock = mock_connect.return_value.cursor.return_value.execute
        execute_mock.assert_has_calls([mock.call(q)], True)
Beispiel #8
0
    def test_batch_size(self, mock_connect, mock_execute):
        customBatchSize = 42
        self.source['__batchSize'] = customBatchSize
        inst = Postgres(self.source, OPTIONS)
        inst.tables = [{'value': 'my_schema.foo_bar'}]

        cursor_return_value = mock_connect.return_value.cursor.return_value
        cursor_return_value.fetchall.return_value = self.mock_recs

        inst.read()
        second_query = mock_execute.call_args_list[1][0][0]
        txt = 'FETCH FORWARD %s' % customBatchSize
        self.assertTrue(second_query.startswith(txt))
Beispiel #9
0
    def test_read(self, mock_connect):
        '''reads a table from the database and validates that each row
        has a __tablename and __schemaname column'''

        inst = Postgres(self.source, OPTIONS)
        inst.tables = [{'value': 'my_schema.foo_bar'}]
        cursor_return_value = mock_connect.return_value.cursor.return_value
        cursor_return_value.fetchall.return_value = self.mock_recs

        rows = inst.read()
        self.assertEqual(len(rows), len(self.mock_recs))
        for x in range(0, len(rows)):
            self.assertEqual(rows[x]['__tablename'], 'foo_bar')
            self.assertEqual(rows[x]['__schemaname'], 'my_schema')
Beispiel #10
0
    def test_no_state_for_empty_results(self, mock_connect, mock_state):
        '''before returning a batch of data, the sources state should be
        reported as well as having the state ID appended to each data object'''

        inst = Postgres(self.source, OPTIONS)
        table_name = 'my_schema.foo_bar'
        inst.tables = [{'value': table_name}]
        result_order = [[], []]
        cursor_return_value = mock_connect.return_value.cursor.return_value
        cursor_return_value.fetchall.side_effect = result_order

        rows = inst.read()

        # State function was called with relevant table name and row count
        mock_state.assert_not_called()
Beispiel #11
0
    def test_recover_from_state(self, mock_connect, mock_execute):
        ''' continues to read a table from the saved state '''

        table_offset = 100
        self.source['state'] = {'my_schema.foo_bar': table_offset}
        inst = Postgres(self.source, OPTIONS)
        inst.tables = [{'value': 'my_schema.foo_bar'}]
        cursor_return_value = mock_connect.return_value.cursor.return_value
        cursor_return_value.fetchall.return_value = self.mock_recs

        inst.read()
        first_query = mock_execute.call_args_list[0][0][0]
        # The query should use `OFFSET` to skip already collected data
        self.assertTrue(first_query.endswith('OFFSET %s' % table_offset))
        # Three records were returned so the loaded count should be OFFSET + 3
        self.assertEqual(inst.loaded, table_offset + len(self.mock_recs))
Beispiel #12
0
    def test_read_end_stream(self, mock_connect):
        '''reads the entire table from the database and validates that the
        stream returns None to indicate the end'''

        inst = Postgres(self.source, OPTIONS)
        inst.tables = [{'value': 'my_schema.foo_bar'}]
        result_order = [self.mock_recs, []]
        cursor_return_value = mock_connect.return_value.cursor.return_value
        cursor_return_value.fetchall.side_effect = result_order

        rows = inst.read()
        self.assertEqual(len(rows), len(self.mock_recs))

        empty = inst.read()
        self.assertEqual(empty, [])
        end = inst.read()
        self.assertEqual(end, None)
Beispiel #13
0
    def test_retry_with_last_values(self, mock_connect, mock_metadata, _):

        mock_metadata.side_effect = lambda *args: [{
            'attname': 'col1',
            'indisunique': True,
            'indisprimary': True
        }, {
            'attname': 'col2',
            'indisunique': True,
            'indisprimary': True
        }]

        inst = Postgres(self.source, OPTIONS)
        inst.tables = [{'value': 'my_schema.foo_bar'}]
        inst.batch_size = 1

        cursor_execute = mock_connect.return_value.cursor.return_value.execute
        cursor_execute.side_effect = [
            lambda *args: None,
            lambda *args: None,
            psycopg2.DatabaseError('TestRetriesError'),
            lambda *args: None,
            lambda *args: None,
        ]

        cursor_return_value = mock_connect.return_value.cursor.return_value
        cursor_return_value.fetchall.return_value = self.mock_recs

        # First read no error
        inst.read()
        # Raise retry error
        inst.read()

        # Extract mock call arguments
        args = mock_connect.return_value.cursor.return_value\
            .execute.call_args_list
        args = [r[0] for r, _ in args]
        args = filter(lambda x: 'DECLARE' in x, args)

        # Second DECLARATION of cursor should start from last row fetched
        self.assertTrue(
            'WHERE ("col1","col2") >= (\'foo3\',\'bar3\')' in args[1])
Beispiel #14
0
    def test_read_from_other_schema(self, mock_connect, mock_metadata, __):

        inst = Postgres(self.source, OPTIONS)
        inst.tables = [{
            'value': 'my_schema.foo_bar'
        }, {
            'value': 'your_schema.bar_foo'
        }]
        cursor_return_value = mock_connect.return_value.cursor.return_value
        mock_data = [self.mock_recs[:1], []] * 2
        cursor_return_value.fetchall.side_effect = mock_data

        expected = [('my_schema', 'foo_bar'), ('my_schema', 'foo_bar'),
                    ('your_schema', 'bar_foo'), ('your_schema', 'bar_foo')]

        for expected_schema, expected_table in expected:
            inst.read()
            _, schema, table = mock_metadata.call_args[0]

            self.assertEqual(schema, expected_schema)
            self.assertEqual(table, expected_table)
Beispiel #15
0
    def test_reports_state(self, mock_connect, mock_state):
        '''before returning a batch of data, the sources state should be
        reported as well as having the state ID appended to each data object'''

        inst = Postgres(self.source, OPTIONS)
        table_name = 'my_schema.foo_bar'
        inst.tables = [{'value': table_name}]
        result_order = [self.mock_recs, []]
        cursor_return_value = mock_connect.return_value.cursor.return_value
        cursor_return_value.fetchall.side_effect = result_order

        rows = inst.read()
        state_id = rows[0]['__state']
        state_obj = dict([(table_name, len(self.mock_recs))])

        msg = 'State ID is not the same in all rows!'
        for row in rows:
            self.assertEqual(row['__state'], state_id, msg)

        # State function was called with relevant table name and row count
        mock_state.assert_called_with(state_id, state_obj)
Beispiel #16
0
    def test_query_with_no_keys(self, mock_connect, _):
        inst = Postgres(self.source, OPTIONS)
        inst.tables = [{'value': 'my_schema.foo_bar'}]

        cursor_return_value = mock_connect.return_value.cursor.return_value
        cursor_return_value.fetchall.side_effect = [
            [],
            [
                {
                    'attname': 'id',
                    'data_type': 'integer'
                },
                {
                    'attname': 'name',
                    'data_type': 'text'
                },
            ],
            [],
        ]

        cursor_execute = mock_connect.return_value.cursor.return_value.execute
        cursor_execute.return_value = lambda *args: None

        inst.read()

        # Extract mock call arguments
        args = mock_connect.return_value.cursor.return_value \
            .execute.call_args_list
        args = [r[0] for r, _ in args]
        args = filter(lambda x: 'DECLARE' in x, args)

        expected = 'DECLARE cur ' \
                   'CURSOR FOR SELECT * FROM "my_schema"."foo_bar" ' \
                   "WHERE (\"inckey\" >= 'incval' AND \"inckey\" <= '100') " \
                   'ORDER BY "id","inckey"'

        self.assertEqual(args[0], expected)
Beispiel #17
0
 def test_connect_error(self, m):
     inst = Postgres(self.source, OPTIONS)
     inst.tables = [{'value': 'schema.foo'}]
     m.side_effect = psycopg2.OperationalError('Mock Error')
     with self.assertRaises(PanoplyException):
         inst.read()
Beispiel #18
0
    def test_read_end_stream(self, mock_connect, mock_execute, mock_metadata,
                             _):
        """reads the entire table from the database and validates that the
        stream returns None to indicate the end"""
        tables = [
            {
                'value': 'public.table1'
            },
            {
                'value': 'public.table2'
            },
            {
                'value': 'public.table3'
            },
        ]

        mock_metadata.side_effect = [
            [{
                'attname': 'col1'
            }],
            [{
                'attname': 'col2'
            }],
            [{
                'attname': 'col3'
            }],
        ]

        inst = Postgres(self.source, OPTIONS)
        inst.tables = tables
        result_order = [
            self.mock_recs, [], self.mock_recs, [], self.mock_recs, []
        ]

        cursor_return_value = mock_connect.return_value.cursor.return_value
        cursor_return_value.fetchall.side_effect = result_order

        # First call to read
        result = inst.read()
        self.assertEqual(len(result), len(self.mock_recs))
        query = mock_execute.call_args_list[0][0][0]
        expected_query = 'FROM "public"."table1" ' \
                         'WHERE ("inckey" >= \'incval\' ' \
                         'AND "inckey" <= \'100\') '\
                         'ORDER BY "col1","inckey"'
        self.assertTrue(expected_query in query)
        query = mock_execute.call_args_list[1][0][0]
        expected_query = 'FETCH FORWARD'
        self.assertTrue(expected_query in query)

        # Second call to read
        result = inst.read()
        self.assertEqual(result, [])
        query = mock_execute.call_args_list[2][0][0]
        expected_query = 'FETCH FORWARD'
        self.assertTrue(expected_query in query)

        # Third call to read
        result = inst.read()
        self.assertEqual(len(result), len(self.mock_recs))
        query = mock_execute.call_args_list[3][0][0]
        expected_query = 'FROM "public"."table2" ' \
                         'WHERE ("inckey" >= \'incval\' ' \
                         'AND "inckey" <= \'100\') '\
                         'ORDER BY "col2","inckey"'
        self.assertTrue(expected_query in query)
        query = mock_execute.call_args_list[4][0][0]
        expected_query = 'FETCH FORWARD'
        self.assertTrue(expected_query in query)

        # Fourth call to read
        result = inst.read()
        self.assertEqual(result, [])
        query = mock_execute.call_args_list[5][0][0]
        expected_query = 'FETCH FORWARD'
        self.assertTrue(expected_query in query)

        # Fifth call to read
        result = inst.read()
        self.assertEqual(len(result), len(self.mock_recs))
        query = mock_execute.call_args_list[6][0][0]
        expected_query = 'FROM "public"."table3" ' \
                         'WHERE ("inckey" >= \'incval\' ' \
                         'AND "inckey" <= \'100\') '\
                         'ORDER BY "col3","inckey"'
        self.assertTrue(expected_query in query)
        query = mock_execute.call_args_list[7][0][0]
        expected_query = 'FETCH FORWARD'
        self.assertTrue(expected_query in query)

        # Sixth call to read
        result = inst.read()
        self.assertEqual(result, [])
        query = mock_execute.call_args_list[8][0][0]
        expected_query = 'FETCH FORWARD'
        self.assertTrue(expected_query in query)

        end = inst.read()
        self.assertEqual(end, None)