Fix DataFrameClient empty tag processing

tzonghao · tzonghao · commit c6334f8cbb13 · 2016-09-09T11:56:27.000-04:00
Empty or None tags should be left out, otherwise it would break line
protocol.
diff --git a/influxdb/_dataframe_client.py b/influxdb/_dataframe_client.py
@@ -279,13 +279,21 @@ def _convert_dataframe_to_lines(self,
         # If tag columns exist, make an array of formatted tag keys and values
         if tag_columns:
             tag_df = dataframe[tag_columns]
+            tag_df = tag_df.fillna('')  # replace NA with empty string
             tag_df = tag_df.sort_index(axis=1)
             tag_df = self._stringify_dataframe(
                 tag_df, numeric_precision, datatype='tag')
-            tags = (',' + (
-                (tag_df.columns.values + '=').tolist() + tag_df)).sum(axis=1)
-            del tag_df
 
+            # prepend tag keys
+            tag_df = tag_df.apply(
+                lambda s: s.apply(
+                    lambda v, l: l + '=' + v if v else None, l=s.name))
+
+            # join tags, but leave out None values
+            tags = tag_df.apply(
+                lambda r: ',' + ','.join(r.dropna()), axis=1)
+
+            del tag_df
         else:
             tags = ''
 
diff --git a/influxdb/tests/dataframe_client_test.py b/influxdb/tests/dataframe_client_test.py
@@ -201,20 +201,29 @@ def test_write_points_from_dataframe_with_tag_cols_and_defaults(self):
     def test_write_points_from_dataframe_with_tag_escaped(self):
         now = pd.Timestamp('1970-01-01 00:00+00:00')
         dataframe = pd.DataFrame(
-            data=[['blue', 1, "1", 1, 1.0, 'hot'],
-                  ['red,green=orange', 0, "2", 2, 2.0, 'cold']],
-            index=[now, now + timedelta(hours=1)],
-            columns=["tag_one", "tag_two", "column_one",
-                     "column_two", "column_three",
-                     "tag_three"])
+            data=[
+                ['blue orange', "1", 1, 'hot=cold'],  # space, equal
+                ['red,green', "2", 2, r'cold\fire'],  # comma, backslash
+                ['some', "2", 2, ''],                 # skip empty
+                ['some', "2", 2, None],               # skip None
+            ],
+            index=pd.period_range(now, freq='H', periods=4),
+            columns=["tag_one", "column_one", "column_two", "tag_three"]
+        )
 
         expected_escaped_tags = (
-            b"foo,tag_one=blue "
+            b"foo,tag_one=blue\\ orange,tag_three=hot\\=cold "
             b"column_one=\"1\",column_two=1i "
             b"0\n"
-            b"foo,tag_one=red\\,green\\=orange "
+            b"foo,tag_one=red\\,green,tag_three=cold\\\\fire "
             b"column_one=\"2\",column_two=2i "
             b"3600000000000\n"
+            b"foo,tag_one=some "
+            b"column_one=\"2\",column_two=2i "
+            b"7200000000000\n"
+            b"foo,tag_one=some "
+            b"column_one=\"2\",column_two=2i "
+            b"10800000000000\n"
         )
 
         with requests_mock.Mocker() as m:
@@ -224,7 +233,7 @@ def test_write_points_from_dataframe_with_tag_escaped(self):
             cli = DataFrameClient(database='db')
             cli.write_points(dataframe, 'foo',
                              field_columns=['column_one', 'column_two'],
-                             tag_columns=['tag_one'])
+                             tag_columns=['tag_one', 'tag_three'])
             self.assertEqual(m.last_request.body, expected_escaped_tags)
 
     def test_write_points_from_dataframe_with_numeric_column_names(self):
