Avoid saving sensitive information in graph.

4 files changed, 56 insertions, 68 deletions

diff --git a/tensorflow/contrib/ignite/kernels/ignite_dataset_ops.cc b/tensorflow/contrib/ignite/kernels/ignite_dataset_ops.cc
index e48fce4ed2..bdaed72387 100644
--- a/tensorflow/contrib/ignite/kernels/ignite_dataset_ops.cc
+++ b/tensorflow/contrib/ignite/kernels/ignite_dataset_ops.cc
@@ -125,35 +125,15 @@ class IgniteDatasetOp : public DatasetOpKernel {
       OP_REQUIRES_OK(ctx,
                      ParseScalarArgument<int32>(ctx, "page_size", &page_size));
 
-    if (env_username)
-      username = string(env_username);
-    else
-      OP_REQUIRES_OK(ctx,
-                     ParseScalarArgument<string>(ctx, "username", &username));
+    if (env_username) username = string(env_username);
 
-    if (env_password)
-      password = string(env_password);
-    else
-      OP_REQUIRES_OK(ctx,
-                     ParseScalarArgument<string>(ctx, "password", &password));
+    if (env_password) password = string(env_password);
 
-    if (env_certfile)
-      certfile = string(env_certfile);
-    else
-      OP_REQUIRES_OK(ctx,
-                     ParseScalarArgument<string>(ctx, "certfile", &certfile));
+    if (env_certfile) certfile = string(env_certfile);
 
-    if (env_keyfile)
-      keyfile = string(env_keyfile);
-    else
-      OP_REQUIRES_OK(ctx,
-                     ParseScalarArgument<string>(ctx, "keyfile", &keyfile));
+    if (env_keyfile) keyfile = string(env_keyfile);
 
-    if (env_cert_password)
-      cert_password = string(env_cert_password);
-    else
-      OP_REQUIRES_OK(ctx, ParseScalarArgument<string>(ctx, "cert_password",
-                                                      &cert_password));
+    if (env_cert_password) cert_password = string(env_cert_password);
 
     const Tensor* schema_tensor;
     OP_REQUIRES_OK(ctx, ctx->input("schema", &schema_tensor));
diff --git a/tensorflow/contrib/ignite/ops/dataset_ops.cc b/tensorflow/contrib/ignite/ops/dataset_ops.cc
index 7d18df11aa..3d6fbe00e6 100644
--- a/tensorflow/contrib/ignite/ops/dataset_ops.cc
+++ b/tensorflow/contrib/ignite/ops/dataset_ops.cc
@@ -26,11 +26,6 @@ REGISTER_OP("IgniteDataset")
     .Input("local: bool")
     .Input("part: int32")
     .Input("page_size: int32")
-    .Input("username: string")
-    .Input("password: string")
-    .Input("certfile: string")
-    .Input("keyfile: string")
-    .Input("cert_password: string")
     .Input("schema: int32")
     .Input("permutation: int32")
     .Output("handle: variant")
@@ -54,11 +49,6 @@ port: Ignite Thin Client Port.
 local: Local flag that defines that data should be fetched from local host only.
 part: Partition data should be fetched from.
 page_size: Page size for Ignite Thin Client.
-username: Username to authenticate via Ignite Thin Client.
-password: Password to authenticate via Ignite Thin Client.
-certfile: SSL certificate to establish SSL connection.
-keyfile: Private key file to establish SSL connection.
-cert_password: SSL certificate password to establish SSL connection.
 schema: Internal structure that defines schema of cache objects.
 permutation: Internal structure that defines permutation of cache objects.
 )doc");
diff --git a/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py b/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py
index c0e24b1c69..7fc9e1fdd1 100644
--- a/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py
+++ b/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py
@@ -732,18 +732,6 @@ class IgniteDataset(Dataset):
     self.part = ops.convert_to_tensor(part, dtype=dtypes.int32, name="part")
     self.page_size = ops.convert_to_tensor(page_size, dtype=dtypes.int32,
                                            name="page_size")
-    self.username = ops.convert_to_tensor("" if username is None else username,
-                                          dtype=dtypes.string, name="username")
-    self.password = ops.convert_to_tensor("" if password is None else password,
-                                          dtype=dtypes.string, name="password")
-    self.certfile = ops.convert_to_tensor("" if certfile is None else certfile,
-                                          dtype=dtypes.string, name="certfile")
-    self.keyfile = ops.convert_to_tensor("" if keyfile is None else keyfile,
-                                         dtype=dtypes.string, name="keyfile")
-    self.cert_password = ops.convert_to_tensor("" if cert_password is None
-                                               else cert_password,
-                                               dtype=dtypes.string,
-                                               name="cert_password")
     self.schema = ops.convert_to_tensor(self.cache_type.to_flat(),
                                         dtype=dtypes.int32, name="schema")
     self.permutation = ops.convert_to_tensor(self.cache_type.to_permutation(),
@@ -753,10 +741,8 @@ class IgniteDataset(Dataset):
   def _as_variant_tensor(self):
     return gen_dataset_ops.ignite_dataset(self.cache_name, self.host,
                                           self.port, self.local, self.part,
-                                          self.page_size, self.username,
-                                          self.password, self.certfile,
-                                          self.keyfile, self.cert_password,
-                                          self.schema, self.permutation)
+                                          self.page_size, self.schema,
+                                          self.permutation)
 
   @property
   def output_classes(self):
