remove SourceType and check; [https://github.com/MarquezProject/marqu… (#104)

* remove SourceType and check; [MarquezProject/marquez-airflow#79]

Signed-off-by: SreeV <[email protected]>

* no value on source type

Signed-off-by: SreeV <[email protected]>

* remove reference to SourceType

Signed-off-by: SreeV <[email protected]>

* no value on source type

Signed-off-by: SreeV <[email protected]>

Author: sreev

marquez_client/client.py

@@ -18,7 +18,7 @@
 
 from marquez_client import errors
 from marquez_client.constants import (DEFAULT_TIMEOUT_MS)
-from marquez_client.models import DatasetType, SourceType, JobType
+from marquez_client.models import (DatasetType, JobType)
 from marquez_client.utils import Utils
 from marquez_client.version import VERSION
 
@@ -75,12 +75,11 @@ def list_namespaces(self, limit=None, offset=None):
     def create_source(self, source_name, source_type, connection_url,
                       description=None):
         Utils.check_name_length(source_name, 'source_name')
-        Utils.is_instance_of(source_type, SourceType)
 
         Utils.is_valid_connection_url(connection_url)
 
         payload = {
-            'type': source_type.value,
+            'type': source_type,
             'connectionUrl': connection_url
         }
 

marquez_client/client_wo.py

@@ -16,7 +16,7 @@
 from six.moves.urllib.parse import quote
 
 from marquez_client.version import VERSION
-from .models import DatasetType, SourceType, JobType
+from .models import (DatasetType, JobType)
 from .utils import Utils
 
 _API_PATH = '/api/v1'
@@ -53,12 +53,11 @@ def create_namespace(self, namespace_name, owner_name, description=None):
     def create_source(self, source_name, source_type, connection_url,
                       description=None):
         Utils.check_name_length(source_name, 'source_name')
-        Utils.is_instance_of(source_type, SourceType)
 
         Utils.is_valid_connection_url(connection_url)
 
         payload = {
-            'type': source_type.value,
+            'type': source_type,
             'connectionUrl': connection_url
         }
 

marquez_client/models.py

@@ -13,15 +13,6 @@
 from enum import Enum
 
 
-class SourceType(Enum):
-    MYSQL = "MYSQL"
-    POSTGRESQL = "POSTGRESQL"
-    REDSHIFT = "REDSHIFT"
-    SNOWFLAKE = "SNOWFLAKE"
-    KAFKA = "KAFKA"
-    REST = "REST"
-
-
 class DatasetType(Enum):
     DB_TABLE = "DB_TABLE"
     STREAM = "STREAM"

marquez_client/utils.py

@@ -17,7 +17,7 @@
 
 from datetime import datetime
 from pyrfc3339 import generate
-from marquez_client.models import DatasetFieldType, DatasetType
+from marquez_client.models import (DatasetFieldType, DatasetType)
 
 log = logging.getLogger(__name__)
 

tests/test_airflow_dag.py

@@ -20,7 +20,7 @@
 import yaml
 
 from marquez_client import Clients
-from marquez_client.models import (SourceType, DatasetType, JobType)
+from marquez_client.models import (DatasetType, JobType)
 
 from marquez_client.utils import Utils
 
@@ -58,7 +58,7 @@ def test_create_dag(self):
             self.client.create_namespace(NAMESPACE, OWNER)
             self.client.create_source(
                 SOURCE,
-                SourceType.POSTGRESQL,
+                'POSTGRESQL',
                 "jdbc:postgresql://localhost:5432/test?user=fred&ssl=true")
             self.client.create_dataset(
                 NAMESPACE, DATASET, DatasetType.DB_TABLE,

tests/test_marquez_client.py

@@ -19,7 +19,7 @@
 import yaml
 
 import marquez_client
-from marquez_client.models import DatasetType, SourceType, JobType, RunState
+from marquez_client.models import (DatasetType, JobType, RunState)
 from marquez_client.utils import Utils
 
 _NAMESPACE = "my-namespace"
@@ -175,7 +175,7 @@ def test_get_dataset(self, mock_get):
     @mock.patch("marquez_client.client.MarquezClient._put")
     def test_create_datasource(self, mock_put):
         source_name = "flight_schedules_db"
-        source_type = SourceType.POSTGRESQL
+        source_type = 'POSTGRESQL'
         source_url = "jdbc:postgresql://localhost:5432/test?" \
                      "user=fred&password=secret&ssl=true"
         description = "PostgreSQL - flight schedules database"
@@ -198,7 +198,7 @@ def test_create_datasource(self, mock_put):
 
         assert response['name'] == source_name
         assert response['connectionUrl'] == source_url
-        assert response['type'] == source_type.value
+        assert response['type'] == source_type
 
     @mock.patch("marquez_client.client.MarquezClient._put")
     def test_create_job(self, mock_put):

tests/test_marquez_write_only_client_file.py

@@ -9,18 +9,17 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import datetime
-import unittest
-
-from marquez_client.models import DatasetType, SourceType, JobType, RunState
-from marquez_client import Clients
 
-import uuid
 import logging
 import logging.config
-import yaml
 import os
+import unittest
+import uuid
 
+import yaml
+
+from marquez_client import Clients
+from marquez_client.models import (DatasetType, JobType)
 from marquez_client.utils import Utils
 
 _NAMESPACE = "my-namespace"
@@ -86,7 +85,7 @@ def test_create_dataset(self):
 
     def test_create_datasource(self):
         source_name = "flight_schedules_db"
-        source_type = SourceType.POSTGRESQL
+        source_type = 'POSTGRESQL'
         source_url = "jdbc:postgresql://localhost:5432/test?" \
                      "user=fred&password=secret&ssl=true"
         description = "PostgreSQL - flight schedules database"

tests/test_marquez_write_only_client_http.py

@@ -9,19 +9,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import datetime
-import unittest
-
-from marquez_client.models import DatasetType, SourceType, JobType, RunState
-from marquez_client import Clients
 
-import uuid
 import logging
 import logging.config
-import yaml
-import mock
 import os
+import unittest
+import uuid
+
+import mock
+import yaml
 
+from marquez_client import Clients
+from marquez_client.models import (DatasetType, JobType)
 from marquez_client.utils import Utils
 
 _NAMESPACE = "my-namespace"
@@ -88,7 +87,7 @@ def test_create_dataset(self, mock_put):
     @mock.patch("marquez_client.http_backend.HttpBackend.put")
     def test_create_datasource(self, mock_put):
         source_name = "flight_schedules_db"
-        source_type = SourceType.POSTGRESQL
+        source_type = 'POSTGRESQL'
         source_url = "jdbc:postgresql://localhost:5432/test?" \
                      "user=fred&password=secret&ssl=true"
         description = "PostgreSQL - flight schedules database"