diff --git a/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py b/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py
index 933e62b804..5d74617690 100644
--- a/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py
+++ b/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py
@@ -35,28 +35,60 @@ class IgniteDatasetTest(test.TestCase):
   """
 
   def test_ignite_dataset_with_plain_client(self):
+    """Test Ignite Dataset with plain client.
+    """
+    self._clear_env()
     ds = IgniteDataset(cache_name="SQL_PUBLIC_TEST_CACHE", port=42300)
-    self.__check_dataset(ds)
+    self._check_dataset(ds)
 
   def test_ignite_dataset_with_ssl_client(self):
-    ds = IgniteDataset(cache_name="SQL_PUBLIC_TEST_CACHE", port=42301,\
-      certfile=os.path.dirname(os.path.realpath(__file__)) +\
-      "/keystore/client.pem", cert_password="123456")
-    self.__check_dataset(ds)
+    """Test Ignite Dataset with ssl client.
+    """
+    self._clear_env()
+    os.environ["IGNITE_DATASET_CERTFILE"] = os.path.dirname(
+        os.path.realpath(__file__)) + "/keystore/client.pem"
+    os.environ["IGNITE_DATASET_CERT_PASSWORD"] = "123456"
+
+    ds = IgniteDataset(cache_name="SQL_PUBLIC_TEST_CACHE", port=42301,
+                       certfile=os.environ["IGNITE_DATASET_CERTFILE"],
+                       cert_password=os.environ["IGNITE_DATASET_CERT_PASSWORD"])
+    self._check_dataset(ds)
 
   def test_ignite_dataset_with_ssl_client_and_auth(self):
-    ds = IgniteDataset(cache_name="SQL_PUBLIC_TEST_CACHE", port=42302,\
-      certfile=os.path.dirname(os.path.realpath(__file__)) +\
-      "/keystore/client.pem", cert_password="123456",\
-      username="ignite", password="ignite")
-    self.__check_dataset(ds)
+    """Test Ignite Dataset with ssl client and authentication.
+    """
+    self._clear_env()
+    os.environ['IGNITE_DATASET_USERNAME'] = "ignite"
+    os.environ['IGNITE_DATASET_PASSWORD'] = "ignite"
+    os.environ['IGNITE_DATASET_CERTFILE'] = os.path.dirname(
+        os.path.realpath(__file__)) + "/keystore/client.pem"
+    os.environ['IGNITE_DATASET_CERT_PASSWORD'] = "123456"
+
+    ds = IgniteDataset(cache_name="SQL_PUBLIC_TEST_CACHE", port=42302,
+                       certfile=os.environ['IGNITE_DATASET_CERTFILE'],
+                       cert_password=os.environ['IGNITE_DATASET_CERT_PASSWORD'],
+                       username=os.environ['IGNITE_DATASET_USERNAME'],
+                       password=os.environ['IGNITE_DATASET_PASSWORD'])
+    self._check_dataset(ds)
+
+  def _clear_env(self):
+    """Clears environment variables used by Ignite Dataset.
+    """
+    if 'IGNITE_DATASET_USERNAME' in os.environ:
+      del os.environ['IGNITE_DATASET_USERNAME']
+    if 'IGNITE_DATASET_PASSWORD' in os.environ:
+      del os.environ['IGNITE_DATASET_PASSWORD']
+    if 'IGNITE_DATASET_CERTFILE' in os.environ:
+      del os.environ['IGNITE_DATASET_CERTFILE']
+    if 'IGNITE_DATASET_CERT_PASSWORD' in os.environ:
+      del os.environ['IGNITE_DATASET_CERT_PASSWORD']
 
-  def __check_dataset(self, dataset):
+  def _check_dataset(self, dataset):
     """Checks that dataset provids correct data.
     """
-    self.assertEquals(tf.int64, dataset.output_types['key'])
-    self.assertEquals(tf.string, dataset.output_types['val']['NAME'])
-    self.assertEquals(tf.int64, dataset.output_types['val']['VAL'])
+    self.assertEqual(tf.int64, dataset.output_types['key'])
+    self.assertEqual(tf.string, dataset.output_types['val']['NAME'])
+    self.assertEqual(tf.int64, dataset.output_types['val']['VAL'])
 
     it = dataset.make_one_shot_iterator()
     ne = it.get_next()
@@ -66,11 +98,11 @@ class IgniteDatasetTest(test.TestCase):
       with self.assertRaises(errors.OutOfRangeError):
         sess.run(ne)
 
-    self.assertEquals({'key': 1, 'val': {'NAME': b'TEST1', 'VAL': 42}},\
+    self.assertEqual({'key': 1, 'val': {'NAME': b'TEST1', 'VAL': 42}},\
       rows[0])
-    self.assertEquals({'key': 2, 'val': {'NAME': b'TEST2', 'VAL': 43}},\
+    self.assertEqual({'key': 2, 'val': {'NAME': b'TEST2', 'VAL': 43}},\
       rows[1])
-    self.assertEquals({'key': 3, 'val': {'NAME': b'TEST3', 'VAL': 44}},\
+    self.assertEqual({'key': 3, 'val': {'NAME': b'TEST3', 'VAL': 44}},\
       rows[2])
 
 if __name__ == "__main__":